feat: Multiplexed sessions - Remove handling of MethodNotImplemented exception from DatabaseSessionManager and add unit tests.

Signed-off-by: Taylor Curran <taylor.curran@improving.com>

Author: currantw

google/cloud/spanner_v1/database_sessions_manager.py

@@ -11,38 +11,40 @@
 # 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.
-import datetime
-import threading
-import time
-import weakref
-
-from google.api_core.exceptions import MethodNotImplemented
+from datetime import timedelta
+from threading import Event, Lock, Thread
+from time import sleep, time
+from weakref import ref
 
+from google.cloud.spanner_v1.session import Session
+from google.cloud.spanner_v1.session_options import TransactionType
 from google.cloud.spanner_v1._opentelemetry_tracing import (
     get_current_span,
     add_span_event,
 )
-from google.cloud.spanner_v1.session import Session
-from google.cloud.spanner_v1.session_options import TransactionType
 
 
 class DatabaseSessionsManager(object):
     """Manages sessions for a Cloud Spanner database.
+
     Sessions can be checked out from the database session manager for a specific
     transaction type using :meth:`get_session`, and returned to the session manager
     using :meth:`put_session`.
-    The sessions returned by the session manager depend on the client's session options (see
-    :class:`~google.cloud.spanner_v1.session_options.SessionOptions`) and the provided session
-    pool (see :class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`).
+
+    The sessions returned by the session manager depend on the client's session options
+    (see :class:`~google.cloud.spanner_v1.session_options.SessionOptions`) and the
+    provided session pool (see :class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`).
+
     :type database: :class:`~google.cloud.spanner_v1.database.Database`
     :param database: The database to manage sessions for.
+
     :type pool: :class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`
     :param pool: The pool to get non-multiplexed sessions from.
     """
 
     # Intervals for the maintenance thread to check and refresh the multiplexed session.
-    _MAINTENANCE_THREAD_POLLING_INTERVAL = datetime.timedelta(minutes=10)
-    _MAINTENANCE_THREAD_REFRESH_INTERVAL = datetime.timedelta(days=7)
+    _MAINTENANCE_THREAD_POLLING_INTERVAL = timedelta(minutes=10)
+    _MAINTENANCE_THREAD_REFRESH_INTERVAL = timedelta(days=7)
 
     def __init__(self, database, pool):
         self._database = database
@@ -56,45 +58,29 @@ def __init__(self, database, pool):
         # so that the thread can terminate if the use of multiplexed session has been
         # disabled for all transactions.
         self._multiplexed_session = None
-        self._multiplexed_session_maintenance_thread = None
-        self._multiplexed_session_lock = threading.Lock()
-        self._is_multiplexed_sessions_disabled_event = threading.Event()
-
-    @property
-    def _logger(self):
-        """The logger used by this database session manager.
-
-        :rtype: :class:`logging.Logger`
-        :returns: The logger.
-        """
-        return self._database.logger
+        self._multiplexed_session_thread = None
+        self._multiplexed_session_lock = Lock()
+        self._multiplexed_session_disabled_event = Event()
 
     def get_session(self, transaction_type: TransactionType) -> Session:
         """Returns a session for the given transaction type from the database session manager.
+
         :rtype: :class:`~google.cloud.spanner_v1.session.Session`
         :returns: a session for the given transaction type.
         """
 
         session_options = self._database.session_options
         use_multiplexed = session_options.use_multiplexed(transaction_type)
 
+        # TODO multiplexed: enable for read/write transactions
         if use_multiplexed and transaction_type == TransactionType.READ_WRITE:
             raise NotImplementedError(
                 f"Multiplexed sessions are not yet supported for {transaction_type} transactions."
             )
 
