feat: use external header file in admin import demo

Author: fbiville

examples/neo4j-admin/src/test/java/org/neo4j/importer/Neo4jAdminExampleIT.java

@@ -21,8 +21,8 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.io.*;
 import java.net.URL;
+import java.nio.file.Files;
 import java.sql.DriverManager;
-import java.sql.SQLException;
 import java.util.*;
 import java.util.function.Function;
 import java.util.stream.Collectors;
@@ -61,17 +61,13 @@ public class Neo4jAdminExampleIT {
 
     private static final String SHARED_FOLDER = "/admin-import/";
 
-    // neo4j-admin database import from Parquet files is an unreleased feature, and we are using a custom docker image
-    // for that reason, until the feature is publicly available. Set this image name using
-    // `NEO4J_PRE_RELEASE_DOCKER_IMAGE` environment variable.
     @Container
-    private static final GenericContainer<?> NEO4J = new Neo4jContainer<>(DockerImageName.parse(
-                            Optional.ofNullable(System.getenv("NEO4J_PRE_RELEASE_DOCKER_IMAGE"))
-                                    .orElseThrow(
-                                            () -> new IllegalArgumentException(
-                                                    "Docker image name is not set through NEO4J_PRE_RELEASE_DOCKER_IMAGE environment variable!")))
-                    .asCompatibleSubstituteFor("neo4j"))
+    private static final GenericContainer<?> NEO4J = new Neo4jContainer<>(
+                    DockerImageName.parse("neo4j:2025.03-enterprise"))
             .withEnv("NEO4J_ACCEPT_LICENSE_AGREEMENT", "yes")
+            .withNeo4jConfig("dbms.integrations.cloud_storage.gs.project_id", "connectors-public")
+            .withNeo4jConfig("server.config.strict_validation.enabled", "false")
+            .withNeo4jConfig("internal.dbms.cloud.storage.gs.host", "https://storage.googleapis.com")
             .withAdminPassword("letmein!")
             .withCreateContainerCmdModifier(cmd -> cmd.withUser("neo4j"))
             .withFileSystemBind(
@@ -106,7 +102,7 @@ void runs_an_offline_import_of_dvd_rental_data_set() throws Exception {
                 var specification = ImportSpecificationDeserializer.deserialize(reader);
                 var sharedFolder = pathFor(SHARED_FOLDER);
                 var neo4jAdmin = new Neo4jAdmin(sharedFolder, driver, TARGET_DATABASE);
-                neo4jAdmin.copyFiles(specification);
+                neo4jAdmin.createHeaderFiles(specification);
                 neo4jAdmin.executeImport(specification, NEO4J);
             }
         }
@@ -338,35 +334,35 @@ private void migrateTemporalProperties() {
             }
         }
 
