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/database.py

@@ -1100,8 +1100,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.
+                   "client_context" (Optional) Client context to use for all requests
+                   made by this transaction.
 
         :rtype: Any
         :returns: The return value of ``func``.
@@ -1526,7 +1526,9 @@ def __enter__(self):
         transaction_type = TransactionType.READ_WRITE
         self._session = self._database.sessions_manager.get_session(transaction_type)
 
-        return MutationGroups(session=self._session, client_context=self._client_context)
+        return MutationGroups(
+            session=self._session, client_context=self._client_context
+        )
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """End ``with`` block."""

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)