From 8115a010513bf332350613c6ce2f5cdcc2abeca4 Mon Sep 17 00:00:00 2001
From: Dmitriy Tverdiakov <11927660+injectives@users.noreply.github.com>
Date: Tue, 24 Jan 2023 09:44:20 +0000
Subject: [PATCH] Add a new QueryTask.ResultFinisher interface (#1367)

This is to make the `Result.keys()` value available as well.
---
 .../main/java/org/neo4j/driver/QueryTask.java | 49 ++++++++++++++++---
 .../driver/internal/InternalQueryTask.java    | 21 +-------
 .../internal/InternalQueryTaskTest.java       | 10 ++--
 3 files changed, 50 insertions(+), 30 deletions(-)

diff --git a/driver/src/main/java/org/neo4j/driver/QueryTask.java b/driver/src/main/java/org/neo4j/driver/QueryTask.java
index 76f071fcf2..255893f787 100644
--- a/driver/src/main/java/org/neo4j/driver/QueryTask.java
+++ b/driver/src/main/java/org/neo4j/driver/QueryTask.java
@@ -18,10 +18,13 @@
  */
 package org.neo4j.driver;
 
+import java.util.List;
 import java.util.Map;
 import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.stream.Collector;
+import java.util.stream.Collectors;
+import org.neo4j.driver.internal.EagerResultValue;
 import org.neo4j.driver.summary.ResultSummary;
 import org.neo4j.driver.util.Experimental;
 
@@ -79,12 +82,13 @@
  *         .execute(mapping(record -> record.get("N").asLong(), maxBy(Long::compare)));
  * }
  * </pre>
