Add support for multiplexed sessions - part 1 (classes)

- Adds `SessionOptions` and `DatabaseSessionManager` class support multiplexed session in future work.
- Update `Client`, `Instance`, `Database`, `Session` and associated `Checkout` classes accordingly.
- Add unit tests for new classes and update those for existing classes.

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

Author: currantw

google/cloud/spanner_dbapi/connection.py

@@ -330,8 +330,13 @@ def _session_checkout(self):
         """
         if self.database is None:
             raise ValueError("Database needs to be passed for this operation")
+
         if not self._session:
-            self._session = self.database._pool.get()
+            self._session = (
+                self.database._session_manager.get_session_for_read_only()
+                if self.read_only
+                else self.database._session_manager.get_session_for_read_write()
+            )
 
         return self._session
 
@@ -344,7 +349,7 @@ def _release_session(self):
             return
         if self.database is None:
             raise ValueError("Database needs to be passed for this operation")
-        self.database._pool.put(self._session)
+        self.database._session_manager.put_session(self._session)
         self._session = None
 
     def transaction_checkout(self):
@@ -400,7 +405,7 @@ def close(self):
             self._transaction.rollback()
 
         if self._own_pool and self.database:
-            self.database._pool.clear()
+            self.database._session_manager._pool.clear()
 
         self.is_closed = True
 

google/cloud/spanner_v1/client.py

@@ -70,6 +70,7 @@
 except ImportError:  # pragma: NO COVER
     HAS_GOOGLE_CLOUD_MONITORING_INSTALLED = False
 
+from google.cloud.spanner_v1.session_options import SessionOptions
 
 _CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__)
 EMULATOR_ENV_VAR = "SPANNER_EMULATOR_HOST"
@@ -171,6 +172,9 @@ class Client(ClientWithProject):
         or :class:`dict`
     :param default_transaction_options: (Optional) Default options to use for all transactions.
 
+    :type session_options: :class:`~google.cloud.spanner_v1.SessionOptions`
+    :param session_options: (Optional) Options for client sessions.
+
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`
     """
@@ -193,6 +197,7 @@ def __init__(
         directed_read_options=None,
         observability_options=None,
         default_transaction_options: Optional[DefaultTransactionOptions] = None,
+        session_options=None,
     ):
         self._emulator_host = _get_spanner_emulator_host()
 
@@ -262,6 +267,8 @@ def __init__(
             )
         self._default_transaction_options = default_transaction_options
 
+        self._session_options = session_options or SessionOptions()
+
     @property
     def credentials(self):
         """Getter for client's credentials.
@@ -525,3 +532,12 @@ def default_transaction_options(
             )
 
         self._default_transaction_options = default_transaction_options
+
+    @property
+    def session_options(self):
+        """Returns the session options for the client.
+
+        :rtype: :class:`~google.cloud.spanner_v1.SessionOptions`
+        :returns: The session options for the client.
+        """
+        return self._session_options

google/cloud/spanner_v1/database.py

@@ -40,6 +40,7 @@
 from google.cloud.spanner_admin_database_v1 import RestoreDatabaseRequest
 from google.cloud.spanner_admin_database_v1 import UpdateDatabaseDdlRequest
 from google.cloud.spanner_admin_database_v1.types import DatabaseDialect
+from google.cloud.spanner_v1.database_sessions_manager import DatabaseSessionsManager
 from google.cloud.spanner_v1.transaction import BatchTransactionId
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import Type
@@ -59,7 +60,6 @@
 from google.cloud.spanner_v1.keyset import KeySet
 from google.cloud.spanner_v1.merged_result_set import MergedResultSet
 from google.cloud.spanner_v1.pool import BurstyPool
-from google.cloud.spanner_v1.pool import SessionCheckout
 from google.cloud.spanner_v1.session import Session
 from google.cloud.spanner_v1.snapshot import _restart_on_unavailable
 from google.cloud.spanner_v1.snapshot import Snapshot
@@ -70,7 +70,6 @@
 from google.cloud.spanner_v1.table import Table
 from google.cloud.spanner_v1._opentelemetry_tracing import (
     add_span_event,
-    get_current_span,
     trace_call,
 )
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
@@ -191,10 +190,10 @@ def __init__(
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
-
-        self._pool = pool
         pool.bind(self)
 
+        self._session_manager = DatabaseSessionsManager(database=self, pool=pool)
+
     @classmethod
     def from_pb(cls, database_pb, instance, pool=None):
         """Creates an instance of this class from a protobuf.
