[ADR-001] Initial draft of bookmark manager

Author: robsdedude

neo4j/_async/bookmark_manager.py

@@ -0,0 +1,86 @@
+# Copyright (c) "Neo4j"
+# Neo4j Sweden AB [https://neo4j.com]
+#
+# This file is part of Neo4j.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from __future__ import annotations
+
+import typing as t
+from collections import defaultdict
+
+from .._async_compat.concurrency import (
+    AsyncCooperativeLock,
+    AsyncLock,
+)
+from .._async_compat.util import AsyncUtil
+from ..api import AsyncBookmarkManager
+
+
+class AsyncNeo4jBookmarkManager(AsyncBookmarkManager):
+    def __init__(self, initial_bookmarks=None, bookmark_supplier=None,
+                 notify_bookmarks=None):
+        super().__init__()
+        self._initial_bookmarks = initial_bookmarks
+        self._bookmark_supplier = bookmark_supplier
+        self._notify_bookmarks = notify_bookmarks
+        self._bookmarks = defaultdict(set)
+        if bookmark_supplier or notify_bookmarks:
+            self._lock = AsyncLock()
+        else:
+            self._lock = AsyncCooperativeLock()
+
+    async def update_bookmarks(
+        self, database: str, previous_bookmarks: t.Iterable[str],
+        new_bookmarks: t.Iterable[str]
+    ) -> None:
+        with self._lock:
+            new_bms = set(new_bookmarks)
+            if not new_bms:
+                return
+            prev_bms = set(previous_bookmarks)
+            curr_bms = self._bookmarks[database]
+            curr_bms.difference_update(prev_bms)
+            curr_bms.update(new_bms)
+
+            if self._notify_bookmarks:
+                await AsyncUtil.callback(
+                    self._notify_bookmarks, database, tuple(curr_bms)
+                )
+
+    async def _get_bookmarks(self, database: str) -> t.Set[str]:
+        bms = self._bookmarks[database]
+        if self._bookmark_supplier:
+            extra_bms = await AsyncUtil.callback(
+                self._bookmark_supplier, database
+            )
+            if extra_bms is not None:
+                bms &= set(extra_bms)
+        return bms
+
+    async def get_bookmarks(self, database: str) -> t.Set[str]:
+        with self._lock:
+            return await self._get_bookmarks(database)
+
+    async def get_all_bookmarks(
+        self, must_included_databases: t.Iterable[str]
+    ) -> t.Set[str]:
+        with self._lock:
+            bms = set()
+            databases = (set(must_included_databases)
+                         | set(self._bookmarks.keys()))
+            for database in databases:
+                bms.update(await self._get_bookmarks(database))
+            return bms

neo4j/_async/driver.py

@@ -43,7 +43,9 @@
 )
 from ..addressing import Address
 from ..api import (
+    AsyncBookmarkManager,
     Auth,
+    BookmarkManager,
     Bookmarks,
     DRIVER_BOLT,
     DRIVER_NEO4J,
@@ -62,9 +64,14 @@
     URI_SCHEME_NEO4J_SECURE,
     URI_SCHEME_NEO4J_SELF_SIGNED_CERTIFICATE,
 )
+from .bookmark_manager import AsyncNeo4jBookmarkManager
 from .work import AsyncSession
 
 
+_T_BmSupplier = t.Callable[[str], t.Union[Bookmarks, t.Awaitable[Bookmarks]]]
+_T_NotifyBm = t.Callable[[str, Bookmarks], t.Union[None, t.Awaitable[None]]]
+
+
 class AsyncGraphDatabase:
     """Accessor for :class:`neo4j.Driver` construction.
     """
