Merge pull request #346 from martin-neotech/4.0-protocol-handlers-io

4.0 protocol handlers io

Author: technige

neo4j/io/__init__.py

@@ -90,10 +90,55 @@ class Bolt:
     the handshake was carried out.
     """
 
-    MAGIC_PREAMBLE = 0x6060B017
+    MAGIC_PREAMBLE = b"\x60\x60\xB0\x17"
 
     PROTOCOL_VERSION = None
 
+    @classmethod
+    def get_handshake(cls):
+        """ Return the supported Bolt versions as bytes.
+        The length is 16 bytes as specified in the Bolt version negotiation.
+        :return: bytes
+        """
+        offered_versions = sorted(cls.protocol_handlers().keys(), reverse=True)[:4]
+        return b"".join(version.to_bytes() for version in offered_versions).ljust(16, b"\x00")
+
+    @classmethod
+    def protocol_handlers(cls, protocol_version=None):
+        """ Return a dictionary of available Bolt protocol handlers,
+        keyed by version tuple. If an explicit protocol version is
+        provided, the dictionary will contain either zero or one items,
+        depending on whether that version is supported. If no protocol
+        version is provided, all available versions will be returned.
+
+        :param protocol_version: tuple identifying a specific protocol
+            version (e.g. (3, 5)) or None
+        :return: dictionary of version tuple to handler class for all
+            relevant and supported protocol versions
+        :raise TypeError: if protocol version is not passed in a tuple
+        """
+
+        # Carry out subclass imports locally to avoid circular
+        # dependency issues.
+        from neo4j.io._bolt3 import Bolt3
+        from neo4j.io._bolt4x0 import Bolt4x0
+
+        handlers = {
+            Bolt3.PROTOCOL_VERSION: Bolt3,
+            Bolt4x0.PROTOCOL_VERSION: Bolt4x0
+        }
+
+        if protocol_version is None:
+            return handlers
+
+        if not isinstance(protocol_version, tuple):
+            raise TypeError("Protocol version must be specified as a tuple")
+
+        if protocol_version in handlers:
+            return {protocol_version: handlers[protocol_version]}
+
+        return {}
+
     @classmethod
     def ping(cls, address, *, timeout=None, **config):
         """ Attempt to establish a Bolt connection, returning the
