Support routing context from bolt routing uri

Author: Zhen

neo4j/addressing.py

@@ -79,6 +79,25 @@ def parse(cls, string, default_port=0):
         """
         return cls.from_uri("//{}".format(string), default_port)
 
+    @classmethod
+    def parse_routing_context(cls, uri):
+        query = urlparse(uri).query
+        if not query:
+            return {}
+
+        context = {}
+        parameters = [x for x in query.split('&') if x]
+        for keyValue in parameters:
+            pair = keyValue.split('=')
+            if len(pair) != 2 or not pair[0] or not pair[1]:
+                raise ValueError("Invalid parameters: '" + keyValue + "' in URI '" + uri + "'.")
+            key = pair[0]
+            value = pair[1]
+            if key in context:
+                raise ValueError("Duplicated query parameters with key '" + key + "' found in URL '" + uri + "'")
+            context[key] = value
+        return context
+
 
 def resolve(socket_address):
     try:

neo4j/util.py

@@ -25,6 +25,25 @@
 from sys import stdout
 
 
+class ServerVersion(object):
+    def __init__(self, product, version_tuple, tags_tuple):
+        self.product = product
+        self.version_tuple = version_tuple
+        self.tags_tuple = tags_tuple
+
+    def at_least_version(self, major, minor):
+        return self.version_tuple >= (major, minor)
+
+    @classmethod
+    def from_str(cls, full_version):
+        if full_version is None:
+            return ServerVersion("Neo4j", (3, 0), ())
+        product, _, tagged_version = full_version.partition("/")
+        tags = tagged_version.split("-")
+        version = map(int, tags[0].split("."))
+        return ServerVersion(product, tuple(version), tuple(tags[1:]))
+
+
 class ColourFormatter(logging.Formatter):
     """ Colour formatter for pretty log output.
     """

neo4j/v1/routing.py

@@ -30,6 +30,7 @@
 from neo4j.v1.exceptions import SessionExpired
 from neo4j.v1.security import SecurityPlan
 from neo4j.v1.session import BoltSession
+from neo4j.util import ServerVersion
 
 
 class RoundRobinSet(MutableSet):
@@ -152,14 +153,23 @@ class RoutingConnectionPool(ConnectionPool):
     """ Connection pool with routing table.
     """
 
-    routing_info_procedure = "dbms.cluster.routing.getServers"
+    call_get_servers = "CALL dbms.cluster.routing.getServers"
+    get_routing_table_param = "context"
+    call_get_routing_table = "CALL dbms.cluster.routing.getRoutingTable({" + get_routing_table_param + "})"
 
-    def __init__(self, connector, initial_address, *routers):
+    def __init__(self, connector, initial_address, routing_context, *routers):
         super(RoutingConnectionPool, self).__init__(connector)
         self.initial_address = initial_address
+        self.routing_context = routing_context
         self.routing_table = RoutingTable(routers)
         self.refresh_lock = Lock()
 
+    def routing_info_procedure(self, connection):
+        if ServerVersion.from_str(connection.server.version).at_least_version(3, 2):
+            return self.call_get_routing_table, {self.get_routing_table_param: self.routing_context}
+        else:
+            return self.call_get_servers, None
+
     def fetch_routing_info(self, address):
         """ Fetch raw routing info from a given router address.
 
@@ -170,8 +180,9 @@ def fetch_routing_info(self, address):
                                    if routing support is broken
         """
         try:
-            with BoltSession(lambda _: self.acquire_direct(address)) as session:
-                return list(session.run("CALL %s" % self.routing_info_procedure))
+            connection = self.acquire_direct(address)
+            with BoltSession(lambda _: connection) as session:
+                return list(session.run(*self.routing_info_procedure(connection)))
         except CypherError as error:
             if error.code == "Neo.ClientError.Procedure.ProcedureNotFound":
                 raise ServiceUnavailable("Server {!r} does not support routing".format(address))
@@ -313,6 +324,7 @@ def __init__(self, uri, **config):
         self.initial_address = initial_address = SocketAddress.from_uri(uri, DEFAULT_PORT)
         self.security_plan = security_plan = SecurityPlan.build(**config)
         self.encrypted = security_plan.encrypted
+        routing_context = SocketAddress.parse_routing_context(uri)
         if not security_plan.routing_compatible:
             # this error message is case-specific as there is only one incompatible
             # scenario right now
@@ -321,7 +333,7 @@ def __init__(self, uri, **config):
         def connector(a):
             return connect(a, security_plan.ssl_context, **config)
 
-        pool = RoutingConnectionPool(connector, initial_address, *resolve(initial_address))
+        pool = RoutingConnectionPool(connector, initial_address, routing_context *resolve(initial_address))
         try:
             pool.update_routing_table()
         except:

test/integration/tools.py

@@ -33,6 +33,7 @@
 from boltkit.controller import WindowsController, UnixController
 
 from neo4j.v1 import GraphDatabase, AuthError
+from neo4j.util import ServerVersion
 
 from test.env import NEO4J_SERVER_PACKAGE, NEO4J_USER, NEO4J_PASSWORD
 
@@ -89,17 +90,11 @@ def server_version_info(cls):
         with GraphDatabase.driver(cls.bolt_uri, auth=cls.auth_token) as driver:
             with driver.session() as session:
                 full_version = session.run("RETURN 1").summary().server.version
-                if full_version is None:
-                    return "Neo4j", (3, 0), ()
-                product, _, tagged_version = full_version.partition("/")
-                tags = tagged_version.split("-")
-                version = map(int, tags[0].split("."))
-                return product, tuple(version), tuple(tags[1:])
+                return ServerVersion.from_str(full_version)
 
     @classmethod
     def at_least_version(cls, major, minor):
-        _, server_version, _ = cls.server_version_info()
-        return server_version >= (major, minor)
+        return cls.server_version_info().at_least_version(major, minor);
 
     @classmethod
     def delete_known_hosts_file(cls):