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

packages/google-cloud-datastore/google/cloud/datastore/_http.py

@@ -61,6 +61,7 @@ def _request(
     database,
     retry=None,
     timeout=None,
+    metadata=None,
 ):
     """Make a request over the Http transport to the Cloud Datastore API.
 
@@ -104,6 +105,9 @@ def _request(
         "User-Agent": user_agent,
         connection_module.CLIENT_INFO_HEADER: user_agent,
     }
+    if metadata:
+        for k, v in metadata:
+            headers[k] = v
     _update_headers(headers, project, database)
     api_url = build_api_url(project, method, base_url)
 
@@ -143,6 +147,7 @@ def _rpc(
     database,
     retry=None,
     timeout=None,
+    metadata=None,
 ):
     """Make a protobuf RPC request.
 
@@ -185,7 +190,7 @@ def _rpc(
     req_data = request_pb._pb.SerializeToString()
     kwargs = _make_retry_timeout_kwargs(retry, timeout)
     response = _request(
-        http, project, method, req_data, base_url, client_info, database, **kwargs
+        http, project, method, req_data, base_url, client_info, database, metadata=metadata, **kwargs
     )
     return response_pb_cls.deserialize(response)
 
@@ -226,7 +231,7 @@ class HTTPDatastoreAPI(object):
     def __init__(self, client):
         self.client = client
 
-    def lookup(self, request, retry=None, timeout=None):
+    def lookup(self, request, retry=None, timeout=None, metadata=None):
         """Perform a ``lookup`` request.
 
         :type request: :class:`_datastore_pb2.LookupRequest` or dict
@@ -256,10 +261,9 @@ def lookup(self, request, retry=None, timeout=None):
             _datastore_pb2.LookupResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
-    def run_query(self, request, retry=None, timeout=None):
+    def run_query(self, request, retry=None, timeout=None, metadata=None):
         """Perform a ``runQuery`` request.
 
         :type request: :class:`_datastore_pb2.BeginTransactionRequest` or dict
@@ -289,10 +293,9 @@ def run_query(self, request, retry=None, timeout=None):
             _datastore_pb2.RunQueryResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
-    def run_aggregation_query(self, request, retry=None, timeout=None):
+    def run_aggregation_query(self, request, retry=None, timeout=None, metadata=None):
         """Perform a ``runAggregationQuery`` request.
 
         :type request: :class:`_datastore_pb2.BeginTransactionRequest` or dict
@@ -324,10 +327,9 @@ def run_aggregation_query(self, request, retry=None, timeout=None):
             _datastore_pb2.RunAggregationQueryResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
-    def begin_transaction(self, request, retry=None, timeout=None):
+    def begin_transaction(self, request, retry=None, timeout=None, metadata=None):
         """Perform a ``beginTransaction`` request.
 
         :type request: :class:`_datastore_pb2.BeginTransactionRequest` or dict
@@ -357,10 +359,9 @@ def begin_transaction(self, request, retry=None, timeout=None):
             _datastore_pb2.BeginTransactionResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
-    def commit(self, request, retry=None, timeout=None):
+    def commit(self, request, retry=None, timeout=None, metadata=None):
         """Perform a ``commit`` request.
 
         :type request: :class:`_datastore_pb2.CommitRequest` or dict
@@ -390,10 +391,9 @@ def commit(self, request, retry=None, timeout=None):
             _datastore_pb2.CommitResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
-    def rollback(self, request, retry=None, timeout=None):
+    def rollback(self, request, retry=None, timeout=None, metadata=None):
         """Perform a ``rollback`` request.
 
         :type request: :class:`_datastore_pb2.RollbackRequest` or dict
@@ -423,10 +423,9 @@ def rollback(self, request, retry=None, timeout=None):
             _datastore_pb2.RollbackResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
-    def allocate_ids(self, request, retry=None, timeout=None):
+    def allocate_ids(self, request, retry=None, timeout=None, metadata=None):
         """Perform an ``allocateIds`` request.
 
         :type request: :class:`_datastore_pb2.AllocateIdsRequest` or dict
@@ -456,10 +455,9 @@ def allocate_ids(self, request, retry=None, timeout=None):
             _datastore_pb2.AllocateIdsResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
-    def reserve_ids(self, request, retry=None, timeout=None):
+    def reserve_ids(self, request, retry=None, timeout=None, metadata=None):
         """Perform an ``reserveIds`` request.
 
         :type request: :class:`_datastore_pb2.ReserveIdsRequest` or dict
@@ -489,8 +487,7 @@ def reserve_ids(self, request, retry=None, timeout=None):
             _datastore_pb2.ReserveIdsResponse,
             database_id,
             retry=retry,
-            timeout=timeout,
-        )
+            timeout=timeout, metadata=metadata)
 
 
 def _update_headers(headers, project_id, database_id=None):