- * If there is a need to access {@link ResultSummary} value, another method option is available:
+ * If there is a need to access {@link Result#keys()} and/or {@link ResultSummary} value, another method option is
+ * available:
  * <pre>
  * {@code
  * import static java.util.stream.Collectors.*;
  *
- * private record ResultValue(Set<Long> values, ResultSummary summary) {}
+ * private record ResultValue(List<String> keys, Set<Long> values, ResultSummary summary) {}
  *
  * var result = driver.queryTask("UNWIND range(0, 5) as N RETURN N")
  *                     .execute(Collectors.mapping(record -> record.get("N").asLong(), toSet()), ResultValue::new);
@@ -118,7 +122,9 @@ public interface QueryTask {
      *
      * @return an instance of result containing all records, keys and result summary
      */
-    EagerResult execute();
+    default EagerResult execute() {
+        return execute(Collectors.toList(), EagerResultValue::new);
+    }
 
     /**
      * Executes query, collects {@link Record} values using the provided {@link Collector} and produces a final result.
@@ -129,22 +135,49 @@ public interface QueryTask {
      * @return the final result value
      */
     default <T> T execute(Collector<Record, ?, T> recordCollector) {
-        return execute(recordCollector, (collectorResult, ignored) -> collectorResult);
+        return execute(recordCollector, (ignoredKeys, collectorResult, ignoredSummary) -> collectorResult);
     }
 
     /**
      * Executes query, collects {@link Record} values using the provided {@link Collector} and produces a final result
      * by invoking the provided {@link BiFunction} with the collected result and {@link ResultSummary} values.
+     * <p>
+     * If any of the arguments throws an exception implementing the
+     * {@link org.neo4j.driver.exceptions.RetryableException} marker interface, the query is retried automatically in
+     * the same way as in the transaction functions. Exceptions not implementing the interface trigger transaction
+     * rollback and are then propagated to the user.
      *
      * @param recordCollector collector instance responsible for processing {@link Record} values and producing a
      *                        collected result, the collector may be used multiple times if query is retried
-     * @param finisherWithSummary function accepting both the collected result and {@link ResultSummary} values to
-     *                            output the final result, the function may be invoked multiple times if query is
-     *                            retried
+     * @param resultFinisher function accepting the {@link Result#keys()}, collected result and {@link ResultSummary}
+     *                       values to output the final result value, the function may be invoked multiple times if
+     *                       query is retried
      * @param <A> the mutable accumulation type of the collector's reduction operation
      * @param <R> the collector's result type
      * @param <T> the final result type
      * @return the final result value
      */
-    <A, R, T> T execute(Collector<Record, A, R> recordCollector, BiFunction<R, ResultSummary, T> finisherWithSummary);
+    <A, R, T> T execute(Collector<Record, A, R> recordCollector, ResultFinisher<R, T> resultFinisher);
+
+    /**
+     * A function accepting the {@link Result#keys()}, collected result and {@link ResultSummary} values to produce a
+     * final result value.
+     *
+     * @param <S> the collected value type
+     * @param <T> the final value type
+     * @since 5.5
+     */
+    @Experimental
+    @FunctionalInterface
+    interface ResultFinisher<S, T> {
+        /**
+         * Accepts the {@link Result#keys()}, collected result and {@link ResultSummary} values to produce the final
+         * result value.
+         * @param value the collected value
+         * @param keys the {@link Result#keys()} value
+         * @param summary the {@link ResultSummary} value
+         * @return the final value
+         */
+        T finish(List<String> keys, S value, ResultSummary summary);
+    }
 }
diff --git a/driver/src/main/java/org/neo4j/driver/internal/InternalQueryTask.java b/driver/src/main/java/org/neo4j/driver/internal/InternalQueryTask.java
index 093e3d1e64..486f126dd0 100644
--- a/driver/src/main/java/org/neo4j/driver/internal/InternalQueryTask.java
+++ b/driver/src/main/java/org/neo4j/driver/internal/InternalQueryTask.java
@@ -20,28 +20,17 @@
 
 import static java.util.Objects.requireNonNull;
 
-import java.util.Collections;
-import java.util.List;
 import java.util.Map;
-import java.util.function.BiFunction;
 import java.util.stream.Collector;
-import java.util.stream.Collectors;
 import org.neo4j.driver.Driver;
-import org.neo4j.driver.EagerResult;
 import org.neo4j.driver.Query;
 import org.neo4j.driver.QueryConfig;
 import org.neo4j.driver.QueryTask;
 import org.neo4j.driver.Record;
 import org.neo4j.driver.SessionConfig;
 import org.neo4j.driver.TransactionCallback;
-import org.neo4j.driver.summary.ResultSummary;
 
 public class InternalQueryTask implements QueryTask {
-    private static final BiFunction<List<Record>, ResultSummary, EagerResult> EAGER_RESULT_FINISHER =
-            (records, summary) -> {
-                var keys = records.stream().findFirst().map(Record::keys).orElseGet(Collections::emptyList);
-                return new EagerResultValue(keys, records, summary);
-            };
     private final Driver driver;
     private final Query query;
     private final QueryConfig config;
@@ -68,13 +57,7 @@ public QueryTask withConfig(QueryConfig config) {
     }
 
     @Override
-    public EagerResult execute() {
-        return execute(Collectors.toList(), EAGER_RESULT_FINISHER);
-    }
-
-    @Override
-    public <A, R, T> T execute(
-            Collector<Record, A, R> recordCollector, BiFunction<R, ResultSummary, T> finisherWithSummary) {
+    public <A, R, T> T execute(Collector<Record, A, R> recordCollector, ResultFinisher<R, T> resultFinisher) {
         var sessionConfigBuilder = SessionConfig.builder();
         config.database().ifPresent(sessionConfigBuilder::withDatabase);
         config.impersonatedUser().ifPresent(sessionConfigBuilder::withImpersonatedUser);
@@ -91,7 +74,7 @@ public <A, R, T> T execute(
                 }
                 var finishedValue = finisher.apply(container);
                 var summary = result.consume();
-                return finisherWithSummary.apply(finishedValue, summary);
+                return resultFinisher.finish(result.keys(), finishedValue, summary);
             };
             return switch (config.routing()) {
                 case WRITERS -> session.executeWrite(txCallback);
diff --git a/driver/src/test/java/org/neo4j/driver/internal/InternalQueryTaskTest.java b/driver/src/test/java/org/neo4j/driver/internal/InternalQueryTaskTest.java
index 13b13bc540..c5de47c688 100644
--- a/driver/src/test/java/org/neo4j/driver/internal/InternalQueryTaskTest.java
+++ b/driver/src/test/java/org/neo4j/driver/internal/InternalQueryTaskTest.java
@@ -26,6 +26,7 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 
+import java.util.List;
 import java.util.Map;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
@@ -39,6 +40,7 @@
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.Query;
 import org.neo4j.driver.QueryConfig;
+import org.neo4j.driver.QueryTask;
 import org.neo4j.driver.Record;
 import org.neo4j.driver.Result;
 import org.neo4j.driver.RoutingControl;
@@ -131,6 +133,8 @@ void shouldExecuteAndReturnResult(RoutingControl routingControl) {
         });
         var result = mock(Result.class);
         given(txContext.run(any(Query.class))).willReturn(result);
+        var keys = List.of("key");
+        given(result.keys()).willReturn(keys);
         given(result.hasNext()).willReturn(true, false);
         var record = mock(Record.class);
         given(result.next()).willReturn(record);
@@ -152,9 +156,9 @@ var record = mock(Record.class);
         Function<Object, String> finisher = mock(Function.class);
         given(finisher.apply(resultContainer)).willReturn(collectorResult);
         given(recordCollector.finisher()).willReturn(finisher);
-        BiFunction<String, ResultSummary, String> finisherWithSummary = mock(BiFunction.class);
+        QueryTask.ResultFinisher<String, String> finisherWithSummary = mock(QueryTask.ResultFinisher.class);
         var expectedExecuteResult = "1";
-        given(finisherWithSummary.apply(any(String.class), any(ResultSummary.class)))
+        given(finisherWithSummary.finish(any(List.class), any(String.class), any(ResultSummary.class)))
                 .willReturn(expectedExecuteResult);
         var queryTask = new InternalQueryTask(driver, query, config).withParameters(params);
 
@@ -181,7 +185,7 @@ var record = mock(Record.class);
         then(accumulator).should().accept(resultContainer, record);
         then(recordCollector).should().finisher();
         then(finisher).should().apply(resultContainer);
-        then(finisherWithSummary).should().apply(collectorResult, summary);
+        then(finisherWithSummary).should().finish(keys, collectorResult, summary);
         assertEquals(expectedExecuteResult, executeResult);
     }
 }