feat(resolver): Send resolved address in ROUTE routing context (#1732)

injectives · web-flow · commit c309efd6a821 · 2026-03-17T15:54:54.000Z
Sending resolved address in ROUTE routing context helps when connecting to DBMS that has `dbms.routing.default_router=SERVER` mode enabled as returned routing table includes the address that the driver could establish connection with.
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/ClusterCompositionProvider.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/ClusterCompositionProvider.java
@@ -22,8 +22,13 @@
 import org.neo4j.driver.Bookmark;
 import org.neo4j.driver.internal.DatabaseName;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.net.ServerAddress;
 
 public interface ClusterCompositionProvider {
     CompletionStage<ClusterComposition> getClusterComposition(
-            Connection connection, DatabaseName databaseName, Bookmark bookmark, String impersonatedUser);
+            Connection connection,
+            ServerAddress address,
+            DatabaseName databaseName,
+            Bookmark bookmark,
+            String impersonatedUser);
 }
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/RediscoveryImpl.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/RediscoveryImpl.java
@@ -281,8 +281,8 @@ private CompletionStage<ClusterComposition> lookupOnRouter(
                 .thenApply(address -> addAndReturn(seenServers, address))
                 .thenCompose(connectionPool::acquire)
                 .thenApply(connection -> ImpersonationUtil.ensureImpersonationSupport(connection, impersonatedUser))
-                .thenCompose(connection ->
-                        provider.getClusterComposition(connection, routingTable.database(), bookmark, impersonatedUser))
+                .thenCompose(connection -> provider.getClusterComposition(
+                        connection, routerAddress, routingTable.database(), bookmark, impersonatedUser))
                 .handle((response, error) -> {
                     Throwable cause = Futures.completionExceptionCause(error);
                     if (cause != null) {
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/RouteMessageRoutingProcedureRunner.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/RouteMessageRoutingProcedureRunner.java
@@ -39,6 +39,7 @@
 import org.neo4j.driver.internal.handlers.RouteMessageResponseHandler;
 import org.neo4j.driver.internal.messaging.request.RouteMessage;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.net.ServerAddress;
 
 /**
  * This implementation of the {@link RoutingProcedureRunner} access the routing procedure
@@ -61,10 +62,21 @@ protected RouteMessageRoutingProcedureRunner(
 
     @Override
     public CompletionStage<RoutingProcedureResponse> run(
-            Connection connection, DatabaseName databaseName, Bookmark bookmark, String impersonatedUser) {
+            Connection connection,
+            ServerAddress address,
+            DatabaseName databaseName,
+            Bookmark bookmark,
+            String impersonatedUser) {
         CompletableFuture<Map<String, Value>> completableFuture = createCompletableFuture.get();
 
         DirectConnection directConnection = toDirectConnection(connection, databaseName, impersonatedUser);
+        Map<String, Value> routingContext = this.routingContext;
+        if (routingContext.containsKey(RoutingContext.ROUTING_ADDRESS_KEY)) {
+            routingContext = new HashMap<>(routingContext);
+            routingContext.put(
+                    RoutingContext.ROUTING_ADDRESS_KEY,
+                    Values.value(String.format("%s:%d", address.host(), address.port())));
+        }
         directConnection.writeAndFlush(
                 new RouteMessage(
                         routingContext, bookmark, databaseName.databaseName().orElse(null), impersonatedUser),
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingContext.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingContext.java
@@ -29,7 +29,7 @@
 
 public class RoutingContext {
     public static final RoutingContext EMPTY = new RoutingContext();
-    private static final String ROUTING_ADDRESS_KEY = "address";
+    public static final String ROUTING_ADDRESS_KEY = "address";
 
     private final Map<String, String> context;
     private final boolean isServerRoutingEnabled;
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureClusterCompositionProvider.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureClusterCompositionProvider.java
@@ -34,6 +34,7 @@
 import org.neo4j.driver.internal.DatabaseName;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.util.Clock;
+import org.neo4j.driver.net.ServerAddress;
 
 public class RoutingProcedureClusterCompositionProvider implements ClusterCompositionProvider {
     private static final String PROTOCOL_ERROR_MESSAGE = "Failed to parse '%s' result received from server due to ";
@@ -64,7 +65,11 @@ public RoutingProcedureClusterCompositionProvider(Clock clock, RoutingContext ro
 
     @Override
     public CompletionStage<ClusterComposition> getClusterComposition(
-            Connection connection, DatabaseName databaseName, Bookmark bookmark, String impersonatedUser) {
+            Connection connection,
+            ServerAddress address,
+            DatabaseName databaseName,
+            Bookmark bookmark,
+            String impersonatedUser) {
         RoutingProcedureRunner runner;
 
         if (supportsRouteMessage(connection)) {
@@ -75,7 +80,8 @@ public CompletionStage<ClusterComposition> getClusterComposition(
             runner = singleDatabaseRoutingProcedureRunner;
         }
 
-        return runner.run(connection, databaseName, bookmark, impersonatedUser).thenApply(this::processRoutingResponse);
+        return runner.run(connection, address, databaseName, bookmark, impersonatedUser)
+                .thenApply(this::processRoutingResponse);
     }
 
     private ClusterComposition processRoutingResponse(RoutingProcedureResponse response) {
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureRunner.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureRunner.java
@@ -22,6 +22,7 @@
 import org.neo4j.driver.Bookmark;
 import org.neo4j.driver.internal.DatabaseName;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.net.ServerAddress;
 
 /**
  * Interface which defines the standard way to get the routing table
@@ -31,11 +32,16 @@ public interface RoutingProcedureRunner {
      * Run the calls to the server
      *
      * @param connection       The connection which will be used to call the server
+     * @param address          The router address to include in routing context
      * @param databaseName     The database name
      * @param bookmark         The bookmark used to query the routing information
      * @param impersonatedUser The impersonated user, should be {@code null} for non-impersonated requests
      * @return The routing table
      */
     CompletionStage<RoutingProcedureResponse> run(
-            Connection connection, DatabaseName databaseName, Bookmark bookmark, String impersonatedUser);
+            Connection connection,
+            ServerAddress address,
+            DatabaseName databaseName,
+            Bookmark bookmark,
+            String impersonatedUser);
 }
diff --git a/driver/src/main/java/org/neo4j/driver/internal/cluster/SingleDatabaseRoutingProcedureRunner.java b/driver/src/main/java/org/neo4j/driver/internal/cluster/SingleDatabaseRoutingProcedureRunner.java
@@ -40,6 +40,7 @@
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.util.Futures;
 import org.neo4j.driver.internal.util.ServerVersion;
+import org.neo4j.driver.net.ServerAddress;
 
 /**
  * This implementation of the {@link RoutingProcedureRunner} works with single database versions of Neo4j calling
@@ -59,7 +60,11 @@ public SingleDatabaseRoutingProcedureRunner(RoutingContext context, Logging logg
 
     @Override
     public CompletionStage<RoutingProcedureResponse> run(
-            Connection connection, DatabaseName databaseName, Bookmark bookmark, String impersonatedUser) {
+            Connection connection,
+            ServerAddress address,
+            DatabaseName databaseName,
+            Bookmark bookmark,
+            String impersonatedUser) {
         DirectConnection delegate = connection(connection);
         Query procedure = procedureQuery(connection.serverVersion(), databaseName);
         BookmarkHolder bookmarkHolder = bookmarkHolder(bookmark);
diff --git a/driver/src/test/java/org/neo4j/driver/internal/cluster/AbstractRoutingProcedureRunnerTest.java b/driver/src/test/java/org/neo4j/driver/internal/cluster/AbstractRoutingProcedureRunnerTest.java
@@ -37,6 +37,7 @@
 import org.neo4j.driver.Record;
 import org.neo4j.driver.exceptions.ClientException;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.net.ServerAddress;
 
 abstract class AbstractRoutingProcedureRunnerTest {
     @Test
@@ -45,7 +46,8 @@ void shouldReturnFailedResponseOnClientException() {
         SingleDatabaseRoutingProcedureRunner runner =
                 singleDatabaseRoutingProcedureRunner(RoutingContext.EMPTY, failedFuture(error));
 
-        RoutingProcedureResponse response = await(runner.run(connection(), defaultDatabase(), empty(), null));
+        RoutingProcedureResponse response =
+                await(runner.run(connection(), ServerAddress.of("localhost", 7687), defaultDatabase(), empty(), null));
 
         assertFalse(response.isSuccess());
         assertEquals(error, response.error());
@@ -57,8 +59,10 @@ void shouldReturnFailedStageOnError() {
         SingleDatabaseRoutingProcedureRunner runner =
                 singleDatabaseRoutingProcedureRunner(RoutingContext.EMPTY, failedFuture(error));
 
-        Exception e =
-                assertThrows(Exception.class, () -> await(runner.run(connection(), defaultDatabase(), empty(), null)));
+        Exception e = assertThrows(
+                Exception.class,
+                () -> await(runner.run(
+                        connection(), ServerAddress.of("localhost", 7687), defaultDatabase(), empty(), null)));
         assertEquals(error, e);
     }
 
@@ -67,7 +71,8 @@ void shouldReleaseConnectionOnSuccess() {
         SingleDatabaseRoutingProcedureRunner runner = singleDatabaseRoutingProcedureRunner(RoutingContext.EMPTY);
 
         Connection connection = connection();
-        RoutingProcedureResponse response = await(runner.run(connection, defaultDatabase(), empty(), null));
+        RoutingProcedureResponse response =
+                await(runner.run(connection, ServerAddress.of("localhost", 7687), defaultDatabase(), empty(), null));
 
         assertTrue(response.isSuccess());
         verify(connection).release();
@@ -81,7 +86,9 @@ void shouldPropagateReleaseError() {
         Connection connection = connection(failedFuture(releaseError));
 
         RuntimeException e = assertThrows(
-                RuntimeException.class, () -> await(runner.run(connection, defaultDatabase(), empty(), null)));
+                RuntimeException.class,
+                () -> await(
+                        runner.run(connection, ServerAddress.of("localhost", 7687), defaultDatabase(), empty(), null)));
         assertEquals(releaseError, e);
         verify(connection).release();
     }
diff --git a/driver/src/test/java/org/neo4j/driver/internal/cluster/MultiDatabasesRoutingProcedureRunnerTest.java b/driver/src/test/java/org/neo4j/driver/internal/cluster/MultiDatabasesRoutingProcedureRunnerTest.java
@@ -51,13 +51,15 @@
 import org.neo4j.driver.internal.BookmarkHolder;
 import org.neo4j.driver.internal.ReadOnlyBookmarkHolder;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.net.ServerAddress;
 
 class MultiDatabasesRoutingProcedureRunnerTest extends AbstractRoutingProcedureRunnerTest {
     @ParameterizedTest
     @ValueSource(strings = {"", SYSTEM_DATABASE_NAME, " this is a db name "})
     void shouldCallGetRoutingTableWithEmptyMapOnSystemDatabaseForDatabase(String db) {
         TestRoutingProcedureRunner runner = new TestRoutingProcedureRunner(RoutingContext.EMPTY);
-        RoutingProcedureResponse response = await(runner.run(connection(), database(db), empty(), null));
+        RoutingProcedureResponse response =
+                await(runner.run(connection(), ServerAddress.of("localhost", 7687), database(db), empty(), null));
 
         assertTrue(response.isSuccess());
         assertEquals(1, response.records().size());
@@ -77,7 +79,8 @@ void shouldCallGetRoutingTableWithParamOnSystemDatabaseForDatabase(String db) {
         RoutingContext context = new RoutingContext(uri);
 
         TestRoutingProcedureRunner runner = new TestRoutingProcedureRunner(context);
-        RoutingProcedureResponse response = await(runner.run(connection(), database(db), empty(), null));
+        RoutingProcedureResponse response =
+                await(runner.run(connection(), ServerAddress.of("localhost", 7687), database(db), empty(), null));
 
         assertTrue(response.isSuccess());
         assertEquals(1, response.records().size());
diff --git a/driver/src/test/java/org/neo4j/driver/internal/cluster/RediscoveryTest.java b/driver/src/test/java/org/neo4j/driver/internal/cluster/RediscoveryTest.java
@@ -77,6 +77,7 @@
 import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.internal.util.FakeClock;
 import org.neo4j.driver.internal.util.ImmediateSchedulingEventExecutor;
+import org.neo4j.driver.net.ServerAddress;
 import org.neo4j.driver.net.ServerAddressResolver;
 
 class RediscoveryTest {
@@ -552,7 +553,11 @@ private static ClusterCompositionProvider compositionProviderMock(
             Map<BoltServerAddress, Object> responsesByAddress) {
         ClusterCompositionProvider provider = mock(ClusterCompositionProvider.class);
         when(provider.getClusterComposition(
-                        any(Connection.class), any(DatabaseName.class), any(InternalBookmark.class), any()))
+                        any(Connection.class),
+                        any(ServerAddress.class),
+                        any(DatabaseName.class),
+                        any(InternalBookmark.class),
+                        any()))
                 .then(invocation -> {
                     Connection connection = invocation.getArgument(0);
                     BoltServerAddress address = connection.serverAddress();
diff --git a/driver/src/test/java/org/neo4j/driver/internal/cluster/RouteMessageRoutingProcedureRunnerTest.java b/driver/src/test/java/org/neo4j/driver/internal/cluster/RouteMessageRoutingProcedureRunnerTest.java
@@ -34,7 +34,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -49,6 +48,7 @@
 import org.neo4j.driver.internal.handlers.RouteMessageResponseHandler;
 import org.neo4j.driver.internal.messaging.request.RouteMessage;
 import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.net.ServerAddress;
 import org.neo4j.driver.util.TestUtil;
 
 class RouteMessageRoutingProcedureRunnerTest {
@@ -77,7 +77,8 @@ void shouldRequestRoutingTableForAllValidInputScenarios(RoutingContext routingCo
         CompletableFuture<Void> releaseConnectionFuture = CompletableFuture.completedFuture(null);
         doReturn(releaseConnectionFuture).when(connection).release();
 
-        RoutingProcedureResponse response = TestUtil.await(runner.run(connection, databaseName, null, null));
+        RoutingProcedureResponse response =
+                TestUtil.await(runner.run(connection, ServerAddress.of("localhost", 7687), databaseName, null, null));
 
         assertNotNull(response);
         assertTrue(response.isSuccess());
@@ -104,8 +105,8 @@ void shouldReturnFailureWhenSomethingHappensGettingTheRoutingTable() {
         CompletableFuture<Void> releaseConnectionFuture = CompletableFuture.completedFuture(null);
         doReturn(releaseConnectionFuture).when(connection).release();
 
-        RoutingProcedureResponse response =
-                TestUtil.await(runner.run(connection, DatabaseNameUtil.defaultDatabase(), null, null));
+        RoutingProcedureResponse response = TestUtil.await(runner.run(
+                connection, ServerAddress.of("localhost", 7687), DatabaseNameUtil.defaultDatabase(), null, null));
 
         assertNotNull(response);
         assertFalse(response.isSuccess());
@@ -124,8 +125,13 @@ private void verifyMessageWasWrittenAndFlushed(
             RoutingContext routingContext,
             Bookmark bookmark,
             DatabaseName databaseName) {
-        Map<String, Value> context = routingContext.toMap().entrySet().stream()
-                .collect(Collectors.toMap(Map.Entry::getKey, entry -> Values.value(entry.getValue())));
+        Map<String, Value> context = new HashMap<>();
+        routingContext.toMap().entrySet().stream()
+                .filter(entry -> !entry.getKey().equals(RoutingContext.ROUTING_ADDRESS_KEY))
+                .forEach(entry -> context.put(entry.getKey(), Values.value(entry.getValue())));
+        if (routingContext.toMap().containsKey(RoutingContext.ROUTING_ADDRESS_KEY)) {
+            context.put(RoutingContext.ROUTING_ADDRESS_KEY, Values.value("localhost:7687"));
+        }
 
         verify(connection)
                 .writeAndFlush(
diff --git a/driver/src/test/java/org/neo4j/driver/internal/cluster/RoutingProcedureClusterCompositionProviderTest.java b/driver/src/test/java/org/neo4j/driver/internal/cluster/RoutingProcedureClusterCompositionProviderTest.java
diff --git a/driver/src/test/java/org/neo4j/driver/internal/cluster/SingleDatabaseRoutingProcedureRunnerTest.java b/driver/src/test/java/org/neo4j/driver/internal/cluster/SingleDatabaseRoutingProcedureRunnerTest.java
diff --git a/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/StartTest.java b/testkit-backend/src/main/java/neo4j/org/testkit/backend/messages/requests/StartTest.java