-        if use_multiplexed:
-            try:
-                session = self._get_multiplexed_session()
-
-            # If multiplexed sessions are not supported, disable
-            # them for all transactions and return a non-multiplexed session.
-            except MethodNotImplemented:
-                self._disable_multiplexed_sessions()
-                session = self._pool.get()
-
-        else:
-            session = self._pool.get()
+        session = (
+            self._get_multiplexed_session() if use_multiplexed else self._pool.get()
+        )
 
         add_span_event(
             get_current_span(),
@@ -106,6 +92,7 @@ def get_session(self, transaction_type: TransactionType) -> Session:
 
     def put_session(self, session: Session) -> None:
         """Returns the session to the database session manager.
+
         :type session: :class:`~google.cloud.spanner_v1.session.Session`
         :param session: The session to return to the database session manager.
         """
@@ -124,12 +111,12 @@ def put_session(self, session: Session) -> None:
 
     def _get_multiplexed_session(self) -> Session:
         """Returns a multiplexed session from the database session manager.
+
         If the multiplexed session is not defined, creates a new multiplexed
         session and starts a maintenance thread to periodically delete and
         recreate it so that it remains valid. Otherwise, simply returns the
         current multiplexed session.
-        :raises MethodNotImplemented:
-            if multiplexed sessions are not supported.
+
         :rtype: :class:`~google.cloud.spanner_v1.session.Session`
         :returns: a multiplexed session.
         """
@@ -138,18 +125,14 @@ def _get_multiplexed_session(self) -> Session:
             if self._multiplexed_session is None:
                 self._multiplexed_session = self._build_multiplexed_session()
 
-                # Build and start a thread to maintain the multiplexed session.
-                self._multiplexed_session_maintenance_thread = (
-                    self._build_maintenance_thread()
-                )
-                self._multiplexed_session_maintenance_thread.start()
+                self._multiplexed_session_thread = self._build_maintenance_thread()
+                self._multiplexed_session_thread.start()
 
             return self._multiplexed_session
 
     def _build_multiplexed_session(self) -> Session:
         """Builds and returns a new multiplexed session for the database session manager.
-        :raises MethodNotImplemented:
-            if multiplexed sessions are not supported.
+
         :rtype: :class:`~google.cloud.spanner_v1.session.Session`
         :returns: a new multiplexed session.
         """
@@ -159,34 +142,34 @@ def _build_multiplexed_session(self) -> Session:
             database_role=self._database.database_role,
             is_multiplexed=True,
         )
-
         session.create()
 
-        self._logger.info("Created multiplexed session.")
+        self._database.logger.info("Created multiplexed session.")
 
         return session
 
     def _disable_multiplexed_sessions(self) -> None:
         """Disables multiplexed sessions for all transactions."""
 
         self._multiplexed_session = None
-        self._is_multiplexed_sessions_disabled_event.set()
-        self._database.session_options.disable_multiplexed(self._logger)
+        self._multiplexed_session_disabled_event.set()
+        self._database.session_options.disable_multiplexed(self._database.logger)
 
-    def _build_maintenance_thread(self) -> threading.Thread:
+    def _build_maintenance_thread(self) -> Thread:
         """Builds and returns a multiplexed session maintenance thread for
         the database session manager. This thread will periodically delete
         and recreate the multiplexed session to ensure that it is always valid.
+
         :rtype: :class:`threading.Thread`
         :returns: a multiplexed session maintenance thread.
         """
 
         # Use a weak reference to the database session manager to avoid
         # creating a circular reference that would prevent the database
         # session manager from being garbage collected.
-        session_manager_ref = weakref.ref(self)
+        session_manager_ref = ref(self)
 
