fix: transaction_tag should be set on BeginTransactionRequest

Author: surbhigarg92

google/cloud/spanner_v1/snapshot.py

@@ -901,13 +901,17 @@ def attempt_tracking_method():
 
         return [partition.partition_token for partition in response.partitions]
 
-    def _begin_transaction(self, mutation: Mutation = None) -> bytes:
+    def _begin_transaction(self, mutation: Mutation = None, transaction_tag: str = None) -> bytes:
         """Begins a transaction on the database.
 
         :type mutation: :class:`~google.cloud.spanner_v1.mutation.Mutation`
         :param mutation: (Optional) Mutation to include in the begin transaction
             request. Required for mutation-only transactions with multiplexed sessions.
 
+        :type transaction_tag: str
+        :param transaction_tag: (Optional) Transaction tag to include in the begin transaction
+            request. Required for multiplexed sessions.
+
         :rtype: bytes
         :returns: identifier for the transaction.
 
@@ -931,6 +935,10 @@ def _begin_transaction(self, mutation: Mutation = None) -> bytes:
                 (_metadata_with_leader_aware_routing(database._route_to_leader_enabled))
             )
 
+        request_options = RequestOptions()
+        if transaction_tag is not None:
+            request_options.transaction_tag = transaction_tag
+
         with trace_call(
             name=f"CloudSpanner.{type(self).__name__}.begin",
             session=session,
@@ -944,6 +952,7 @@ def wrapped_method():
                 begin_transaction_request = BeginTransactionRequest(
                     session=session.name,
                     options=self._build_transaction_selector_pb().begin,
+                    request_options=request_options,
                     mutation_key=mutation,
                 )
                 begin_transaction_method = functools.partial(

google/cloud/spanner_v1/transaction.py

@@ -714,7 +714,7 @@ def _begin_transaction(self, mutation: Mutation = None) -> bytes:
         if self.rolled_back:
             raise ValueError("Transaction is already rolled back")
 
-        return super(Transaction, self)._begin_transaction(mutation=mutation)
+        return super(Transaction, self)._begin_transaction(mutation=mutation, transaction_tag=self.transaction_tag)
 
     def _begin_mutations_only_transaction(self) -> None:
         """Begins a mutations-only transaction on the database."""

tests/unit/test_transaction.py

@@ -463,10 +463,14 @@ def _commit_helper(
         if mutations is not None:
             self.assertEqual(transaction._transaction_id, TRANSACTION_ID)
 
+            request_options = RequestOptions()
+            request_options.transaction_tag = TRANSACTION_TAG
+
             expected_begin_transaction_request = BeginTransactionRequest(
                 session=session.name,
                 options=TransactionOptions(read_write=TransactionOptions.ReadWrite()),
                 mutation_key=expected_begin_mutation,
+                request_options=request_options,
             )
 
             expected_begin_metadata = base_metadata.copy()