Adding multidb feature detection Driver#supportsMultiDb

For a single driver, we will open a connection with the init server and then check if the server version is 4.0+ and talking with driver using bolt v4.
For a routing driver, we will go through all addresses that are resolved from the init seed server. Then we will check with the first successfully connected server if the server version is 4.0+ and talking with driver using protocol bolt v4.

Author: Zhen Li

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

@@ -194,4 +194,17 @@ public interface Driver extends AutoCloseable
      * @return a {@link CompletionStage completion stage} that represents the asynchronous verification.
      */
     CompletionStage<Void> verifyConnectivityAsync();
+
+    /**
+     * Returns true if the server or cluster the driver connects to supports multi-databases, otherwise false.
+     * @return true if the server or cluster the driver connects to supports multi-databases, otherwise false.
+     */
+    boolean supportsMultiDb();
+
+    /**
+     * Asynchronous check if the server or cluster the driver connects to supports multi-databases.
+     * @return a {@link CompletionStage completion stage} that returns true if the server or cluster
+     * the driver connects to supports multi-databases, otherwise false.
+     */
+    CompletionStage<Boolean> supportsMultiDbAsync();
 }

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

@@ -26,7 +26,7 @@
 import org.neo4j.driver.internal.spi.ConnectionPool;
 import org.neo4j.driver.internal.spi.ConnectionProvider;
 
-import static org.neo4j.driver.internal.async.ImmutableConnectionContext.simple;
+import static org.neo4j.driver.internal.messaging.request.MultiDatabaseUtil.supportsMultiDatabase;
 
 /**
  * Simple {@link ConnectionProvider connection provider} that obtains connections form the given pool only for
@@ -46,15 +46,13 @@ public class DirectConnectionProvider implements ConnectionProvider
     @Override
     public CompletionStage<Connection> acquireConnection( ConnectionContext context )
     {
-        return connectionPool.acquire( address ).thenApply( connection -> new DirectConnection( connection, context.databaseName(), context.mode() ) );
+        return acquireConnection().thenApply( connection -> new DirectConnection( connection, context.databaseName(), context.mode() ) );
     }
 
     @Override
     public CompletionStage<Void> verifyConnectivity()
     {
-        // We verify the connection by establishing a connection with the remote server specified by the address.
-        // Connection context will be ignored as no query is run in this connection and the connection is released immediately.
-        return acquireConnection( simple() ).thenCompose( Connection::release );
+        return acquireConnection().thenCompose( Connection::release );
     }
 
     @Override
@@ -63,8 +61,26 @@ public CompletionStage<Void> close()
         return connectionPool.close();
     }
 
+    @Override
+    public CompletionStage<Boolean> supportsMultiDbAsync()
+    {
+        return acquireConnection().thenCompose( conn -> {
+            boolean supportsMultiDatabase = supportsMultiDatabase( conn );
+            return conn.release().thenApply( ignored -> supportsMultiDatabase );
+        } );
+    }
+
     public BoltServerAddress getAddress()
     {
         return address;
     }
+
+    /**
+     * Used only for grabbing a connection with the server after hello message.
+     * This connection cannot be directly used for running any queries as it is missing necessary connection context
+     */
+    private CompletionStage<Connection> acquireConnection()
+    {
+        return connectionPool.acquire( address );
+    }
 }

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

@@ -141,6 +141,18 @@ public CompletionStage<Void> verifyConnectivityAsync()
         return sessionFactory.verifyConnectivity();
     }
 
+    @Override
+    public boolean supportsMultiDb()
+    {
+        return Futures.blockingGet( supportsMultiDbAsync() );
+    }
+
+    @Override
+    public CompletionStage<Boolean> supportsMultiDbAsync()
+    {
+        return sessionFactory.supportsMultiDbAsync();
+    }
+
     @Override
     public void verifyConnectivity()
     {

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

@@ -30,4 +30,6 @@ public interface SessionFactory
     CompletionStage<Void> verifyConnectivity();
 
     CompletionStage<Void> close();
+
+    CompletionStage<Boolean> supportsMultiDbAsync();
 }

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

@@ -79,6 +79,12 @@ public CompletionStage<Void> close()
         return connectionProvider.close();
     }
 
+    @Override
+    public CompletionStage<Boolean> supportsMultiDbAsync()
+    {
+        return connectionProvider.supportsMultiDbAsync();
+    }
+
     /**
      * Get the underlying connection provider.
      * <p>

driver/src/main/java/org/neo4j/driver/internal/cluster/Rediscovery.java

@@ -18,12 +18,16 @@
  */
 package org.neo4j.driver.internal.cluster;
 
