Plumb test_context_manager

Author: odeke-em

google/cloud/spanner_v1/_helpers.py

@@ -566,6 +566,10 @@ 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)

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
@@ -222,13 +223,24 @@ def commit(
             trace_attributes,
             observability_options=observability_options,
         ):
-            method = functools.partial(
-                api.commit,
-                request=request,
-                metadata=metadata,
-            )
+            attempt = AtomicCounter(0)
+            next_nth_request = database._next_nth_request
+
+            def wrapped_method(*args, **kwargs):
+                all_metadata = database.metadata_with_request_id(
+                    next_nth_request,
+                    attempt.increment(),
+                    metadata,
+                )
+                method = functools.partial(
+                    api.commit,
+                    request=request,
+                    metadata=all_metadata,
+                )
+                return method(*args, **kwargs)
+
             response = _retry(
-                method,
+                wrapped_method,
                 allowed_exceptions={InternalServerError: _check_rst_stream_error},
             )
         self.committed = response.commit_timestamp
@@ -341,13 +353,24 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             trace_attributes,
             observability_options=observability_options,
         ):
-            method = functools.partial(
-                api.batch_write,
-                request=request,
-                metadata=metadata,
-            )
+            attempt = AtomicCounter(0)
+            next_nth_request = database._next_nth_request
+
+            def wrapped_method(*args, **kwargs):
+                all_metadata = database.metadata_with_request_id(
+                    next_nth_request,
+                    attempt.increment(),
+                    metadata,
+                )
+                method = functools.partial(
+                    api.batch_write,
+                    request=request,
+                    metadata=all_metadata,
+                )
+                return method(*args, **kwargs)
+
             response = _retry(
-                method,
+                wrapped_method,
                 allowed_exceptions={InternalServerError: _check_rst_stream_error},
             )
         self.committed = True

google/cloud/spanner_v1/client.py

@@ -25,7 +25,6 @@
 """
 import grpc
 import os
-import sys
 import warnings
 
 from google.api_core.gapic_v1 import client_info
@@ -208,7 +207,6 @@ def __init__(
 
     @property
     def _next_nth_request(self):
-        print("next_nth_request called by", sys._getframe().f_back.f_code.co_name)
         return self._nth_request.increment()
 
     @property

google/cloud/spanner_v1/database.py

@@ -804,7 +804,9 @@ def wrapped_method(*args, **kwargs):
 
     @property
     def _next_nth_request(self):
-        return self._instance._client._next_nth_request
+        if self._instance and self._instance._client:
+            return self._instance._client._next_nth_request
+        return 1
 
     @property
     def _nth_client_id(self):

tests/mockserver_tests/mock_server_test_base.py

@@ -20,6 +20,7 @@
     start_mock_server,
     SpannerServicer,
 )
+from google.cloud.spanner_v1.client import Client
 import google.cloud.spanner_v1.types.type as spanner_type
 import google.cloud.spanner_v1.types.result_set as result_set
 from google.api_core.client_options import ClientOptions
@@ -133,6 +134,7 @@ def teardown_class(cls):
         if MockServerTestBase.server is not None:
             MockServerTestBase.server.stop(grace=None)
             MockServerTestBase.server = None
+        Client.NTH_CLIENT.reset()
 
     def setup_method(self, *args, **kwargs):
         self._client = None

tests/unit/test_batch.py

@@ -20,6 +20,10 @@
     enrich_with_otel_scope,
 )
 from google.cloud.spanner_v1 import RequestOptions
+from google.cloud.spanner_v1._helpers import (
+    _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"]
@@ -435,6 +439,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}.1.1.1",
+                ),
             ],
         )
         self.assertEqual(request_options, RequestOptions())
@@ -622,6 +630,32 @@ def session_id(self):
 class _Database(object):
     name = "testing"
     _route_to_leader_enabled = True
+    NTH_CLIENT = 1
+
+    def __init__(self):
+        self._nth_request = 0
+
+    @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,
+            self._channel_id,
+            nth_request,
+            nth_attempt,
+            prior_metadata,
+        )
+
+    @property
+    def _channel_id(self):
+        return 1
 
 
 class _FauxSpannerAPI:

tests/unit/test_database.py

@@ -3251,11 +3251,15 @@ def __init__(self, name, instance=None):
 
     @property
     def _next_nth_request(self):
-        return self._instance._client._next_nth_request
+        if self._instance and self._instance._client:
+            return self._instance._client._next_nth_request
+        return 1
 
     @property
     def _nth_client_id(self):
-        return self._instance._client._nth_client_id
+        if self._instance and self._instance._client:
+            return self._instance._client._nth_client_id
+        return 1
 
     def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
         return _metadata_with_request_id(