From eb03c32ffd0411390078f8ab1bf8bd07b187105c Mon Sep 17 00:00:00 2001 From: Surbhi Garg Date: Thu, 11 Dec 2025 11:24:37 +0530 Subject: [PATCH 1/2] fix: transaction_tag should be set on BeginTransactionRequest --- google/cloud/spanner_v1/snapshot.py | 23 +++++++++++++++++++---- google/cloud/spanner_v1/transaction.py | 4 +++- tests/unit/test_session.py | 5 ++++- tests/unit/test_transaction.py | 4 ++++ 4 files changed, 30 insertions(+), 6 deletions(-) diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py index 46b0f5af8d..89cbc9fe88 100644 --- a/google/cloud/spanner_v1/snapshot.py +++ b/google/cloud/spanner_v1/snapshot.py @@ -901,13 +901,19 @@ 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. + :rtype: bytes :returns: identifier for the transaction. @@ -931,6 +937,17 @@ def _begin_transaction(self, mutation: Mutation = None) -> bytes: (_metadata_with_leader_aware_routing(database._route_to_leader_enabled)) ) + begin_request_kwargs = { + "session": session.name, + "options": self._build_transaction_selector_pb().begin, + "mutation_key": mutation, + } + + if transaction_tag: + begin_request_kwargs["request_options"] = RequestOptions( + transaction_tag=transaction_tag + ) + with trace_call( name=f"CloudSpanner.{type(self).__name__}.begin", session=session, @@ -942,9 +959,7 @@ def _begin_transaction(self, mutation: Mutation = None) -> bytes: def wrapped_method(): begin_transaction_request = BeginTransactionRequest( - session=session.name, - options=self._build_transaction_selector_pb().begin, - mutation_key=mutation, + **begin_request_kwargs ) begin_transaction_method = functools.partial( api.begin_transaction, diff --git a/google/cloud/spanner_v1/transaction.py b/google/cloud/spanner_v1/transaction.py index b9e14a0040..de8b421840 100644 --- a/google/cloud/spanner_v1/transaction.py +++ b/google/cloud/spanner_v1/transaction.py @@ -714,7 +714,9 @@ 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.""" diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index bfbd6edd5e..8026c50c24 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -2005,9 +2005,12 @@ def unit_of_work(txn, *args, **kw): self.assertEqual(kw, {"some_arg": "def"}) expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite()) + expected_request_options = RequestOptions(transaction_tag=transaction_tag) gax_api.begin_transaction.assert_called_once_with( request=BeginTransactionRequest( - session=self.SESSION_NAME, options=expected_options + session=self.SESSION_NAME, + options=expected_options, + request_options=expected_request_options, ), metadata=[ ("google-cloud-resource-prefix", database.name), diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py index 39656cb8d1..8f432876db 100644 --- a/tests/unit/test_transaction.py +++ b/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() From dd5ad6da07f2f5617d33fb3e4f5b7c88eafb4833 Mon Sep 17 00:00:00 2001 From: Surbhi Garg Date: Mon, 15 Dec 2025 10:06:06 +0530 Subject: [PATCH 2/2] review comments --- tests/unit/test_transaction.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/unit/test_transaction.py b/tests/unit/test_transaction.py index 8f432876db..510251656e 100644 --- a/tests/unit/test_transaction.py +++ b/tests/unit/test_transaction.py @@ -463,14 +463,11 @@ 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, + request_options=RequestOptions(transaction_tag=TRANSACTION_TAG), ) expected_begin_metadata = base_metadata.copy()