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))); * } * - * 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: *
  * {@code
  * import static java.util.stream.Collectors.*;
  *
- * private record ResultValue(Set values, ResultSummary summary) {}
+ * private record ResultValue(List keys, Set 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 execute(Collector 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.
+     * 

+ * 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 the mutable accumulation type of the collector's reduction operation * @param the collector's result type * @param the final result type * @return the final result value */ - T execute(Collector recordCollector, BiFunction finisherWithSummary); + T execute(Collector recordCollector, ResultFinisher resultFinisher); + + /** + * A function accepting the {@link Result#keys()}, collected result and {@link ResultSummary} values to produce a + * final result value. + * + * @param the collected value type + * @param the final value type + * @since 5.5 + */ + @Experimental + @FunctionalInterface + interface ResultFinisher { + /** + * 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 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, 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 T execute( - Collector recordCollector, BiFunction finisherWithSummary) { + public T execute(Collector recordCollector, ResultFinisher resultFinisher) { var sessionConfigBuilder = SessionConfig.builder(); config.database().ifPresent(sessionConfigBuilder::withDatabase); config.impersonatedUser().ifPresent(sessionConfigBuilder::withImpersonatedUser); @@ -91,7 +74,7 @@ public 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 finisher = mock(Function.class); given(finisher.apply(resultContainer)).willReturn(collectorResult); given(recordCollector.finisher()).willReturn(finisher); - BiFunction finisherWithSummary = mock(BiFunction.class); + QueryTask.ResultFinisher 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); } }