diff --git a/nutkit/frontend/__init__.py b/nutkit/frontend/__init__.py index 3919b2ab1..f36ec6234 100644 --- a/nutkit/frontend/__init__.py +++ b/nutkit/frontend/__init__.py @@ -1,3 +1,6 @@ -from .config import Neo4jBookmarkManagerConfig +from .bookmark_manager import ( + BookmarkManager, + Neo4jBookmarkManagerConfig, +) from .driver import Driver from .exceptions import ApplicationCodeError diff --git a/nutkit/frontend/bookmark_manager.py b/nutkit/frontend/bookmark_manager.py new file mode 100644 index 000000000..98c9dfefd --- /dev/null +++ b/nutkit/frontend/bookmark_manager.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import ( + Any, + Callable, + ClassVar, + Dict, + List, + Optional, +) + +from nutkit.backend import Backend + +from .. import protocol + + +@dataclass +class Neo4jBookmarkManagerConfig: + initial_bookmarks: Optional[Dict[str, List[str]]] = None + bookmarks_supplier: Optional[Callable[[str], List[str]]] = None + bookmarks_consumer: Optional[Callable[[str, List[str]], None]] = None + + +@dataclass +class BookmarkManager: + _registry: ClassVar[Dict[Any, BookmarkManager]] = {} + + def __init__(self, backend: Backend, config: Neo4jBookmarkManagerConfig): + self._backend = backend + self.config = config + + req = protocol.NewBookmarkManager( + config.initial_bookmarks, + config.bookmarks_supplier is not None, + config.bookmarks_consumer is not None + ) + res = backend.send_and_receive(req) + if not isinstance(res, protocol.BookmarkManager): + raise Exception("Should be BookmarkManager but was %s" % res) + + self._bookmark_manager = res + self._registry[self._bookmark_manager.id] = self + + @property + def id(self): + return self._bookmark_manager.id + + @classmethod + def process_callbacks(cls, request): + if isinstance(request, protocol.BookmarksSupplierRequest): + if request.bookmark_manager_id not in cls._registry: + raise Exception( + "Backend provided unknown Bookmark manager id: " + f"{request.bookmark_manager_id} not found" + ) + + manager = cls._registry[request.bookmark_manager_id] + supply = manager.config.bookmarks_supplier + if supply is not None: + bookmarks = supply(request.database) + return protocol.BookmarksSupplierCompleted(request.id, + bookmarks) + if isinstance(request, protocol.BookmarksConsumerRequest): + if request.bookmark_manager_id not in cls._registry: + raise ValueError( + "Backend provided unknown Bookmark manager id: " + f"{request.bookmark_manager_id} not found" + ) + manager = cls._registry[request.bookmark_manager_id] + consume = manager.config.bookmarks_consumer + if consume is not None: + consume(request.database, request.bookmarks) + return protocol.BookmarksConsumerCompleted(request.id) + + def close(self, hooks=None): + res = self._backend.send_and_receive( + protocol.BookmarkManagerClose(self.id), + hooks=hooks + ) + if not isinstance(res, protocol.BookmarkManager): + raise Exception("Should be BookmarkManager but was %s" % res) + del self._registry[self.id] diff --git a/nutkit/frontend/config.py b/nutkit/frontend/config.py deleted file mode 100644 index 728fd117b..000000000 --- a/nutkit/frontend/config.py +++ /dev/null @@ -1,29 +0,0 @@ -from dataclasses import dataclass -from typing import ( - Callable, - Dict, - List, - Optional, -) - - -@dataclass -class Neo4jBookmarkManagerConfig: - initial_bookmarks: Optional[Dict[str, List[str]]] = None - bookmarks_supplier: Optional[Callable[[str], List[str]]] = None - bookmarks_consumer: Optional[Callable[[str, List[str]], None]] = None - - -def from_bookmark_manager_config_to_protocol( - config: Optional[Neo4jBookmarkManagerConfig] -) -> Optional[Dict]: - if config is not None: - return { - "initialBookmarks": config.initial_bookmarks, - "bookmarksSupplierRegistered": - config.bookmarks_supplier is not None, - "bookmarksConsumerRegistered": - config.bookmarks_consumer is not None, - } - - return None diff --git a/nutkit/frontend/driver.py b/nutkit/frontend/driver.py index aec57d0c1..6a39eb378 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -1,8 +1,5 @@ -from typing import Optional - from .. import protocol -from .config import from_bookmark_manager_config_to_protocol -from .config import Neo4jBookmarkManagerConfig as BMMConfig +from .bookmark_manager import BookmarkManager from .session import Session @@ -13,15 +10,11 @@ def __init__(self, backend, uri, auth_token, user_agent=None, max_tx_retry_time_ms=None, encrypted=None, trusted_certificates=None, liveness_check_timeout_ms=None, max_connection_pool_size=None, - connection_acquisition_timeout_ms=None, - bookmark_manager_config: Optional[BMMConfig] = None): + connection_acquisition_timeout_ms=None): self._backend = backend self._resolver_fn = resolver_fn self._domain_name_resolver_fn = domain_name_resolver_fn - self._bookmark_manager_config = bookmark_manager_config - bookmark_manager = from_bookmark_manager_config_to_protocol( - bookmark_manager_config - ) + req = protocol.NewDriver( uri, auth_token, userAgent=user_agent, resolverRegistered=resolver_fn is not None, @@ -32,7 +25,6 @@ def __init__(self, backend, uri, auth_token, user_agent=None, liveness_check_timeout_ms=liveness_check_timeout_ms, max_connection_pool_size=max_connection_pool_size, connection_acquisition_timeout_ms=connection_acquisition_timeout_ms, # noqa: E501 - bookmark_manager=bookmark_manager ) res = backend.send_and_receive(req) if not isinstance(res, protocol.Driver): @@ -59,29 +51,11 @@ def receive(self, timeout=None, hooks=None, *, allow_resolution): hooks=hooks ) continue - if self._bookmark_manager_config is not None: - supply = self._bookmark_manager_config.bookmarks_supplier - if supply is not None: - if isinstance(res, protocol.BookmarksSupplierRequest): - bookmarks = supply(res.database) - self._backend.send( - protocol.BookmarksSupplierCompleted( - res.id, - bookmarks - ), - hooks=hooks - ) - continue - consume = self._bookmark_manager_config.bookmarks_consumer - if consume is not None: - if isinstance(res, protocol.BookmarksConsumerRequest): - consume(res.database, res.bookmarks) - self._backend.send( - protocol.BookmarksConsumerCompleted( - res.id - ) - ) - continue + bookmark_manager_response = BookmarkManager.process_callbacks(res) + if bookmark_manager_response: + self._backend.send(bookmark_manager_response, hooks=hooks) + continue + return res def send(self, req, hooks=None): @@ -128,16 +102,17 @@ def close(self): def session(self, access_mode, bookmarks=None, database=None, fetch_size=None, impersonated_user=None, - ignore_bookmark_manager=None): + bookmark_manager=None): req = protocol.NewSession( self._driver.id, access_mode, bookmarks=bookmarks, database=database, fetchSize=fetch_size, impersonatedUser=impersonated_user, - ignore_bookmark_manager=ignore_bookmark_manager + bookmark_manager=bookmark_manager ) res = self.send_and_receive(req, allow_resolution=False) if not isinstance(res, protocol.Session): raise Exception("Should be session") + return Session(self, res) def resolve(self, address): diff --git a/nutkit/protocol/requests.py b/nutkit/protocol/requests.py index 6c957b629..8519eecc7 100644 --- a/nutkit/protocol/requests.py +++ b/nutkit/protocol/requests.py @@ -71,8 +71,7 @@ def __init__( fetchSize=None, maxTxRetryTimeMs=None, encrypted=None, trustedCertificates=None, liveness_check_timeout_ms=None, max_connection_pool_size=None, - connection_acquisition_timeout_ms=None, - bookmark_manager=None + connection_acquisition_timeout_ms=None ): # Neo4j URI to connect to self.uri = uri @@ -88,8 +87,6 @@ def __init__( self.livenessCheckTimeoutMs = liveness_check_timeout_ms self.maxConnectionPoolSize = max_connection_pool_size self.connectionAcquisitionTimeoutMs = connection_acquisition_timeout_ms - if bookmark_manager is not None: - self.bookmarkManager = bookmark_manager # (bool) whether to enable or disable encryption # field missing in message: use driver default (should be False) if encrypted is not None: @@ -202,8 +199,8 @@ class BookmarksSupplierCompleted: Pushes bookmarks for a given database to the Bookmark Manager. """ - def __init__(self, requestId, bookmarks): - self.requestId = requestId + def __init__(self, request_id, bookmarks): + self.requestId = request_id self.bookmarks = bookmarks @@ -214,8 +211,35 @@ class BookmarksConsumerCompleted: Signal the method call has finished """ - def __init__(self, requestId): - self.requestId = requestId + def __init__(self, request_id): + self.requestId = request_id + + +class NewBookmarkManager: + """Instantiates a bookmark manager by calling the default factory. + + Backend should respond with a BookmarkManager response. + """ + + def __init__(self, initial_bookmarks, + bookmarks_supplier_registered, bookmarks_consumer_registered): + self.initialBookmarks = initial_bookmarks + self.bookmarksSupplierRegistered = bookmarks_supplier_registered + self.bookmarksConsumerRegistered = bookmarks_consumer_registered + + +class BookmarkManagerClose: + """Destroy the bookmark manager in the backend and free the resources. + + The driver-provided BookmarkManager implementation does not have a close + method. This message is an instruction solely for the backend to be able to + destroy the bookmark manager object when done testing it to free resources. + + Backend should respond with a BookmarkManager response. + """ + + def __init__(self, id): + self.id = id class DomainNameResolutionCompleted: @@ -258,7 +282,7 @@ class NewSession: def __init__(self, driverId, accessMode, bookmarks=None, database=None, fetchSize=None, impersonatedUser=None, - ignore_bookmark_manager=None): + bookmark_manager=None): # Id of driver on backend that session should be created on self.driverId = driverId # Session accessmode: 'r' for read access and 'w' for write access. @@ -268,8 +292,8 @@ def __init__(self, driverId, accessMode, bookmarks=None, self.database = database self.fetchSize = fetchSize self.impersonatedUser = impersonatedUser - if ignore_bookmark_manager is not None: - self.ignoreBookmarkManager = ignore_bookmark_manager + if bookmark_manager is not None: + self.bookmarkManagerId = bookmark_manager.id class SessionClose: diff --git a/nutkit/protocol/responses.py b/nutkit/protocol/responses.py index f7b33d414..f982a2723 100644 --- a/nutkit/protocol/responses.py +++ b/nutkit/protocol/responses.py @@ -71,6 +71,13 @@ def __init__(self, id, address): self.address = address +class BookmarkManager: + """Represents a created Bookmark Manager.""" + + def __init__(self, id): + self.id = id + + class BookmarksSupplierRequest: """ Represents a bookmark supplier in the Bookmark Manager. @@ -81,9 +88,10 @@ class BookmarksSupplierRequest: bookmarks for all databases. """ - def __init__(self, id, database=None): + def __init__(self, id, bookmarkManagerId, database=None): # id of the callback request self.id = id + self.bookmark_manager_id = bookmarkManagerId self.database = database @@ -95,9 +103,10 @@ class BookmarksConsumerRequest: to send the new bookmark set for a given database. """ - def __init__(self, id, database, bookmarks): + def __init__(self, id, bookmarkManagerId, database, bookmarks): # id of the callback request self.id = id + self.bookmark_manager_id = bookmarkManagerId self.database = database self.bookmarks = bookmarks diff --git a/tests/stub/driver_parameters/scripts/transaction_chaining.script b/tests/stub/driver_parameters/scripts/transaction_chaining.script index 2e0e53729..03dfffca2 100644 --- a/tests/stub/driver_parameters/scripts/transaction_chaining.script +++ b/tests/stub/driver_parameters/scripts/transaction_chaining.script @@ -40,6 +40,15 @@ A: HELLO {"{}": "*"} S: SUCCESS {"type": "r"} C: COMMIT S: SUCCESS {"bookmark": "bm3"} + ---- + C: BEGIN {"tx_metadata": {"return_bookmark": "bm4"}, "[bookmarks]": "*", "[db]": "*"} + S: SUCCESS {} + C: RUN "RETURN 1 as n" {} {} + S: SUCCESS {"fields": ["n"]} + C: PULL {"n": 1000} + S: SUCCESS {"type": "r"} + C: COMMIT + S: SUCCESS {"bookmark": "bm4"} ---- C: BEGIN {"tx_metadata": {"order": "adb"}, "[bookmarks]": "*", "[db]": "*"} S: SUCCESS {} diff --git a/tests/stub/driver_parameters/test_bookmark_manager.py b/tests/stub/driver_parameters/test_bookmark_manager.py index de5994dcc..6097e9244 100644 --- a/tests/stub/driver_parameters/test_bookmark_manager.py +++ b/tests/stub/driver_parameters/test_bookmark_manager.py @@ -2,6 +2,7 @@ import re from nutkit.frontend import ( + BookmarkManager, Driver, Neo4jBookmarkManagerConfig, ) @@ -21,6 +22,7 @@ def setUp(self): self._server = StubServer(9010) self._router = StubServer(9000) self._driver = None + self._bookmark_managers = [] def tearDown(self): self._server.reset() @@ -28,6 +30,9 @@ def tearDown(self): if self._driver: self._driver.close() + for bookmark_manager in self._bookmark_managers: + bookmark_manager.close() + self._bookmark_managers.clear() return super().tearDown() @@ -36,16 +41,17 @@ def test_should_keep_track_of_session_run(self): self._start_server(self._server, "session_run_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) list(s1.run("RETURN BOOKMARK bm1")) s1.close() - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) list(s2.run("RETURN BOOKMARK bm2")) s2.close() - s3 = self._driver.session("w") + s3 = self._driver.session("w", bookmark_manager=manager) list(s3.run("RETURN BOOKMARK bm3")) s3.close() @@ -69,20 +75,21 @@ def test_should_keep_track_of_tx_in_sequence(self): self._start_server(self._server, "transaction_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() s2.close() - s3 = self._driver.session("w") + s3 = self._driver.session("w", bookmark_manager=manager) tx3 = s3.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx3.run("RETURN 1 as n")) tx3.commit() @@ -107,20 +114,21 @@ def test_should_not_replace_bookmarks_with_empty_bookmarks(self): self._start_server(self._server, "transaction_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("r") + s2 = self._driver.session("r", bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "empty"}) list(tx2.run("RETURN 1 as n")) tx2.commit() s2.close() - s3 = self._driver.session("w") + s3 = self._driver.session("w", bookmark_manager=manager) tx3 = s3.begin_transaction(tx_meta={"return_bookmark": "bm3"}) list(tx3.run("RETURN 1 as n")) tx3.commit() @@ -145,12 +153,13 @@ def test_should_keep_track_of_tx_in_parallel(self): self._start_server(self._server, "transaction_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) @@ -159,7 +168,7 @@ def test_should_keep_track_of_tx_in_parallel(self): tx2.commit() s2.close() - s3 = self._driver.session("w") + s3 = self._driver.session("w", bookmark_manager=manager) tx3 = s3.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx3.run("RETURN 1 as n")) tx3.commit() @@ -182,26 +191,31 @@ def test_should_manage_explicity_session_bookmarks(self): self._start_server(self._server, "transaction_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w", bookmarks=[]) + s2 = self._driver.session("w", bookmarks=[], bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() s2.close() - s3 = self._driver.session("w", bookmarks=["bm2", "unmanaged"]) + s3 = self._driver.session( + "w", + bookmarks=["bm2", "unmanaged"], + bookmark_manager=manager + ) tx3 = s3.begin_transaction(tx_meta={"return_bookmark": "bm3"}) list(tx3.run("RETURN 1 as n")) tx3.commit() s3.close() - s4 = self._driver.session("w") + s4 = self._driver.session("w", bookmark_manager=manager) tx4 = s4.begin_transaction(tx_meta={"return_bookmark": "bm3"}) list(tx4.run("RETURN 1 as n")) tx4.commit() @@ -225,19 +239,20 @@ def test_should_manage_explicity_session_bookmarks(self): bookmarks=["bm3"] ) - def test_should_be_able_to_ignore_bookmark_manager_in_a_sesssion(self): + def test_should_ignore_bookmark_manager_not_set_in_a_sesssion(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w", ignore_bookmark_manager=True) + s2 = self._driver.session("w") tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() @@ -245,7 +260,6 @@ def test_should_be_able_to_ignore_bookmark_manager_in_a_sesssion(self): s3 = self._driver.session( "w", - ignore_bookmark_manager=True, bookmarks=["unmanaged"] ) tx3 = s3.begin_transaction(tx_meta={"return_bookmark": "bm3"}) @@ -253,7 +267,7 @@ def test_should_be_able_to_ignore_bookmark_manager_in_a_sesssion(self): tx3.commit() s3.close() - s4 = self._driver.session("w") + s4 = self._driver.session("w", bookmark_manager=manager) tx4 = s4.begin_transaction(tx_meta={"return_bookmark": "bm3"}) list(tx4.run("RETURN 1 as n")) tx4.commit() @@ -280,17 +294,27 @@ def test_should_use_initial_bookmark_set_in_the_fist_tx(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - self._driver = self._new_driver(Neo4jBookmarkManagerConfig( - initial_bookmarks={"neo4j": ["fist_bm"]} - )) + self._driver, manager = self._new_driver_and_bookmark_manager( + Neo4jBookmarkManagerConfig( + initial_bookmarks={"neo4j": ["fist_bm"]} + ) + ) - s1 = self._driver.session("w", database="neo4j") + s1 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w", database="neo4j") + s2 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() @@ -313,17 +337,27 @@ def test_should_send_all_db_bookmarks_and_update_only_relevant(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - self._driver = self._new_driver(Neo4jBookmarkManagerConfig( - initial_bookmarks={"neo4j": ["fist_bm"], "adb": ["adb:bm1"]} - )) + self._driver, manager = self._new_driver_and_bookmark_manager( + Neo4jBookmarkManagerConfig( + initial_bookmarks={"neo4j": ["fist_bm"], "adb": ["adb:bm1"]} + ) + ) - s1 = self._driver.session("w", database="neo4j") + s1 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w", database="neo4j") + s2 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() @@ -347,26 +381,43 @@ def test_should_handle_database_redirection_in_tx(self): self._start_server(self._server, "transaction_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w", database="neo4j") + s1 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w", database="neo4j") + s2 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx2 = s2.begin_transaction(tx_meta={"order": "adb"}) list(tx2.run("USE adb RETURN 1 as n")) tx2.commit() s2.close() - s3 = self._driver.session("w", database="neo4j") + s3 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx3 = s3.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx3.run("RETURN 1 as n")) tx3.commit() s3.close() - s4 = self._driver.session("w", database="neo4j") + s4 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx4 = s4.begin_transaction(tx_meta={"return_bookmark": "bm3"}) list(tx4.run("RETURN 1 as n")) tx4.commit() @@ -397,20 +448,21 @@ def test_should_handle_database_redirection_in_session_run(self): self._start_server(self._server, "session_run_chaining.script") self._driver = self._new_driver() + manager = self._new_bookmark_manager() - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) list(s1.run("RETURN BOOKMARK bm1")) s1.close() - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) list(s2.run("USE adb RETURN BOOKMARK adb:bm4")) s2.close() - s3 = self._driver.session("w") + s3 = self._driver.session("w", bookmark_manager=manager) list(s3.run("RETURN BOOKMARK bm2")) s3.close() - s4 = self._driver.session("w") + s4 = self._driver.session("w", bookmark_manager=manager) list(s4.run("RETURN BOOKMARK bm3")) s4.close() @@ -436,17 +488,19 @@ def test_should_resolve_database_name_with_system_bookmarks(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - self._driver = self._new_driver(Neo4jBookmarkManagerConfig( - initial_bookmarks={"system": ["sys:bm1"]} - )) + self._driver, manager = self._new_driver_and_bookmark_manager( + Neo4jBookmarkManagerConfig( + initial_bookmarks={"system": ["sys:bm1"]} + ) + ) - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() @@ -471,18 +525,6 @@ def test_should_resolve_database_name_with_system_bookmarks(self): bookmarks=["sys:bm1", "bm1"] ) - def _new_driver(self, bookmark_manager_config=None): - if bookmark_manager_config is None: - bookmark_manager_config = Neo4jBookmarkManagerConfig() - uri = "neo4j://%s" % self._router.address - auth = types.AuthorizationToken("basic", principal="neo4j", - credentials="pass") - return Driver( - self._backend, - uri, auth, - bookmark_manager_config=bookmark_manager_config - ) - def test_should_call_bookmark_supplier_for_all_get_bookmarks_calls(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") @@ -494,7 +536,7 @@ def get_bookmarks(db): get_bookmarks_calls.append([db]) return [] - self._driver = self._new_driver( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -503,13 +545,17 @@ def get_bookmarks(db): ) ) - s1 = self._driver.session("w", database="neo4j") + s1 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() @@ -558,7 +604,7 @@ def get_bookmarks(db): return system_bookmarks + neo4j_bookmarks return bookmarks.get(db, []) - self._driver = self._new_driver( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -567,13 +613,17 @@ def get_bookmarks(db): ) ) - s1 = self._driver.session("w", database="neo4j") + s1 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"return_bookmark": "bm2"}) list(tx2.run("RETURN 1 as n")) tx2.commit() @@ -603,7 +653,7 @@ def test_should_call_bookmarks_consumer_when_new_bookmarks_arrive(self): def bookmarks_consumer(db, bookmarks): bookmarks_consumer_calls.append([db, bookmarks]) - self._driver = self._new_driver( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -612,13 +662,17 @@ def bookmarks_consumer(db, bookmarks): ) ) - s1 = self._driver.session("w", database="neo4j") + s1 = self._driver.session( + "w", + database="neo4j", + bookmark_manager=manager + ) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() s1.close() - s2 = self._driver.session("w") + s2 = self._driver.session("w", bookmark_manager=manager) tx2 = s2.begin_transaction(tx_meta={"order": "adb"}) list(tx2.run("USE adb RETURN 1 as n")) tx2.commit() @@ -642,7 +696,7 @@ def test_should_call_bookmarks_consumer_for_default_db(self): def bookmarks_consumer(db, bookmarks): bookmarks_consumer_calls.append([db, bookmarks]) - self._driver = self._new_driver( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -651,7 +705,7 @@ def bookmarks_consumer(db, bookmarks): ) ) - s1 = self._driver.session("w") + s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) list(tx1.run("RETURN 1 as n")) tx1.commit() @@ -671,6 +725,69 @@ def bookmarks_consumer(db, bookmarks): ] ]) + def test_multiple_bookmark_manager(self): + self._start_server(self._router, "router_with_db_name.script") + self._start_server(self._server, "transaction_chaining.script") + + self._driver = self._new_driver() + + manager1 = self._new_bookmark_manager( + Neo4jBookmarkManagerConfig( + initial_bookmarks={"neo4j": ["manager_01_initial_bm"]} + ) + ) + + manager2 = self._new_bookmark_manager( + Neo4jBookmarkManagerConfig( + initial_bookmarks={"neo4j": ["manager_02_initial_bm"]} + ) + ) + + manager1_s1 = self._driver.session("w", bookmark_manager=manager1) + tx1 = manager1_s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) + list(tx1.run("RETURN 1 as n")) + tx1.commit() + manager1_s1.close() + + manager2_s1 = self._driver.session("w", bookmark_manager=manager2) + tx2 = manager2_s1.begin_transaction(tx_meta={"return_bookmark": "bm2"}) + list(tx2.run("RETURN 1 as n")) + tx2.commit() + manager2_s1.close() + + manager2_s2 = self._driver.session("w", bookmark_manager=manager2) + tx3 = manager2_s2.begin_transaction(tx_meta={"return_bookmark": "bm3"}) + list(tx3.run("RETURN 1 as n")) + tx3.commit() + manager2_s2.close() + + manager1_s2 = self._driver.session("w", bookmark_manager=manager1) + tx4 = manager1_s2.begin_transaction(tx_meta={"return_bookmark": "bm4"}) + list(tx4.run("RETURN 1 as n")) + tx4.commit() + manager1_s2.close() + + self._server.reset() + begin_requests = self._server.get_requests("BEGIN") + + self.assertEqual(len(begin_requests), 4) + self.assert_begin( + begin_requests[0], + bookmarks=["manager_01_initial_bm"] + ) + self.assert_begin( + begin_requests[1], + bookmarks=["manager_02_initial_bm"] + ) + self.assert_begin( + begin_requests[2], + bookmarks=["bm2"] + ) + self.assert_begin( + begin_requests[3], + bookmarks=["bm1"] + ) + def _start_server(self, server, script): server.start(self.script_path(script), vars_={"#HOST#": self._router.host}) @@ -733,3 +850,28 @@ def supports_minimal_bookmarks(self): return self.driver_supports_features( types.Feature.OPT_MINIMAL_BOOKMARKS_SET ) + + def _new_driver(self): + uri = "neo4j://%s" % self._router.address + auth = types.AuthorizationToken("basic", principal="neo4j", + credentials="pass") + driver = Driver( + self._backend, + uri, auth + ) + return driver + + def _new_bookmark_manager(self, bookmark_manager_config=None): + if bookmark_manager_config is None: + bookmark_manager_config = Neo4jBookmarkManagerConfig() + bookmark_manager = BookmarkManager( + self._backend, + bookmark_manager_config + ) + self._bookmark_managers.append(bookmark_manager) + return bookmark_manager + + def _new_driver_and_bookmark_manager(self, bookmark_manager_config=None): + bookmark_manager = self._new_bookmark_manager(bookmark_manager_config) + driver = self._new_driver() + return driver, bookmark_manager