Use the supportMultiDb feature detection result, to choose the database to connect.

Author: Zhen Li

driver/src/main/java/org/neo4j/driver/internal/async/ImmutableConnectionContext.java

@@ -24,14 +24,16 @@
 import org.neo4j.driver.internal.spi.Connection;
 
 import static org.neo4j.driver.internal.DatabaseNameUtil.defaultDatabase;
+import static org.neo4j.driver.internal.DatabaseNameUtil.systemDatabase;
 import static org.neo4j.driver.internal.InternalBookmark.empty;
 
 /**
  * A {@link Connection} shall fulfil this {@link ImmutableConnectionContext} when acquired from a connection provider.
  */
 public class ImmutableConnectionContext implements ConnectionContext
 {
-    private static final ConnectionContext SIMPLE = new ImmutableConnectionContext( defaultDatabase(), empty(), AccessMode.READ );
+    private static final ConnectionContext SINGLE_DB_CONTEXT = new ImmutableConnectionContext( defaultDatabase(), empty(), AccessMode.READ );
+    private static final ConnectionContext MULTI_DB_CONTEXT = new ImmutableConnectionContext( systemDatabase(), empty(), AccessMode.READ );
 
     private final DatabaseName databaseName;
     private final AccessMode mode;
@@ -65,10 +67,10 @@ public Bookmark rediscoveryBookmark()
     /**
      * A simple context is used to test connectivity with a remote server/cluster.
      * As long as there is a read only service, the connection shall be established successfully.
-     * This context should be applicable for both bolt v4 and bolt v3 routing table rediscovery.
+     * Depending on whether multidb is supported or not, this method returns different context for routing table discovery.
      */
-    public static ConnectionContext simple()
+    public static ConnectionContext simple( boolean supportsMultiDb )
     {
-        return SIMPLE;
+        return supportsMultiDb ? MULTI_DB_CONTEXT : SINGLE_DB_CONTEXT;
     }
 }

driver/src/main/java/org/neo4j/driver/internal/cluster/loadbalancing/LoadBalancer.java

