feat: Add support for multiplexed sessions

Author: rahul2393

.github/workflows/integration-tests-against-emulator-with-multiplexed-session.yaml

@@ -0,0 +1,33 @@
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+name: Run Spanner integration tests against emulator
+jobs:
+  system-tests:
+    runs-on: ubuntu-latest
+
+    services:
+      emulator:
+        image: gcr.io/cloud-spanner-emulator/emulator:latest
+        ports:
+          - 9010:9010
+          - 9020:9020
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+      - name: Setup Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: 3.8
+      - name: Install nox
+        run: python -m pip install nox
+      - name: Run system tests
+        run: nox -s system
+        env:
+          SPANNER_EMULATOR_HOST: localhost:9010
+          GOOGLE_CLOUD_PROJECT: emulator-test-project
+          GOOGLE_CLOUD_TESTS_CREATE_SPANNER_INSTANCE: true
+          GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS: true

.kokoro/presubmit/integration-multiplexed-sessions-enabled.cfg

@@ -0,0 +1,17 @@
+# Format: //devtools/kokoro/config/proto/build.proto
+
+# Only run a subset of all nox sessions
+env_vars: {
+    key: "NOX_SESSION"
+    value: "unit-3.8 unit-3.12 system-3.8"
+}
+
+env_vars: {
+    key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS"
+    value: "true"
+}
+
+env_vars: {
+    key: "GOOGLE_CLOUD_SPANNER_MULTIPLEXED_SESSIONS_PARTITIONED_OPS"
+    value: "true"
+}

google/cloud/spanner_dbapi/connection.py

@@ -28,6 +28,7 @@
 from google.cloud.spanner_dbapi.transaction_helper import TransactionRetryHelper
 from google.cloud.spanner_dbapi.cursor import Cursor
 from google.cloud.spanner_v1 import RequestOptions, TransactionOptions
+from google.cloud.spanner_v1.session_options import TransactionType
 from google.cloud.spanner_v1.snapshot import Snapshot
 
 from google.cloud.spanner_dbapi.exceptions import (
@@ -357,7 +358,12 @@ 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()
+            transaction_type = (
+                TransactionType.READ_ONLY
+                if self.read_only
+                else TransactionType.READ_WRITE
+            )
+            self._session = self.database._session_manager.get_session(transaction_type)
 
         return self._session
 
@@ -370,7 +376,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):
@@ -432,7 +438,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

@@ -60,6 +60,7 @@
 from google.cloud.spanner_v1.metrics.metrics_exporter import (
     CloudMonitoringMetricsExporter,
 )
+from google.cloud.spanner_v1.session_options import SessionOptions
 
 try:
     from opentelemetry import metrics
