Router and test_routing with boltkit

Author: technige

neo4j/v1/exceptions.py

@@ -51,3 +51,8 @@ class TransactionError(Exception):
 class ResultError(Exception):
     """ Raised when an error occurs while consuming a result.
     """
+
+
+class ServiceUnavailable(Exception):
+    """ Raised when no database service is available.
+    """

neo4j/v1/session.py

@@ -40,7 +40,8 @@
 from .constants import DEFAULT_PORT, ENCRYPTION_DEFAULT, TRUST_DEFAULT, TRUST_SIGNED_CERTIFICATES, \
     TRUST_ON_FIRST_USE, READ_ACCESS, TRUST_SYSTEM_CA_SIGNED_CERTIFICATES, \
     TRUST_ALL_CERTIFICATES, TRUST_CUSTOM_CA_SIGNED_CERTIFICATES
-from .exceptions import CypherError, ProtocolError, ResultError, TransactionError
+from .exceptions import CypherError, ProtocolError, ResultError, TransactionError, \
+    ServiceUnavailable
 from .ssl_compat import SSL_AVAILABLE, SSLContext, PROTOCOL_SSLv23, OP_NO_SSLv2, CERT_REQUIRED
 from .summary import ResultSummary
 from .types import hydrated
@@ -196,53 +197,88 @@ def session(self, access_mode=None):
         return Session(self.pool.acquire(self.address))
 
 