@@ -103,14 +103,15 @@ public CompletionStage<Connection> acquireConnection( ConnectionContext context
     @Override
     public CompletionStage<Void> verifyConnectivity()
     {
-        return routingTables.refreshRoutingTable( simple() ).handle( ( ignored, error ) -> {
+        return this.supportsMultiDbAsync().thenCompose( supports -> routingTables.refreshRoutingTable( simple( supports ) ) ).handle( ( ignored, error ) -> {
             if ( error != null )
             {
                 Throwable cause = Futures.completionExceptionCause( error );
                 if ( cause instanceof ServiceUnavailableException )
                 {
                     throw Futures.asCompletionException( new ServiceUnavailableException(
-                            "Unable to connect to database, ensure the database is running and that there is a working network connection to it.", cause ) );
+                            "Unable to connect to database management service, ensure the database is running and that there is a working network connection to it.",
+                            cause ) );
                 }
                 throw Futures.asCompletionException( cause );
             }

driver/src/test/java/org/neo4j/driver/internal/cluster/RoutingTableHandlerTest.java

@@ -115,7 +115,7 @@ void acquireShouldUpdateRoutingTableWhenKnownRoutingTableIsStale()
 
         RoutingTableHandler handler = newRoutingTableHandler( routingTable, rediscovery, connectionPool );
 
-        assertNotNull( await( handler.refreshRoutingTable( simple() ) ) );
+        assertNotNull( await( handler.refreshRoutingTable( simple( false ) ) ) );
 
         verify( rediscovery ).lookupClusterComposition( eq ( routingTable ) , eq ( connectionPool ), any() );
         assertArrayEquals( new BoltServerAddress[]{reader1, reader2}, routingTable.readers().toArray() );
@@ -188,7 +188,7 @@ public void removeAged()
 
         RoutingTableHandler handler = newRoutingTableHandler( routingTable, rediscovery, connectionPool, registry );
 
-        RoutingTable actual = await( handler.refreshRoutingTable( simple() ) );
+        RoutingTable actual = await( handler.refreshRoutingTable( simple( false ) ) );
         assertEquals( routingTable, actual );
 
         verify( connectionPool ).retainAll( new HashSet<>( asList( A, B, C ) ) );
@@ -208,7 +208,7 @@ void shouldRemoveRoutingTableHandlerIfFailedToLookup() throws Throwable
         // When
 
         RoutingTableHandler handler = newRoutingTableHandler( routingTable, rediscovery, connectionPool, registry );
-        assertThrows( RuntimeException.class, () -> await( handler.refreshRoutingTable( simple() ) ) );
+        assertThrows( RuntimeException.class, () -> await( handler.refreshRoutingTable( simple( false ) ) ) );
 
         // Then
         verify( registry ).remove( defaultDatabase() );

driver/src/test/java/org/neo4j/driver/internal/cluster/loadbalancing/LoadBalancerTest.java

@@ -29,6 +29,7 @@
 import java.util.LinkedHashSet;
 import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Function;
 
 import org.neo4j.driver.AccessMode;
 import org.neo4j.driver.exceptions.ServiceUnavailableException;
@@ -161,7 +162,7 @@ void shouldSelectLeastConnectedAddress()
         Set<BoltServerAddress> seenAddresses = new HashSet<>();
         for ( int i = 0; i < 10; i++ )
         {
-            Connection connection = await( loadBalancer.acquireConnection( simple() ) );
+            Connection connection = await( loadBalancer.acquireConnection( newBoltV4ConnectionContext() ) );
             seenAddresses.add( connection.serverAddress() );
         }
 
@@ -185,7 +186,7 @@ void shouldRoundRobinWhenNoActiveConnections()
         Set<BoltServerAddress> seenAddresses = new HashSet<>();
         for ( int i = 0; i < 10; i++ )
         {
-            Connection connection = await( loadBalancer.acquireConnection( simple() ) );
+            Connection connection = await( loadBalancer.acquireConnection( newBoltV4ConnectionContext() ) );
             seenAddresses.add( connection.serverAddress() );
         }
 
@@ -204,7 +205,7 @@ void shouldTryMultipleServersAfterRediscovery()
 
         LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTable );
 
-        Connection connection = await( loadBalancer.acquireConnection( simple() ) );
+        Connection connection = await( loadBalancer.acquireConnection( newBoltV4ConnectionContext() ) );
 
         assertNotNull( connection );
         assertEquals( B, connection.serverAddress() );
@@ -257,21 +258,93 @@ void shouldSuccessOnFirstSuccessfulServer() throws Throwable
         assertTrue( await( loadBalancer.supportsMultiDbAsync() ) );
     }
 
+    @Test
+    void shouldThrowModifiedErrorWhenSupportMultiDbTestFails() throws Throwable
+    {
+        Set<BoltServerAddress> unavailableAddresses = asOrderedSet( A, B );
+        ConnectionPool connectionPool = newConnectionPoolMockWithFailures( unavailableAddresses );
+
+        Rediscovery rediscovery = mock( Rediscovery.class );
+        when( rediscovery.resolve() ).thenReturn( Arrays.asList( A, B ) );
+
+        LoadBalancer loadBalancer = newLoadBalancer( connectionPool, rediscovery );
+
+        ServiceUnavailableException exception = assertThrows( ServiceUnavailableException.class, () -> await( loadBalancer.verifyConnectivity() ) );
+        assertThat( exception.getMessage(), startsWith( "Unable to connect to database management service," ) );
+    }
+
+    @Test
+    void shouldThrowModifiedErrorWhenRefreshRoutingTableFails() throws Throwable
+    {
+        ConnectionPool connectionPool = newConnectionPoolMock();
+
+        Rediscovery rediscovery = mock( Rediscovery.class );
+        when( rediscovery.resolve() ).thenReturn( Arrays.asList( A, B ) );
+
+        RoutingTableRegistry routingTables = mock( RoutingTableRegistry.class );
+        when( routingTables.refreshRoutingTable( any( ConnectionContext.class ) ) ).thenThrow( new ServiceUnavailableException( "boooo" ) );
+
+        LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTables, rediscovery );
+
+        ServiceUnavailableException exception = assertThrows( ServiceUnavailableException.class, () -> await( loadBalancer.verifyConnectivity() ) );
+        assertThat( exception.getMessage(), startsWith( "Unable to connect to database management service," ) );
+        verify( routingTables ).refreshRoutingTable( any( ConnectionContext.class ) );
+    }
+
+    @Test
+    void shouldThrowOriginalErrorWhenRefreshRoutingTableFails() throws Throwable
+    {
+        ConnectionPool connectionPool = newConnectionPoolMock();
+
+        Rediscovery rediscovery = mock( Rediscovery.class );
+        when( rediscovery.resolve() ).thenReturn( Arrays.asList( A, B ) );
+
+        RoutingTableRegistry routingTables = mock( RoutingTableRegistry.class );
+        when( routingTables.refreshRoutingTable( any( ConnectionContext.class ) ) ).thenThrow( new SecurityException( "boo" ) );
+
+        LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTables, rediscovery );
+
+        SecurityException exception = assertThrows( SecurityException.class, () -> await( loadBalancer.verifyConnectivity() ) );
+        assertThat( exception.getMessage(), startsWith( "boo" ) );
+        verify( routingTables ).refreshRoutingTable( any( ConnectionContext.class ) );
+    }
+
+    @Test
+    void shouldReturnSuccessVerifyConnectivity() throws Throwable
+    {
+        ConnectionPool connectionPool = newConnectionPoolMock();
+
+        Rediscovery rediscovery = mock( Rediscovery.class );
+        when( rediscovery.resolve() ).thenReturn( Arrays.asList( A, B ) );
+
+        RoutingTableRegistry routingTables = mock( RoutingTableRegistry.class );
+        when( routingTables.refreshRoutingTable( any( ConnectionContext.class ) ) ).thenReturn( Futures.completedWithNull() );
+
+        LoadBalancer loadBalancer = newLoadBalancer( connectionPool, routingTables, rediscovery );
+
+        await( loadBalancer.verifyConnectivity() );
+        verify( routingTables ).refreshRoutingTable( any( ConnectionContext.class ) );
+    }
+
     private static ConnectionPool newConnectionPoolMock()
     {
         return newConnectionPoolMockWithFailures( emptySet() );
     }
 