packages/google-cloud-datastore/google/cloud/datastore/aggregation.py

@@ -488,8 +488,16 @@ def _next_page(self):
                 )  # copy for testability
                 request["aggregation_query"] = new_query_pb
 
+            is_transaction = (
+                transaction_id is not None or new_transaction_options is not None
+            )
+            metadata = (
+                self.client._transaction_rpc_metadata
+                if is_transaction
+                else self.client._rpc_metadata
+            )
             response_pb = self.client._datastore_api.run_aggregation_query(
-                request=request.copy(), **kwargs
+                request=request.copy(), metadata=metadata, **kwargs
             )
             # capture explain metrics if present in response
             # should only be present in last response, and only if explain_options was set

packages/google-cloud-datastore/google/cloud/datastore/batch.py

@@ -300,8 +300,10 @@ def _commit(self, retry, timeout):
         """
         if self._id is None:
             mode = _datastore_pb2.CommitRequest.Mode.NON_TRANSACTIONAL
+            metadata = self._client._rpc_metadata
         else:
             mode = _datastore_pb2.CommitRequest.Mode.TRANSACTIONAL
+            metadata = self._client._transaction_rpc_metadata
 
         kwargs = {}
 
@@ -322,6 +324,7 @@ def _commit(self, retry, timeout):
 
         commit_response_pb = self._client._datastore_api.commit(
             request=request,
+            metadata=metadata,
             **kwargs,
         )
 

packages/google-cloud-datastore/google/cloud/datastore/client.py

@@ -127,6 +127,7 @@ def _extended_lookup(
     timeout=None,
     read_time=None,
     database=None,
+    metadata=None,
 ):
     """Repeat lookup until all keys found (unless stop requested).
 
