Introduce an override AuthToken support to ExecutableQuery (neo4j#1532)

* Introduce an override AuthToken support to ExecutableQuery

This update adds support for specifying an override `AuthToken` for `ExecutableQuery`. This feature requires Bolt 5.1 or greater.

* Update visibility

Author: injectives

driver/clirr-ignored-differences.xml

@@ -597,4 +597,10 @@
         <differenceType>8001</differenceType>
     </difference>
 
+    <difference>
+        <className>org/neo4j/driver/ExecutableQuery</className>
+        <differenceType>7012</differenceType>
+        <method>org.neo4j.driver.ExecutableQuery withAuthToken(org.neo4j.driver.AuthToken)</method>
+    </difference>
+
 </differences>

driver/src/main/java/org/neo4j/driver/Driver.java

@@ -19,6 +19,7 @@
 import java.util.concurrent.CompletionStage;
 import org.neo4j.driver.async.AsyncSession;
 import org.neo4j.driver.exceptions.ClientException;
+import org.neo4j.driver.exceptions.UnsupportedFeatureException;
 import org.neo4j.driver.reactive.ReactiveSession;
 import org.neo4j.driver.reactive.RxSession;
 import org.neo4j.driver.types.TypeSystem;
@@ -143,7 +144,7 @@ default <T extends BaseSession> T session(Class<T> sessionClass) {
      * Instantiate a new session of a supported type with the supplied {@link AuthToken}.
      * <p>
      * This method allows creating a session with a different {@link AuthToken} to the one used on the driver level.
-     * The minimum Bolt protocol version is 5.1. An {@link IllegalStateException} will be emitted on session interaction
+     * The minimum Bolt protocol version is 5.1. An {@link UnsupportedFeatureException} will be emitted on session interaction
      * for previous Bolt versions.
      * <p>
      * Supported types are:
@@ -214,7 +215,7 @@ default <T extends BaseSession> T session(Class<T> sessionClass, SessionConfig s
      * {@link AuthToken}.
      * <p>
      * This method allows creating a session with a different {@link AuthToken} to the one used on the driver level.
-     * The minimum Bolt protocol version is 5.1. An {@link IllegalStateException} will be emitted on session interaction
+     * The minimum Bolt protocol version is 5.1. An {@link UnsupportedFeatureException} will be emitted on session interaction
      * for previous Bolt versions.
      * <p>
      * Supported types are:

driver/src/main/java/org/neo4j/driver/ExecutableQuery.java

@@ -22,6 +22,7 @@
 import java.util.function.Consumer;
 import java.util.stream.Collector;
 import java.util.stream.Collectors;
+import org.neo4j.driver.exceptions.UnsupportedFeatureException;
 import org.neo4j.driver.internal.EagerResultValue;
 import org.neo4j.driver.summary.ResultSummary;
 
@@ -97,7 +98,7 @@ public interface ExecutableQuery {
     /**
      * Sets query parameters.
      *
-     * @param parameters parameters map, must not be {@code null}
+     * @param parameters parameters map, must not be {@literal null}
      * @return a new executable query
      */
     ExecutableQuery withParameters(Map<String, Object> parameters);
@@ -107,11 +108,27 @@ public interface ExecutableQuery {
      * <p>
      * By default, {@link ExecutableQuery} has {@link QueryConfig#defaultConfig()} value.
      *
-     * @param config query config, must not be {@code null}
+     * @param config query config, must not be {@literal null}
      * @return a new executable query
      */
     ExecutableQuery withConfig(QueryConfig config);
 
+    /**
+     * Sets an {@link AuthToken} to be used for this query.
+     * <p>
+     * The default value is {@literal null}.
+     * <p>
+     * The minimum Bolt protocol version for this feature is 5.1. An {@link UnsupportedFeatureException} will be emitted on
+     * query execution for previous Bolt versions.
+     *
+     * @param authToken the {@link AuthToken} for this query or {@literal null} to use the driver default
+     * @return a new executable query
+     * @since 5.18
+     */
+    default ExecutableQuery withAuthToken(AuthToken authToken) {
+        throw new UnsupportedFeatureException("Session AuthToken is not supported.");
+    }
+
     /**
      * Executes query, collects all results eagerly and returns a result.
      *

driver/src/main/java/org/neo4j/driver/internal/InternalDriver.java

@@ -87,7 +87,7 @@ public class InternalDriver implements Driver {
 
     @Override
     public ExecutableQuery executableQuery(String query) {
-        return new InternalExecutableQuery(this, new Query(query), QueryConfig.defaultConfig());
+        return new InternalExecutableQuery(this, new Query(query), QueryConfig.defaultConfig(), null);
     }
 
     @Override

driver/src/main/java/org/neo4j/driver/internal/InternalExecutableQuery.java

@@ -21,12 +21,14 @@
 import java.util.Map;
 import java.util.stream.Collector;
 import org.neo4j.driver.AccessMode;
+import org.neo4j.driver.AuthToken;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.ExecutableQuery;
 import org.neo4j.driver.Query;
 import org.neo4j.driver.QueryConfig;
 import org.neo4j.driver.Record;
 import org.neo4j.driver.RoutingControl;
+import org.neo4j.driver.Session;
 import org.neo4j.driver.SessionConfig;
 import org.neo4j.driver.TransactionCallback;
 import org.neo4j.driver.TransactionConfig;
@@ -36,26 +38,33 @@ public class InternalExecutableQuery implements ExecutableQuery {
     private final Driver driver;
     private final Query query;
     private final QueryConfig config;
+    private final AuthToken authToken;
 
-    public InternalExecutableQuery(Driver driver, Query query, QueryConfig config) {
+    public InternalExecutableQuery(Driver driver, Query query, QueryConfig config, AuthToken authToken) {
         requireNonNull(driver, "driver must not be null");
         requireNonNull(query, "query must not be null");
         requireNonNull(config, "config must not be null");
         this.driver = driver;
         this.query = query;
         this.config = config;
+        this.authToken = authToken;
     }
 
     @Override
     public ExecutableQuery withParameters(Map<String, Object> parameters) {
         requireNonNull(parameters, "parameters must not be null");
-        return new InternalExecutableQuery(driver, query.withParameters(parameters), config);
+        return new InternalExecutableQuery(driver, query.withParameters(parameters), config, authToken);
     }
 
     @Override
     public ExecutableQuery withConfig(QueryConfig config) {
         requireNonNull(config, "config must not be null");
-        return new InternalExecutableQuery(driver, query, config);
+        return new InternalExecutableQuery(driver, query, config, authToken);
+    }
+
+    @Override
+    public ExecutableQuery withAuthToken(AuthToken authToken) {
+        return new InternalExecutableQuery(driver, query, config, authToken);
     }
 
     @Override
@@ -68,7 +77,7 @@ public <A, R, T> T execute(Collector<Record, A, R> recordCollector, ResultFinish
         var supplier = recordCollector.supplier();
         var accumulator = recordCollector.accumulator();
         var finisher = recordCollector.finisher();
-        try (var session = (InternalSession) driver.session(sessionConfigBuilder.build())) {
+        try (var session = (InternalSession) driver.session(Session.class, sessionConfigBuilder.build(), authToken)) {
             TransactionCallback<T> txCallback = tx -> {
                 var result = tx.run(query);
                 var container = supplier.get();
@@ -89,22 +98,27 @@ public <A, R, T> T execute(Collector<Record, A, R> recordCollector, ResultFinish
     }
 
     // For testing only
-    public Driver driver() {
+    Driver driver() {
         return driver;
     }
 
     // For testing only
-    public String query() {
+    String query() {
         return query.text();
     }
 
     // For testing only
-    public Map<String, Object> parameters() {
+    Map<String, Object> parameters() {
         return query.parameters().asMap();
     }
 
     // For testing only
-    public QueryConfig config() {
+    QueryConfig config() {
         return config;
     }
+
+    // For testing only
+    AuthToken authToken() {
+        return authToken;
+    }
 }

driver/src/test/java/org/neo4j/driver/internal/InternalExecutableQueryTest.java

@@ -17,6 +17,7 @@
 package org.neo4j.driver.internal;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.eq;
@@ -35,6 +36,7 @@
 import org.junit.jupiter.params.provider.MethodSource;
 import org.mockito.ArgumentCaptor;
 import org.neo4j.driver.AccessMode;
+import org.neo4j.driver.AuthTokens;
 import org.neo4j.driver.BookmarkManager;
 import org.neo4j.driver.Driver;
 import org.neo4j.driver.ExecutableQuery;
@@ -43,6 +45,7 @@
 import org.neo4j.driver.Record;
 import org.neo4j.driver.Result;
 import org.neo4j.driver.RoutingControl;
+import org.neo4j.driver.Session;
 import org.neo4j.driver.SessionConfig;
 import org.neo4j.driver.TransactionCallback;
 import org.neo4j.driver.TransactionConfig;
@@ -55,27 +58,27 @@ class InternalExecutableQueryTest {
     void shouldNotAcceptNullDriverOnInstantiation() {
         assertThrows(
                 NullPointerException.class,
-                () -> new InternalExecutableQuery(null, new Query("string"), QueryConfig.defaultConfig()));
+                () -> new InternalExecutableQuery(null, new Query("string"), QueryConfig.defaultConfig(), null));
     }
 
     @Test
     void shouldNotAcceptNullQueryOnInstantiation() {
         assertThrows(
                 NullPointerException.class,
-                () -> new InternalExecutableQuery(mock(Driver.class), null, QueryConfig.defaultConfig()));
+                () -> new InternalExecutableQuery(mock(Driver.class), null, QueryConfig.defaultConfig(), null));
     }
 
     @Test
     void shouldNotAcceptNullConfigOnInstantiation() {
         assertThrows(
                 NullPointerException.class,
-                () -> new InternalExecutableQuery(mock(Driver.class), new Query("string"), null));
+                () -> new InternalExecutableQuery(mock(Driver.class), new Query("string"), null, null));
     }
 
     @Test
     void shouldNotAcceptNullParameters() {
         var executableQuery =
-                new InternalExecutableQuery(mock(Driver.class), new Query("string"), QueryConfig.defaultConfig());
+                new InternalExecutableQuery(mock(Driver.class), new Query("string"), QueryConfig.defaultConfig(), null);
         assertThrows(NullPointerException.class, () -> executableQuery.withParameters(null));
     }
 
@@ -84,7 +87,7 @@ void shouldUpdateParameters() {
         // GIVEN
         var query = new Query("string");
         var params = Map.<String, Object>of("$param", "value");
-        var executableQuery = new InternalExecutableQuery(mock(Driver.class), query, QueryConfig.defaultConfig());
+        var executableQuery = new InternalExecutableQuery(mock(Driver.class), query, QueryConfig.defaultConfig(), null);
 
         // WHEN
         executableQuery = (InternalExecutableQuery) executableQuery.withParameters(params);
@@ -96,15 +99,15 @@ void shouldUpdateParameters() {
     @Test
     void shouldNotAcceptNullConfig() {
         var executableQuery =
-                new InternalExecutableQuery(mock(Driver.class), new Query("string"), QueryConfig.defaultConfig());
+                new InternalExecutableQuery(mock(Driver.class), new Query("string"), QueryConfig.defaultConfig(), null);
         assertThrows(NullPointerException.class, () -> executableQuery.withConfig(null));
     }
 
     @Test
     void shouldUpdateConfig() {
         // GIVEN
         var query = new Query("string");
-        var executableQuery = new InternalExecutableQuery(mock(Driver.class), query, QueryConfig.defaultConfig());
+        var executableQuery = new InternalExecutableQuery(mock(Driver.class), query, QueryConfig.defaultConfig(), null);
         var config = QueryConfig.builder().withDatabase("database").build();
 
         // WHEN
@@ -127,7 +130,8 @@ void shouldExecuteAndReturnResult(RoutingControl routingControl) {
         var bookmarkManager = mock(BookmarkManager.class);
         given(driver.executableQueryBookmarkManager()).willReturn(bookmarkManager);
         var session = mock(InternalSession.class);
-        given(driver.session(any(SessionConfig.class))).willReturn(session);
+        given(driver.session(eq(Session.class), any(SessionConfig.class), eq(null)))
+                .willReturn(session);
         var txContext = mock(TransactionContext.class);
         var accessMode = routingControl.equals(RoutingControl.WRITE) ? AccessMode.WRITE : AccessMode.READ;
         given(session.execute(
@@ -169,14 +173,14 @@ var record = mock(Record.class);
         var expectedExecuteResult = "1";
         given(finisherWithSummary.finish(any(List.class), any(String.class), any(ResultSummary.class)))
                 .willReturn(expectedExecuteResult);
-        var executableQuery = new InternalExecutableQuery(driver, query, config).withParameters(params);
+        var executableQuery = new InternalExecutableQuery(driver, query, config, null).withParameters(params);
 
         // WHEN
         var executeResult = executableQuery.execute(recordCollector, finisherWithSummary);
 
         // THEN
         var sessionConfigCapture = ArgumentCaptor.forClass(SessionConfig.class);
-        then(driver).should().session(sessionConfigCapture.capture());
+        then(driver).should().session(eq(Session.class), sessionConfigCapture.capture(), eq(null));
         var sessionConfig = sessionConfigCapture.getValue();
         @SuppressWarnings("OptionalGetWithoutIsPresent")
         var expectedSessionConfig = SessionConfig.builder()
@@ -205,4 +209,25 @@ var record = mock(Record.class);
         then(finisherWithSummary).should().finish(keys, collectorResult, summary);
         assertEquals(expectedExecuteResult, executeResult);
     }
+
+    @Test
+    void shouldAllowNullAuthToken() {
+        var executableQuery =
+                new InternalExecutableQuery(mock(Driver.class), new Query("string"), QueryConfig.defaultConfig(), null);
+
+        executableQuery.withAuthToken(null);
+
+        assertNull(executableQuery.authToken());
+    }
+
+    @Test
+    void shouldUpdateAuthToken() {
+        var executableQuery =
+                new InternalExecutableQuery(mock(Driver.class), new Query("string"), QueryConfig.defaultConfig(), null);
+        var authToken = AuthTokens.basic("user", "password");
+
+        executableQuery = (InternalExecutableQuery) executableQuery.withAuthToken(authToken);
+
+        assertEquals(authToken, executableQuery.authToken());
+    }
 }

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/ExecuteQuery.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.CompletionStage;
 import lombok.Getter;
 import lombok.Setter;
+import neo4j.org.testkit.backend.AuthTokenUtil;
 import neo4j.org.testkit.backend.TestkitState;
 import neo4j.org.testkit.backend.messages.requests.deserializer.TestkitCypherParamDeserializer;
 import neo4j.org.testkit.backend.messages.responses.EagerResult;
@@ -73,10 +74,15 @@ public TestkitResponse process(TestkitState testkitState) {
 
         Optional.ofNullable(data.getConfig().getTxMeta()).ifPresent(configBuilder::withMetadata);
 
+        var authToken = data.getConfig().getAuthorizationToken() != null
+                ? AuthTokenUtil.parseAuthToken(data.getConfig().getAuthorizationToken())
+                : null;
+
         var params = data.getParams() != null ? data.getParams() : Collections.<String, Object>emptyMap();
         var eagerResult = driver.executableQuery(data.getCypher())
                 .withParameters(params)
                 .withConfig(configBuilder.build())
+                .withAuthToken(authToken)
                 .execute();
 
         return EagerResult.builder()
@@ -135,5 +141,7 @@ public static class QueryConfigData {
 
         @JsonDeserialize(using = TestkitCypherParamDeserializer.class)
         private Map<String, Serializable> txMeta;
+
+        private AuthorizationToken authorizationToken;
     }
 }

testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/GetFeatures.java

@@ -78,6 +78,7 @@ public class GetFeatures implements TestkitRequest {
             "Optimization:ResultListFetchAll",
             "Feature:API:Result.Single",
             "Feature:API:Driver.ExecuteQuery",
+            "Feature:API:Driver.ExecuteQuery:WithAuth",
             "Feature:API:Driver.VerifyAuthentication",
             "Optimization:ExecuteQueryPipelining"));