Adding tests for Bolt 5.4 and telemetry (#594)

Co-authored-by: Antonio Barcélos <[email protected]>

Author: robsdedude

boltstub/bolt_protocol.py

@@ -526,3 +526,20 @@ class Bolt5x3Protocol(Bolt5x2Protocol):
     equivalent_versions = set()
 
     server_agent = "Neo4j/5.9.0"
+
+
+class Bolt5x4Protocol(Bolt5x3Protocol):
+    protocol_version = (5, 4)
+    version_aliases = set()
+    # allow the server to negotiate other bolt versions
+    equivalent_versions = set()
+
+    server_agent = "Neo4j/5.14.0"
+
+    messages = {
+        "C": {
+            **Bolt5x3Protocol.messages["C"],
+            b"\x54": "TELEMETRY",
+        },
+        "S": Bolt5x3Protocol.messages["S"],
+    }

nutkit/frontend/__init__.py

@@ -10,3 +10,4 @@
 from .driver import Driver
 from .exceptions import ApplicationCodeError
 from .fake_time import FakeTime
+from .session import Session

nutkit/frontend/driver.py

@@ -17,7 +17,8 @@ def __init__(self, backend, uri, auth_token, user_agent=None,
                  max_connection_pool_size=None,
                  connection_acquisition_timeout_ms=None,
                  notifications_min_severity=None,
-                 notifications_disabled_categories=None):
+                 notifications_disabled_categories=None,
+                 telemetry_disabled=None):
         self._backend = backend
         self._resolver_fn = resolver_fn
         self._domain_name_resolver_fn = domain_name_resolver_fn
@@ -48,7 +49,8 @@ def __init__(self, backend, uri, auth_token, user_agent=None,
             max_connection_pool_size=max_connection_pool_size,
             connection_acquisition_timeout_ms=connection_acquisition_timeout_ms,  # noqa: E501
             notifications_min_severity=notifications_min_severity,
-            notifications_disabled_categories=notifications_disabled_categories
+            notifications_disabled_categories=notifications_disabled_categories,  # noqa: E501
+            telemetry_disabled=telemetry_disabled
         )
         res = backend.send_and_receive(req)
         if not isinstance(res, protocol.Driver):

nutkit/protocol/feature.py

@@ -107,6 +107,8 @@ class Feature(Enum):
     BOLT_5_2 = "Feature:Bolt:5.2"
     # The driver supports Bolt protocol version 5.3
     BOLT_5_3 = "Feature:Bolt:5.3"
+    # The driver supports Bolt protocol version 5.4
+    BOLT_5_4 = "Feature:Bolt:5.4"
     # The driver supports patching DateTimes to use UTC for Bolt 4.3 and 4.4
     BOLT_PATCH_UTC = "Feature:Bolt:Patch:UTC"
     # The driver supports impersonation

nutkit/protocol/requests.py

@@ -74,7 +74,8 @@ def __init__(
         liveness_check_timeout_ms=None, max_connection_pool_size=None,
         connection_acquisition_timeout_ms=None,
         notifications_min_severity=None,
-        notifications_disabled_categories=None
+        notifications_disabled_categories=None,
+        telemetry_disabled=None
     ):
         # Neo4j URI to connect to
         self.uri = uri
