feat(firestore): add x-goog-is-transaction routing header

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):

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,
         )
 

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,
         )
 

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,
         )
 

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,
         )
 

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,
         )
 

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()

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:

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
 

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,
         )
 

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,
         )
 

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

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,
         )
 

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()

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,
         )
 

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():

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,
         )
 

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,
         )
 

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,
     )
 

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,
     )
 

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