feat: Propagate client_context in Session and update tests

- Update Session.transaction to accept client_context.
- Update unit tests to support client_context propagation.
- Update mock objects in tests to match the expected attribute hierarchy.
- Clean up nested imports in test files.

Author: aseering

google/cloud/spanner_v1/session.py

@@ -472,9 +472,14 @@ def batch(self):
 
         return Batch(self)
 
-    def transaction(self) -> Transaction:
+    def transaction(self, client_context=None) -> Transaction:
         """Create a transaction to perform a set of reads with shared staleness.
 
+        :type client_context: :class:`~google.cloud.spanner_v1.types.ClientContext`
+            or :class:`dict`
+        :param client_context: (Optional) Client context to use for all requests made
+                               by this transaction.
+
         :rtype: :class:`~google.cloud.spanner_v1.transaction.Transaction`
         :returns: a transaction bound to this session
 
@@ -483,7 +488,7 @@ def transaction(self) -> Transaction:
         if self._session_id is None:
             raise ValueError("Session has not been created.")
 
-        return Transaction(self)
+        return Transaction(self, client_context=client_context)
 
     def run_in_transaction(self, func, *args, **kw):
         """Perform a unit of work in a transaction, retrying on abort.
@@ -512,6 +517,8 @@ def run_in_transaction(self, func, *args, **kw):
                    the DDL option `allow_txn_exclusion` being false or unset.
                    "isolation_level" sets the isolation level for the transaction.
                    "read_lock_mode" sets the read lock mode for the transaction.
+                   "client_context" (Optional) Client context to use for all requests made
+                                    by this transaction.
 
         :rtype: Any
         :returns: The return value of ``func``.
@@ -529,6 +536,7 @@ def run_in_transaction(self, func, *args, **kw):
         )
         isolation_level = kw.pop("isolation_level", None)
         read_lock_mode = kw.pop("read_lock_mode", None)
+        client_context = kw.pop("client_context", None)
 
         database = self._database
         log_commit_stats = database.log_commit_stats
@@ -554,7 +562,7 @@ def run_in_transaction(self, func, *args, **kw):
             previous_transaction_id: Optional[bytes] = None
 
             while True:
-                txn = self.transaction()
+                txn = self.transaction(client_context=client_context)
                 txn.transaction_tag = transaction_tag
                 txn.exclude_txn_from_change_streams = exclude_txn_from_change_streams
                 txn.isolation_level = isolation_level

tests/unit/spanner_dbapi/test_connection.py

@@ -872,6 +872,7 @@ class _Client(object):
     def __init__(self, project="project_id"):
         self.project = project
         self.project_name = "projects/" + self.project
+        self._client_context = None
 
     def instance(self, instance_id="instance_id"):
         return _Instance(name=instance_id, client=self)

tests/unit/test_backup.py

@@ -679,6 +679,7 @@ class _Client(object):
     def __init__(self, project=TestBackup.PROJECT_ID):
         self.project = project
         self.project_name = "projects/" + self.project
+        self._client_context = None
 
 
 class _Instance(object):

tests/unit/test_batch.py

@@ -806,6 +806,9 @@ class _Database(object):
 
     def __init__(self, enable_end_to_end_tracing=False):
         self.name = "testing"
+        self._instance = mock.Mock()
+        self._instance._client = mock.Mock()
+        self._instance._client._client_context = None
         self._route_to_leader_enabled = True
         if enable_end_to_end_tracing:
             self.observability_options = dict(enable_end_to_end_tracing=True)

tests/unit/test_database.py

@@ -30,6 +30,7 @@
     RequestOptions,
     DirectedReadOptions,
     DefaultTransactionOptions,
+    ExecuteSqlRequest,
 )
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
@@ -2599,6 +2600,7 @@ def test__get_snapshot_new_wo_staleness(self):
             exact_staleness=None,
             multi_use=True,
             transaction_id=None,
+            client_context=None,
         )
         snapshot.begin.assert_called_once_with()
 
@@ -2614,6 +2616,7 @@ def test__get_snapshot_w_read_timestamp(self):
             exact_staleness=None,
             multi_use=True,
             transaction_id=None,
+            client_context=None,
         )
         snapshot.begin.assert_called_once_with()
 
@@ -2629,6 +2632,7 @@ def test__get_snapshot_w_exact_staleness(self):
             exact_staleness=duration,
             multi_use=True,
             transaction_id=None,
+            client_context=None,
         )
         snapshot.begin.assert_called_once_with()
 