@@ -757,12 +802,18 @@ def _handshake(s, resolved_address):
     """
     local_port = s.getsockname()[1]
 
-    # Send details of the protocol versions supported
-    supported_versions = [3, 0, 0, 0]
-    handshake = [Bolt.MAGIC_PREAMBLE] + supported_versions
-    log.debug("[#%04X]  C: <MAGIC> 0x%08X", local_port, Bolt.MAGIC_PREAMBLE)
-    log.debug("[#%04X]  C: <HANDSHAKE> 0x%08X 0x%08X 0x%08X 0x%08X", local_port, *supported_versions)
-    data = b"".join(struct_pack(">I", num) for num in handshake)
+    # TODO: Optimize logging code
+    handshake = Bolt.get_handshake()
+    import struct
+    handshake = struct.unpack(">16B", handshake)
+    handshake = [handshake[i:i + 4] for i in range(0, len(handshake), 4)]
+
+    supported_versions = [("0x%02X%02X%02X%02X" % (vx[0], vx[1], vx[2], vx[3])) for vx in handshake]
+
+    log.debug("[#%04X]  C: <MAGIC> 0x%08X", local_port, int.from_bytes(Bolt.MAGIC_PREAMBLE, byteorder="big"))
+    log.debug("[#%04X]  C: <HANDSHAKE> %s %s %s %s", local_port, *supported_versions)
+
+    data = Bolt.MAGIC_PREAMBLE + Bolt.get_handshake()
     s.sendall(data)
 
     # Handle the handshake response
@@ -796,17 +847,7 @@ def _handshake(s, resolved_address):
                                  "(looks like HTTP)".format(resolved_address))
     agreed_version = data[-1], data[-2]
     log.debug("[#%04X]  S: <HANDSHAKE> 0x%06X%02X", local_port, agreed_version[1], agreed_version[0])
-    if agreed_version == (0, 0):
-        log.debug("[#%04X]  C: <CLOSE>", local_port)
-        s.shutdown(SHUT_RDWR)
-        s.close()
-    elif agreed_version in ((3, 0),):
-        return s, agreed_version
-    else:
-        log.debug("[#%04X]  S: <CLOSE>", local_port)
-        s.close()
-        raise ProtocolError("Unknown Bolt protocol version: "
-                            "{}".format(agreed_version))
+    return s, agreed_version
 
 
 def connect(address, *, timeout=None, config):

tests/stub/conftest.py

@@ -116,3 +116,16 @@ def __init__(self, *servers):
 @fixture
 def script():
     return lambda *paths: path_join(dirname(__file__), "scripts", *paths)
+
+
+@fixture
+def driver_info():
+    """ Base class for test cases that integrate with a server.
+    """
+    return {
+        "uri": "bolt://localhost:7687",
+        "bolt_routing_uri": "bolt+routing://localhost:7687",
+        "user": "test",
+        "password": "test",
+        "auth_token": ("test", "test")
+    }

tests/stub/scripts/v4x0/empty.script

@@ -0,0 +1,4 @@
+!: BOLT 4
+!: AUTO HELLO
+!: AUTO GOODBYE
+!: AUTO RESET

tests/stub/test_directdriver.py

@@ -19,42 +19,63 @@
 # limitations under the License.
 
 
+import pytest
+
 from neo4j.exceptions import ServiceUnavailable
 
-from neo4j import GraphDatabase, BoltDriver
+from neo4j import (
+    GraphDatabase,
+    BoltDriver,
+)
+
+from tests.stub.conftest import (
+    StubCluster,
+)
 
-from tests.stub.conftest import StubTestCase, StubCluster
+# python -m pytest tests/stub/test_directdriver.py
 
 
-class BoltDriverTestCase(StubTestCase):
+@pytest.mark.parametrize(
+    "test_script",
+    [
+        "v3/empty.script",
+        "v4x0/empty.script",
+    ]
+)
+def test_bolt_uri_constructs_bolt_driver(driver_info, test_script):
+    # python -m pytest tests/stub/test_directdriver.py -s -v -k test_bolt_uri_constructs_bolt_driver
+    with StubCluster(test_script):
+        uri = "bolt://127.0.0.1:9001"
+        with GraphDatabase.driver(uri, auth=driver_info["auth_token"]) as driver:
+            assert isinstance(driver, BoltDriver)
 
-    def test_bolt_uri_constructs_bolt_driver(self):
-        with StubCluster("v3/empty.script"):
-            uri = "bolt://127.0.0.1:9001"
-            with GraphDatabase.driver(uri, auth=self.auth_token) as driver:
-                assert isinstance(driver, BoltDriver)
 
-    def test_direct_disconnect_on_run(self):
-        with StubCluster("v3/disconnect_on_run.script"):
-            uri = "bolt://127.0.0.1:9001"
-            with GraphDatabase.driver(uri, auth=self.auth_token) as driver:
-                with self.assertRaises(ServiceUnavailable):
-                    with driver.session() as session:
-                        session.run("RETURN $x", {"x": 1}).consume()
+def test_direct_disconnect_on_run(driver_info):
+    # python -m pytest tests/stub/test_directdriver.py -s -v -k test_direct_disconnect_on_run
+    with StubCluster("v3/disconnect_on_run.script"):
+        uri = "bolt://127.0.0.1:9001"
+        with GraphDatabase.driver(uri, auth=driver_info["auth_token"]) as driver:
+            with pytest.raises(ServiceUnavailable):
+                with driver.session() as session:
+                    session.run("RETURN $x", {"x": 1}).consume()
 
-    def test_direct_disconnect_on_pull_all(self):
-        with StubCluster("v3/disconnect_on_pull_all.script"):
-            uri = "bolt://127.0.0.1:9001"
-            with GraphDatabase.driver(uri, auth=self.auth_token) as driver:
-                with self.assertRaises(ServiceUnavailable):
-                    with driver.session() as session:
-                        session.run("RETURN $x", {"x": 1}).consume()
 
-    def test_direct_session_close_after_server_close(self):
-        with StubCluster("v3/disconnect_after_init.script"):
-            uri = "bolt://127.0.0.1:9001"
-            with GraphDatabase.driver(uri, auth=self.auth_token, max_retry_time=0,
-                                      acquire_timeout=3, user_agent="test") as driver:
+def test_direct_disconnect_on_pull_all(driver_info):
+    # python -m pytest tests/stub/test_directdriver.py -s -v -k test_direct_disconnect_on_pull_all
+    with StubCluster("v3/disconnect_on_pull_all.script"):
+        uri = "bolt://127.0.0.1:9001"
+        with GraphDatabase.driver(uri, auth=driver_info["auth_token"]) as driver:
+            with pytest.raises(ServiceUnavailable):
                 with driver.session() as session:
-                    with self.assertRaises(ServiceUnavailable):
-                        session.write_transaction(lambda tx: tx.run("CREATE (a:Item)"))
+                    session.run("RETURN $x", {"x": 1}).consume()
+
+
+def test_direct_session_close_after_server_close(driver_info):
+    # python -m pytest tests/stub/test_directdriver.py -s -v -k test_direct_session_close_after_server_close
+    with StubCluster("v3/disconnect_after_init.script"):
+        uri = "bolt://127.0.0.1:9001"
+        with GraphDatabase.driver(uri, auth=driver_info["auth_token"], max_retry_time=0,
+                                  acquire_timeout=3, user_agent="test") as driver:
+            with driver.session() as session:
+                with pytest.raises(ServiceUnavailable):
+                    session.write_transaction(lambda tx: tx.run("CREATE (a:Item)"))

tests/unit/io/test_class_bolt.py

@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+# Copyright (c) 2002-2020 "Neo4j,"
+# Neo4j Sweden AB [http://neo4j.com]
+#
+# This file is part of Neo4j.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import pytest
+from neo4j.io import Bolt
+
+# python -m pytest tests/unit/io/test_class_bolt.py -s -v
+
+
+def test_class_method_protocol_handlers():
+    # python -m pytest tests/unit/io/test_class_bolt.py -s -v -k test_class_method_protocol_handlers
+    protocol_handlers = Bolt.protocol_handlers()
+    assert len(protocol_handlers) == 2
+
+
+@pytest.mark.parametrize(
+    "test_input, expected",
+    [
+        ((0, 0), 0),
+        ((4, 0), 1),
+    ]
+)
+def test_class_method_protocol_handlers_with_protocol_version(test_input, expected):
+    # python -m pytest tests/unit/io/test_class_bolt.py -s -v -k test_class_method_protocol_handlers_with_protocol_version
+    protocol_handlers = Bolt.protocol_handlers(protocol_version=test_input)
+    assert len(protocol_handlers) == expected
+
+
+def test_class_method_protocol_handlers_with_invalid_protocol_version():
+    # python -m pytest tests/unit/io/test_class_bolt.py -s -v -k test_class_method_protocol_handlers_with_invalid_protocol_version
+    with pytest.raises(TypeError):
+        Bolt.protocol_handlers(protocol_version=2)
+
+
+def test_class_method_get_handshake():
+    # python -m pytest tests/unit/io/test_class_bolt.py -s -v -k test_class_method_get_handshake
+    handshake = Bolt.get_handshake()
+    assert handshake == b"\x00\x00\x00\x04\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x00"
+
+
+def test_magic_preamble():
+    # python -m pytest tests/unit/io/test_class_bolt.py -s -v -k test_magic_preamble
+    preamble = 0x6060B017
+    preamble_bytes = preamble.to_bytes(4, byteorder="big")
+    assert Bolt.MAGIC_PREAMBLE == preamble_bytes