bolt 4.1 now adds routing context to the hello message (#439)

Author: martin-neotech

neo4j/io/__init__.py

@@ -177,12 +177,13 @@ def ping(cls, address, *, timeout=None, **config):
             return protocol_version
 
     @classmethod
-    def open(cls, address, *, auth=None, timeout=None, **pool_config):
+    def open(cls, address, *, auth=None, timeout=None, routing_context=None, **pool_config):
         """ Open a new Bolt connection to a given server address.
 
         :param address:
         :param auth:
-        :param timeout: The connection timeout
+        :param timeout: the connection timeout in seconds
+        :param routing_context: dict containing routing context
         :param pool_config:
         :return:
         :raise BoltHandshakeError: raised if the Bolt Protocol can not negotiate a protocol version.
@@ -200,15 +201,15 @@ def open(cls, address, *, auth=None, timeout=None, **pool_config):
         if pool_config.protocol_version == (3, 0):
             # Carry out Bolt subclass imports locally to avoid circular dependency issues.
             from neo4j.io._bolt3 import Bolt3
-            connection = Bolt3(address, s, pool_config.max_connection_lifetime, auth=auth, user_agent=pool_config.user_agent)
+            connection = Bolt3(address, s, pool_config.max_connection_lifetime, auth=auth, user_agent=pool_config.user_agent, routing_context=routing_context)
         elif pool_config.protocol_version == (4, 0):
             # Carry out Bolt subclass imports locally to avoid circular dependency issues.
             from neo4j.io._bolt4x0 import Bolt4x0
-            connection = Bolt4x0(address, s, pool_config.max_connection_lifetime, auth=auth, user_agent=pool_config.user_agent)
+            connection = Bolt4x0(address, s, pool_config.max_connection_lifetime, auth=auth, user_agent=pool_config.user_agent, routing_context=routing_context)
         elif pool_config.protocol_version == (4, 1):
             # Carry out Bolt subclass imports locally to avoid circular dependency issues.
             from neo4j.io._bolt4x1 import Bolt4x1
-            connection = Bolt4x1(address, s, pool_config.max_connection_lifetime, auth=auth, user_agent=pool_config.user_agent)
+            connection = Bolt4x1(address, s, pool_config.max_connection_lifetime, auth=auth, user_agent=pool_config.user_agent, routing_context=routing_context)
         else:
             log.debug("[#%04X]  S: <CLOSE>", s.getpeername()[1])
             s.shutdown(SHUT_RDWR)
@@ -499,27 +500,35 @@ def close(self):
 class BoltPool(IOPool):
 
     @classmethod
-    def open(cls, address, *, auth, pool_config, workspace_config):
+    def open(cls, address, *, auth, pool_config, workspace_config, routing_context=None):
         """Create a new BoltPool
 
         :param address:
         :param auth:
         :param pool_config:
         :param workspace_config:
+        :param routing_context:
         :return: BoltPool
         """
 
+        if routing_context is None:
+            routing_context = {}
+        elif "address" in routing_context:
+            raise ConfigurationError("The key 'address' is reserved for routing context.")
+        routing_context["address"] = str(address)
+
         def opener(addr, timeout):
-            return Bolt.open(addr, auth=auth, timeout=timeout, **pool_config)
+            return Bolt.open(addr, auth=auth, timeout=timeout, routing_context=routing_context, **pool_config)
 
-        pool = cls(opener, pool_config, workspace_config, address)
+        pool = cls(opener, pool_config, workspace_config, routing_context, address)
         seeds = [pool.acquire() for _ in range(pool_config.init_size)]
         pool.release(*seeds)
         return pool
 
-    def __init__(self, opener, pool_config, workspace_config, address):
+    def __init__(self, opener, pool_config, workspace_config, routing_context, address):
         super(BoltPool, self).__init__(opener, pool_config, workspace_config)
         self.address = address
+        self.routing_context = routing_context
 
     def __repr__(self):
         return "<{} address={!r}>".format(self.__class__.__name__, self.address)
@@ -545,10 +554,17 @@ def open(cls, *addresses, auth, pool_config, workspace_config, routing_context=N
         :return: Neo4jPool
         """
 
+        address = addresses[0]
+        if routing_context is None:
+            routing_context = {}
+        elif "address" in routing_context:
+            raise ConfigurationError("The key 'address' is reserved for routing context.")
+        routing_context["address"] = str(address)
+
         def opener(addr, timeout):
-            return Bolt.open(addr, auth=auth, timeout=timeout, **pool_config)
+            return Bolt.open(addr, auth=auth, timeout=timeout, routing_context=routing_context, **pool_config)
 
-        pool = cls(opener, pool_config, workspace_config, routing_context, addresses)
+        pool = cls(opener, pool_config, workspace_config, routing_context, address)
 
         try:
             pool.update_routing_table(database=workspace_config.database)
@@ -558,7 +574,7 @@ def opener(addr, timeout):
         else:
             return pool
 
-    def __init__(self, opener, pool_config, workspace_config, routing_context, addresses):
+    def __init__(self, opener, pool_config, workspace_config, routing_context, address):
         """
 
         :param opener:
@@ -569,15 +585,10 @@ def __init__(self, opener, pool_config, workspace_config, routing_context, addre
         """
         super(Neo4jPool, self).__init__(opener, pool_config, workspace_config)
         # Each database have a routing table, the default database is a special case.
-        log.debug("[#0000]  C: <NEO4J POOL> routing addresses %r", addresses)
-        self.init_address = addresses[0]
-        self.routing_tables = {workspace_config.database: RoutingTable(database=workspace_config.database, routers=addresses)}
+        log.debug("[#0000]  C: <NEO4J POOL> routing address %r", address)
+        self.address = address
+        self.routing_tables = {workspace_config.database: RoutingTable(database=workspace_config.database, routers=[address])}
         self.routing_context = routing_context
-        if self.routing_context is None:
-            self.routing_context = {}
-        elif "address" in self.routing_context:
-            raise ConfigurationError("The key 'address' is reserved for routing context.")
-        self.routing_context["address"] = str(self.init_address)
         self.refresh_lock = Lock()
 
     def __repr__(self):
@@ -621,7 +632,7 @@ def fetch_routing_info(self, *, address, timeout, database):
         :param address: router address
         :param timeout: seconds
         :param database: the data base name to get routing table for
-        :param init_address: the address by which the client initially contacted the server as a hint for inclusion in the returned routing table.
+        :param address: the address by which the client initially contacted the server as a hint for inclusion in the returned routing table.
 
         :return: list of routing records or
                  None if no connection could be established

neo4j/io/_bolt3.py

@@ -76,7 +76,7 @@ class Bolt3(Bolt):
     #: The pool of which this connection is a member
     pool = None
 
-    def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=None, user_agent=None):
+    def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=None, user_agent=None, routing_context=None):
         self.unresolved_address = unresolved_address
         self.socket = sock
         self.server_info = ServerInfo(Address(sock.getpeername()), Bolt3.PROTOCOL_VERSION)