@@ -3540,6 +3544,7 @@ def __init__(
         self.directed_read_options = directed_read_options
         self.default_transaction_options = default_transaction_options
         self.observability_options = observability_options
+        self._client_context = None
         self._nth_client_id = _Client.NTH_CLIENT.increment()
         self._nth_request = AtomicCounter()
 
@@ -3589,6 +3594,13 @@ class _Database(object):
     def __init__(self, name, instance=None):
         self.name = name
         self.database_id = name.rsplit("/", 1)[1]
+        if instance is None:
+            instance = mock.Mock()
+            instance._client = mock.Mock()
+            instance._client._client_context = None
+            instance._client._query_options = ExecuteSqlRequest.QueryOptions(
+                optimizer_version="1"
+            )
         self._instance = instance
         from logging import Logger
 

tests/unit/test_instance.py

@@ -1023,6 +1023,7 @@ def __init__(self, project, timeout_seconds=None):
         self.route_to_leader_enabled = True
         self.directed_read_options = None
         self.default_transaction_options = DefaultTransactionOptions()
+        self._client_context = None
 
     def copy(self):
         from copy import deepcopy

tests/unit/test_pool.py

@@ -19,7 +19,21 @@
 from datetime import datetime, timedelta
 
 import mock
+from google.cloud.spanner_v1 import pool as MUT
 from google.cloud.spanner_v1 import _opentelemetry_tracing
+from google.cloud.spanner_v1 import ExecuteSqlRequest
+from google.cloud.spanner_v1 import BatchCreateSessionsResponse
+from google.cloud.spanner_v1 import Session
+from google.cloud.spanner_v1 import SpannerClient
+from google.cloud.spanner_v1.database import Database
+from google.cloud.spanner_v1.pool import AbstractSessionPool
+from google.cloud.spanner_v1.pool import SessionCheckout
+from google.cloud.spanner_v1.pool import FixedSizePool
+from google.cloud.spanner_v1.pool import BurstyPool
+from google.cloud.spanner_v1.pool import PingingPool
+from google.cloud.spanner_v1.transaction import Transaction
+from google.cloud.exceptions import NotFound
+from google.cloud._testing import _Monkey
 from google.cloud.spanner_v1._helpers import (
     _metadata_with_request_id,
     _metadata_with_request_id_and_req_id,
@@ -40,20 +54,17 @@
 
 
 def _make_database(name="name"):
-    from google.cloud.spanner_v1.database import Database
 
     return mock.create_autospec(Database, instance=True)
 
 
 def _make_session():
-    from google.cloud.spanner_v1.database import Session
 
     return mock.create_autospec(Session, instance=True)
 
 
 class TestAbstractSessionPool(unittest.TestCase):
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.pool import AbstractSessionPool
 
         return AbstractSessionPool
 
@@ -129,7 +140,6 @@ def test__new_session_w_database_role(self):
         self.assertEqual(new_session.database_role, database_role)
 
     def test_session_wo_kwargs(self):
-        from google.cloud.spanner_v1.pool import SessionCheckout
 
         pool = self._make_one()
         checkout = pool.session()
@@ -139,7 +149,6 @@ def test_session_wo_kwargs(self):
         self.assertEqual(checkout._kwargs, {})
 
     def test_session_w_kwargs(self):
-        from google.cloud.spanner_v1.pool import SessionCheckout
 
         pool = self._make_one()
         checkout = pool.session(foo="bar")
@@ -164,7 +173,6 @@ class TestFixedSizePool(OpenTelemetryBase):
     enrich_with_otel_scope(BASE_ATTRIBUTES)
 
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.pool import FixedSizePool
 
         return FixedSizePool
 
@@ -559,7 +567,6 @@ class TestBurstyPool(OpenTelemetryBase):
     enrich_with_otel_scope(BASE_ATTRIBUTES)
 
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.pool import BurstyPool
 
         return BurstyPool
 
@@ -850,7 +857,6 @@ class TestPingingPool(OpenTelemetryBase):
     enrich_with_otel_scope(BASE_ATTRIBUTES)
 
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.pool import PingingPool
 
         return PingingPool
 
@@ -946,8 +952,6 @@ def test_get_hit_no_ping(self, mock_region):
     )
     def test_get_hit_w_ping(self, mock_region):
         import datetime
-        from google.cloud._testing import _Monkey
-        from google.cloud.spanner_v1 import pool as MUT
 
         pool = self._make_one(size=4)
         database = _Database("name")
@@ -974,8 +978,6 @@ def test_get_hit_w_ping(self, mock_region):
     )
     def test_get_hit_w_ping_expired(self, mock_region):
         import datetime