@@ -94,6 +101,8 @@ def driver(
             ssl_context: ssl.SSLContext = ...,
             user_agent: str = ...,
             keep_alive: bool = ...,
+            bookmark_manager: t.Union[AsyncBookmarkManager,
+                                      BookmarkManager, None] = ...,
 
             # undocumented/unsupported options
             # they may be change or removed any time without prior notice
@@ -208,6 +217,36 @@ def driver(cls, uri, *, auth=None, **config) -> AsyncDriver:
             return cls.neo4j_driver(parsed.netloc, auth=auth,
                                     routing_context=routing_context, **config)
 
+    @classmethod
+    def bookmark_manager(
+        cls, initial_bookmarks: Bookmarks = None,
+        bookmark_supplier: _T_BmSupplier = None,
+        notify_bookmarks: _T_NotifyBm = None
+    ) -> AsyncBookmarkManager:
+        """Create a default :class:`AsyncBookmarkManager`.
+
+        :param initial_bookmarks:
+            The initial set of bookmarks. The default bookmark manager will
+            seed the set of bookmarks for each database with this value.
+        :param bookmark_supplier:
+            Function which will be called every time the default bookmark
+            manager's method :meth:`.AsyncBookmarkManager.get_bookmarks`
+            gets called. The result of ``bookmark_supplier`` will be
+            concatenated with the internal set of bookmarks and used to
+            configure the session in creation.
+        :param notify_bookmarks:
+            Function which will be called whenever the set of bookmarks
+            handled by the bookmark manager gets updated with the new
+            internal bookmark set.
+
+        :returns: A default implementation of :class:`AsyncBookmarkManager`.
+        """
+        return AsyncNeo4jBookmarkManager(
+            initial_bookmarks=initial_bookmarks,
+            bookmark_supplier=bookmark_supplier,
+            notify_bookmarks=notify_bookmarks
+        )
+
     @classmethod
     def bolt_driver(cls, target, *, auth=None, **config):
         """ Create a driver for direct Bolt server access that uses
@@ -353,6 +392,8 @@ def session(
             initial_retry_delay: float = ...,
             retry_delay_multiplier: float = ...,
             retry_delay_jitter_factor: float = ...,
+            bookmark_manager: t.Union[AsyncBookmarkManager,
+                                      BookmarkManager, None] = ...,
         ) -> AsyncSession:
             ...
 
@@ -394,6 +435,8 @@ async def verify_connectivity(
             initial_retry_delay: float = ...,
             retry_delay_multiplier: float = ...,
             retry_delay_jitter_factor: float = ...,
+            bookmark_manager: t.Union[AsyncBookmarkManager,
+                                      BookmarkManager, None] = ...,
         ) -> None:
             ...
 
@@ -456,6 +499,8 @@ async def get_server_info(
             initial_retry_delay: float = ...,
             retry_delay_multiplier: float = ...,
             retry_delay_jitter_factor: float = ...,
+            bookmark_manager: t.Union[AsyncBookmarkManager,
+                                      BookmarkManager, None] = ...,
         ) -> ServerInfo:
             ...
 

neo4j/_async/work/session.py

@@ -100,9 +100,12 @@ class AsyncSession(AsyncWorkspace):
     _state_failed = False
 
     def __init__(self, pool, session_config):
-        super().__init__(pool, session_config)
         assert isinstance(session_config, SessionConfig)
-        self._bookmarks = self._prepare_bookmarks(session_config.bookmarks)
+        super().__init__(pool, session_config)
+        if session_config.bookmarks is not None:
+            self._bookmarks = self._prepare_bookmarks(session_config.bookmarks)
+        else:
+            self._bookmark_manager = session_config.bookmark_manager
 
     async def __aenter__(self) -> AsyncSession:
         return self
@@ -147,10 +150,6 @@ async def _disconnect(self, sync=False):
             self._handle_cancellation(message="_disconnect")
             raise
 
-    def _collect_bookmark(self, bookmark):
-        if bookmark:
-            self._bookmarks = bookmark,
-
     def _handle_cancellation(self, message="General"):
         self._transaction = None
         self._auto_result = None
@@ -165,7 +164,7 @@ def _handle_cancellation(self, message="General"):
 
     async def _result_closed(self):
         if self._auto_result:
-            self._collect_bookmark(self._auto_result._bookmark)
+            await self._update_bookmark(self._auto_result._bookmark)
             self._auto_result = None
             await self._disconnect()
 
@@ -196,7 +195,7 @@ async def close(self) -> None:
                 if self._state_failed is False:
                     try:
                         await self._auto_result.consume()
-                        self._collect_bookmark(self._auto_result._bookmark)
+                        await self._update_bookmark(self._auto_result._bookmark)
                     except Exception as error:
                         # TODO: Investigate potential non graceful close states
                         self._auto_result = None
@@ -336,7 +335,7 @@ async def last_bookmark(self) -> t.Optional[str]:
             await self._auto_result.consume()
 
         if self._transaction and self._transaction._closed:
-            self._collect_bookmark(self._transaction._bookmark)
+            await self._update_bookmark(self._transaction._bookmark)
             self._transaction = None
 
         if self._bookmarks:
@@ -377,14 +376,14 @@ async def last_bookmarks(self) -> Bookmarks:
             await self._auto_result.consume()
 
         if self._transaction and self._transaction._closed():
-            self._collect_bookmark(self._transaction._bookmark)
+            await self._update_bookmark(self._transaction._bookmark)
             self._transaction = None
 
         return Bookmarks.from_raw_values(self._bookmarks)
 
     async def _transaction_closed_handler(self):
         if self._transaction:
-            self._collect_bookmark(self._transaction._bookmark)
+            await self._update_bookmark(self._transaction._bookmark)
             self._transaction = None
             await self._disconnect()
 

neo4j/_async/work/workspace.py

@@ -20,6 +20,7 @@
 
 import asyncio
 
+from ..._async_compat.util import AsyncUtil
 from ..._conf import WorkspaceConfig
 from ..._deadline import Deadline
 from ..._meta import (
@@ -44,7 +45,9 @@ def __init__(self, pool, config):
         self._connection_access_mode = None
         # Sessions are supposed to cache the database on which to operate.
         self._cached_database = False
-        self._bookmarks = None
+        self._bookmarks = ()
+        self._bookmark_manager = None
+        self._last_from_bookmark_manager = None
         # Workspace has been closed.
         self._closed = False
 
@@ -77,6 +80,40 @@ def _set_cached_database(self, database):
         self._cached_database = True
         self._config.database = database
 
+    async def _get_bookmarks(self, database):
+        if self._bookmark_manager is not None:
+            self._bookmarks = tuple(
+                await AsyncUtil.callback(
+                    self._bookmark_manager.get_bookmarks, database
+                )
+            )
+        return self._bookmarks
+
+    async def _get_all_bookmarks(self, must_included_databases):
+        if self._bookmark_manager is not None:
+            self._bookmarks = tuple(
+                await AsyncUtil.callback(
+                    self._bookmark_manager.get_all_bookmarks,
+                    must_included_databases
+                )
+            )
+        return self._bookmarks
+
+    async def _update_bookmarks(self, database, new_bookmarks):
+        if not new_bookmarks:
+            return
+        previous_bookmarks = self._bookmarks
+        self._bookmarks = new_bookmarks
+        if self._bookmark_manager is None:
+            return
+        await self._bookmark_manager.update_bookmarks(
+            database, previous_bookmarks, new_bookmarks
+        )
+
+    async def _update_bookmark(self, bookmark):
+        if bookmark:
+            await self._update_bookmarks(self._config.database, (bookmark,))
+
     async def _connect(self, access_mode, **acquire_kwargs):
         timeout = Deadline(self._config.session_connection_timeout)
         if self._connection:

neo4j/_conf.py

@@ -429,6 +429,10 @@ class WorkspaceConfig(Config):
     impersonated_user = None
     # Note that you need appropriate permissions to do so.
 
+    #: Bookmark Manager
+    bookmark_manager = None
+    # Specify the bookmark manager to be used for sessions by default.
+
 
 class SessionConfig(WorkspaceConfig):
     """ Session configuration.