Add feature toggle for enabling arrow database import

Co-Authored-By: Paul Horn <[email protected]>

Author: s1ck

core/src/main/java/org/neo4j/gds/utils/GdsFeatureToggles.java

@@ -34,7 +34,8 @@ public enum GdsFeatureToggles {
     USE_BIT_ID_MAP(true),
     USE_SHARDED_ID_MAP(false),
     USE_UNCOMPRESSED_ADJACENCY_LIST(false),
-    USE_REORDERED_ADJACENCY_LIST(false);
+    USE_REORDERED_ADJACENCY_LIST(false),
+    ENABLE_ARROW_DATABASE_IMPORT(true);
 
     public boolean isEnabled() {
         return current.get();

proc/misc/src/main/java/org/neo4j/gds/FeatureToggleProc.java

@@ -152,6 +152,21 @@ public Stream<FeatureState> resetUseReorderedAdjacencyList() {
         return Stream.of(new FeatureState(GdsFeatureToggles.USE_REORDERED_ADJACENCY_LIST.isEnabled()));
     }
 
+    @Internal
+    @Procedure("gds.features.enableArrowDatabaseImport")
+    @Description("Enables support for importing Neo4j databases via the GDS Arrow Flight Server.")
+    public void enableArrowDatabaseImport(@Name(value = "useReorderedAdjacencyList") boolean useReorderedAdjacencyList) {
+        GdsFeatureToggles.ENABLE_ARROW_DATABASE_IMPORT.toggle(useReorderedAdjacencyList);
+    }
+
+    @Internal
+    @Procedure("gds.features.enableArrowDatabaseImport.reset")
+    @Description("Sets the default behaviour for enabling Neo4j database import via the GDS Arrow Flight Server. That value is returned.")
+    public Stream<FeatureState> resetEnableArrowDatabaseImport() {
+        GdsFeatureToggles.ENABLE_ARROW_DATABASE_IMPORT.reset();
+        return Stream.of(new FeatureState(GdsFeatureToggles.ENABLE_ARROW_DATABASE_IMPORT.isEnabled()));
+    }
+
     @Internal
     @Procedure("gds.features.maxArrayLengthShift.reset")
     @Description("Set the value of the max array size before paging to the default. That value is returned.")

proc/misc/src/test/java/org/neo4j/gds/FeatureToggleProcTest.java

@@ -31,6 +31,8 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.neo4j.gds.utils.GdsFeatureToggles.ENABLE_ARROW_DATABASE_IMPORT;
 import static org.neo4j.gds.utils.GdsFeatureToggles.SKIP_ORPHANS;
 import static org.neo4j.gds.utils.GdsFeatureToggles.USE_PARALLEL_PROPERTY_VALUE_INDEX;
 import static org.neo4j.gds.utils.GdsFeatureToggles.USE_PARTITIONED_SCAN;
@@ -165,6 +167,25 @@ void resetUseReorderedAdjacencyList() {
         assertFalse(USE_REORDERED_ADJACENCY_LIST.isEnabled());
     }
 
+    @Test
+    void toggleEnableArrowDatabaseImport() {
+        var enableArrowDatabaseImport = ENABLE_ARROW_DATABASE_IMPORT.isEnabled();
+        runQuery("CALL gds.features.enableArrowDatabaseImport($value)", Map.of("value", !enableArrowDatabaseImport));
+        assertEquals(!enableArrowDatabaseImport, ENABLE_ARROW_DATABASE_IMPORT.isEnabled());
+        runQuery("CALL gds.features.enableArrowDatabaseImport($value)", Map.of("value", enableArrowDatabaseImport));
+        assertEquals(enableArrowDatabaseImport, ENABLE_ARROW_DATABASE_IMPORT.isEnabled());
+    }
+
+    @Test
+    void resetEnableArrowDatabaseImport() {
+        ENABLE_ARROW_DATABASE_IMPORT.reset();
+        assertCypherResult(
+            "CALL gds.features.enableArrowDatabaseImport.reset()",
+            List.of(Map.of("enabled", true))
+        );
+        assertTrue(ENABLE_ARROW_DATABASE_IMPORT.isEnabled());
+    }
+
     @Test
     void toggleMaxArrayLengthShift() {
         var maxArrayLengthShift = GdsFeatureToggles.MAX_ARRAY_LENGTH_SHIFT.get();