-        from google.cloud._testing import _Monkey
-        from google.cloud.spanner_v1 import pool as MUT
 
         pool = self._make_one(size=4)
         database = _Database("name")
@@ -1097,8 +1099,6 @@ def test_spans_put_full(self, mock_region):
     )
     def test_put_non_full(self, mock_region):
         import datetime
-        from google.cloud._testing import _Monkey
-        from google.cloud.spanner_v1 import pool as MUT
 
         pool = self._make_one(size=1)
         session_queue = pool._sessions = _Queue()
@@ -1172,8 +1172,6 @@ def test_ping_oldest_fresh(self, mock_region):
     )
     def test_ping_oldest_stale_but_exists(self, mock_region):
         import datetime
-        from google.cloud._testing import _Monkey
-        from google.cloud.spanner_v1 import pool as MUT
 
         pool = self._make_one(size=1)
         database = _Database("name")
@@ -1193,8 +1191,6 @@ def test_ping_oldest_stale_but_exists(self, mock_region):
     )
     def test_ping_oldest_stale_and_not_exists(self, mock_region):
         import datetime
-        from google.cloud._testing import _Monkey
-        from google.cloud.spanner_v1 import pool as MUT
 
         pool = self._make_one(size=1)
         database = _Database("name")
@@ -1257,7 +1253,6 @@ def test_spans_get_and_leave_empty_pool(self, mock_region):
 
 class TestSessionCheckout(unittest.TestCase):
     def _getTargetClass(self):
-        from google.cloud.spanner_v1.pool import SessionCheckout
 
         return SessionCheckout
 
@@ -1314,7 +1309,6 @@ def test_context_manager_w_kwargs(self):
 
 
 def _make_transaction(*args, **kw):
-    from google.cloud.spanner_v1.transaction import Transaction
 
     txn = mock.create_autospec(Transaction)(*args, **kw)
     txn.committed = None
@@ -1352,14 +1346,12 @@ def exists(self):
         return self._exists
 
     def ping(self):
-        from google.cloud.exceptions import NotFound
 
         self._pinged = True
         if not self._exists:
             raise NotFound("expired session")
 
     def delete(self):
-        from google.cloud.exceptions import NotFound
 
         self._deleted = True
         if not self._exists:
@@ -1391,9 +1383,6 @@ def mock_batch_create_sessions(
             metadata=[],
             labels={},
         ):
-            from google.cloud.spanner_v1 import BatchCreateSessionsResponse
-            from google.cloud.spanner_v1 import Session
-
             database_role = request.session_template.creator_role if request else None
             if request.session_count < 2:
                 response = BatchCreateSessionsResponse(
@@ -1408,10 +1397,15 @@ def mock_batch_create_sessions(
                 )
             return response
 
-        from google.cloud.spanner_v1 import SpannerClient
-
         self.spanner_api = mock.create_autospec(SpannerClient, instance=True)
         self.spanner_api.batch_create_sessions.side_effect = mock_batch_create_sessions
+        self._instance = mock.Mock()
+        self._instance._client = mock.Mock()
+        self._instance._client._client_context = None
+        self._instance._client.spanner_api = self.spanner_api
+        self._instance._client._query_options = ExecuteSqlRequest.QueryOptions(
+            optimizer_version="1"
+        )
 
     @property
     def database_role(self):

tests/unit/test_session.py

@@ -194,6 +194,9 @@ def _make_database(
         database.database_role = database_role
         database._route_to_leader_enabled = True
         database.default_transaction_options = default_transaction_options
+        database._instance = mock.Mock()
+        database._instance._client = mock.Mock()
+        database._instance._client._client_context = None
         inject_into_mock_database(database)
 
         return database

tests/unit/test_snapshot.py

@@ -2182,6 +2182,7 @@ def __init__(self):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
         self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
+        self._client_context = None
         self._nth_client_id = _Client.NTH_CLIENT.increment()
         self._nth_request = AtomicCounter()
 

tests/unit/test_spanner.py

@@ -1280,6 +1280,7 @@ def __init__(self):
         self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
         self.directed_read_options = None
         self.default_transaction_options = DefaultTransactionOptions()
+        self._client_context = None
         self._nth_client_id = _Client.NTH_CLIENT.increment()
         self._nth_request = AtomicCounter()