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 with multiplexed sessions
+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_v1/database.py

@@ -62,6 +62,8 @@
 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
+from google.cloud.spanner_v1.database_sessions_manager import DatabaseSessionsManager
 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
@@ -199,6 +201,10 @@ def __init__(
 
         self._pool = pool
         pool.bind(self)
+        
+        # Initialize session options and sessions manager for multiplexed session support
+        self.session_options = SessionOptions()
+        self._sessions_manager = DatabaseSessionsManager(self, pool)
 
     @classmethod
     def from_pb(cls, database_pb, instance, pool=None):
@@ -759,7 +765,11 @@ def execute_pdml():
                 "CloudSpanner.Database.execute_partitioned_pdml",
                 observability_options=self.observability_options,
             ) as span, MetricsCapture():
-                with SessionCheckout(self._pool) as session:
+                from google.cloud.spanner_v1.session_options import TransactionType
+                
+                # Use sessions manager for partitioned DML operations
+                session = self._sessions_manager.get_session(TransactionType.PARTITIONED)
+                try:
                     add_span_event(span, "Starting BeginTransaction")
                     txn = api.begin_transaction(
                         session=session.name,
@@ -802,6 +812,8 @@ def execute_pdml():
                     list(result_set)  # consume all partials
 
                     return result_set.stats.row_count_lower_bound
+                finally:
+                    self._sessions_manager.put_session(session)
 
         return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
 
@@ -1240,6 +1252,15 @@ def observability_options(self):
         opts["db_name"] = self.name
         return opts
 
+    @property
+    def sessions_manager(self):
+        """Returns the database sessions manager.
+        
+        :rtype: :class:`~google.cloud.spanner_v1.database_sessions_manager.DatabaseSessionsManager`
+        :returns: The sessions manager for this database.
+        """
+        return self._sessions_manager
+
 
 class BatchCheckout(object):
     """Context manager for using a batch from a database.
@@ -1290,8 +1311,12 @@ def __init__(
 
     def __enter__(self):
         """Begin ``with`` block."""
+        from google.cloud.spanner_v1.session_options import TransactionType
+        
         current_span = get_current_span()
-        session = self._session = self._database._pool.get()
+        session = self._session = self._database.sessions_manager.get_session(
+            TransactionType.READ_WRITE
+        )
         add_span_event(current_span, "Using session", {"id": session.session_id})
         batch = self._batch = Batch(session)
         if self._request_options.transaction_tag:
@@ -1316,7 +1341,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)
+            self._database.sessions_manager.put_session(self._session)
             current_span = get_current_span()
             add_span_event(
                 current_span,
@@ -1344,7 +1369,11 @@ def __init__(self, database):
 
     def __enter__(self):
         """Begin ``with`` block."""
-        session = self._session = self._database._pool.get()
+        from google.cloud.spanner_v1.session_options import TransactionType
+        
+        session = self._session = self._database.sessions_manager.get_session(
+            TransactionType.READ_WRITE
+        )
         return MutationGroups(session)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -1355,7 +1384,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             if not self._session.exists():
                 self._session = self._database._pool._new_session()
                 self._session.create()
-        self._database._pool.put(self._session)
+        self._database.sessions_manager.put_session(self._session)
 
 
 class SnapshotCheckout(object):
@@ -1383,7 +1412,11 @@ def __init__(self, database, **kw):
 
     def __enter__(self):
         """Begin ``with`` block."""
-        session = self._session = self._database._pool.get()
+        from google.cloud.spanner_v1.session_options import TransactionType
+        
+        session = self._session = self._database.sessions_manager.get_session(
+            TransactionType.READ_ONLY
+        )
         return Snapshot(session, **self._kw)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -1394,7 +1427,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             if not self._session.exists():
                 self._session = self._database._pool._new_session()
                 self._session.create()
-        self._database._pool.put(self._session)
+        self._database.sessions_manager.put_session(self._session)
 
 
 class BatchSnapshot(object):
@@ -1474,10 +1507,13 @@ def _get_session(self):
            all partitions have been processed.
         """
         if self._session is None:
-            session = self._session = self._database.session()
-            if self._session_id is None:
-                session.create()
-            else:
+            from google.cloud.spanner_v1.session_options import TransactionType
+            
+            # Use sessions manager for partition operations
+            session = self._session = self._database.sessions_manager.get_session(
+                TransactionType.PARTITIONED
+            )
+            if self._session_id is not None:
                 session._session_id = self._session_id
         return self._session