ADR 032: Advertised Address

nutkit/protocol/feature.py

@@ -220,6 +220,10 @@ class Feature(Enum):
     CONF_HINT_CON_RECV_TIMEOUT = "ConfHint:connection.recv_timeout_seconds"
 
     # === BACKEND FEATURES FOR TESTING ===
+    # The backend/driver offers a way to configure a driver with a custom DNS
+    # resolver. This configuration option is for testing purposes and might
+    # not be exposed to the user.
+    BACKEND_DNS_RESOLVER = "Backend:DNSResolver"
     # The backend understands the FakeTimeInstall, FakeTimeUninstall and
     # FakeTimeTick protocol messages and provides a way to mock the system
     # time. This is mainly used for testing various timeouts.

tests/stub/advertised_address/scripts/advertised_address_direct.script

@@ -0,0 +1,24 @@
+!: BOLT 5.8
+
+A: HELLO {"{}": "*"}
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"}
+
+*: RESET
+
+C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [1]
+S: SUCCESS {}
+
+*: RESET
+
+C: RUN "RETURN 2 AS n" {"{}": "*"} {"{}": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [2]
+S: SUCCESS {}
+
+*: RESET
+A: GOODBYE

tests/stub/advertised_address/scripts/advertised_address_direct_warm.script

@@ -0,0 +1,44 @@
+!: BOLT 5.8
+
+A: HELLO {"{}": "*"}
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "something_wild:1234"}
+
+*: RESET
+
+C: RUN "RETURN 1 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [1]
+S: SUCCESS {}
+
+*: RESET
+
+A: LOGOFF
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "#HOST#:#PORT#"}
+
+*: RESET
+
+C: RUN "RETURN 2 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [2]
+S: SUCCESS {}
+
+*: RESET
+
+A: LOGOFF
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "certainly-unusable:7688"}
+
+*: RESET
+
+C: RUN "RETURN 3 AS n" {"{}": "*"} {"mode": "r", "[db]": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [3]
+S: SUCCESS {}
+
+*: RESET
+A: GOODBYE

tests/stub/advertised_address/scripts/advertised_address_routing.script

@@ -0,0 +1,21 @@
+!: BOLT 5.8
+
+A: HELLO {"{}": "*"}
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"}
+
+*: RESET
+
+C: ROUTE {"{}": "*"} {"[]": "*"} {"{}": "*"}
+S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["doesnt_exist2:#PORT#"], "role":"READ"}, {"addresses": ["#ADVERTISED_HOST#:#PORT#"], "role":"WRITE"}]}}
+
+*: RESET
+
+C: RUN "RETURN 1 AS n" {"{}": "*"} {"{}": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [1]
+S: SUCCESS {}
+
+*: RESET
+A: GOODBYE

tests/stub/advertised_address/scripts/advertised_address_routing_warm.script

@@ -0,0 +1,49 @@
+!: BOLT 5.8
+
+A: HELLO {"{}": "*"}
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "#HOST#:#PORT#"}
+
+*: RESET
+
+C: ROUTE {"{}": "*"} {"[]": "*"} {"{}": "*"}
+S: SUCCESS { "rt": { "ttl": 1000, "servers": [{"addresses": ["doesnt_exist:9001"], "role":"ROUTE"}, {"addresses": ["#ADVERTISED_HOST#:#PORT#"], "role":"READ"}, {"addresses": ["#HOST#:#PORT#"], "role":"WRITE"}]}}
+
+*: RESET
+
+C: RUN "RETURN 1 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [1]
+S: SUCCESS {}
+
+*: RESET
+
+A: LOGOFF
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"}
+
+*: RESET
+
+C: RUN "RETURN 2 AS n" {"{}": "*"} {"[mode]": "w", "[db]": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [2]
+S: SUCCESS {}
+
+*: RESET
+
+A: LOGOFF
+C: LOGON {"{}": "*"}
+S: SUCCESS {"advertised_address": "#ADVERTISED_HOST#:#PORT#"}
+
+*: RESET
+
+C: RUN "RETURN 3 AS n" {"{}": "*"} {"mode": "r", "[db]": "*"}
+S: SUCCESS {"fields": ["n"]}
+C: PULL {"{}": "*"}
+S: RECORD [3]
+S: SUCCESS {}
+
+*: RESET
+A: GOODBYE

tests/stub/advertised_address/test_advertised_address.py

@@ -0,0 +1,198 @@
+from abc import ABC
+from collections import deque
+from contextlib import contextmanager
+
+import nutkit.protocol as types
+from nutkit.frontend import Driver
+from tests.shared import (
+    driver_feature,
+    TestkitTestCase,
+)
+from tests.stub.shared import StubServer
+
+_FAKE_ADDRESS = "banana.example.com"
+_FAKE_ADVERTISED_ADDRESS = "cucumber.example.com"
+
+
+class _AdvertisedAddressTestCase(TestkitTestCase, ABC):
+    @contextmanager
+    def server(self, script, vars_=None):
+        server = StubServer(9001)
+        vars_ = {
+            "#ADVERTISED_HOST#": f"{_FAKE_ADVERTISED_ADDRESS}",
+            "#PORT#": server.port,
+            "#HOST#": server.host,
+            **(vars_ or {}),
+        }
+        server.start(path=self.script_path(script),
+                     vars_=vars_)
+        try:
+            yield server
+        except Exception:
+            server.reset()
+            raise
+
+        server.done()
+
+    @contextmanager
+    def driver(self, server, routing=True, dns_resolver=None, **kwargs):
+        auth = types.AuthorizationToken("bearer", credentials="foo")
+        scheme = "neo4j" if routing else "bolt"
+        uri = f"{scheme}://{_FAKE_ADDRESS}:{server.port}"
+        driver = Driver(
+            self._backend, uri, auth, domain_name_resolver_fn=dns_resolver,
+            **kwargs,
+        )
+        try:
+            yield driver
+        finally:
+            driver.close()
+
+    @contextmanager
+    def session(self, driver, access_mode="w", session_config=None):
+        if session_config is None:
+            session_config = {}
+        session = driver.session(access_mode, **session_config)
+        try:
+            yield session
+        finally:
+            session.close()
+
+    @staticmethod
+    def _make_dns_resolver(*expected_resolved_pairs):
+        dns_expectations = deque(expected_resolved_pairs)
+
+        def dns_resolver(name):
+            nonlocal dns_expectations
+            expectation, result = dns_expectations.popleft()
+            parts = name.rsplit(":", 1)
+            sep = port = ""
+            if len(parts) == 2:
+                name, port = parts
+                sep = ":"
+            print(parts, expectation)
+            assert name == expectation
+            return [sep.join((host, port)) for host in result]
+
+        return dns_resolver
+
+
+class TestAdvertisedAddress(_AdvertisedAddressTestCase):
+    required_features = (
+        types.Feature.BOLT_5_8,
+    )
+
+    def test_routing_driver_reuses_connection_according_to_advertised_address(
+        self,
+    ):
+        with self.server("advertised_address_routing.script") as server:
+            self._test_reuses_connection(
+                server,
+                driver_kwargs={"routing": True},
+            )
+
+    def test_direct_driver_reuses_connection_regardless_of_advertised_address(
+        self
+    ):
+        with self.server("advertised_address_direct.script") as server:
+            self._test_reuses_connection(
+                server,
+                driver_kwargs={
+                    "routing": False,
+                    "max_connection_pool_size": 1,
+                },
+                repetitions=2,
+            )
+
+    @driver_feature(types.Feature.BACKEND_DNS_RESOLVER)
+    def _test_reuses_connection(
+        self,
+        server,
+        *,
+        driver_kwargs=None,
+        repetitions=1,
+    ):
+        if driver_kwargs is None:
+            driver_kwargs = {}
+
+        dns_resolver = self._make_dns_resolver((_FAKE_ADDRESS, [server.host]))
+
+        with self.driver(
+            server,
+            dns_resolver=dns_resolver,
+            **driver_kwargs
+        ) as driver:
+            for i in range(repetitions):
+                with self.session(driver) as session:
+                    list(session.run(f"RETURN {i + 1} AS n"))
+
+    @driver_feature(
+        types.Feature.BACKEND_DNS_RESOLVER,
+        types.Feature.API_SESSION_AUTH_CONFIG,
+    )
+    def test_warm_routing_driver_reuses_connection_according_to_advertised_address(  # noqa: E501
+        self,
+    ):
+        with self.server("advertised_address_routing_warm.script") as server:
+            self._test_reuses_warm_connection(
+                server,
+                driver_kwargs={
+                    "routing": True,
+                    "max_connection_pool_size": 1,
+                }
+            )
+
+    @driver_feature(
+        types.Feature.BACKEND_DNS_RESOLVER,
+        types.Feature.API_SESSION_AUTH_CONFIG,
+    )
+    def test_warm_direct_driver_reuses_connection_according_to_advertised_address(  # noqa: E501
+        self,
+    ):
+        with self.server("advertised_address_direct_warm.script") as server:
+            self._test_reuses_warm_connection(
+                server,
+                driver_kwargs={
+                    "routing": False,
+                    "max_connection_pool_size": 1,
+                }
+            )
+
+    @driver_feature(types.Feature.BACKEND_DNS_RESOLVER)
+    def _test_reuses_warm_connection(
+        self,
+        server,
+        *,
+        driver_kwargs=None,
+    ):
+        if driver_kwargs is None:
+            driver_kwargs = {}
+
+        dns_resolver = self._make_dns_resolver((_FAKE_ADDRESS, [server.host]))
+
+        with self.driver(
+            server,
+            dns_resolver=dns_resolver,
+            **driver_kwargs
+        ) as driver:
+            with self.session(
+                driver,
+                session_config={"database": "neo4j"},
+            ) as session:
+                list(session.run("RETURN 1 AS n"))
+
+            # Using session auth to cause LOGOFF/LOGON
+            # during which the server will change its advertised address.
+            auth = types.AuthorizationToken("bearer", credentials="bar")
+            with self.session(
+                driver,
+                session_config={"database": "neo4j", "auth_token": auth}
+            ) as session:
+                list(session.run("RETURN 2 AS n"))
+
+            with self.session(
+                driver,
+                access_mode="r",
+                session_config={"database": "neo4j"},
+            ) as session:
+                list(session.run("RETURN 3 AS n"))

tests/stub/routing/test_routing_v5x0.py

@@ -2412,6 +2412,7 @@ def test_should_ignore_system_bookmark_when_getting_rt_for_multi_db(self):
         self.assertEqual([1], sequence)
         self.assertEqual(["foo:6678"], last_bookmarks)
 
+    @driver_feature(types.Feature.BACKEND_DNS_RESOLVER)
     def _test_should_request_rt_from_all_initial_routers_until_successful(
         self, failure_script
     ):