@@ -172,6 +173,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`
     """
@@ -196,6 +200,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()
 
@@ -268,6 +273,7 @@ def __init__(
         self._default_transaction_options = default_transaction_options
         self._nth_client_id = Client.NTH_CLIENT.increment()
         self._nth_request = AtomicCounter(0)
+        self._session_options = session_options or SessionOptions()
 
     @property
     def _next_nth_request(self):
@@ -373,6 +379,14 @@ def default_transaction_options(self):
         """
         return self._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
+
     @property
     def directed_read_options(self):
         """Getter for directed_read_options.

google/cloud/spanner_v1/database.py

@@ -57,11 +57,11 @@
 )
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.batch import MutationGroups
+from google.cloud.spanner_v1.database_sessions_manager import DatabaseSessionsManager
 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.session_options import SessionOptions, TransactionType
 from google.cloud.spanner_v1.snapshot import _restart_on_unavailable
 from google.cloud.spanner_v1.snapshot import Snapshot
 from google.cloud.spanner_v1.streamed import StreamedResultSet
@@ -196,9 +196,9 @@ 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):
@@ -462,6 +462,14 @@ def spanner_api(self):
 
         return self._spanner_api
 
+    @property
+    def session_options(self) -> SessionOptions:
+        """Session options for the database.
+        :rtype: :class:`~google.cloud.spanner_v1.session_options.SessionOptions`
+        :returns: the session options
+        """
+        return self._instance._client.session_options
+
     def metadata_with_request_id(
         self, nth_request, nth_attempt, prior_metadata=[], span=None
     ):
@@ -759,18 +767,31 @@ 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, TransactionType.PARTITIONED) as session:
                     add_span_event(span, "Starting BeginTransaction")
-                    txn = api.begin_transaction(
-                        session=session.name,
-                        options=txn_options,
-                        metadata=self.metadata_with_request_id(
-                            self._next_nth_request,
-                            1,
-                            metadata,
-                            span,
-                        ),
-                    )
+                    try:
+                        txn = api.begin_transaction(
+                            session=session.name,
+                            options=txn_options,
+                            metadata=self.metadata_with_request_id(
+                                self._next_nth_request,
+                                1,
+                                metadata,
+                                span,
+                            ),
+                        )
+                    # If partitioned DML is not supported with multiplexed sessions,
+                    # disable multiplexed sessions for partitioned transactions before
+                    # re-raising the error.
+                    except NotImplementedError as exc:
+                        if (
+                                "Transaction type partitioned_dml not supported with multiplexed sessions"
+                                in str(exc)
+                        ):
+                            self.session_options.disable_multiplexed(
+                                self.logger, TransactionType.PARTITIONED
+                            )
+                        raise exc
 
                     txn_selector = TransactionSelector(id=txn.id)
 
@@ -792,6 +813,7 @@ def execute_pdml():
                         method=method,
                         trace_name="CloudSpanner.ExecuteStreamingSql",
                         request=request,
+                        session=session,
                         metadata=metadata,
                         transaction_selector=txn_selector,
                         observability_options=self.observability_options,
@@ -817,23 +839,6 @@ def _nth_client_id(self):
             return self._instance._client._nth_client_id
         return 0
 
-    def session(self, labels=None, database_role=None):
-        """Factory to create a session for this database.
-
-        :type labels: dict (str -> str) or None
-        :param labels: (Optional) user-assigned labels for the session.
-
-        :type database_role: str
-        :param database_role: (Optional) user-assigned database_role for the session.
-
-        :rtype: :class:`~google.cloud.spanner_v1.session.Session`
-        :returns: a session bound to this database.
-        """
-        # 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)
-
     def snapshot(self, **kw):
         """Return an object which wraps a snapshot.
 
@@ -995,7 +1000,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
@@ -1241,6 +1246,50 @@ 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,  # type: ignore
+        transaction_type: TransactionType = TransactionType.READ_WRITE,
+    ):
+        # Move import here to avoid circular import
+        from google.cloud.spanner_v1.database import Database
+        if not isinstance(database, Database):
+            raise TypeError(
+                "{class_name} must receive an instance of {expected_class_name}. Received: {actual_class_name}".format(
+                    class_name=self.__class__.__name__,
+                    expected_class_name=Database.__name__,
+                    actual_class_name=database.__class__.__name__,
+                )
+            )
+
+        if not isinstance(transaction_type, TransactionType):
+            raise TypeError(
+                "{class_name} must receive an instance of {expected_class_name}. Received: {actual_class_name}".format(
+                    class_name=self.__class__.__name__,
+                    expected_class_name=TransactionType.__name__,
+                    actual_class_name=transaction_type.__class__.__name__,
+                )
+            )
+
+        self._database = database
+        self._transaction_type = transaction_type
+
+    def __enter__(self):
+        session_manager = self._database._session_manager
+        self._session = session_manager.get_session(self._transaction_type)
+        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.
 
@@ -1929,3 +1978,4 @@ def _retry_on_aborted(func, retry_config):
     """
     retry = retry_config.with_predicate(if_exception_type(Aborted))
     return retry(func)
+