From 07cd587b814a73f260b0fa729867432adbeb4571 Mon Sep 17 00:00:00 2001 From: Antonio Barcelos Date: Wed, 24 Aug 2022 16:51:20 +0200 Subject: [PATCH 01/12] Moving BookmarkManager configuration to `driver.session` --- nutkit/frontend/__init__.py | 6 +- nutkit/frontend/bookmark_manager.py | 41 ++++ nutkit/frontend/config.py | 29 --- nutkit/frontend/driver.py | 49 ++-- nutkit/protocol/requests.py | 27 ++- nutkit/protocol/responses.py | 13 +- .../test_bookmark_manager.py | 209 +++++++++++------- 7 files changed, 237 insertions(+), 137 deletions(-) create mode 100644 nutkit/frontend/bookmark_manager.py delete mode 100644 nutkit/frontend/config.py diff --git a/nutkit/frontend/__init__.py b/nutkit/frontend/__init__.py index 3919b2ab1..629175e0d 100644 --- a/nutkit/frontend/__init__.py +++ b/nutkit/frontend/__init__.py @@ -1,3 +1,7 @@ -from .config import Neo4jBookmarkManagerConfig +from .bookmark_manager import ( + BookmarkManager, + create_bookmark_manager, + 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..18ee7f472 --- /dev/null +++ b/nutkit/frontend/bookmark_manager.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass +from typing import ( + Callable, + 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: + config: Neo4jBookmarkManagerConfig + id: int + + +def create_bookmark_manager(backend: Backend, + config: Neo4jBookmarkManagerConfig): + 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) + + return BookmarkManager( + config, + res.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..7336d1044 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -1,8 +1,8 @@ from typing import Optional +from nutkit.protocol.responses import BookmarkManager + from .. import protocol -from .config import from_bookmark_manager_config_to_protocol -from .config import Neo4jBookmarkManagerConfig as BMMConfig from .session import Session @@ -13,15 +13,12 @@ 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 - ) + self._bookmarks_managers = {} + req = protocol.NewDriver( uri, auth_token, userAgent=user_agent, resolverRegistered=resolver_fn is not None, @@ -32,7 +29,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 +55,36 @@ 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): + if isinstance(res, protocol.BookmarksSupplierRequest): + if res.bookmarkManagerId in self._bookmarks_managers: + manager = self._bookmarks_managers[res.bookmarkManagerId] + supply = manager.config.bookmarks_supplier + if supply is not None: bookmarks = supply(res.database) self._backend.send( protocol.BookmarksSupplierCompleted( res.id, + res.bookmarkManagerId, 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) + if isinstance(res, protocol.BookmarksConsumerRequest): + if res.bookmarkManagerId in self._bookmarks_managers: + manager = self._bookmarks_managers[res.bookmarkManagerId] + consume = manager.config.bookmarks_consumer + if consume is not None: + bookmarks = consume(res.database, res.bookmarks) self._backend.send( protocol.BookmarksConsumerCompleted( - res.id - ) + res.id, + res.bookmarkManagerId + ), + hooks=hooks ) continue + return res def send(self, req, hooks=None): @@ -128,16 +131,20 @@ def close(self): def session(self, access_mode, bookmarks=None, database=None, fetch_size=None, impersonated_user=None, - ignore_bookmark_manager=None): + bookmark_manager: Optional[BookmarkManager] = 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") + + if bookmark_manager is not None: + self._bookmarks_managers[bookmark_manager.id] = bookmark_manager + return Session(self, res) def resolve(self, address): diff --git a/nutkit/protocol/requests.py b/nutkit/protocol/requests.py index 6c957b629..464c41029 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,9 @@ class BookmarksSupplierCompleted: Pushes bookmarks for a given database to the Bookmark Manager. """ - def __init__(self, requestId, bookmarks): + def __init__(self, requestId, bookmarkManagerId, bookmarks): self.requestId = requestId + self.bookmarkManagerId = bookmarkManagerId self.bookmarks = bookmarks @@ -214,8 +212,19 @@ class BookmarksConsumerCompleted: Signal the method call has finished """ - def __init__(self, requestId): + def __init__(self, requestId, bookmarkManagerId): self.requestId = requestId + self.bookmarkManagerId = bookmarkManagerId + + +class NewBookmarkManager: + """Instantiates a bookmark manager by calling the default factory.""" + + def __init__(self, initialBookmarks, + bookmarksSupplierRegistered, bookmarksConsumerRegistered): + self.initialBookmarks = initialBookmarks + self.bookmarksSupplierRegistered = bookmarksSupplierRegistered + self.bookmarksConsumerRegistered = bookmarksConsumerRegistered class DomainNameResolutionCompleted: @@ -258,7 +267,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 +277,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..d6dfecae4 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.bookmarkManagerId = 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.bookmarkManagerId = bookmarkManagerId self.database = database self.bookmarks = bookmarks diff --git a/tests/stub/driver_parameters/test_bookmark_manager.py b/tests/stub/driver_parameters/test_bookmark_manager.py index de5994dcc..4ecdc1672 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 ( + create_bookmark_manager, Driver, Neo4jBookmarkManagerConfig, ) @@ -35,17 +36,17 @@ def test_should_keep_track_of_session_run(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "session_run_chaining.script") - self._driver = self._new_driver() + (self._driver, manager) = self._new_driver_and_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() @@ -68,21 +69,21 @@ def test_should_keep_track_of_tx_in_sequence(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - self._driver = self._new_driver() + (self._driver, manager) = self._new_driver_and_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() @@ -106,21 +107,21 @@ def test_should_not_replace_bookmarks_with_empty_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() + (self._driver, manager) = self._new_driver_and_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() @@ -144,13 +145,13 @@ def test_should_keep_track_of_tx_in_parallel(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - self._driver = self._new_driver() + (self._driver, manager) = self._new_driver_and_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 +160,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() @@ -181,27 +182,31 @@ def test_should_manage_explicity_session_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() + (self._driver, manager) = self._new_driver_and_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 +230,19 @@ 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() + (self._driver, manager) = self._new_driver_and_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 +250,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 +257,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 +284,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 +327,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() @@ -346,27 +370,43 @@ def test_should_handle_database_redirection_in_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() + (self._driver, manager) = self._new_driver_and_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() @@ -396,21 +436,21 @@ def test_should_handle_database_redirection_in_session_run(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "session_run_chaining.script") - self._driver = self._new_driver() + (self._driver, manager) = self._new_driver_and_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 +476,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 +513,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 +524,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 +533,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 +592,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 +601,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 +641,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 +650,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 +684,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 +693,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() @@ -733,3 +775,20 @@ def supports_minimal_bookmarks(self): return self.driver_supports_features( types.Feature.OPT_MINIMAL_BOOKMARKS_SET ) + + def _new_driver_and_bookmark_manager(self, bookmark_manager_config=None): + if bookmark_manager_config is None: + bookmark_manager_config = Neo4jBookmarkManagerConfig() + bookmark_manager = create_bookmark_manager( + self._backend, + bookmark_manager_config + ) + + uri = "neo4j://%s" % self._router.address + auth = types.AuthorizationToken("basic", principal="neo4j", + credentials="pass") + driver = Driver( + self._backend, + uri, auth + ) + return (driver, bookmark_manager) From eb6a7bc0dad85cf6498be569938d3c9294b2574f Mon Sep 17 00:00:00 2001 From: Antonio Barcelos Date: Thu, 25 Aug 2022 15:55:52 +0200 Subject: [PATCH 02/12] Test Multiple Bookmark Managers Co-Authored-By: Florent Biville --- .../scripts/transaction_chaining.script | 9 ++ .../test_bookmark_manager.py | 85 +++++++++++++++++-- 2 files changed, 86 insertions(+), 8 deletions(-) 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 4ecdc1672..dea14b74b 100644 --- a/tests/stub/driver_parameters/test_bookmark_manager.py +++ b/tests/stub/driver_parameters/test_bookmark_manager.py @@ -713,6 +713,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}) @@ -776,14 +839,7 @@ def supports_minimal_bookmarks(self): types.Feature.OPT_MINIMAL_BOOKMARKS_SET ) - def _new_driver_and_bookmark_manager(self, bookmark_manager_config=None): - if bookmark_manager_config is None: - bookmark_manager_config = Neo4jBookmarkManagerConfig() - bookmark_manager = create_bookmark_manager( - self._backend, - bookmark_manager_config - ) - + def _new_driver(self): uri = "neo4j://%s" % self._router.address auth = types.AuthorizationToken("basic", principal="neo4j", credentials="pass") @@ -791,4 +847,17 @@ def _new_driver_and_bookmark_manager(self, bookmark_manager_config=None): 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() + return create_bookmark_manager( + self._backend, + bookmark_manager_config + ) + + 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) From eb180bebfe641eaaa1648e0648a71a362fef27dd Mon Sep 17 00:00:00 2001 From: Antonio Barcelos Date: Thu, 25 Aug 2022 16:00:43 +0200 Subject: [PATCH 03/12] Small clean-up Co-Authored-By: Florent Biville --- .../test_bookmark_manager.py | 24 ++++++++++++------- 1 file changed, 16 insertions(+), 8 deletions(-) diff --git a/tests/stub/driver_parameters/test_bookmark_manager.py b/tests/stub/driver_parameters/test_bookmark_manager.py index dea14b74b..f39a81049 100644 --- a/tests/stub/driver_parameters/test_bookmark_manager.py +++ b/tests/stub/driver_parameters/test_bookmark_manager.py @@ -36,7 +36,8 @@ def test_should_keep_track_of_session_run(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "session_run_chaining.script") - (self._driver, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session("w", bookmark_manager=manager) list(s1.run("RETURN BOOKMARK bm1")) @@ -69,7 +70,8 @@ def test_should_keep_track_of_tx_in_sequence(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - (self._driver, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) @@ -107,7 +109,8 @@ def test_should_not_replace_bookmarks_with_empty_bookmarks(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - (self._driver, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) @@ -145,7 +148,8 @@ def test_should_keep_track_of_tx_in_parallel(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - (self._driver, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) @@ -182,7 +186,8 @@ def test_should_manage_explicity_session_bookmarks(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - (self._driver, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) @@ -234,7 +239,8 @@ 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, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session("w", bookmark_manager=manager) tx1 = s1.begin_transaction(tx_meta={"return_bookmark": "bm1"}) @@ -370,7 +376,8 @@ def test_should_handle_database_redirection_in_tx(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "transaction_chaining.script") - (self._driver, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session( "w", @@ -436,7 +443,8 @@ def test_should_handle_database_redirection_in_session_run(self): self._start_server(self._router, "router_with_db_name.script") self._start_server(self._server, "session_run_chaining.script") - (self._driver, manager) = self._new_driver_and_bookmark_manager() + self._driver = self._new_driver() + manager = self._new_bookmark_manager() s1 = self._driver.session("w", bookmark_manager=manager) list(s1.run("RETURN BOOKMARK bm1")) From b9bd74b915afa3cd60ea78d9449c4121d96fa60b Mon Sep 17 00:00:00 2001 From: Antonio Barcelos Date: Mon, 29 Aug 2022 09:31:58 +0200 Subject: [PATCH 04/12] Store BookmarkManager before create session --- nutkit/frontend/driver.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nutkit/frontend/driver.py b/nutkit/frontend/driver.py index 7336d1044..c9f881916 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -132,6 +132,9 @@ def close(self): def session(self, access_mode, bookmarks=None, database=None, fetch_size=None, impersonated_user=None, bookmark_manager: Optional[BookmarkManager] = None): + if bookmark_manager is not None: + self._bookmarks_managers[bookmark_manager.id] = bookmark_manager + req = protocol.NewSession( self._driver.id, access_mode, bookmarks=bookmarks, database=database, fetchSize=fetch_size, @@ -142,9 +145,6 @@ def session(self, access_mode, bookmarks=None, database=None, if not isinstance(res, protocol.Session): raise Exception("Should be session") - if bookmark_manager is not None: - self._bookmarks_managers[bookmark_manager.id] = bookmark_manager - return Session(self, res) def resolve(self, address): From dd32b8b59ab68683cd8dc7bf7fd4c39b98080010 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Antonio=20Barc=C3=A9los?= Date: Mon, 29 Aug 2022 12:12:44 +0200 Subject: [PATCH 05/12] Apply suggestions from code review Co-authored-by: Robsdedude --- .../driver_parameters/test_bookmark_manager.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/stub/driver_parameters/test_bookmark_manager.py b/tests/stub/driver_parameters/test_bookmark_manager.py index f39a81049..0de9a6597 100644 --- a/tests/stub/driver_parameters/test_bookmark_manager.py +++ b/tests/stub/driver_parameters/test_bookmark_manager.py @@ -290,7 +290,7 @@ 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, manager) = self._new_driver_and_bookmark_manager( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={"neo4j": ["fist_bm"]} ) @@ -333,7 +333,7 @@ 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, manager) = self._new_driver_and_bookmark_manager( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={"neo4j": ["fist_bm"], "adb": ["adb:bm1"]} ) @@ -484,7 +484,7 @@ 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, manager) = self._new_driver_and_bookmark_manager( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={"system": ["sys:bm1"]} ) @@ -532,7 +532,7 @@ def get_bookmarks(db): get_bookmarks_calls.append([db]) return [] - (self._driver, manager) = self._new_driver_and_bookmark_manager( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -600,7 +600,7 @@ def get_bookmarks(db): return system_bookmarks + neo4j_bookmarks return bookmarks.get(db, []) - (self._driver, manager) = self._new_driver_and_bookmark_manager( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -649,7 +649,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, manager) = self._new_driver_and_bookmark_manager( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -868,4 +868,4 @@ def _new_bookmark_manager(self, bookmark_manager_config=None): 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) + return driver, bookmark_manager From c978f7d517dfa96fa653bad48e7742fb87569555 Mon Sep 17 00:00:00 2001 From: Antonio Barcelos Date: Mon, 29 Aug 2022 12:17:58 +0200 Subject: [PATCH 06/12] bookmarkManagerId -> bookmark_manager_id --- nutkit/frontend/driver.py | 12 ++++++------ nutkit/protocol/responses.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/nutkit/frontend/driver.py b/nutkit/frontend/driver.py index c9f881916..dc348aac9 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -56,30 +56,30 @@ def receive(self, timeout=None, hooks=None, *, allow_resolution): ) continue if isinstance(res, protocol.BookmarksSupplierRequest): - if res.bookmarkManagerId in self._bookmarks_managers: - manager = self._bookmarks_managers[res.bookmarkManagerId] + if res.bookmark_manager_id in self._bookmarks_managers: + manager = self._bookmarks_managers[res.bookmark_manager_id] supply = manager.config.bookmarks_supplier if supply is not None: bookmarks = supply(res.database) self._backend.send( protocol.BookmarksSupplierCompleted( res.id, - res.bookmarkManagerId, + res.bookmark_manager_id, bookmarks ), hooks=hooks ) continue if isinstance(res, protocol.BookmarksConsumerRequest): - if res.bookmarkManagerId in self._bookmarks_managers: - manager = self._bookmarks_managers[res.bookmarkManagerId] + if res.bookmark_manager_id in self._bookmarks_managers: + manager = self._bookmarks_managers[res.bookmark_manager_id] consume = manager.config.bookmarks_consumer if consume is not None: bookmarks = consume(res.database, res.bookmarks) self._backend.send( protocol.BookmarksConsumerCompleted( res.id, - res.bookmarkManagerId + res.bookmark_manager_id ), hooks=hooks ) diff --git a/nutkit/protocol/responses.py b/nutkit/protocol/responses.py index d6dfecae4..f982a2723 100644 --- a/nutkit/protocol/responses.py +++ b/nutkit/protocol/responses.py @@ -91,7 +91,7 @@ class BookmarksSupplierRequest: def __init__(self, id, bookmarkManagerId, database=None): # id of the callback request self.id = id - self.bookmarkManagerId = bookmarkManagerId + self.bookmark_manager_id = bookmarkManagerId self.database = database @@ -106,7 +106,7 @@ class BookmarksConsumerRequest: def __init__(self, id, bookmarkManagerId, database, bookmarks): # id of the callback request self.id = id - self.bookmarkManagerId = bookmarkManagerId + self.bookmark_manager_id = bookmarkManagerId self.database = database self.bookmarks = bookmarks From 798a66a11ec7a1a76c69c596704f919d84cf3d68 Mon Sep 17 00:00:00 2001 From: Rouven Bauer Date: Mon, 29 Aug 2022 12:21:55 +0200 Subject: [PATCH 07/12] Protocol parameter/attribute name style --- nutkit/frontend/driver.py | 12 +++--------- nutkit/protocol/requests.py | 25 +++++++++++++------------ 2 files changed, 16 insertions(+), 21 deletions(-) diff --git a/nutkit/frontend/driver.py b/nutkit/frontend/driver.py index dc348aac9..e54e75101 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -62,11 +62,8 @@ def receive(self, timeout=None, hooks=None, *, allow_resolution): if supply is not None: bookmarks = supply(res.database) self._backend.send( - protocol.BookmarksSupplierCompleted( - res.id, - res.bookmark_manager_id, - bookmarks - ), + protocol.BookmarksSupplierCompleted(res.id, + bookmarks), hooks=hooks ) continue @@ -77,10 +74,7 @@ def receive(self, timeout=None, hooks=None, *, allow_resolution): if consume is not None: bookmarks = consume(res.database, res.bookmarks) self._backend.send( - protocol.BookmarksConsumerCompleted( - res.id, - res.bookmark_manager_id - ), + protocol.BookmarksConsumerCompleted(res.id), hooks=hooks ) continue diff --git a/nutkit/protocol/requests.py b/nutkit/protocol/requests.py index 464c41029..51bf8d9bc 100644 --- a/nutkit/protocol/requests.py +++ b/nutkit/protocol/requests.py @@ -199,9 +199,8 @@ class BookmarksSupplierCompleted: Pushes bookmarks for a given database to the Bookmark Manager. """ - def __init__(self, requestId, bookmarkManagerId, bookmarks): - self.requestId = requestId - self.bookmarkManagerId = bookmarkManagerId + def __init__(self, request_id, bookmarks): + self.requestId = request_id self.bookmarks = bookmarks @@ -212,19 +211,21 @@ class BookmarksConsumerCompleted: Signal the method call has finished """ - def __init__(self, requestId, bookmarkManagerId): - self.requestId = requestId - self.bookmarkManagerId = bookmarkManagerId + def __init__(self, request_id): + self.requestId = request_id class NewBookmarkManager: - """Instantiates a bookmark manager by calling the default factory.""" + """Instantiates a bookmark manager by calling the default factory. + + Backend should respond with a BookmarkManager response. + """ - def __init__(self, initialBookmarks, - bookmarksSupplierRegistered, bookmarksConsumerRegistered): - self.initialBookmarks = initialBookmarks - self.bookmarksSupplierRegistered = bookmarksSupplierRegistered - self.bookmarksConsumerRegistered = bookmarksConsumerRegistered + 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 DomainNameResolutionCompleted: From 6f583d75bc54d1623872594a997616ef54bdee26 Mon Sep 17 00:00:00 2001 From: Rouven Bauer Date: Mon, 29 Aug 2022 12:23:26 +0200 Subject: [PATCH 08/12] BMM consume callback has not return value --- nutkit/frontend/driver.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nutkit/frontend/driver.py b/nutkit/frontend/driver.py index e54e75101..a9e79c7f9 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -72,7 +72,7 @@ def receive(self, timeout=None, hooks=None, *, allow_resolution): manager = self._bookmarks_managers[res.bookmark_manager_id] consume = manager.config.bookmarks_consumer if consume is not None: - bookmarks = consume(res.database, res.bookmarks) + consume(res.database, res.bookmarks) self._backend.send( protocol.BookmarksConsumerCompleted(res.id), hooks=hooks From 87d8708459862a0126f8090a855567aca594096c Mon Sep 17 00:00:00 2001 From: Antonio Barcelos Date: Mon, 29 Aug 2022 12:29:53 +0200 Subject: [PATCH 09/12] Remove type hint --- nutkit/frontend/driver.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/nutkit/frontend/driver.py b/nutkit/frontend/driver.py index a9e79c7f9..92c21e542 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -1,7 +1,3 @@ -from typing import Optional - -from nutkit.protocol.responses import BookmarkManager - from .. import protocol from .session import Session @@ -125,7 +121,7 @@ def close(self): def session(self, access_mode, bookmarks=None, database=None, fetch_size=None, impersonated_user=None, - bookmark_manager: Optional[BookmarkManager] = None): + bookmark_manager=None): if bookmark_manager is not None: self._bookmarks_managers[bookmark_manager.id] = bookmark_manager From da40310e0398d9af2f694a8c16aae25f1e994f4b Mon Sep 17 00:00:00 2001 From: Rouven Bauer Date: Mon, 29 Aug 2022 13:22:54 +0200 Subject: [PATCH 10/12] Add BMM close message + refactor BMM callback loop --- nutkit/frontend/__init__.py | 1 - nutkit/frontend/bookmark_manager.py | 80 ++++++++++++++----- nutkit/frontend/driver.py | 31 ++----- nutkit/protocol/requests.py | 10 +++ .../test_bookmark_manager.py | 12 ++- 5 files changed, 85 insertions(+), 49 deletions(-) diff --git a/nutkit/frontend/__init__.py b/nutkit/frontend/__init__.py index 629175e0d..f36ec6234 100644 --- a/nutkit/frontend/__init__.py +++ b/nutkit/frontend/__init__.py @@ -1,6 +1,5 @@ from .bookmark_manager import ( BookmarkManager, - create_bookmark_manager, Neo4jBookmarkManagerConfig, ) from .driver import Driver diff --git a/nutkit/frontend/bookmark_manager.py b/nutkit/frontend/bookmark_manager.py index 18ee7f472..98c9dfefd 100644 --- a/nutkit/frontend/bookmark_manager.py +++ b/nutkit/frontend/bookmark_manager.py @@ -1,6 +1,10 @@ +from __future__ import annotations + from dataclasses import dataclass from typing import ( + Any, Callable, + ClassVar, Dict, List, Optional, @@ -20,22 +24,60 @@ class Neo4jBookmarkManagerConfig: @dataclass class BookmarkManager: - config: Neo4jBookmarkManagerConfig - id: int - - -def create_bookmark_manager(backend: Backend, - config: Neo4jBookmarkManagerConfig): - 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) - - return BookmarkManager( - config, - res.id - ) + _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/driver.py b/nutkit/frontend/driver.py index 92c21e542..a068565f0 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -1,4 +1,5 @@ from .. import protocol +from .bookmark_manager import BookmarkManager from .session import Session @@ -51,29 +52,10 @@ def receive(self, timeout=None, hooks=None, *, allow_resolution): hooks=hooks ) continue - if isinstance(res, protocol.BookmarksSupplierRequest): - if res.bookmark_manager_id in self._bookmarks_managers: - manager = self._bookmarks_managers[res.bookmark_manager_id] - supply = manager.config.bookmarks_supplier - if supply is not None: - bookmarks = supply(res.database) - self._backend.send( - protocol.BookmarksSupplierCompleted(res.id, - bookmarks), - hooks=hooks - ) - continue - if isinstance(res, protocol.BookmarksConsumerRequest): - if res.bookmark_manager_id in self._bookmarks_managers: - manager = self._bookmarks_managers[res.bookmark_manager_id] - consume = manager.config.bookmarks_consumer - if consume is not None: - consume(res.database, res.bookmarks) - self._backend.send( - protocol.BookmarksConsumerCompleted(res.id), - hooks=hooks - ) - continue + bookmark_manager_response = BookmarkManager.process_callbacks(res) + if bookmark_manager_response: + self._backend.send(bookmark_manager_response, hooks=hooks) + continue return res @@ -122,9 +104,6 @@ def close(self): def session(self, access_mode, bookmarks=None, database=None, fetch_size=None, impersonated_user=None, bookmark_manager=None): - if bookmark_manager is not None: - self._bookmarks_managers[bookmark_manager.id] = bookmark_manager - req = protocol.NewSession( self._driver.id, access_mode, bookmarks=bookmarks, database=database, fetchSize=fetch_size, diff --git a/nutkit/protocol/requests.py b/nutkit/protocol/requests.py index 51bf8d9bc..6897f1041 100644 --- a/nutkit/protocol/requests.py +++ b/nutkit/protocol/requests.py @@ -228,6 +228,16 @@ def __init__(self, initial_bookmarks, self.bookmarksConsumerRegistered = bookmarks_consumer_registered +class BookmarkManagerClose: + """Destroy the bookmark manager in the backend and free the resources. + + Backend should respond with a BookmarkManager response. + """ + + def __init__(self, id): + self.id = id + + class DomainNameResolutionCompleted: """ Results of a DNS resolution. diff --git a/tests/stub/driver_parameters/test_bookmark_manager.py b/tests/stub/driver_parameters/test_bookmark_manager.py index 0de9a6597..6097e9244 100644 --- a/tests/stub/driver_parameters/test_bookmark_manager.py +++ b/tests/stub/driver_parameters/test_bookmark_manager.py @@ -2,7 +2,7 @@ import re from nutkit.frontend import ( - create_bookmark_manager, + BookmarkManager, Driver, Neo4jBookmarkManagerConfig, ) @@ -22,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() @@ -29,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() @@ -692,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, manager) = self._new_driver_and_bookmark_manager( + self._driver, manager = self._new_driver_and_bookmark_manager( Neo4jBookmarkManagerConfig( initial_bookmarks={ "adb": adb_bookmarks @@ -860,10 +864,12 @@ def _new_driver(self): def _new_bookmark_manager(self, bookmark_manager_config=None): if bookmark_manager_config is None: bookmark_manager_config = Neo4jBookmarkManagerConfig() - return create_bookmark_manager( + 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) From 5db8305b3228e7ea2f587d1a99726cf1dc1e2251 Mon Sep 17 00:00:00 2001 From: Rouven Bauer Date: Mon, 29 Aug 2022 15:24:58 +0200 Subject: [PATCH 11/12] code clean-up --- nutkit/frontend/driver.py | 1 - 1 file changed, 1 deletion(-) diff --git a/nutkit/frontend/driver.py b/nutkit/frontend/driver.py index a068565f0..6a39eb378 100644 --- a/nutkit/frontend/driver.py +++ b/nutkit/frontend/driver.py @@ -14,7 +14,6 @@ def __init__(self, backend, uri, auth_token, user_agent=None, self._backend = backend self._resolver_fn = resolver_fn self._domain_name_resolver_fn = domain_name_resolver_fn - self._bookmarks_managers = {} req = protocol.NewDriver( uri, auth_token, userAgent=user_agent, From c403cd5c11f5f7d9407829118adea7a53df1a599 Mon Sep 17 00:00:00 2001 From: Rouven Bauer Date: Mon, 29 Aug 2022 15:28:16 +0200 Subject: [PATCH 12/12] Clarify docs for BookmarkManagerClose message It's a backend-only message and does not correspond to any driver API. --- nutkit/protocol/requests.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/nutkit/protocol/requests.py b/nutkit/protocol/requests.py index 6897f1041..8519eecc7 100644 --- a/nutkit/protocol/requests.py +++ b/nutkit/protocol/requests.py @@ -231,6 +231,10 @@ def __init__(self, initial_bookmarks, 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. """