[olbYUBeG] Added meta-data to batches in periodic iterate (#610)

core/src/main/java/apoc/periodic/Periodic.java

@@ -43,6 +43,7 @@
 import org.neo4j.graphdb.schema.ConstraintDefinition;
 import org.neo4j.graphdb.schema.IndexDefinition;
 import org.neo4j.graphdb.schema.Schema;
+import org.neo4j.kernel.impl.coreapi.InternalTransaction;
 import org.neo4j.logging.Log;
 import org.neo4j.procedure.*;
 
@@ -307,6 +308,13 @@ public Stream<BatchAndTotalResult> iterate(
                 "retries", 0)); // todo sleep/delay or push to end of batch to try again or immediate ?
         int failedParams = Util.toInteger(config.getOrDefault("failedParams", -1));
 
+        final Map<String, Object> metaData;
+        if(tx instanceof InternalTransaction iTx){
+            metaData = iTx.kernelTransaction().getMetaData();
+        } else {
+            metaData = Map.of();
+        }
+
         BatchMode batchMode = BatchMode.fromConfig(config);
         Map<String, Object> params = (Map<String, Object>) config.getOrDefault("params", Collections.emptyMap());
 
@@ -333,6 +341,9 @@ public Stream<BatchAndTotalResult> iterate(
                     retries,
                     result,
                     (tx, p) -> {
+                        if(tx instanceof InternalTransaction iTx){
+                            iTx.setMetaData(metaData);
+                        }
                         final Result r = tx.execute(innerStatement, merge(params, p));
                         Iterators.count(r); // XXX: consume all results
                         return r.getQueryStatistics();

it/src/test/java/apoc/it/core/PeriodicIterateTest.java

@@ -0,0 +1,58 @@
+package apoc.it.core;
+
+import apoc.util.Neo4jContainerExtension;
+import apoc.util.TestContainerUtil;
+import apoc.util.TestUtil;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.neo4j.driver.Session;
+import org.neo4j.driver.TransactionConfig;
+
+import java.util.List;
+import java.util.Map;
+
+import static apoc.util.TestContainerUtil.createDB;
+import static apoc.util.TestContainerUtil.dockerImageForNeo4j;
+import static org.junit.Assert.*;
+import static org.junit.Assert.fail;
+
+public class PeriodicIterateTest {
+
+    // The Query Log is only accessible on Enterprise
+    @Test
+    public void check_metadata_in_batches() {
+        try {
+            Neo4jContainerExtension neo4jContainer = createDB(
+                    TestContainerUtil.Neo4jVersion.ENTERPRISE, List.of(TestContainerUtil.ApocPackage.CORE), !TestUtil.isRunningInCI())
+                    .withNeo4jConfig("dbms.transaction.timeout", "60s");
+
+            neo4jContainer.start();
+
+
+            Session session = neo4jContainer.getSession();
+            session.run(
+                    "CALL apoc.periodic.iterate(\"MATCH (p:Person) RETURN p\"," +
+                            "\"SET p.name='test'\"," +
+                            "{batchSize:1, parallel:false})",
+                    TransactionConfig.builder()
+                            .withMetadata(Map.of("shouldAppear", "inBatches"))
+                            .build()
+                    ).stream().count();
+            var queryLogs = neo4jContainer.queryLogs();
+            assertTrue(queryLogs.contains("SET p.name='test' - {_batch: [], _count: 0} - runtime=pipelined - {shouldAppear: 'inBatches'}"));
+            session.close();
+            neo4jContainer.close();
+        } catch (Exception ex) {
+            // if Testcontainers wasn't able to retrieve the docker image we ignore the test
+            if (TestContainerUtil.isDockerImageAvailable(ex)) {
+                ex.printStackTrace();
+                fail("Should not have thrown exception when trying to start Neo4j: " + ex);
+            } else if (!TestUtil.isRunningInCI()) {
+                fail("The docker image " + dockerImageForNeo4j(TestContainerUtil.Neo4jVersion.ENTERPRISE)
+                        + " could not be loaded. Check whether it's available locally / in the CI. Exception:"
+                        + ex);
+            }
+        }
+    }
+}

test-utils/src/main/java/apoc/util/Neo4jContainerExtension.java

@@ -21,6 +21,7 @@
 import static apoc.util.TestContainerUtil.Neo4jVersion;
 import static apoc.util.TestContainerUtil.Neo4jVersion.ENTERPRISE;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -109,6 +110,7 @@ public void dumpLogs() {
                 System.err.println(execInContainer("cat", "logs/debug.log").toString());
                 System.err.println(execInContainer("cat", "logs/http.log").toString());
                 System.err.println(execInContainer("cat", "logs/security.log").toString());
+                System.err.println(execInContainer("cat", "logs/query.log").toString());
             } else {
                 System.err.println("Failed to find logs");
             }
@@ -117,6 +119,10 @@ public void dumpLogs() {
         }
     }
 
+    public String queryLogs() throws IOException, InterruptedException {
+       return execInContainer("cat", "logs/query.log").toString();
+    }
+
     private void executeScript(String filePath) {
         InputStream resource = Thread.currentThread().getContextClassLoader().getResourceAsStream(filePath);
         if (resource == null) {