-class RoutingDriver(Driver):
-    """ A :class:`.RoutingDriver` is created from a `bolt+routing` URI.
+def parse_address(address):
+    """ Convert an address string to a tuple.
     """
+    host, _, port = address.partition(":")
+    return host, int(port)
 
-    def __init__(self, address, **config):
-        self.address = address
-        self.security_plan = security_plan = SecurityPlan.build(address, **config)
-        self.encrypted = security_plan.encrypted
-        if not security_plan.routing_compatible:
-            # this error message is case-specific as there is only one incompatible
-            # scenario right now
-            raise RuntimeError("TRUST_ON_FIRST_USE is not compatible with routing")
-        Driver.__init__(self, lambda a: connect(a, security_plan.ssl_context, **config))
-        self._lock = Lock()
-        self._expiry_time = None
-        self._routers = RoundRobinSet([address])
-        self._readers = RoundRobinSet()
-        self._writers = RoundRobinSet()
-        self.discover()
+
+class Router(object):
+
+    timer = monotonic
+
+    def __init__(self, pool, initial_address):
+        self.pool = pool
+        self.lock = Lock()
+        self.expiry_time = None
+        self.routers = RoundRobinSet([initial_address])
+        self.readers = RoundRobinSet()
+        self.writers = RoundRobinSet()
+
+    def stale(self):
+        expired = self.expiry_time is None or self.expiry_time <= self.timer()
+        return expired or len(self.routers) <= 1 or not self.readers or not self.writers
 
     def discover(self):
-        with self._lock:
-            for router in list(self._routers):
+        with self.lock:
+            if not self.routers:
+                raise ServiceUnavailable("No routers available")
+            for router in list(self.routers):
                 session = Session(self.pool.acquire(router))
                 try:
                     record = session.run("CALL dbms.cluster.routing.getServers").single()
+                except CypherError as error:
+                    if error.code == "Neo.ClientError.Procedure.ProcedureNotFound":
+                        raise ServiceUnavailable("Server does not support routing")
+                    raise
                 except ResultError:
                     raise RuntimeError("TODO")
-                new_expiry_time = monotonic() + record["ttl"]
+                new_expiry_time = self.timer() + record["ttl"]
                 servers = record["servers"]
                 new_routers = [s["addresses"] for s in servers if s["role"] == "ROUTE"][0]
                 new_readers = [s["addresses"] for s in servers if s["role"] == "READ"][0]
                 new_writers = [s["addresses"] for s in servers if s["role"] == "WRITE"][0]
                 if new_routers and new_readers and new_writers:
-                    self._expiry_time = new_expiry_time
-                    self._routers.replace(new_routers)
-                    self._readers.replace(new_readers)
-                    self._writers.replace(new_writers)
-                else:
-                    raise RuntimeError("TODO")
+                    self.expiry_time = new_expiry_time
+                    self.routers.replace(map(parse_address, new_routers))
+                    self.readers.replace(map(parse_address, new_readers))
+                    self.writers.replace(map(parse_address, new_writers))
+                    return
+            raise ServiceUnavailable("Unable to establish routing information")
+
+    def acquire_read_connection(self):
+        if self.stale():
+            self.discover()
+        return self.pool.acquire(next(self.readers))
+
+    def acquire_write_connection(self):
+        if self.stale():
+            self.discover()
+        return self.pool.acquire(next(self.writers))
+
+
+class RoutingDriver(Driver):
+    """ A :class:`.RoutingDriver` is created from a `bolt+routing` URI.
+    """
+
+    def __init__(self, address, **config):
+        self.security_plan = security_plan = SecurityPlan.build(address, **config)
+        self.encrypted = security_plan.encrypted
+        if not security_plan.routing_compatible:
+            # this error message is case-specific as there is only one incompatible
+            # scenario right now
+            raise RuntimeError("TRUST_ON_FIRST_USE is not compatible with routing")
+        Driver.__init__(self, lambda a: connect(a, security_plan.ssl_context, **config))
+        self.router = Router(self.pool, address)
+        self.router.discover()
 
     def session(self, access_mode=None):
         if access_mode == READ_ACCESS:
-            address = next(self._readers)
+            connection = self.router.acquire_read_connection()
         else:
-            address = next(self._writers)
-        return Session(self.pool.acquire(address))
+            connection = self.router.acquire_write_connection()
+        return Session(connection)
 
 
 class StatementResult(object):

test/resources/discover_servers.script

@@ -4,5 +4,5 @@
 C: RUN "CALL dbms.cluster.routing.getServers" {}
    PULL_ALL
 S: SUCCESS {"fields": ["ttl", "servers"]}
-   RECORD [9223372036854775807, [{"addresses": ["127.0.0.1:9001"],"role": "WRITE"}, {"addresses": ["127.0.0.1:9002","127.0.0.1:9003"], "role": "READ"},{"addresses": ["127.0.0.1:9001","127.0.0.1:9002","127.0.0.1:9003"], "role": "ROUTE"}]]
+   RECORD [300, [{"addresses": ["127.0.0.1:9001"],"role": "WRITE"}, {"addresses": ["127.0.0.1:9002","127.0.0.1:9003"], "role": "READ"},{"addresses": ["127.0.0.1:9001","127.0.0.1:9002","127.0.0.1:9003"], "role": "ROUTE"}]]
    SUCCESS {}

test/resources/return_1.script

@@ -0,0 +1,13 @@
+!: AUTO INIT
+!: AUTO RESET
+
+C: RUN "CALL dbms.cluster.routing.getServers" {}
+   PULL_ALL
+S: SUCCESS {"fields": ["ttl", "servers"]}
+   RECORD [300, [{"addresses": ["127.0.0.1:9001"],"role": "WRITE"}, {"addresses": ["127.0.0.1:9001"], "role": "READ"},{"addresses": ["127.0.0.1:9001", "127.0.0.1:9002"], "role": "ROUTE"}]]
+   SUCCESS {}
+C: RUN "RETURN $x" {"x": 1}
+   PULL_ALL
+S: SUCCESS {"fields": ["x"]}
+   RECORD [1]
+   SUCCESS {}

test/test_routing.py

@@ -19,39 +19,49 @@
 # limitations under the License.
 
 
-# from os.path import join as join_path, normpath
-# from subprocess import check_call
-from unittest import TestCase
+from neo4j.v1 import basic_auth, GraphDatabase, Router, ConnectionPool, connect, READ_ACCESS
+from test.util import ServerTestCase
 
-from neo4j.v1 import basic_auth, GraphDatabase, RoutingDriver, READ_ACCESS
 
+class RouterTestCase(ServerTestCase):
 
-# from os.path import dirname
+    def setUp(self):
+        self.pool = ConnectionPool(lambda a: connect(a, auth=basic_auth("neo4j", "password")))
 
-# TEST = normpath(dirname(__file__))
-# TEST_RESOURCES = join_path(TEST, "resources")
-# BOLT_ROUTING_URI = "bolt+routing://127.0.0.1:9001"
-# AUTH_TOKEN = basic_auth("neotest", "neotest")
+    def test_router_is_initially_stale(self):
+        router = Router(self.pool, ("127.0.0.1", 7687))
+        assert router.stale()
 
+    def test_discovery(self):
+        self.start_stub_server(9001, "discover_servers.script")
+        router = Router(self.pool, ("127.0.0.1", 9001))
+        router.timer = lambda: 0
+        router.discover()
+        assert router.expiry_time == 300
+        assert router.routers == {'127.0.0.1:9001', '127.0.0.1:9002', '127.0.0.1:9003'}
+        assert router.readers == {'127.0.0.1:9002', '127.0.0.1:9003'}
+        assert router.writers == {'127.0.0.1:9001'}
 
-class LocalClusterIntegrationTestCase(TestCase):
+
+class LocalClusterIntegrationTestCase(ServerTestCase):
+
+    def test_should_discover_servers_on_driver_construction(self):
+        self.start_stub_server(9001, "discover_servers.script")
+        uri = "bolt+routing://127.0.0.1:9001"
+        driver = GraphDatabase.driver(uri, auth=basic_auth("neo4j", "password"), encrypted=False)
+        assert driver.router.routers == {('127.0.0.1', 9001), ('127.0.0.1', 9002), ('127.0.0.1', 9003)}
+        assert driver.router.readers == {('127.0.0.1', 9002), ('127.0.0.1', 9003)}
+        assert driver.router.writers == {('127.0.0.1', 9001)}
 
     def test_should_be_able_to_run_cypher(self):
-        uri = "bolt+routing://ec2-54-78-203-70.eu-west-1.compute.amazonaws.com:26000"
-        driver = GraphDatabase.driver(uri, auth=basic_auth("neo4j", "password"))
+        self.start_stub_server(9001, "return_1.script")
+        uri = "bolt+routing://127.0.0.1:9001"
+        driver = GraphDatabase.driver(uri, auth=basic_auth("neo4j", "password"), encrypted=False)
         try:
             with driver.session(READ_ACCESS) as session:
-                result = session.run("UNWIND range(1, 3) AS n RETURN n")
+                result = session.run("RETURN $x", {"x": 1})
                 for record in result:
-                    print(record)
-                print(result.summary.metadata)
-                print(session.connection.address)
+                    assert record["x"] == 1
+                assert session.connection.address == ('127.0.0.1', 9001)
         finally:
             driver.close()
-
-    def test_should_discover_servers_on_driver_construction(self):
-        uri = "bolt+routing://ec2-54-78-203-70.eu-west-1.compute.amazonaws.com:26000"
-        driver = GraphDatabase.driver(uri, auth=basic_auth("neo4j", "password"))
-        print(driver._routers)
-        print(driver._readers)
-        print(driver._writers)

test/util.py

@@ -21,9 +21,10 @@
 
 import functools
 from os import getenv, remove, rename
-from os.path import isfile
+from os.path import isfile, dirname, join as path_join
 from socket import create_connection
 from subprocess import check_call, CalledProcessError
+from threading import Thread
 from time import sleep
 from unittest import TestCase
 
@@ -96,3 +97,18 @@ def tearDown(self):
             if isfile(self.known_hosts):
                 remove(self.known_hosts)
             rename(self.known_hosts_backup, self.known_hosts)
+
+    def start_stub_server(self, port, script):
+        StubServer(port, script).start()
+        sleep(0.5)
+
+
+class StubServer(Thread):
+
+    def __init__(self, port, script):
+        super(StubServer, self).__init__()
+        self.port = port
+        self.script = path_join(dirname(__file__), "resources", script)
+
+    def run(self):
+        check_call(["boltstub", str(self.port), self.script])