+import java.util.List;
 import java.util.concurrent.CompletionStage;
 
 import org.neo4j.driver.Bookmark;
+import org.neo4j.driver.internal.BoltServerAddress;
 import org.neo4j.driver.internal.spi.ConnectionPool;
 
 public interface Rediscovery
 {
     CompletionStage<ClusterComposition> lookupClusterComposition( RoutingTable routingTable, ConnectionPool connectionPool, Bookmark bookmark );
+
+    List<BoltServerAddress> resolve();
 }

driver/src/main/java/org/neo4j/driver/internal/cluster/RediscoveryImpl.java

@@ -198,7 +198,7 @@ private CompletionStage<ClusterComposition> lookupOnInitialRouter( RoutingTable
         List<BoltServerAddress> addresses;
         try
         {
-            addresses = resolve( initialRouter );
+            addresses = resolve();
         }
         catch ( Throwable error )
         {
@@ -256,9 +256,10 @@ private ClusterComposition handleRoutingProcedureError( Throwable error, Routing
         return null;
     }
 
-    private List<BoltServerAddress> resolve( BoltServerAddress address )
+    @Override
+    public List<BoltServerAddress> resolve()
     {
-        return resolver.resolve( address )
+        return resolver.resolve( initialRouter )
                 .stream()
                 .flatMap( resolved -> resolveAll( BoltServerAddress.from( resolved ) ) )
                 .collect( toList() ); // collect to list to preserve the order

driver/src/main/java/org/neo4j/driver/internal/cluster/RoutingProcedureClusterCompositionProvider.java

@@ -23,16 +23,16 @@
 import java.util.concurrent.CompletionStage;
 
 import org.neo4j.driver.Bookmark;
-import org.neo4j.driver.Record;
 import org.neo4j.driver.Query;
+import org.neo4j.driver.Record;
 import org.neo4j.driver.exceptions.ProtocolException;
 import org.neo4j.driver.exceptions.value.ValueException;
 import org.neo4j.driver.internal.DatabaseName;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.internal.util.Clock;
-import org.neo4j.driver.internal.util.ServerVersion;
 
 import static java.lang.String.format;
+import static org.neo4j.driver.internal.messaging.request.MultiDatabaseUtil.supportsMultiDatabase;
 
 public class RoutingProcedureClusterCompositionProvider implements ClusterCompositionProvider
 {
@@ -59,7 +59,7 @@ public RoutingProcedureClusterCompositionProvider( Clock clock, RoutingContext r
     public CompletionStage<ClusterComposition> getClusterComposition( Connection connection, DatabaseName databaseName, Bookmark bookmark )
     {
         RoutingProcedureRunner runner;
-        if ( connection.serverVersion().greaterThanOrEqual( ServerVersion.v4_0_0 ) )
+        if ( supportsMultiDatabase( connection ) )
         {
             runner = multiDatabaseRoutingProcedureRunner;
         }

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

@@ -20,8 +20,10 @@
 
 import io.netty.util.concurrent.EventExecutorGroup;
 
+import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionStage;
+import java.util.function.Supplier;
 
 import org.neo4j.driver.AccessMode;
 import org.neo4j.driver.Logger;
@@ -49,6 +51,10 @@
 
 import static java.lang.String.format;
 import static org.neo4j.driver.internal.async.ImmutableConnectionContext.simple;
+import static org.neo4j.driver.internal.messaging.request.MultiDatabaseUtil.supportsMultiDatabase;
+import static org.neo4j.driver.internal.util.Futures.completedWithNull;
+import static org.neo4j.driver.internal.util.Futures.failedFuture;
+import static org.neo4j.driver.internal.util.Futures.onErrorContinue;
 
 public class LoadBalancer implements ConnectionProvider
 {
@@ -58,20 +64,29 @@ public class LoadBalancer implements ConnectionProvider
     private final LoadBalancingStrategy loadBalancingStrategy;
     private final EventExecutorGroup eventExecutorGroup;
     private final Logger log;
+    private final Rediscovery rediscovery;
 
     public LoadBalancer( BoltServerAddress initialRouter, RoutingSettings settings, ConnectionPool connectionPool,
             EventExecutorGroup eventExecutorGroup, Clock clock, Logging logging,
             LoadBalancingStrategy loadBalancingStrategy, ServerAddressResolver resolver )
     {
-        this( connectionPool, createRoutingTables( connectionPool, eventExecutorGroup, initialRouter, resolver, settings, clock, logging ),
-                loadBalancerLogger( logging ), loadBalancingStrategy, eventExecutorGroup );
+        this( connectionPool, createRediscovery( eventExecutorGroup, initialRouter, resolver, settings, clock, logging ), settings, loadBalancingStrategy,
+                eventExecutorGroup, clock, loadBalancerLogger( logging ) );
     }
 
-    LoadBalancer( ConnectionPool connectionPool, RoutingTableRegistry routingTables, Logger log, LoadBalancingStrategy loadBalancingStrategy,
-            EventExecutorGroup eventExecutorGroup )
+    private LoadBalancer( ConnectionPool connectionPool, Rediscovery rediscovery, RoutingSettings settings, LoadBalancingStrategy loadBalancingStrategy,
+            EventExecutorGroup eventExecutorGroup, Clock clock, Logger log )
+    {
+        this( connectionPool, createRoutingTables( connectionPool, rediscovery, settings, clock, log ), rediscovery, loadBalancingStrategy, eventExecutorGroup,
+                log );
+    }
+
+    LoadBalancer( ConnectionPool connectionPool, RoutingTableRegistry routingTables, Rediscovery rediscovery, LoadBalancingStrategy loadBalancingStrategy,
+            EventExecutorGroup eventExecutorGroup, Logger log )
     {
         this.connectionPool = connectionPool;
         this.routingTables = routingTables;
+        this.rediscovery = rediscovery;
         this.loadBalancingStrategy = loadBalancingStrategy;
         this.eventExecutorGroup = eventExecutorGroup;
         this.log = log;
@@ -109,6 +124,38 @@ public CompletionStage<Void> close()
         return connectionPool.close();
     }
 
+    @Override
+    public CompletionStage<Boolean> supportsMultiDbAsync()
+    {
+        List<BoltServerAddress> addresses;
+
+        try
+        {
+            addresses = rediscovery.resolve();
+        }
+        catch ( Throwable error )
+        {
+            return failedFuture( error );
+        }
+
+        CompletableFuture<Boolean> result = completedWithNull();
+        Throwable baseError = new ServiceUnavailableException( "Failed to perform multi-databases feature detection with the following servers: " + addresses );
+
+        for ( BoltServerAddress address : addresses )
+        {
+            result = onErrorContinue( result, baseError, () -> supportsMultiDbAsync( address ) );
+        }
+        return onErrorContinue( result, baseError, (Supplier<CompletableFuture<Boolean>>) () -> failedFuture( baseError ) );
+    }
+
+    private CompletionStage<Boolean> supportsMultiDbAsync( BoltServerAddress address )
+    {
+        return connectionPool.acquire( address ).thenCompose( conn -> {
+            boolean supportsMultiDatabase = supportsMultiDatabase( conn );
+            return conn.release().thenApply( ignored -> supportsMultiDatabase );
+        } );
+    }
+
     private CompletionStage<Connection> acquire( AccessMode mode, RoutingTable routingTable )
     {
         AddressSet addresses = addressSet( mode, routingTable );
@@ -181,17 +228,16 @@ private BoltServerAddress selectAddress( AccessMode mode, AddressSet servers )
         }
     }
 
-    private static RoutingTableRegistry createRoutingTables( ConnectionPool connectionPool, EventExecutorGroup eventExecutorGroup, BoltServerAddress initialRouter,
-            ServerAddressResolver resolver, RoutingSettings settings, Clock clock, Logging logging )
+    private static RoutingTableRegistry createRoutingTables( ConnectionPool connectionPool, Rediscovery rediscovery, RoutingSettings settings, Clock clock,
+            Logger log )
     {
-        Logger log = loadBalancerLogger( logging );
-        Rediscovery rediscovery = createRediscovery( eventExecutorGroup, initialRouter, resolver, settings, clock, log );
         return new RoutingTableRegistryImpl( connectionPool, rediscovery, clock, log, settings.routingTablePurgeDelayMs() );
     }
 
     private static Rediscovery createRediscovery( EventExecutorGroup eventExecutorGroup, BoltServerAddress initialRouter, ServerAddressResolver resolver,
-            RoutingSettings settings, Clock clock, Logger log )
+            RoutingSettings settings, Clock clock, Logging logging )
     {
+        Logger log = loadBalancerLogger( logging );
         ClusterCompositionProvider clusterCompositionProvider = new RoutingProcedureClusterCompositionProvider( clock, settings.routingContext() );
         return new RediscoveryImpl( initialRouter, settings, clusterCompositionProvider, eventExecutorGroup, resolver, log );
     }

driver/src/main/java/org/neo4j/driver/internal/messaging/request/MultiDatabaseUtil.java

@@ -20,6 +20,8 @@
 
 import org.neo4j.driver.exceptions.ClientException;
 import org.neo4j.driver.internal.DatabaseName;
+import org.neo4j.driver.internal.spi.Connection;
+import org.neo4j.driver.internal.util.ServerVersion;
 
 public final class MultiDatabaseUtil
 {
@@ -31,4 +33,9 @@ public static void assertEmptyDatabaseName( DatabaseName databaseName, int boltV
                     "Database name: '%s'", boltVersion, databaseName.description() ) );
         }
     }
+
+    public static boolean supportsMultiDatabase( Connection connection )
+    {
+        return connection.serverVersion().greaterThanOrEqual( ServerVersion.v4_0_0 ) && connection.protocol().version() >= 4;
+    }
 }

driver/src/main/java/org/neo4j/driver/internal/spi/ConnectionProvider.java

@@ -36,4 +36,6 @@ public interface ConnectionProvider
     CompletionStage<Void> verifyConnectivity();
 
     CompletionStage<Void> close();
+
+    CompletionStage<Boolean> supportsMultiDbAsync();
 }

driver/src/main/java/org/neo4j/driver/internal/util/Futures.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.Future;
 import java.util.function.BiConsumer;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 
 import org.neo4j.driver.internal.async.connection.EventLoopGroupFactory;
 
@@ -220,6 +221,35 @@ else if ( error2 != null )
         }
     }
 