@@ -90,6 +90,7 @@ def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=No
         self.supports_multiple_results = False
         self.supports_multiple_databases = False
         self._is_reset = True
+        self.routing_context = routing_context
 
         # Determine the user agent
         if user_agent:

neo4j/io/_bolt4x0.py

@@ -75,7 +75,7 @@ class Bolt4x0(Bolt):
     #: The pool of which this connection is a member
     pool = None
 
-    def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=None, user_agent=None):
+    def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=None, user_agent=None, routing_context=None):
         self.unresolved_address = unresolved_address
         self.socket = sock
         self.server_info = ServerInfo(Address(sock.getpeername()), Bolt4x0.PROTOCOL_VERSION)
@@ -89,6 +89,7 @@ def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=No
         self.supports_multiple_results = True
         self.supports_multiple_databases = True
         self._is_reset = True
+        self.routing_context = routing_context
 
         # Determine the user agent
         if user_agent:

neo4j/io/_bolt4x1.py

@@ -75,7 +75,7 @@ class Bolt4x1(Bolt):
     #: The pool of which this connection is a member
     pool = None
 
-    def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=None, user_agent=None):
+    def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=None, user_agent=None, routing_context=None):
         self.unresolved_address = unresolved_address
         self.socket = sock
         self.server_info = ServerInfo(Address(sock.getpeername()), Bolt4x1.PROTOCOL_VERSION)
