Add support of passing routing parameters via driver config

Author: Zhen

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

@@ -33,17 +33,16 @@
 
 public class GetServersProcedureClusterCompositionProvider implements ClusterCompositionProvider
 {
-    private final String GET_SERVERS = "dbms.cluster.routing.getServers";
-    private final String PROTOCOL_ERROR_MESSAGE = "Failed to parse '" + GET_SERVERS +
-                                                         "' result received from server due to ";
+
+    private final String PROTOCOL_ERROR_MESSAGE = "Failed to parse `%s' result received from server due to ";
 
     private final Clock clock;
     private final Logger log;
     private final GetServersProcedureRunner getServersRunner;
 
-    public GetServersProcedureClusterCompositionProvider( Clock clock, Logger log )
+    public GetServersProcedureClusterCompositionProvider( Clock clock, Logger log, RoutingSettings settings )
     {
-        this( clock, log, new GetServersProcedureRunner() );
+        this( clock, log, new GetServersProcedureRunner( settings.routingParameters ) );
     }
 
     GetServersProcedureClusterCompositionProvider( Clock clock, Logger log, GetServersProcedureRunner getServersRunner )
@@ -67,7 +66,9 @@ public ClusterCompositionResponse getClusterComposition( Connection connection )
         {
             return new ClusterCompositionResponse.Failure( new ServiceUnavailableException( format(
                     "Failed to call '%s' procedure on server. " +
-                    "Please make sure that there is a Neo4j 3.1+ causal cluster up running.", GET_SERVERS ), e ) );
+                    "Please make sure that there is a Neo4j 3.1+ causal cluster up running.",
+                    getServersRunner.procedureCalled() ), e
+            ) );
         }
 
         log.info( "Got getServers response: %s", records );
@@ -77,7 +78,8 @@ public ClusterCompositionResponse getClusterComposition( Connection connection )
         if ( records.size() != 1 )
         {
             return new ClusterCompositionResponse.Failure( new ProtocolException( format(
-                    "%srecords received '%s' is too few or too many.", PROTOCOL_ERROR_MESSAGE,
+                    PROTOCOL_ERROR_MESSAGE +
+                    "records received '%s' is too few or too many.", getServersRunner.procedureCalled(),
                     records.size() ) ) );
         }
 
@@ -90,14 +92,16 @@ public ClusterCompositionResponse getClusterComposition( Connection connection )
         catch ( ValueException e )
         {
             return new ClusterCompositionResponse.Failure( new ProtocolException( format(
-                    "%sunparsable record received.", PROTOCOL_ERROR_MESSAGE ), e ) );
+                    PROTOCOL_ERROR_MESSAGE +
+                    "unparsable record received.", getServersRunner.procedureCalled() ), e ) );
         }
 
         // the cluster result is not a legal reply
         if ( !cluster.hasRoutersAndReaders() )
         {
             return new ClusterCompositionResponse.Failure( new ProtocolException( format(
-                    "%sno router or reader found in response.", PROTOCOL_ERROR_MESSAGE ) ) );
+                    PROTOCOL_ERROR_MESSAGE +
+                    "no router or reader found in response.", getServersRunner.procedureCalled() ) ) );
         }
 
         // all good

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

@@ -22,17 +22,49 @@
 import java.util.List;
 
 import org.neo4j.driver.internal.NetworkSession;
-import org.neo4j.driver.internal.SessionResourcesHandler;
 import org.neo4j.driver.internal.spi.Connection;
 import org.neo4j.driver.v1.Record;
 import org.neo4j.driver.v1.Statement;
