Include batch*

Author: odeke-em

google/cloud/spanner_v1/batch.py

@@ -26,6 +26,7 @@
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
     _merge_Transaction_Options,
+    AtomicCounter,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1 import RequestOptions
@@ -249,17 +250,25 @@ def commit(
             observability_options=observability_options,
             metadata=metadata,
         ), MetricsCapture():
-            method = functools.partial(
-                api.commit,
-                request=request,
-                metadata=metadata,
-            )
+            attempt = AtomicCounter()
+            nth_request = getattr(database, "_next_nth_request", 0)
+
+            def wrapped_method(*args, **kwargs):
+                method = functools.partial(
+                    api.commit,
+                    request=request,
+                    metadata=database.metadata_with_request_id(
+                        nth_request, attempt.increment(), metadata
+                    ),
+                )
+                return method(*args, **kwargs)
+
             deadline = time.time() + kwargs.get(
                 "timeout_secs", DEFAULT_RETRY_TIMEOUT_SECS
             )
             default_retry_delay = kwargs.get("default_retry_delay", None)
             response = _retry_on_aborted_exception(
-                method,
+                wrapped_method,
                 deadline=deadline,
                 default_retry_delay=default_retry_delay,
             )

tests/unit/test_batch.py

@@ -38,8 +38,10 @@
 from google.rpc.status_pb2 import Status
 
 from google.cloud.spanner_v1._helpers import (
+    AtomicCounter,
     _metadata_with_request_id,
 )
+from google.cloud.spanner_v1.request_id_header import REQ_RAND_PROCESS_ID
 
 TABLE_NAME = "citizens"
 COLUMNS = ["email", "first_name", "last_name", "age"]
@@ -252,6 +254,10 @@ def test_commit_ok(self):
             [
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1",
+                ),
             ],
         )
         self.assertEqual(request_options, RequestOptions())
@@ -346,6 +352,10 @@ def _test_commit_with_options(
             [
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1",
+                ),
             ],
         )
         self.assertEqual(actual_request_options, expected_request_options)
@@ -456,6 +466,10 @@ def test_context_mgr_success(self):
             [
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1",
+                ),
             ],
         )
         self.assertEqual(request_options, RequestOptions())
@@ -582,12 +596,6 @@ def _test_batch_write_with_request_options(
                 "traceparent is missing in metadata",
             )
 
-        # expected_metadata.append(
-        #     (
-        #         "x-goog-spanner-request-id",
-        #         f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.1.1.1",
-        #     )
-        # )
         # Remove traceparent from actual metadata for comparison
         filtered_metadata = [item for item in metadata if item[0] != "traceparent"]
 
@@ -647,7 +655,7 @@ def session_id(self):
 class _Database(object):
     name = "testing"
     _route_to_leader_enabled = True
-    NTH_CLIENT = 1
+    NTH_CLIENT_ID = AtomicCounter()
 
     def __init__(self, enable_end_to_end_tracing=False):
         self.name = "testing"
@@ -656,16 +664,13 @@ def __init__(self, enable_end_to_end_tracing=False):
             self.observability_options = dict(enable_end_to_end_tracing=True)
         self.default_transaction_options = DefaultTransactionOptions()
         self._nth_request = 0
+        self._nth_client_id = _Database.NTH_CLIENT_ID.increment()
 
     @property
     def _next_nth_request(self):
         self._nth_request += 1
         return self._nth_request
 
-    @property
-    def _nth_client_id(self):
-        return 1
-
     def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
         return _metadata_with_request_id(
             self._nth_client_id,

tests/unit/test_database.py

@@ -2063,6 +2063,10 @@ def test_context_mgr_success(self):
             metadata=[
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1",
+                ),
             ],
         )
 
@@ -2110,6 +2114,10 @@ def test_context_mgr_w_commit_stats_success(self):
             metadata=[
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1",
+                ),
             ],
         )
 
@@ -2155,6 +2163,10 @@ def test_context_mgr_w_aborted_commit_status(self):
             metadata=[
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{database._nth_client_id}.{database._channel_id}.1.1",
+                ),
             ],
         )