From bb55feff2b8666411b77746f0d7f8954388001e8 Mon Sep 17 00:00:00 2001 From: Nikhil Chennupati Date: Fri, 16 Oct 2026 23:09:44 -0700 Subject: [PATCH 1/3] feat(firestore): add x-goog-is-transaction routing header Add x-goog-is-transaction routing header in Python SDK to differentiate transactional requests and non-transactional requests from firestore. Having a routing header for transactions helps in routing grpc requests at GFE. --- .../google/cloud/firestore_v1/_helpers.py | 13 +++-- .../google/cloud/firestore_v1/aggregation.py | 6 ++- .../cloud/firestore_v1/async_aggregation.py | 6 ++- .../google/cloud/firestore_v1/async_client.py | 6 ++- .../cloud/firestore_v1/async_document.py | 6 ++- .../google/cloud/firestore_v1/async_query.py | 6 ++- .../cloud/firestore_v1/async_transaction.py | 6 +-- .../cloud/firestore_v1/async_vector_query.py | 6 ++- .../google/cloud/firestore_v1/base_client.py | 23 +++++++++ .../google/cloud/firestore_v1/client.py | 6 ++- .../google/cloud/firestore_v1/document.py | 6 ++- .../cloud/firestore_v1/pipeline_result.py | 20 +++++++- .../google/cloud/firestore_v1/query.py | 6 ++- .../google/cloud/firestore_v1/transaction.py | 6 +-- .../google/cloud/firestore_v1/vector_query.py | 6 ++- .../tests/unit/v1/test__helpers.py | 17 ++++++- .../tests/unit/v1/test_aggregation.py | 10 ++-- .../tests/unit/v1/test_async_aggregation.py | 4 +- .../tests/unit/v1/test_async_client.py | 6 ++- .../tests/unit/v1/test_async_document.py | 6 ++- .../tests/unit/v1/test_async_pipeline.py | 2 + .../tests/unit/v1/test_async_query.py | 4 +- .../tests/unit/v1/test_async_transaction.py | 50 +++++++++---------- .../tests/unit/v1/test_async_vector_query.py | 12 ++--- .../tests/unit/v1/test_base_client.py | 33 ++++++++++++ .../tests/unit/v1/test_client.py | 6 ++- .../tests/unit/v1/test_document.py | 6 ++- .../tests/unit/v1/test_pipeline.py | 2 + .../tests/unit/v1/test_query.py | 10 ++-- .../tests/unit/v1/test_transaction.py | 50 +++++++++---------- .../tests/unit/v1/test_vector_query.py | 12 ++--- 31 files changed, 260 insertions(+), 98 deletions(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py index 53fc98573814..b642ff74d749 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/_helpers.py @@ -1114,16 +1114,23 @@ def get_transaction_id(transaction, read_operation=True) -> Union[bytes, None]: return transaction.id -def metadata_with_prefix(prefix: str, **kw) -> List[Tuple[str, str]]: - """Create RPC metadata containing a prefix. +def metadata_with_prefix( + prefix: str, is_transaction: bool = False, **kw +) -> List[Tuple[str, str]]: + """Create RPC metadata containing a prefix and transaction routing header. Args: prefix (str): appropriate resource path. + is_transaction (bool): Optional. Whether the RPC is part of a + transaction. Defaults to False. Returns: List[Tuple[str, str]]: RPC metadata with supplied prefix """ - return [("google-cloud-resource-prefix", prefix)] + return [ + ("google-cloud-resource-prefix", prefix), + ("x-goog-is-transaction", "true" if is_transaction else "false"), + ] class WriteOption(object): diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/aggregation.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/aggregation.py index 496fd487d971..7739fb663a7c 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/aggregation.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/aggregation.py @@ -120,7 +120,11 @@ def _get_stream_iterator( return self._client._firestore_api.run_aggregation_query( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_aggregation.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_aggregation.py index 8c4a2f46dd15..52b7a78d4918 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_aggregation.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_aggregation.py @@ -159,7 +159,11 @@ async def _make_stream( response_iterator = await self._client._firestore_api.run_aggregation_query( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py index 5ae112956f9b..6d0ed64c847e 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_client.py @@ -296,7 +296,11 @@ async def get_all( response_iterator = await self._firestore_api.batch_get_documents( request=request, - metadata=self._rpc_metadata, + metadata=( + self._transaction_rpc_metadata + if transaction is not None + else self._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_document.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_document.py index 8c2a15e44d09..393d20df57d5 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_document.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_document.py @@ -376,7 +376,11 @@ async def get( response_iter = await self._client._firestore_api.batch_get_documents( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_query.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_query.py index 0495d46cdb8b..b3061b655541 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_query.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_query.py @@ -413,7 +413,11 @@ async def _make_stream( response_iterator = await self._client._firestore_api.run_query( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_transaction.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_transaction.py index 4971c1ea023e..26ca8019d9ce 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_transaction.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_transaction.py @@ -111,7 +111,7 @@ async def _begin(self, retry_id: bytes | None = None) -> None: "database": self._client._database_string, "options": self._options_protobuf(retry_id), }, - metadata=self._client._rpc_metadata, + metadata=self._client._transaction_rpc_metadata, ) self._id = transaction_response.transaction @@ -132,7 +132,7 @@ async def _rollback(self) -> None: "database": self._client._database_string, "transaction": self._id, }, - metadata=self._client._rpc_metadata, + metadata=self._client._transaction_rpc_metadata, ) finally: # clean up, even if rollback fails @@ -159,7 +159,7 @@ async def _commit(self) -> list: "writes": self._write_pbs, "transaction": self._id, }, - metadata=self._client._rpc_metadata, + metadata=self._client._transaction_rpc_metadata, ) self._clean_up() diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_vector_query.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_vector_query.py index 0db958ac9e83..43647c8f6d99 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/async_vector_query.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/async_vector_query.py @@ -154,7 +154,11 @@ async def _make_stream( response_iterator = await self._client._firestore_api.run_query( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) async for response in response_iterator: diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py index 95166266bef2..da9de9daf9e2 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/base_client.py @@ -124,6 +124,7 @@ class BaseClient(ClientWithProject): _firestore_api_internal = None _database_string_internal = None _rpc_metadata_internal = None + _transaction_rpc_metadata_internal = None def __init__( self, @@ -292,6 +293,28 @@ def _rpc_metadata(self): return self._rpc_metadata_internal + @property + def _transaction_rpc_metadata(self): + """The transactional RPC metadata for this client's associated database. + + Returns: + Sequence[Tuple(str, str)]: Transactional RPC metadata with resource + prefix for the database associated with this client. + """ + if self._transaction_rpc_metadata_internal is None: + self._transaction_rpc_metadata_internal = _helpers.metadata_with_prefix( + self._database_string, + is_transaction=True, + ) + + if self._emulator_host is not None: + # The emulator requires additional metadata to be set. + self._transaction_rpc_metadata_internal.append( + ("authorization", "Bearer owner") + ) + + return self._transaction_rpc_metadata_internal + def collection(self, *collection_path) -> BaseCollectionReference: raise NotImplementedError diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py index 9f35fb3f59d5..11e11a919f81 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/client.py @@ -262,7 +262,11 @@ def get_all( response_iterator = self._firestore_api.batch_get_documents( request=request, - metadata=self._rpc_metadata, + metadata=( + self._transaction_rpc_metadata + if transaction is not None + else self._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/document.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/document.py index 30c91974ab3a..3d5266c7aa96 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/document.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/document.py @@ -412,7 +412,11 @@ def get( response_iter = self._client._firestore_api.batch_get_documents( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py index eb25b1f6b24c..e76d78160c8b 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py @@ -278,7 +278,15 @@ def __iter__(self) -> Iterator[T]: raise RuntimeError(f"{self.__class__.__name__} can only be iterated once") self._started = True request = self._build_request() - stream = self._client._firestore_api.execute_pipeline(request) + metadata = ( + self._client._transaction_rpc_metadata + if self.transaction is not None + else self._client._rpc_metadata + ) + stream = self._client._firestore_api.execute_pipeline( + request, + metadata=metadata, + ) for response in stream: yield from self._process_response(response) @@ -293,7 +301,15 @@ async def __aiter__(self) -> AsyncIterator[T]: raise RuntimeError(f"{self.__class__.__name__} can only be iterated once") self._started = True request = self._build_request() - stream = await self._client._firestore_api.execute_pipeline(request) + metadata = ( + self._client._transaction_rpc_metadata + if self.transaction is not None + else self._client._rpc_metadata + ) + stream = await self._client._firestore_api.execute_pipeline( + request, + metadata=metadata, + ) async for response in stream: for result in self._process_response(response): yield result diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/query.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/query.py index 89b8fdd9e2bc..5da6afef9399 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/query.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/query.py @@ -268,7 +268,11 @@ def _get_stream_iterator( response_iterator = self._client._firestore_api.run_query( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/transaction.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/transaction.py index 1d2cde079b48..5420a297f4df 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/transaction.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/transaction.py @@ -96,7 +96,7 @@ def _begin(self, retry_id: bytes | None = None) -> None: "database": self._client._database_string, "options": self._options_protobuf(retry_id), }, - metadata=self._client._rpc_metadata, + metadata=self._client._transaction_rpc_metadata, ) self._id = transaction_response.transaction @@ -117,7 +117,7 @@ def _rollback(self) -> None: "database": self._client._database_string, "transaction": self._id, }, - metadata=self._client._rpc_metadata, + metadata=self._client._transaction_rpc_metadata, ) finally: # clean up, even if rollback fails @@ -144,7 +144,7 @@ def _commit(self) -> list: "writes": self._write_pbs, "transaction": self._id, }, - metadata=self._client._rpc_metadata, + metadata=self._client._transaction_rpc_metadata, ) self._clean_up() diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/vector_query.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/vector_query.py index 9d26ac0ae21e..7376ae9892f0 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/vector_query.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/vector_query.py @@ -113,7 +113,11 @@ def _get_stream_iterator(self, transaction, retry, timeout, explain_options=None response_iterator = self._client._firestore_api.run_query( request=request, - metadata=self._client._rpc_metadata, + metadata=( + self._client._transaction_rpc_metadata + if transaction is not None + else self._client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py b/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py index fb2ac42a0b5d..db110faea3f3 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test__helpers.py @@ -2505,7 +2505,22 @@ def test_metadata_with_prefix(): database_string = "projects/prahj/databases/dee-bee" metadata = metadata_with_prefix(database_string) - assert metadata == [("google-cloud-resource-prefix", database_string)] + assert metadata == [ + ("google-cloud-resource-prefix", database_string), + ("x-goog-is-transaction", "false"), + ] + + +def test_metadata_with_prefix_is_transaction(): + from google.cloud.firestore_v1._helpers import metadata_with_prefix + + database_string = "projects/prahj/databases/dee-bee" + metadata = metadata_with_prefix(database_string, is_transaction=True) + + assert metadata == [ + ("google-cloud-resource-prefix", database_string), + ("x-goog-is-transaction", "true"), + ] def test_writeoption_modify_write(): diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_aggregation.py b/packages/google-cloud-firestore/tests/unit/v1/test_aggregation.py index 35c8b5aea312..212c3edacce3 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_aggregation.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_aggregation.py @@ -691,7 +691,7 @@ def test_aggregation_query_get_transaction(): "structured_aggregation_query": aggregation_query._to_protobuf(), "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -797,7 +797,11 @@ def _stream_w_exception(*_args, **_kw): assert calls[0] == mock.call( request=expected_request, - metadata=client._rpc_metadata, + metadata=( + client._transaction_rpc_metadata + if transaction is not None + else client._rpc_metadata + ), **kwargs, ) @@ -1017,7 +1021,7 @@ def test_aggregation_from_query(): "structured_aggregation_query": aggregation_query._to_protobuf(), "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_aggregation.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_aggregation.py index 9dbd7ba64edd..10f9a974a212 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_aggregation.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_aggregation.py @@ -518,7 +518,7 @@ async def test_async_aggregation_query_get_transaction(): "structured_aggregation_query": aggregation_query._to_protobuf(), "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -579,7 +579,7 @@ async def test_async_aggregation_from_query(): "structured_aggregation_query": aggregation_query._to_protobuf(), "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_client.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_client.py index 3aeef8f9f9c0..a667af14191e 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_client.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_client.py @@ -337,7 +337,11 @@ async def _get_all_helper( client._firestore_api.batch_get_documents.assert_called_once_with( request=expected_request, - metadata=client._rpc_metadata, + metadata=( + client._transaction_rpc_metadata + if txn_id is not None + else client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_document.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_document.py index 38fe9482b654..2e2a5c5c65f3 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_document.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_document.py @@ -491,7 +491,11 @@ def WhichOneof(val): firestore_api.batch_get_documents.assert_called_once_with( request=expected_request, - metadata=client._rpc_metadata, + metadata=( + client._transaction_rpc_metadata + if use_transaction + else client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py index 5f9e7f2749ab..2ccfcc80b0f4 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_pipeline.py @@ -142,6 +142,7 @@ async def test_async_pipeline_stream_empty(): assert isinstance(request, ExecutePipelineRequest) assert request.structured_pipeline == ppl_1._to_pb() assert request.database == "projects/A/databases/B" + assert mock_rpc.call_args[1]["metadata"] == client._rpc_metadata @pytest.mark.asyncio @@ -327,6 +328,7 @@ async def test_async_pipeline_stream_with_transaction(): assert request.structured_pipeline == ppl_1._to_pb() assert request.database == "projects/A/databases/B" assert request.transaction == b"123" + assert mock_rpc.call_args[1]["metadata"] == client._transaction_rpc_metadata @pytest.mark.asyncio diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_query.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_query.py index 7bc98147fee6..ef4d1aed1fb5 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_query.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_query.py @@ -499,7 +499,7 @@ async def test_asyncquery_stream_with_transaction(): "structured_query": query._to_protobuf(), "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -550,7 +550,7 @@ async def test_asyncquery_stream_with_transaction_and_read_time(): "transaction": txn_id, "read_time": read_time, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_transaction.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_transaction.py index a09dd1c5b9ef..36667746a470 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_transaction.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_transaction.py @@ -106,7 +106,7 @@ async def test_asynctransaction__begin(): # Verify the called mock. firestore_api.begin_transaction.assert_called_once_with( request={"database": client._database_string, "options": None}, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -148,7 +148,7 @@ async def test_asynctransaction__rollback(): # Verify the called mock. firestore_api.rollback.assert_called_once_with( request={"database": client._database_string, "transaction": txn_id}, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -193,7 +193,7 @@ async def test_asynctransaction__rollback_failure(): # Verify the called mock. firestore_api.rollback.assert_called_once_with( request={"database": client._database_string, "transaction": txn_id}, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -243,7 +243,7 @@ async def test_asynctransaction__commit(): "writes": write_pbs, "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -294,7 +294,7 @@ async def test_asynctransaction__commit_failure(): "writes": write_pbs, "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -470,7 +470,7 @@ async def _get_w_query_helper( # Verify the mock call. firestore_api.run_query.assert_called_once_with( request=request, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) await returned_generator.aclose() @@ -540,7 +540,7 @@ async def test_asynctransactional__pre_commit_success(): firestore_api = transaction._client._firestore_api firestore_api.begin_transaction.assert_called_once_with( request={"database": transaction._client._database_string, "options": None}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_not_called() firestore_api.commit.assert_not_called() @@ -575,7 +575,7 @@ async def test_asynctransactional__pre_commit_retry_id_already_set_success(): "database": transaction._client._database_string, "options": options_, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_not_called() firestore_api.commit.assert_not_called() @@ -600,7 +600,7 @@ async def test_asynctransactional___call__success_first_attempt(): firestore_api = transaction._client._firestore_api firestore_api.begin_transaction.assert_called_once_with( request={"database": transaction._client._database_string, "options": None}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_not_called() firestore_api.commit.assert_called_once_with( @@ -609,7 +609,7 @@ async def test_asynctransactional___call__success_first_attempt(): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -652,18 +652,18 @@ async def test_asynctransactional___call__success_second_attempt(): expected_calls = [ mock.call( request={"database": db_str, "options": None}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ), mock.call( request={"database": db_str, "options": options_}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ), ] assert firestore_api.begin_transaction.mock_calls == expected_calls firestore_api.rollback.assert_not_called() commit_call = mock.call( request={"database": db_str, "writes": [], "transaction": txn_id}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) assert firestore_api.commit.mock_calls == [commit_call, commit_call] @@ -716,7 +716,7 @@ async def test_asynctransactional___call__failure_max_attempts(max_attempts): "database": transaction._client._database_string, "options": None if i == 0 else options_, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) for i in range(max_attempts) ] @@ -728,14 +728,14 @@ async def test_asynctransactional___call__failure_max_attempts(max_attempts): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -779,14 +779,14 @@ async def test_asynctransactional___call__failure_readonly(max_attempts): read_only=common.TransactionOptions.ReadOnly() ), }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.commit.assert_called_once_with( request={ @@ -794,7 +794,7 @@ async def test_asynctransactional___call__failure_readonly(max_attempts): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -834,14 +834,14 @@ async def test_asynctransactional___call__failure_with_non_retryable(max_attempt "database": transaction._client._database_string, "options": None, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.commit.assert_called_once_with( request={ @@ -849,7 +849,7 @@ async def test_asynctransactional___call__failure_with_non_retryable(max_attempt "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -893,14 +893,14 @@ async def test_asynctransactional___call__failure_with_rollback_failure(): "database": transaction._client._database_string, "options": None, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.commit.assert_called_once_with( request={ @@ -908,7 +908,7 @@ async def test_asynctransactional___call__failure_with_rollback_failure(): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_async_vector_query.py b/packages/google-cloud-firestore/tests/unit/v1/test_async_vector_query.py index fe9df07b0597..5aa96706fce1 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_async_vector_query.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_async_vector_query.py @@ -137,7 +137,7 @@ async def _async_vector_query_get_helper( firestore_api.run_query.assert_called_once_with( request=request, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -285,7 +285,7 @@ async def test_async_vector_query_with_filter(distance_measure, expected_distanc "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -359,7 +359,7 @@ async def test_async_vector_query_with_distance_result_field( "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -433,7 +433,7 @@ async def test_async_vector_query_with_distance_threshold( "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -509,7 +509,7 @@ async def test_vector_query_collection_group(distance_measure, expected_distance "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -625,7 +625,7 @@ async def _async_vector_query_stream_helper( firestore_api.run_query.assert_called_once_with( request=request, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py b/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py index 3762af7815eb..54e3a1c8659d 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_base_client.py @@ -217,6 +217,7 @@ def test_baseclient___rpc_metadata_property(): assert client._rpc_metadata == [ ("google-cloud-resource-prefix", client._database_string), + ("x-goog-is-transaction", "false"), ] @@ -233,6 +234,38 @@ def test_baseclient__rpc_metadata_property_with_emulator(): assert client._rpc_metadata == [ ("google-cloud-resource-prefix", client._database_string), + ("x-goog-is-transaction", "false"), + ("authorization", "Bearer owner"), + ] + + +def test_baseclient___transaction_rpc_metadata_property(): + credentials = _make_credentials() + database = "quanta" + client = _make_base_client( + project=PROJECT, credentials=credentials, database=database + ) + + assert client._transaction_rpc_metadata == [ + ("google-cloud-resource-prefix", client._database_string), + ("x-goog-is-transaction", "true"), + ] + + +def test_baseclient__transaction_rpc_metadata_property_with_emulator(): + emulator_host = "localhost:8081" + with mock.patch("os.getenv") as getenv: + getenv.return_value = emulator_host + + credentials = _make_credentials() + database = "quanta" + client = _make_base_client( + project=PROJECT, credentials=credentials, database=database + ) + + assert client._transaction_rpc_metadata == [ + ("google-cloud-resource-prefix", client._database_string), + ("x-goog-is-transaction", "true"), ("authorization", "Bearer owner"), ] diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_client.py b/packages/google-cloud-firestore/tests/unit/v1/test_client.py index 9d0199f9258b..771cafbcf74e 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_client.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_client.py @@ -430,7 +430,11 @@ def _get_all_helper( client._firestore_api.batch_get_documents.assert_called_once_with( request=expected_request, - metadata=client._rpc_metadata, + metadata=( + client._transaction_rpc_metadata + if txn_id is not None + else client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_document.py b/packages/google-cloud-firestore/tests/unit/v1/test_document.py index 96e21cb6c320..573d22956757 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_document.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_document.py @@ -485,7 +485,11 @@ def WhichOneof(val): firestore_api.batch_get_documents.assert_called_once_with( request=expected_request, - metadata=client._rpc_metadata, + metadata=( + client._transaction_rpc_metadata + if use_transaction + else client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py index fa2f19ae109a..0546fdccd2e9 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_pipeline.py @@ -147,6 +147,7 @@ def test_pipeline_stream_empty(): assert isinstance(request, ExecutePipelineRequest) assert request.structured_pipeline == ppl_1._to_pb() assert request.database == "projects/A/databases/B" + assert mock_rpc.call_args[1]["metadata"] == client._rpc_metadata def test_pipeline_stream_no_doc_ref(): @@ -320,6 +321,7 @@ def test_pipeline_stream_with_transaction(): assert request.structured_pipeline == ppl_1._to_pb() assert request.database == "projects/A/databases/B" assert request.transaction == b"123" + assert mock_rpc.call_args[1]["metadata"] == client._transaction_rpc_metadata def test_pipeline_stream_with_read_time(): diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_query.py b/packages/google-cloud-firestore/tests/unit/v1/test_query.py index 7eaeef61ba78..026fb5e63ede 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_query.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_query.py @@ -495,7 +495,7 @@ def test_query_stream_with_transaction(database): "structured_query": query._to_protobuf(), "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -546,7 +546,7 @@ def test_query_stream_with_transaction_and_read_time(database): "transaction": txn_id, "read_time": read_time, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -848,7 +848,11 @@ def _stream_w_exception(*_args, **_kw): assert calls[0] == mock.call( request=expected_request, - metadata=client._rpc_metadata, + metadata=( + client._transaction_rpc_metadata + if transaction is not None + else client._rpc_metadata + ), **kwargs, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_transaction.py b/packages/google-cloud-firestore/tests/unit/v1/test_transaction.py index 7a93615b3ea4..29dac84e5fd0 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_transaction.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_transaction.py @@ -109,7 +109,7 @@ def test_transaction__begin(database): # Verify the called mock. firestore_api.begin_transaction.assert_called_once_with( request={"database": client._database_string, "options": None}, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -155,7 +155,7 @@ def test_transaction__rollback(database): # Verify the called mock. firestore_api.rollback.assert_called_once_with( request={"database": client._database_string, "transaction": txn_id}, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -205,7 +205,7 @@ def test_transaction__rollback_failure(database): # Verify the called mock. firestore_api.rollback.assert_called_once_with( request={"database": client._database_string, "transaction": txn_id}, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -258,7 +258,7 @@ def test_transaction__commit(database): "writes": write_pbs, "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -312,7 +312,7 @@ def test_transaction__commit_failure(database): "writes": write_pbs, "transaction": txn_id, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, ) @@ -490,7 +490,7 @@ def _transaction_get_w_query_helper( # Verify the mock call. firestore_api.run_query.assert_called_once_with( request=request, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -559,7 +559,7 @@ def test__transactional__pre_commit_success(database): firestore_api = transaction._client._firestore_api firestore_api.begin_transaction.assert_called_once_with( request={"database": transaction._client._database_string, "options": None}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_not_called() firestore_api.commit.assert_not_called() @@ -594,7 +594,7 @@ def test__transactional__pre_commit_retry_id_already_set_success(database): "database": transaction._client._database_string, "options": options_, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_not_called() firestore_api.commit.assert_not_called() @@ -619,7 +619,7 @@ def test__transactional___call__success_first_attempt(database): firestore_api = transaction._client._firestore_api firestore_api.begin_transaction.assert_called_once_with( request={"database": transaction._client._database_string, "options": None}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_not_called() firestore_api.commit.assert_called_once_with( @@ -628,7 +628,7 @@ def test__transactional___call__success_first_attempt(database): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -671,18 +671,18 @@ def test__transactional___call__success_second_attempt(database): expected_calls = [ mock.call( request={"database": db_str, "options": None}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ), mock.call( request={"database": db_str, "options": options_}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ), ] assert firestore_api.begin_transaction.mock_calls == expected_calls firestore_api.rollback.assert_not_called() commit_call = mock.call( request={"database": db_str, "writes": [], "transaction": txn_id}, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) assert firestore_api.commit.mock_calls == [commit_call, commit_call] @@ -737,7 +737,7 @@ def test_transactional___call__failure_max_attempts(database, max_attempts): "database": transaction._client._database_string, "options": None if i == 0 else options_, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) for i in range(max_attempts) ] @@ -749,14 +749,14 @@ def test_transactional___call__failure_max_attempts(database, max_attempts): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -802,14 +802,14 @@ def test_transactional___call__failure_readonly(database, max_attempts): read_only=common.TransactionOptions.ReadOnly() ), }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.commit.assert_called_once_with( request={ @@ -817,7 +817,7 @@ def test_transactional___call__failure_readonly(database, max_attempts): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -859,14 +859,14 @@ def test_transactional___call__failure_with_non_retryable(database, max_attempts "database": transaction._client._database_string, "options": None, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.commit.assert_called_once_with( request={ @@ -874,7 +874,7 @@ def test_transactional___call__failure_with_non_retryable(database, max_attempts "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) @@ -918,14 +918,14 @@ def test_transactional___call__failure_with_rollback_failure(database): "database": transaction._client._database_string, "options": None, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.rollback.assert_called_once_with( request={ "database": transaction._client._database_string, "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) firestore_api.commit.assert_called_once_with( request={ @@ -933,7 +933,7 @@ def test_transactional___call__failure_with_rollback_failure(database): "writes": [], "transaction": txn_id, }, - metadata=transaction._client._rpc_metadata, + metadata=transaction._client._transaction_rpc_metadata, ) diff --git a/packages/google-cloud-firestore/tests/unit/v1/test_vector_query.py b/packages/google-cloud-firestore/tests/unit/v1/test_vector_query.py index ad88478c8310..1bed8b2ce0a3 100644 --- a/packages/google-cloud-firestore/tests/unit/v1/test_vector_query.py +++ b/packages/google-cloud-firestore/tests/unit/v1/test_vector_query.py @@ -216,7 +216,7 @@ def _vector_query_get_helper(distance_measure, expected_distance, explain_option expected_request["explain_options"] = explain_options._to_dict() firestore_api.run_query.assert_called_once_with( request=expected_request, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -311,7 +311,7 @@ def test_vector_query_with_distance_result_field(distance_measure, expected_dist "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -377,7 +377,7 @@ def test_vector_query_with_distance_threshold(distance_measure, expected_distanc "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -453,7 +453,7 @@ def test_vector_query_with_filter(distance_measure, expected_distance): "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -528,7 +528,7 @@ def test_vector_query_collection_group(distance_measure, expected_distance): "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) @@ -590,7 +590,7 @@ def test_vector_query_list_as_query_vector(): "structured_query": expected_pb, "transaction": _TXN_ID, }, - metadata=client._rpc_metadata, + metadata=client._transaction_rpc_metadata, **kwargs, ) From f8f17c7613e6b0fefb11a060bd3f1f6ab7ad129d Mon Sep 17 00:00:00 2001 From: Nikhil Date: Sat, 17 Oct 2026 00:51:22 -0700 Subject: [PATCH 2/3] Update packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../google/cloud/firestore_v1/pipeline_result.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py index e76d78160c8b..236e6b56f028 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py @@ -280,7 +280,7 @@ def __iter__(self) -> Iterator[T]: request = self._build_request() metadata = ( self._client._transaction_rpc_metadata - if self.transaction is not None + if self._transaction is not None else self._client._rpc_metadata ) stream = self._client._firestore_api.execute_pipeline( From d61c3ad59de217e4cf744a8509d04b8d7233421d Mon Sep 17 00:00:00 2001 From: Nikhil Date: Sat, 17 Oct 2026 00:51:43 -0700 Subject: [PATCH 3/3] Update packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../google/cloud/firestore_v1/pipeline_result.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py index 236e6b56f028..a4e7d58dd4f7 100644 --- a/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py +++ b/packages/google-cloud-firestore/google/cloud/firestore_v1/pipeline_result.py @@ -303,7 +303,7 @@ async def __aiter__(self) -> AsyncIterator[T]: request = self._build_request() metadata = ( self._client._transaction_rpc_metadata - if self.transaction is not None + if self._transaction is not None else self._client._rpc_metadata ) stream = await self._client._firestore_api.execute_pipeline(