+import org.neo4j.driver.v1.Value;
+
+import static org.neo4j.driver.internal.SessionResourcesHandler.NO_OP;
+import static org.neo4j.driver.internal.cluster.ServerVersion.v3_2_0;
+import static org.neo4j.driver.internal.cluster.ServerVersion.version;
 
 public class GetServersProcedureRunner
 {
-    private static final String CALL_GET_SERVERS = "CALL dbms.cluster.routing.getServers";
+    private static final String GET_SERVERS = "dbms.cluster.routing.getServers";
+    private static final String GET_SERVERS_V2 = "dbms.cluster.routing.getServersV2";
+
+    private final Value routingParameters;
+    private Statement procedureCalled;
+
+    public GetServersProcedureRunner( Value parameters )
+    {
+        this.routingParameters = parameters;
+    }
 
     public List<Record> run( Connection connection )
     {
-        return NetworkSession.run( connection, new Statement( CALL_GET_SERVERS ), SessionResourcesHandler.NO_OP ).list();
+        if( version( connection.server().version() ).greaterThanOrEqual( v3_2_0 ) )
+        {
+            procedureCalled = new Statement( "CALL " + GET_SERVERS_V2, routingParameters );
+        }
+        else
+        {
+            procedureCalled = new Statement("CALL " + GET_SERVERS );
+        }
+
+        return runProcedure( connection, procedureCalled );
+    }
+
+    List<Record> runProcedure( Connection connection, Statement procedure )
+    {
+        return NetworkSession.run( connection, procedure, NO_OP ).list();
+    }
+
+    Statement procedureCalled()
+    {
+        return procedureCalled;
     }
 }

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

@@ -156,7 +156,8 @@ private RoundRobinAddressSet addressSetFor( AccessMode mode )
     private static Rediscovery createRediscovery( BoltServerAddress initialRouter, RoutingSettings settings,
             Clock clock, Logger log )
     {
-        ClusterCompositionProvider clusterComposition = new GetServersProcedureClusterCompositionProvider( clock, log );
+        ClusterCompositionProvider clusterComposition =
+                new GetServersProcedureClusterCompositionProvider( clock, log, settings );
         return new Rediscovery( initialRouter, settings, clock, log, clusterComposition );
     }
 }

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

@@ -18,14 +18,21 @@
  */
 package org.neo4j.driver.internal.cluster;
 
+import java.util.Map;
+
+import org.neo4j.driver.v1.Value;
+import org.neo4j.driver.v1.Values;
+
 public class RoutingSettings
 {
     final int maxRoutingFailures;
     final long retryTimeoutDelay;
+    final Value routingParameters;
 
-    public RoutingSettings( int maxRoutingFailures, long retryTimeoutDelay )
+    public RoutingSettings( int maxRoutingFailures, long retryTimeoutDelay, Map<String, String> routingParameters )
     {
         this.maxRoutingFailures = maxRoutingFailures;
         this.retryTimeoutDelay = retryTimeoutDelay;
+        this.routingParameters = Values.value( routingParameters );
     }
 }

driver/src/test/java/org/neo4j/driver/v1/util/ServerVersion.java renamed to driver/src/main/java/org/neo4j/driver/internal/cluster/ServerVersion.java

@@ -16,7 +16,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package org.neo4j.driver.v1.util;
+package org.neo4j.driver.internal.cluster;
 
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
@@ -43,6 +43,7 @@ private ServerVersion( int major, int minor, int patch )
     }
     public static final ServerVersion v3_1_0 = new ServerVersion(3, 1, 0);
     public static final ServerVersion v3_0_0 = new ServerVersion(3, 0, 0);