@@ -206,6 +207,15 @@ def _extended_lookup(
     read_options = helpers.get_read_options(
         eventual, transaction_id, read_time, new_transaction_options
     )
+    if metadata is None:
+        is_transaction = (
+            transaction_id is not None or new_transaction_options is not None
+        )
+        metadata = (
+            (("x-goog-is-transaction", "true"),)
+            if is_transaction
+            else (("x-goog-is-transaction", "false"),)
+        )
     loop_num = 0
     while loop_num < _MAX_LOOPS:  # loop against possible deferred.
         loop_num += 1
@@ -217,6 +227,7 @@ def _extended_lookup(
         helpers.set_database_id_to_request(request, database)
         lookup_response = datastore_api.lookup(
             request=request,
+            metadata=metadata,
             **kwargs,
         )
 
@@ -352,6 +363,9 @@ def __init__(
                     api_endpoint = client_options.api_endpoint
             self._base_url = api_endpoint
 
+        self._rpc_metadata = (("x-goog-is-transaction", "false"),)
+        self._transaction_rpc_metadata = (("x-goog-is-transaction", "true"),)
+
     @staticmethod
     def _determine_default(project):
         """Helper:  override default project detection."""
@@ -574,6 +588,13 @@ def get_multi(
         if transaction is None:
             transaction = self.current_transaction
 
+        transaction_id, new_transaction_options = helpers.get_transaction_options(
+            transaction
+        )
+        is_transaction = (
+            transaction_id is not None or new_transaction_options is not None
+        )
+
         entity_pbs = _extended_lookup(
             datastore_api=self._datastore_api,
             project=self.project,
@@ -586,6 +607,11 @@ def get_multi(
             timeout=timeout,
             read_time=read_time,
             database=self.database,
+            metadata=(
+                self._transaction_rpc_metadata
+                if is_transaction
+                else self._rpc_metadata
+            ),
         )
 
         if missing is not None:
@@ -775,6 +801,7 @@ def allocate_ids(self, incomplete_key, num_ids, retry=None, timeout=None):
         helpers.set_database_id_to_request(request, self.database)
         response_pb = self._datastore_api.allocate_ids(
             request=request,
+            metadata=self._rpc_metadata,
             **kwargs,
         )
         allocated_ids = [
@@ -1020,6 +1047,7 @@ def reserve_ids_sequential(self, complete_key, num_ids, retry=None, timeout=None
         helpers.set_database_id_to_request(request, self.database)
         self._datastore_api.reserve_ids(
             request=request,
+            metadata=self._rpc_metadata,
             **kwargs,
         )
         return None
@@ -1074,6 +1102,7 @@ def reserve_ids_multi(self, complete_keys, retry=None, timeout=None):
 
         self._datastore_api.reserve_ids(
             request=request,
+            metadata=self._rpc_metadata,
             **kwargs,
         )
 

packages/google-cloud-datastore/google/cloud/datastore/query.py

@@ -833,8 +833,16 @@ def _next_page(self):
                 new_query_pb.offset -= response_pb.batch.skipped_results
                 request["query"] = new_query_pb
 
+            is_transaction = (
+                transaction_id is not None or new_transaction_options is not None
+            )
+            metadata = (
+                self.client._transaction_rpc_metadata
+                if is_transaction
+                else self.client._rpc_metadata
+            )
             response_pb = self.client._datastore_api.run_query(
-                request=request.copy(), **kwargs
+                request=request.copy(), metadata=metadata, **kwargs
             )
             # capture explain metrics if present in response
             # should only be present in last response, and only if explain_options was set

packages/google-cloud-datastore/google/cloud/datastore/transaction.py

@@ -243,7 +243,9 @@ def begin(self, retry=None, timeout=None):
 
         try:
             response_pb = self._client._datastore_api.begin_transaction(
-                request=request, **kwargs
+                request=request,
+                metadata=self._client._transaction_rpc_metadata,
+                **kwargs
             )
             self._id = response_pb.transaction
         except:  # noqa: E722 do not use bare except, specify exception instead
@@ -299,7 +301,11 @@ def rollback(self, retry=None, timeout=None):
             }
 
             set_database_id_to_request(request, self._client.database)
-            self._client._datastore_api.rollback(request=request, **kwargs)
+            self._client._datastore_api.rollback(
+                request=request,
+                metadata=self._client._transaction_rpc_metadata,
+                **kwargs
+            )
         finally:
             super(Transaction, self).rollback()
             # Clear our own ID in case this gets accidentally reused.

packages/google-cloud-datastore/tests/unit/test__http.py

@@ -212,7 +212,7 @@ def _rpc_helper(retry=None, timeout=None, database=None):
             request_pb,
             datastore_pb2.BeginTransactionResponse,
             database,
-            **kwargs,
+            **kwargs, metadata=None,
         )
 
     assert result == response_pb._pb
@@ -225,7 +225,7 @@ def _rpc_helper(retry=None, timeout=None, database=None):
         base_url,
         client_info,
         database,
-        **kwargs,
+        **kwargs, metadata=None,
     )
 
 

packages/google-cloud-datastore/tests/unit/test_aggregation.py

@@ -694,6 +694,7 @@ def _next_page_helper(txn_id=None, retry=None, timeout=None, database_id=None):
     set_database_id_to_request(expected_request, database_id)
     expected_call = mock.call(
         request=expected_request,
+        metadata=(("x-goog-is-transaction", "true" if txn_id else "false"),),
         **kwargs,
     )
     assert ds_api.run_aggregation_query.call_args_list == (
@@ -794,6 +795,7 @@ def test_transaction_id_populated(database_id, aggregation_type, aggregation_arg
     read_options = request["read_options"]
     # ensure transaction ID is populated
     assert read_options.transaction == client.current_transaction.id
+    assert mock_gapic.call_args[1]["metadata"] == (("x-goog-is-transaction", "true"),)
 
 
 @pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True)
@@ -846,6 +848,34 @@ def test_transaction_begin_later(database_id, aggregation_type, aggregation_args
     # ensure new_transaction is populated
     assert not read_options.transaction
     assert read_options.new_transaction == transaction._options
+    assert mock_gapic.call_args[1]["metadata"] == (("x-goog-is-transaction", "true"),)
+
+
+@pytest.mark.parametrize("database_id", [None, "somedb"], indirect=True)
+def test_transaction_unstarted(database_id):
+    import mock
+
+    transaction = mock.Mock(id=None, _begin_later=False, spec=["id", "_begin_later"])
+    mock_datastore_api = mock.Mock()
+    mock_gapic = mock_datastore_api.run_aggregation_query
+    mock_gapic.return_value = _make_aggregation_query_response([])
+    client = _Client(
+        None,
+        datastore_api=mock_datastore_api,
+        database=database_id,
+        transaction=transaction,
+    )
+
+    query = _make_query(client)
+    aggregation_query = _make_aggregation_query(client=client, query=query)
+    aggregation_query.count()
+    list(aggregation_query.fetch())
+    assert mock_gapic.call_count == 1
+    request = mock_gapic.call_args[1]["request"]
+    read_options = request["read_options"]
+    assert not read_options.transaction
+    assert not read_options._pb.HasField("new_transaction")
+    assert mock_gapic.call_args[1]["metadata"] == (("x-goog-is-transaction", "false"),)
 
 
 class _Client(object):
@@ -862,6 +892,8 @@ def __init__(
         self.database = database
         self._datastore_api = datastore_api
         self.namespace = namespace
+        self._rpc_metadata = (('x-goog-is-transaction', 'false'),)
+        self._transaction_rpc_metadata = (('x-goog-is-transaction', 'true'),)
         self._transaction = transaction
 
     @property