@@ -89,6 +89,7 @@ def __init__(self, unresolved_address, sock, max_connection_lifetime, *, auth=No
         self.supports_multiple_results = True
         self.supports_multiple_databases = True
         self._is_reset = True
+        self.routing_context = routing_context
 
         # Determine the user agent
         if user_agent:
@@ -135,6 +136,7 @@ def local_port(self):
     def hello(self):
         headers = {"user_agent": self.user_agent}
         headers.update(self.auth_dict)
+        headers["routing"] = self.routing_context
         logged_headers = dict(headers)
         if "credentials" in logged_headers:
             logged_headers["credentials"] = "*******"

tests/stub/conftest.py

@@ -21,6 +21,7 @@
 
 import subprocess
 import os
+import time
 
 from threading import Thread
 from time import sleep
@@ -44,18 +45,20 @@ def __init__(self, port, script):
     def run(self):
         self._process = subprocess.Popen(["python", "-m", "boltkit", "stub", "-v", "-l", ":{}".format(str(self.port)), "-t", "10", self.script], stdout=subprocess.PIPE)
         # Need verbose for this to work
-        line = self._process.stdout.readline()
+        line = self._process.stdout.readline().decode("utf-8")
+        log.debug("started stub server {}".format(self.port))
+        log.debug(line.strip("\n"))
 
     def wait(self):
-        try:
-            returncode = self._process.wait(2)
-            if returncode != 0:
-                log.debug("stubserver return code {}".format(returncode))
-                log.debug("check for miss match in script")
-            return returncode == 0
-        except subprocess.TimeoutExpired:
-            log.debug("stubserver timeout!")
-            return False
+        while True:
+            return_code = self._process.poll()
+            if return_code is not None:
+                line = self._process.stdout.readline().decode("utf-8")
+                if line == "":
+                    break
+                log.debug(line.strip("\n"))
+
+        return True
 
     def kill(self):
         # Kill process if not already dead
@@ -129,6 +132,7 @@ class DefaultBoltStubService(BoltStubService):
 class StubCluster(StubCluster):
 
     def __init__(self, *servers):
+        print("")
         scripts = [os.path.join(os.path.dirname(__file__), "scripts", server) for server in servers]
 
         bss = DefaultBoltStubService.load(*scripts)

tests/stub/scripts/v3/empty.script

@@ -2,3 +2,4 @@
 !: AUTO HELLO
 !: AUTO GOODBYE
 !: AUTO RESET
+!: PORT 9001

tests/stub/scripts/v3/empty_explicit_hello_goodbye.script

@@ -1,4 +1,5 @@
 !: BOLT 3
+!: PORT 9001
 
 C: HELLO {"user_agent": "test", "scheme": "basic", "principal": "test", "credentials": "test"}
 S: SUCCESS {"server": "Neo4j/3.5.0", "connection_id": "123e4567-e89b-12d3-a456-426655440000"}

tests/stub/scripts/v4x0/empty.script

@@ -1,4 +1,5 @@
 !: BOLT 4
 !: AUTO HELLO
 !: AUTO GOODBYE
-!: AUTO RESET
+!: AUTO RESET
+!: PORT 9001

tests/stub/scripts/v4x0/empty_explicit_hello_goodbye.script

@@ -1,4 +1,5 @@
 !: BOLT 4
+!: PORT 9001
 
 C: HELLO {"user_agent": "test", "scheme": "basic", "principal": "test", "credentials": "test"}
 S: SUCCESS {"server": "Neo4j/4.0.0", "connection_id": "123e4567-e89b-12d3-a456-426655440000"}

tests/stub/scripts/v4x1/empty_explicit_hello_goodbye.script

@@ -1,6 +1,7 @@
 !: BOLT 4.1
+!: PORT 9001
 
-C: HELLO {"user_agent": "test", "scheme": "basic", "principal": "test", "credentials": "test"}
+C: HELLO {"user_agent": "test", "scheme": "basic", "principal": "test", "credentials": "test", "routing": {"address": "localhost:9001"}}
 S: SUCCESS {"server": "Neo4j/4.1.0", "connection_id": "123e4567-e89b-12d3-a456-426655440000"}
 C: GOODBYE
 S: <EXIT>