Added RoutingSession to call get routing table

Author: Zhen

neo4j/v1/routing.py

@@ -26,7 +26,7 @@
 from neo4j.bolt import ConnectionPool, ServiceUnavailable, ProtocolError, DEFAULT_PORT, connect
 from neo4j.compat.collections import MutableSet, OrderedDict
 from neo4j.exceptions import CypherError
-from neo4j.v1.api import Driver, READ_ACCESS, WRITE_ACCESS
+from neo4j.v1.api import Driver, READ_ACCESS, WRITE_ACCESS, fix_statement, fix_parameters
 from neo4j.v1.exceptions import SessionExpired
 from neo4j.v1.security import SecurityPlan
 from neo4j.v1.session import BoltSession
@@ -150,14 +150,28 @@ def update(self, new_routing_table):
         self.ttl = new_routing_table.ttl
 
 
-class RoutingConnectionPool(ConnectionPool):
-    """ Connection pool with routing table.
-    """
+class RoutingSession(BoltSession):
 
     call_get_servers = "CALL dbms.cluster.routing.getServers"
     get_routing_table_param = "context"
     call_get_routing_table = "CALL dbms.cluster.routing.getRoutingTable({%s})" % get_routing_table_param
 
+    def routing_info_procedure(self, routing_context):
+        if ServerVersion.from_str(self._connection.server.version).at_least_version(3, 2):
+            return self.call_get_routing_table, {self.get_routing_table_param: routing_context}
+        else:
+            return self.call_get_servers, {}
+
+    def __run__(self, ignored, routing_context):
+        # the statement is ignored as it will be get routing table procedure call.
+        statement, parameters = self.routing_info_procedure(routing_context)
+        return self._run(fix_statement(statement), fix_parameters(parameters))
+
+
+class RoutingConnectionPool(ConnectionPool):
+    """ Connection pool with routing table.
+    """
+
     def __init__(self, connector, initial_address, routing_context, *routers):
         super(RoutingConnectionPool, self).__init__(connector)
         self.initial_address = initial_address
@@ -166,12 +180,6 @@ def __init__(self, connector, initial_address, routing_context, *routers):
         self.missing_writer = False
         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, {}
-
     def fetch_routing_info(self, address):
         """ Fetch raw routing info from a given router address.
 
@@ -182,15 +190,8 @@ def fetch_routing_info(self, address):
                                    if routing support is broken
         """
         try:
-            connections = [None]
-
-            def connector(_):
-                connection = self.acquire_direct(address)
-                connections[0] = connection
-                return connection
-
-            with BoltSession(lambda _: connector) as session:
-                return list(session.run(*self.routing_info_procedure(connections[0])))
+            with RoutingSession(lambda _: self.acquire_direct(address)) as session:
+                return list(session.run("ignored", self.routing_context))
         except CypherError as error:
             if error.code == "Neo.ClientError.Procedure.ProcedureNotFound":
                 raise ServiceUnavailable("Server {!r} does not support routing".format(address))

neo4j/v1/session.py

@@ -34,7 +34,7 @@ class BoltSession(Session):
     :param bookmark:
     """
 
-    def __run__(self, statement, parameters):
+    def _run(self, statement, parameters):
         assert isinstance(statement, unicode)
         assert isinstance(parameters, dict)
 
@@ -52,6 +52,9 @@ def __run__(self, statement, parameters):
 
         return result
 
+    def __run__(self, statement, parameters):
+        return self._run(statement, parameters)
+
     def __begin__(self):
         return self.__run__(u"BEGIN", {"bookmark": self._bookmark} if self._bookmark else {})