-    private static ConnectionPool newConnectionPoolMockWithFailures(
-            Set<BoltServerAddress> unavailableAddresses )
+    private static ConnectionPool newConnectionPoolMockWithFailures( Set<BoltServerAddress> unavailableAddresses )
+    {
+        return newConnectionPoolMockWithFailures( unavailableAddresses, address -> new ServiceUnavailableException( address + " is unavailable!" ) );
+    }
+
+    private static ConnectionPool newConnectionPoolMockWithFailures( Set<BoltServerAddress> unavailableAddresses, Function<BoltServerAddress, Throwable> errorAction )
     {
         ConnectionPool pool = mock( ConnectionPool.class );
         when( pool.acquire( any( BoltServerAddress.class ) ) ).then( invocation ->
         {
             BoltServerAddress requestedAddress = invocation.getArgument( 0 );
             if ( unavailableAddresses.contains( requestedAddress ) )
             {
-                return Futures.failedFuture( new ServiceUnavailableException( requestedAddress + " is unavailable!" ) );
+                return Futures.failedFuture( errorAction.apply( requestedAddress ) );
             }
 
             return completedFuture( newBoltV4Connection( requestedAddress ) );
@@ -289,6 +362,11 @@ private static Connection newBoltV4Connection( BoltServerAddress address )
         return connection;
     }
 
+    private static ConnectionContext newBoltV4ConnectionContext()
+    {
+        return simple( true );
+    }
+
     private static LoadBalancer newLoadBalancer( ConnectionPool connectionPool, RoutingTable routingTable )
     {
         // Used only in testing
@@ -305,6 +383,12 @@ private static LoadBalancer newLoadBalancer( ConnectionPool connectionPool, Redi
     {
         // Used only in testing
         RoutingTableRegistry routingTables = mock( RoutingTableRegistry.class );
+        return newLoadBalancer( connectionPool, routingTables, rediscovery );
+    }
+
+    private static LoadBalancer newLoadBalancer( ConnectionPool connectionPool, RoutingTableRegistry routingTables, Rediscovery rediscovery )
+    {
+        // Used only in testing
         return new LoadBalancer( connectionPool, routingTables, rediscovery, new LeastConnectedLoadBalancingStrategy( connectionPool, DEV_NULL_LOGGING ),
                 GlobalEventExecutor.INSTANCE, DEV_NULL_LOGGER );
     }