Remove duplication between Bolt4xN classes

Author: technige

neo4j/io/__init__.py

@@ -134,8 +134,7 @@ def protocol_handlers(cls, protocol_version=None):
 
         # Carry out Bolt subclass imports locally to avoid circular dependency issues.
         from neo4j.io._bolt3 import Bolt3
-        from neo4j.io._bolt4x0 import Bolt4x0
-        from neo4j.io._bolt4x1 import Bolt4x1
+        from neo4j.io._bolt4 import Bolt4x0, Bolt4x1
 
         handlers = {
             Bolt3.PROTOCOL_VERSION: Bolt3,
@@ -204,11 +203,11 @@ def open(cls, address, *, auth=None, timeout=None, routing_context=None, **pool_
             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
+            from neo4j.io._bolt4 import Bolt4x0
             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
+            from neo4j.io._bolt4 import Bolt4x1
             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])
@@ -673,8 +672,7 @@ def fetch_routing_info(self, *, address, timeout, database):
 
         # Carry out Bolt subclass imports locally to avoid circular dependency issues.
         from neo4j.io._bolt3 import Bolt3
-        from neo4j.io._bolt4x0 import Bolt4x0
-        from neo4j.io._bolt4x1 import Bolt4x1
+        from neo4j.io._bolt4 import Bolt4x0, Bolt4x1
 
         from neo4j.api import (
             SYSTEM_DATABASE,

neo4j/io/_bolt3.py

@@ -63,6 +63,10 @@
 
 
 class Bolt3(Bolt):
+    """ Protocol handler for Bolt 3.
+
+    This is supported by Neo4j versions 3.5, 4.0 and 4.1.
+    """
 
     PROTOCOL_VERSION = Version(3, 0)
 
@@ -81,7 +85,7 @@ class Bolt3(Bolt):
     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)
+        self.server_info = ServerInfo(Address(sock.getpeername()), self.PROTOCOL_VERSION)
         self.outbox = Outbox()
         self.inbox = Inbox(self.socket, on_error=self._set_defunct)
         self.packer = Packer(self.outbox)
@@ -136,8 +140,13 @@ def local_port(self):
         except IOError:
             return 0
 
+    def get_base_headers(self):
+        return {
+            "user_agent": self.user_agent,
+        }
+
     def hello(self):
-        headers = {"user_agent": self.user_agent}
+        headers = self.get_base_headers()
         headers.update(self.auth_dict)
         logged_headers = dict(headers)
         if "credentials" in logged_headers:

neo4j/io/_bolt4x0.py renamed to neo4j/io/_bolt4.py

@@ -63,6 +63,10 @@
 
 
 class Bolt4x0(Bolt):
+    """ Protocol handler for Bolt 4.0.
+
+    This is supported by Neo4j versions 4.0 and 4.1.
+    """
 
     PROTOCOL_VERSION = Version(4, 0)
 
@@ -81,7 +85,7 @@ class Bolt4x0(Bolt):
     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)
+        self.server_info = ServerInfo(Address(sock.getpeername()), self.PROTOCOL_VERSION)
         self.outbox = Outbox()
         self.inbox = Inbox(self.socket, on_error=self._set_defunct)
         self.packer = Packer(self.outbox)
@@ -136,8 +140,13 @@ def local_port(self):
         except IOError:
             return 0
 
+    def get_base_headers(self):
+        return {
+            "user_agent": self.user_agent,
+        }
+
     def hello(self):
-        headers = {"user_agent": self.user_agent}
+        headers = self.get_base_headers()
         headers.update(self.auth_dict)
         logged_headers = dict(headers)
         if "credentials" in logged_headers:
@@ -445,3 +454,23 @@ def closed(self):
 
     def defunct(self):
         return self._defunct
+
+
+class Bolt4x1(Bolt4x0):
+    """ Protocol handler for Bolt 4.1.
+
+    This is supported by Neo4j version 4.1.
+    """
+
+    PROTOCOL_VERSION = Version(4, 1)
+
+    def get_base_headers(self):
+        """ Bolt 4.1 passes the routing context, originally taken from
+        the URI, into the connection initialisation message. This
+        enables server-side routing to propagate the same behaviour
+        through its driver.
+        """
+        return {
+            "user_agent": self.user_agent,
+            "routing": self.routing_context,
+        }