-        public void copyFiles(ImportSpecification specification) throws Exception {
+        public void createHeaderFiles(ImportSpecification specification) throws Exception {
             Map<String, Source> indexedSources =
                     specification.getSources().stream().collect(Collectors.toMap(Source::getName, Function.identity()));
             Map<String, NodeTarget> indexedNodes = specification.getTargets().getNodes().stream()
                     .collect(Collectors.toMap(Target::getName, Function.identity()));
             for (Target target : specification.getTargets().getAll()) {
                 switch (target) {
-                    case NodeTarget nodeTarget -> copyFile(indexedSources, nodeTarget);
+                    case NodeTarget nodeTarget -> createHeaderFile(indexedSources, nodeTarget);
                     case RelationshipTarget relationshipTarget ->
-                        copyFile(indexedSources, indexedNodes, relationshipTarget);
+                        createHeaderFile(indexedSources, indexedNodes, relationshipTarget);
                     default -> throw new RuntimeException("unsupported target type: %s".formatted(target.getClass()));
                 }
             }
         }
 
-        private void copyFile(Map<String, Source> sources, NodeTarget nodeTarget) throws Exception {
+        private void createHeaderFile(Map<String, Source> sources, NodeTarget nodeTarget) throws Exception {
             var source = sources.get(nodeTarget.getSource());
             assertThat(source).isInstanceOf(ParquetSource.class);
-            File parquetFile = new File(sharedFolder, fileName(nodeTarget));
+            File parquetFile = new File(sharedFolder, headerFileName(nodeTarget));
             List<String> fields = readFieldNames(source);
             Map<String, String> fieldMappings = computeFieldMappings(fields, nodeTarget);
 
-            copyParquetSource((ParquetSource) source, parquetFile, fieldMappings);
+            createHeaderFile(parquetFile, fieldMappings);
         }
 
-        private void copyFile(
+        private void createHeaderFile(
                 Map<String, Source> sources, Map<String, NodeTarget> nodes, RelationshipTarget relationshipTarget)
                 throws Exception {
-            File parquetFile = new File(sharedFolder, fileName(relationshipTarget));
+            File parquetFile = new File(sharedFolder, headerFileName(relationshipTarget));
 
             var source = sources.get(relationshipTarget.getSource());
             assertThat(source).isInstanceOf(ParquetSource.class);
@@ -377,26 +373,14 @@ private void copyFile(
             Map<String, String> fieldMappings =
                     computeFieldMappings(fields, relationshipTarget, startNodeTarget, endNodeTarget);
 
-            copyParquetSource((ParquetSource) source, parquetFile, fieldMappings);
+            createHeaderFile(parquetFile, fieldMappings);
         }
 
-        // 🐤
-        private void copyParquetSource(ParquetSource source, File targetFile, Map<String, String> fieldMappings)
-                throws SQLException {
-            var renamedColumns = String.join(
-                    ", ",
-                    fieldMappings.entrySet().stream()
-                            .map(e -> String.format("%s AS \"%s\"", e.getKey(), e.getValue()))
-                            .toList());
-
-            try (var connection = DriverManager.getConnection("jdbc:duckdb:");
-                    var statement = connection.prepareStatement(String.format(
-                            "COPY (SELECT %s FROM read_parquet($1)) TO '%s' (FORMAT 'parquet', CODEC 'zstd')",
-                            renamedColumns, targetFile.getAbsolutePath()))) {
+        private void createHeaderFile(File targetFile, Map<String, String> fieldMappings) throws IOException {
 
-                statement.setString(1, source.uri());
-                statement.execute();
-            }
+            var sortedKeys = fieldMappings.keySet().stream().sorted().toList();
+            var values = sortedKeys.stream().map(fieldMappings::get).collect(Collectors.joining(","));
+            Files.writeString(targetFile.toPath(), values + "\n" + String.join(",", sortedKeys));
         }
 
         private static String[] importCommand(ImportSpecification specification, String database) {
@@ -409,14 +393,16 @@ private static String[] importCommand(ImportSpecification specification, String
                 command.append(" --nodes=");
                 command.append(String.join(":", nodeTarget.getLabels()));
                 command.append("=");
-                command.append("/import/%s".formatted(fileName(nodeTarget)));
+                command.append("/import/%s,".formatted(headerFileName(nodeTarget)));
+                command.append("%s".formatted(sourceUri(specification, nodeTarget)));
             }
 
             for (RelationshipTarget relationshipTarget : targets.getRelationships()) {
                 command.append(" --relationships=");
                 command.append(relationshipTarget.getType());
                 command.append("=");
-                command.append("/import/%s".formatted(fileName(relationshipTarget)));
+                command.append("/import/%s,".formatted(headerFileName(relationshipTarget)));
+                command.append("%s".formatted(sourceUri(specification, relationshipTarget)));
             }
 
             return command.toString().split(" ");
@@ -493,8 +479,19 @@ private static Map<String, String> indexByField(List<PropertyMapping> properties
             return result;
         }
 
-        private static String fileName(Target target) {
-            return "%s.parquet".formatted(target.getName());
+        private static String headerFileName(Target target) {
+            return "%s_header.csv".formatted(target.getName());
+        }
+
+        private static String sourceUri(ImportSpecification specification, Target target) {
+            var maybeSource = specification.getSources().stream()
+                    .filter(src -> src.getName().equals(target.getSource()))
+                    .findFirst();
+            assertThat(maybeSource).isPresent();
+            var rawSource = maybeSource.get();
+            assertThat(rawSource).isInstanceOf(ParquetSource.class);
+            var source = (ParquetSource) rawSource;
+            return source.uri();
         }
 
         private static String idSpaceFor(NodeTarget nodeTarget) {

examples/neo4j-admin/src/test/resources/specs/dvd_rental.yaml

@@ -2,28 +2,28 @@ version: "1"
 sources:
   - name: actor
     type: parquet
-    uri: gs://connectors-dev/dvdrental/actor.parquet
+    uri: gs://connectors-public/dvdrental/actor.parquet
   - name: category
     type: parquet
-    uri: gs://connectors-dev/dvdrental/category.parquet
+    uri: gs://connectors-public/dvdrental/category.parquet
   - name: customer
     type: parquet
-    uri: gs://connectors-dev/dvdrental/customer.parquet
+    uri: gs://connectors-public/dvdrental/customer.parquet
   - name: film_actor
     type: parquet
-    uri: gs://connectors-dev/dvdrental/film_actor.parquet
+    uri: gs://connectors-public/dvdrental/film_actor.parquet
   - name: film_category
     type: parquet
-    uri: gs://connectors-dev/dvdrental/film_category.parquet
+    uri: gs://connectors-public/dvdrental/film_category.parquet
   - name: film
     type: parquet
-    uri: gs://connectors-dev/dvdrental/film.parquet
+    uri: gs://connectors-public/dvdrental/film.parquet
   - name: inventory
     type: parquet
-    uri: gs://connectors-dev/dvdrental/inventory.parquet
+    uri: gs://connectors-public/dvdrental/inventory.parquet
   - name: rental
     type: parquet
-    uri: gs://connectors-dev/dvdrental/rental.parquet
+    uri: gs://connectors-public/dvdrental/rental.parquet
 targets:
   nodes:
     - source: actor