feat(x-goog-spanner-request-id): implement Request-ID

google/cloud/spanner_v1/_helpers.py

@@ -32,10 +32,11 @@
 from google.cloud.spanner_v1 import TypeCode
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import JsonObject
-from google.cloud.spanner_v1.request_id_header import with_request_id
+from google.cloud.spanner_v1.request_id_header import REQ_ID_HEADER_KEY, with_request_id
 from google.rpc.error_details_pb2 import RetryInfo
 
 import random
+from typing import Callable
 
 # Validation error messages
 NUMERIC_MAX_SCALE_ERR_MSG = (
@@ -641,6 +642,40 @@ def __radd__(self, n):
         """
         return self.__add__(n)
 
+    def reset(self):
+        with self.__lock:
+            self.__value = 0
+
 
 def _metadata_with_request_id(*args, **kwargs):
     return with_request_id(*args, **kwargs)
+
+
+class InterceptingHeaderInjector:
+    def __init__(self, original_callable: Callable):
+        self._original_callable = original_callable
+
+    def __call__(self, *args, **kwargs):
+        metadata = kwargs.get("metadata", [])
+        # Find all the headers that match the x-goog-spanner-request-id
+        # header an on each retry increment the value.
+        all_metadata = []
+        for key, value in metadata:
+            if key is REQ_ID_HEADER_KEY:
+                # Otherwise now increment the count for the attempt number.
+                splits = value.split(".")
+                attempt_plus_one = int(splits[-1]) + 1
+                splits[-1] = str(attempt_plus_one)
+                value_before = value
+                value = ".".join(splits)
+                print("incrementing value on retry from", value_before, "to", value)
+
+            all_metadata.append(
+                (
+                    key,
+                    value,
+                )
+            )
+
+        kwargs["metadata"] = all_metadata
+        return self._original_callable(*args, **kwargs)

google/cloud/spanner_v1/batch.py

@@ -25,6 +25,7 @@
 from google.cloud.spanner_v1._helpers import (
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
+    AtomicCounter,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1 import RequestOptions
@@ -227,11 +228,20 @@ def commit(
             trace_attributes,
             observability_options=observability_options,
         ):
+            attempt = AtomicCounter(0)
+            next_nth_request = database._next_nth_request
+
+            all_metadata = database.metadata_with_request_id(
+                next_nth_request,
+                attempt.increment(),
+                metadata,
+            )
             method = functools.partial(
                 api.commit,
                 request=request,
-                metadata=metadata,
+                metadata=all_metadata,
             )
+
             deadline = time.time() + kwargs.get(
                 "timeout_secs", DEFAULT_RETRY_TIMEOUT_SECS
             )
@@ -349,11 +359,18 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             trace_attributes,
             observability_options=observability_options,
         ):
+            next_nth_request = database._next_nth_request
+            all_metadata = database.metadata_with_request_id(
+                next_nth_request,
+                0,
+                metadata,
+            )
             method = functools.partial(
                 api.batch_write,
                 request=request,
-                metadata=metadata,
+                metadata=all_metadata,
             )
+
             response = _retry(
                 method,
                 allowed_exceptions={

google/cloud/spanner_v1/client.py

@@ -48,6 +48,7 @@
 from google.cloud.spanner_v1._helpers import _merge_query_options
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
 from google.cloud.spanner_v1.instance import Instance
+from google.cloud.spanner_v1._helpers import AtomicCounter
 
 _CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__)
 EMULATOR_ENV_VAR = "SPANNER_EMULATOR_HOST"
@@ -147,6 +148,8 @@ class Client(ClientWithProject):
     SCOPE = (SPANNER_ADMIN_SCOPE,)
     """The scopes required for Google Cloud Spanner."""
 
+    NTH_CLIENT = AtomicCounter()
+
     def __init__(
         self,
         project=None,
@@ -199,6 +202,12 @@ def __init__(
         self._route_to_leader_enabled = route_to_leader_enabled
         self._directed_read_options = directed_read_options
         self._observability_options = observability_options
+        self._nth_client_id = Client.NTH_CLIENT.increment()
+        self._nth_request = AtomicCounter(0)
+
+    @property
+    def _next_nth_request(self):
+        return self._nth_request.increment()
 
     @property
     def credentials(self):