@@ -96,6 +97,8 @@ def __init__(
             self.notificationsMinSeverity = notifications_min_severity
         if notifications_disabled_categories is not None:
             self.notificationsDisabledCategories = notifications_disabled_categories  # noqa: E501
+        if telemetry_disabled is not None:
+            self.telemetryDisabled = telemetry_disabled
         # (bool) whether to enable or disable encryption
         # field missing in message: use driver default (should be False)
         if encrypted is not None:

tests/stub/driver_parameters/telemetry/__init__.py

@@ -0,0 +1,54 @@
+!: BOLT 5.4
+
+C: HELLO {"{}": "*"}
+IF: #SERVER_TELEMETRY_ENABLED#
+    S: SUCCESS {"server": "Neo4j/5.14.0", "connection_id": "bolt-1", "hints": {"telemetry.enabled": true}}
+ELSE:
+    S: SUCCESS {"server": "Neo4j/5.14.0", "connection_id": "bolt-1"}
+A: LOGON {"{}": "*"}
+*: RESET
+
+{{
+    {{
+        {{
+            # transaction function
+            C: TELEMETRY {"Z": "0"}
+            ----
+            # unmanaged transaction
+            C: TELEMETRY {"Z": "1"}
+            ----
+            # driver level
+            C: TELEMETRY {"Z": "3"}
+        }}
+        C: BEGIN {"{}": "*"}
+        S: SUCCESS {}
+        S: SUCCESS {}
+    ----
+        C: BEGIN {"{}": "*"}
+        S: SUCCESS {}
+    }}
+    C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"}
+    S: SUCCESS {"fields": ["n"]}
+    C: PULL {"{}": "*"}
+    S: RECORD [1]
+    S: SUCCESS {"type": "r"}
+    C: COMMIT
+    S: SUCCESS {}
+----
+    # auto commit
+    {{
+        C: TELEMETRY {"Z": "2"}
+        C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"}
+        S: SUCCESS {}
+        S: SUCCESS {"fields": ["n"]}
+    ----
+        C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"}
+        S: SUCCESS {"fields": ["n"]}
+    }}
+    C: PULL {"{}": "*"}
+    S: RECORD [1]
+    S: SUCCESS {"type": "r"}
+}}
+
+*: RESET
+?: GOODBYE

tests/stub/driver_parameters/telemetry/scripts/reader.script

@@ -0,0 +1,50 @@
+!: BOLT 5.4
+!: ALLOW RESTART
+!: PY connection_count = 0
+
+PY: connection_count += 1
+C: HELLO {"{}": "*"}
+S: SUCCESS {"server": "Neo4j/5.14.0", "connection_id": "bolt-1", "hints": {"telemetry.enabled": true}}
+A: LOGON {"{}": "*"}
+*: RESET
+
+IF: connection_count == 1
+{{
+    C: TELEMETRY {"Z": "*"}
+    S: <EXIT>
+    # oopsies... dead connection... time for a retry
+}}
+ELIF: connection_count == 2
+{{
+    # The driver is expected to send the TELEMETRY message again.
+    # This time, it'll be successful.
+    C: TELEMETRY {"Z": "*"}
+    C: BEGIN {"{}": "*"}
+    S: SUCCESS {}
+    S: FAILURE {"code": "Neo.TransientError.General.PleaseRetry", "message": "something went wrong™"}
+
+    {?
+        C: RUN {"U": "*"} {"{}": "*"} {"{}": "*"}
+        S: IGNORED
+    ?}
+    {?
+        C: PULL {"{}": "*"}
+        S: IGNORED
+    ?}
+
+    +: RESET
+
+    # The previous TELEMETRY message was successful, so the driver should not send it again.
+    C: BEGIN {"{}": "*"}
+    S: SUCCESS {}
+    C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"}
+    S: SUCCESS {"fields": ["n"]}
+    C: PULL {"{}": "*"}
+    S: RECORD [1]
+    S: SUCCESS {"type": "r"}
+    C: COMMIT
+    S: SUCCESS {}
+}}
+
+*: RESET
+?: GOODBYE

tests/stub/driver_parameters/telemetry/scripts/reader_telemetry_retry.script

@@ -0,0 +1,17 @@
+!: BOLT 5.4
+!: AUTO RESET
+!: ALLOW RESTART
+
+C: HELLO {"{}": "*"}
+IF: #SERVER_TELEMETRY_ENABLED#
+    S: SUCCESS {"server": "Neo4j/5.14.0", "connection_id": "bolt-2", "hints": {"telemetry.enabled": true}}
+ELSE:
+    S: SUCCESS {"server": "Neo4j/5.14.0", "connection_id": "bolt-2"}
+A: LOGON {"{}": "*"}
+*: RESET
+{+
+    C: ROUTE "*" "*" "*"
+    S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["#HOST#:9010"], "role":"ROUTE"}, {"addresses": ["#HOST#:9000"], "role":"READ"}, {"addresses": ["#HOST#:9000"], "role":"WRITE"}]}}
+    *: RESET
++}
+?: GOODBYE