@@ -708,7 +707,7 @@ def execute_pdml():
                 "CloudSpanner.Database.execute_partitioned_pdml",
                 observability_options=self.observability_options,
             ) as span, MetricsCapture():
-                with SessionCheckout(self._pool) as session:
+                with SessionCheckout(self) as session:
                     add_span_event(span, "Starting BeginTransaction")
                     txn = api.begin_transaction(
                         session=session.name, options=txn_options, metadata=metadata
@@ -759,6 +758,7 @@ def session(self, labels=None, database_role=None):
         """
         # If role is specified in param, then that role is used
         # instead.
+
         role = database_role or self._database_role
         return Session(self, labels=labels, database_role=role)
 
@@ -923,7 +923,7 @@ def run_in_transaction(self, func, *args, **kw):
             # Check out a session and run the function in a transaction; once
             # done, flip the sanity check bit back.
             try:
-                with SessionCheckout(self._pool) as session:
+                with SessionCheckout(self) as session:
                     return session.run_in_transaction(func, *args, **kw)
             finally:
                 self._local.transaction_running = False
@@ -1160,6 +1160,29 @@ def observability_options(self):
         return opts
 
 
+class SessionCheckout(object):
+    """Context manager for using a session from a database.
+
+    :type database: :class:`~google.cloud.spanner_v1.database.Database`
+    :param database: database to use the session from
+    """
+
+    _session = None  # Not checked out until '__enter__'.
+
+    def __init__(self, database):
+        if not isinstance(database, Database):
+            raise TypeError(f"database must be an instance of {Database.__name__}")
+
+        self._database = database
+
+    def __enter__(self):
+        self._session = self._database._session_manager.get_session_for_read_write()
+        return self._session
+
+    def __exit__(self, *ignored):
+        self._database._session_manager.put_session(self._session)
+
+
 class BatchCheckout(object):
     """Context manager for using a batch from a database.
 
@@ -1194,6 +1217,9 @@ def __init__(
         isolation_level=TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
         **kw,
     ):
+        if not isinstance(database, Database):
+            raise TypeError(f"database must be an instance of {Database.__name__}")
+
         self._database = database
         self._session = self._batch = None
         if request_options is None:
@@ -1209,9 +1235,9 @@ def __init__(
 
     def __enter__(self):
         """Begin ``with`` block."""
-        current_span = get_current_span()
-        session = self._session = self._database._pool.get()
-        add_span_event(current_span, "Using session", {"id": session.session_id})
+        session = (
+            self._session
+        ) = self._database._session_manager.get_session_for_read_only()
         batch = self._batch = Batch(session)
         if self._request_options.transaction_tag:
             batch.transaction_tag = self._request_options.transaction_tag
@@ -1235,13 +1261,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     "CommitStats: {}".format(self._batch.commit_stats),
                     extra={"commit_stats": self._batch.commit_stats},
                 )
-            self._database._pool.put(self._session)
-            current_span = get_current_span()
-            add_span_event(
-                current_span,
-                "Returned session to pool",
-                {"id": self._session.session_id},
-            )
+            self._database._session_manager.put_session(self._session)
 
 
 class MutationGroupsCheckout(object):
@@ -1258,23 +1278,22 @@ class MutationGroupsCheckout(object):
     """
 
     def __init__(self, database):
+        if not isinstance(database, Database):
+            raise TypeError(f"database must be an instance of {Database.__name__}")
+
         self._database = database
         self._session = None
 
     def __enter__(self):
         """Begin ``with`` block."""
-        session = self._session = self._database._pool.get()
+        session = (
+            self._session
+        ) = self._database._session_manager.get_session_for_read_write()
         return MutationGroups(session)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """End ``with`` block."""
-        if isinstance(exc_val, NotFound):
-            # If NotFound exception occurs inside the with block
-            # then we validate if the session still exists.
-            if not self._session.exists():
-                self._session = self._database._pool._new_session()
-                self._session.create()
-        self._database._pool.put(self._session)
+        self._database._session_manager.put_session(self._session)
 
 
 class SnapshotCheckout(object):
@@ -1296,24 +1315,23 @@ class SnapshotCheckout(object):
     """
 
     def __init__(self, database, **kw):
+        if not isinstance(database, Database):
+            raise TypeError(f"database must be an instance of {Database.__name__}")
+
         self._database = database
         self._session = None
         self._kw = kw
 
     def __enter__(self):
         """Begin ``with`` block."""
-        session = self._session = self._database._pool.get()
+        session = (
+            self._session
+        ) = self._database._session_manager.get_session_for_read_only()
         return Snapshot(session, **self._kw)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """End ``with`` block."""
-        if isinstance(exc_val, NotFound):
-            # If NotFound exception occurs inside the with block
-            # then we validate if the session still exists.
-            if not self._session.exists():
-                self._session = self._database._pool._new_session()
-                self._session.create()
-        self._database._pool.put(self._session)
+        self._database._session_manager.put_session(self._session)
 
 
 class BatchSnapshot(object):
@@ -1392,6 +1410,7 @@ def _get_session(self):
            Caller is responsible for cleaning up the session after
            all partitions have been processed.
         """
+
         if self._session is None:
             session = self._session = self._database.session()
             if self._session_id is None:

google/cloud/spanner_v1/database_sessions_manager.py

@@ -0,0 +1,94 @@
+# Copyright 2025 Google LLC All rights reserved.
+#
+# 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
+#
+#     http://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 google.cloud.spanner_v1.session import Session
+
+
+class DatabaseSessionsManager(object):
+    """Manages sessions for a Cloud Spanner database.
+
+    Sessions can be checked out from the session manager using :meth:`get_session_for_read_only`,
+    :meth:`get_session_for_partitioned`, and :meth:`get_session_for_read_write`, 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
+    (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.SessionPool`
+    :param pool: The pool to get non-multiplexed sessions from.
+    """
+
+    def __init__(self, database, pool):
+        self._database = database
+        self._pool = pool
+
+    def get_session_for_read_only(self) -> Session:
+        """Checks out and returns a session for read-only transactions from the session manager.
+
+        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
+        :returns: a session for read-only transactions.
+        """
+
+        if (
+            self._database._instance._client.session_options.use_multiplexed_for_partitioned
+        ):
+            raise NotImplementedError(
+                "Multiplexed sessions are not yet supported for read-only transactions."
+            )
+
+        return self._pool.get()
+
+    def get_session_for_partitioned(self) -> Session:
+        """Checks out and returns a session for partitioned transactions from the session manager.
+
+        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
+        :returns: a session for partitioned transactions.
+        """
+
+        if (
+            self._database._instance._client.session_options.use_multiplexed_for_partitioned
+        ):
+            raise NotImplementedError(
+                "Multiplexed sessions are not yet supported for partitioned transactions."
+            )
+
+        return self._pool.get()
+
+    def get_session_for_read_write(self) -> Session:
+        """Checks out and returns a session for read/write transactions from the session manager.
+
+        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
+        :returns: a session for read-write transactions.
+        """
+
+        if (
+            self._database._instance._client.session_options.use_multiplexed_for_read_write
+        ):
+            raise NotImplementedError(
+                "Multiplexed sessions are not yet supported for read/write transactions."
+            )
+
+        return self._pool.get()
+
+    def put_session(self, session: Session) -> None:
+        """Returns the session to the database session manager."""
+
+        if session.is_multiplexed:
+            raise NotImplementedError("Multiplexed sessions are not yet supported.")
+
+        self._pool.put(session)