+    public static final ServerVersion v3_2_0 = new ServerVersion(3, 2, 0);
 
     public static ServerVersion version( Driver driver )
     {

driver/src/main/java/org/neo4j/driver/internal/summary/InternalServerInfo.java

@@ -26,7 +26,7 @@ public class InternalServerInfo implements ServerInfo
     private final BoltServerAddress address;
     private final String version;
 
-    public InternalServerInfo(BoltServerAddress address, String version)
+    public InternalServerInfo( BoltServerAddress address, String version )
     {
         this.address = address;
         this.version = version;

driver/src/main/java/org/neo4j/driver/v1/Config.java

@@ -19,6 +19,7 @@
 package org.neo4j.driver.v1;
 
 import java.io.File;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 
@@ -74,6 +75,8 @@ public class Config
     private final int connectionTimeoutMillis;
     private final RetrySettings retrySettings;
 
+    private final Map<String, String> routingParameters;
+
     private Config( ConfigBuilder builder)
     {
         this.logging = builder.logging;
@@ -88,6 +91,8 @@ private Config( ConfigBuilder builder)
         this.routingRetryDelayMillis = builder.routingRetryDelayMillis;
         this.connectionTimeoutMillis = builder.connectionTimeoutMillis;
         this.retrySettings = builder.retrySettings;
+
+        this.routingParameters = builder.routingParameters;
     }
 
     /**
@@ -182,7 +187,7 @@ public static Config defaultConfig()
 
     RoutingSettings routingSettings()
     {
-        return new RoutingSettings( routingFailureLimit, routingRetryDelayMillis );
+        return new RoutingSettings( routingFailureLimit, routingRetryDelayMillis, routingParameters );
     }
 
     RetrySettings retrySettings()
@@ -205,6 +210,7 @@ public static class ConfigBuilder
         private long routingRetryDelayMillis = TimeUnit.SECONDS.toMillis( 5 );
         private int connectionTimeoutMillis = (int) TimeUnit.SECONDS.toMillis( 5 );
         private RetrySettings retrySettings = RetrySettings.DEFAULT;
+        private Map<String,String> routingParameters = null;
 
         private ConfigBuilder() {}
 
@@ -473,6 +479,21 @@ public ConfigBuilder withMaxTransactionRetryTime( long value, TimeUnit unit )
             return this;
         }
 
+        /**
+         * Specify routing parameters that would be passed to server in getServers Procedure call for customized
+         * routing table reply.
+         * This parameter is only valid for the routing driver, a.k.a. the driver created use bolt+routing in URI
+         * scheme with 3.2+ Neo4j Casual Cluster servers.
+         * @param parameters The parameters to pass to getServers Procedure
+         * @since 1.3
+         * @return this builder
+         */
+        public ConfigBuilder withRoutingParameters( Map<String, String> parameters )
+        {
+            this.routingParameters = parameters;
+            return this;
+        }
+
         /**
          * Create a config instance from this builder.
          * @return a {@link Config} instance

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

@@ -139,7 +139,7 @@ private Driver createDriver( DriverFactory driverFactory )
     private Driver createDriver( DriverFactory driverFactory, Config config )
     {
         AuthToken auth = AuthTokens.none();
-        RoutingSettings routingSettings = new RoutingSettings( 42, 42 );
+        RoutingSettings routingSettings = new RoutingSettings( 42, 42, null );
         return driverFactory.newInstance( uri, auth, routingSettings, RetrySettings.DEFAULT, config );
     }
 

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

@@ -873,7 +873,7 @@ private static Driver newDriverWithFixedRetries( String uriString, int retries )
     private static Driver newDriver( String uriString, DriverFactory driverFactory )
     {
         URI uri = URI.create( uriString );
-        RoutingSettings routingConf = new RoutingSettings( 1, 1 );
+        RoutingSettings routingConf = new RoutingSettings( 1, 1, null );
         AuthToken auth = AuthTokens.none();
         return driverFactory.newInstance( uri, auth, routingConf, RetrySettings.DEFAULT, config );
     }

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

@@ -344,7 +344,7 @@ private final Driver driverWithServers( long ttl, Map<String,Object>... serverIn
 
     private Driver driverWithPool( ConnectionPool pool )
     {
-        RoutingSettings settings = new RoutingSettings( 10, 5_000 );
+        RoutingSettings settings = new RoutingSettings( 10, 5_000, null );
         ConnectionProvider connectionProvider = new LoadBalancer( SEED, settings, pool, clock, logging );
         Config config = Config.build().withLogging( logging ).toConfig();
         SessionFactory sessionFactory = new NetworkSessionWithAddressFactory( connectionProvider, config );