+    /**
+     * Given a future, if the future completes successfully then return a new completed future with the completed value.
+     * Otherwise if the future completes with an error, then this method first saves the error in the error recorder, and then continues with the onErrorAction.
+     * @param future the future.
+     * @param errorRecorder saves error if the given future completes with an error.
+     * @param onErrorAction continues the future with this action if the future completes with an error.
+     * @param <T> type
+     * @return a new completed future with the same completed value if the given future completes successfully, otherwise continues with the onErrorAction.
+     */
+    @SuppressWarnings( "ThrowableNotThrown" )
+    public static <T> CompletableFuture<T> onErrorContinue( CompletableFuture<T> future, Throwable errorRecorder,
+            Supplier<? extends CompletionStage<T>> onErrorAction )
+    {
+        return future.exceptionally( error -> {
+            Futures.combineErrors( errorRecorder, error );
+            return null;
+        } ).thenCompose( value -> {
+            if ( value != null )
+            {
+                return completedFuture( value );
+            }
+            else
+            {
+                return onErrorAction.get();
+            }
+        } );
+    }
+
+
     private static void safeRun( Runnable runnable )
     {
         try

driver/src/test/java/org/neo4j/driver/integration/RoutingDriverBoltKitTest.java

@@ -1155,6 +1155,34 @@ void shouldRevertToInitialRouterIfKnownRouterThrowsProtocolErrors() throws Excep
         }
     }
 
+    @Test
+    void shouldServerWithBoltV4SupportMultiDb() throws Throwable
+    {
+        StubServer server = StubServer.start( "support_multidb_v4.script", 9001 );
+        try ( Driver driver = GraphDatabase.driver( "neo4j://localhost:9001", INSECURE_CONFIG ) )
+        {
+            assertTrue( driver.supportsMultiDb() );
+        }
+        finally
+        {
+            assertEquals( 0, server.exitStatus() );
+        }
+    }
+
+    @Test
+    void shouldServerWithBoltV3NotSupportMultiDb() throws Throwable
+    {
+        StubServer server = StubServer.start( "support_multidb_v3.script", 9001 );
+        try ( Driver driver = GraphDatabase.driver( "neo4j://localhost:9001", INSECURE_CONFIG ) )
+        {
+            assertFalse( driver.supportsMultiDb() );
+        }
+        finally
+        {
+            assertEquals( 0, server.exitStatus() );
+        }
+    }
+
     private static Driver newDriverWithSleeplessClock( String uriString, Config config )
     {
         DriverFactory driverFactory = new DriverFactoryWithClock( new SleeplessClock() );