-        return threading.Thread(
+        return Thread(
             target=self._maintain_multiplexed_session,
             name=f"maintenance-multiplexed-session-{self._multiplexed_session.name}",
             args=[session_manager_ref],
@@ -196,10 +179,12 @@ def _build_maintenance_thread(self) -> threading.Thread:
     @staticmethod
     def _maintain_multiplexed_session(session_manager_ref) -> None:
         """Maintains the multiplexed session for the database session manager.
+
         This method will delete and recreate the referenced database session manager's
         multiplexed session to ensure that it is always valid. The method will run until
         the database session manager is deleted, the multiplexed session is deleted, or
         building a multiplexed session fails.
+
         :type session_manager_ref: :class:`_weakref.ReferenceType`
         :param session_manager_ref: A weak reference to the database session manager.
         """
@@ -215,7 +200,7 @@ def _maintain_multiplexed_session(session_manager_ref) -> None:
             session_manager._MAINTENANCE_THREAD_REFRESH_INTERVAL.total_seconds()
         )
 
-        session_created_time = time.time()
+        session_created_time = time()
 
         while True:
             # Terminate the thread is the database session manager has been deleted.
@@ -224,26 +209,18 @@ def _maintain_multiplexed_session(session_manager_ref) -> None:
                 return
 
             # Terminate the thread if the use of multiplexed sessions has been disabled.
-            if session_manager._is_multiplexed_sessions_disabled_event.is_set():
+            if session_manager._multiplexed_session_disabled_event.is_set():
                 return
 
             # Wait for until the refresh interval has elapsed.
-            if time.time() - session_created_time < refresh_interval_seconds:
-                time.sleep(polling_interval_seconds)
+            if time() - session_created_time < refresh_interval_seconds:
+                sleep(polling_interval_seconds)
                 continue
 
             with session_manager._multiplexed_session_lock:
                 session_manager._multiplexed_session.delete()
+                session_manager._multiplexed_session = (
+                    session_manager._build_multiplexed_session()
+                )
 
-                try:
-                    session_manager._multiplexed_session = (
-                        session_manager._build_multiplexed_session()
-                    )
-
-                # Disable multiplexed sessions for all transactions and terminate
-                # the thread if building a multiplexed session fails.
-                except MethodNotImplemented:
-                    session_manager._disable_multiplexed_sessions()
-                    return
-
-            session_created_time = time.time()
+            session_created_time = time()

tests/_builders.py

@@ -14,9 +14,75 @@
 
 from mock import create_autospec
 
+# Default values used to populate required or expected attributes.
+# Tests should not depend on them: if a test requires a specific
+# identifier or name, it should set it explicitly.
+_PROJECT_ID = "default-project-id"
+_INSTANCE_ID = "default-instance-id"
+_DATABASE_ID = "default-database-id"
+
 
 def build_logger():
     """Builds and returns a logger for testing."""
     from logging import Logger
 
     return create_autospec(Logger, instance=True)
+
+
+# Client objects
+# --------------
+
+
+def build_client(**kwargs):
+    """Builds and returns a client for testing using the given arguments.
+    If a required argument is not provided, a default value will be used."""
+    from google.cloud.spanner_v1 import Client
+
+    if "project" not in kwargs:
+        kwargs["project"] = _PROJECT_ID
+
+    return Client(**kwargs)
+
+
+def build_database(**kwargs):
+    """Builds and returns a database for testing using the given arguments.
+    If a required argument is not provided, a default value will be used.."""
+    from google.cloud.spanner_v1.database import Database
+
+    if "database_id" not in kwargs:
+        kwargs["database_id"] = _DATABASE_ID
+
+    if "logger" not in kwargs:
+        kwargs["logger"] = build_logger()
+
+    if "instance" not in kwargs or isinstance(kwargs["instance"], dict):
+        instance_args = kwargs.pop("instance", {})
+        kwargs["instance"] = build_instance(**instance_args)
+
+    database = Database(**kwargs)
+    database._spanner_api = build_spanner_api()
+
+    return database
+
+
+def build_instance(**kwargs):
+    """Builds and returns an instance for testing using the given arguments.
+    If a required argument is not provided, a default value will be used."""
+    from google.cloud.spanner_v1.instance import Instance
+
+    if "instance_id" not in kwargs:
+        kwargs["instance_id"] = _INSTANCE_ID
+
+    if "client" not in kwargs or isinstance(kwargs["client"], dict):
+        client_args = kwargs.pop("client", {})
+        kwargs["client"] = build_client(**client_args)
+
+    return Instance(**kwargs)
+
+
+def build_spanner_api():
+    """Builds and returns a mock Spanner Client API for testing using the given arguments.
+    Commonly used methods are mocked to return default values."""
+    from google.cloud.spanner_v1 import SpannerClient
+
+    return create_autospec(SpannerClient, instance=True)