chore(x-goog-spanner-request-id): more updates for batch_write + mockserver tests

This change plumbs in some x-goog-spanner-request-id updates
for batch_write and some tests too.

Updates #1261

Author: odeke-em

google/cloud/spanner_v1/_helpers.py

@@ -19,6 +19,7 @@
 import math
 import time
 import base64
+import inspect
 import threading
 
 from google.protobuf.struct_pb2 import ListValue
@@ -33,7 +34,7 @@
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import JsonObject, Interval
 from google.cloud.spanner_v1 import TransactionOptions
-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
 
 try:
@@ -587,6 +588,20 @@ def _check_rst_stream_error(exc):
         raise
 
 
+def _check_unavailable(exc):
+    resumable_error = (
+        any(
+            resumable_message in exc.message
+            for resumable_message in (
+                "INTERNAL",
+                "Service unavailable",
+            )
+        ),
+    )
+    if not resumable_error:
+        raise
+
+
 def _metadata_with_leader_aware_routing(value, **kw):
     """Create RPC metadata containing a leader aware routing header
 

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
@@ -385,13 +386,22 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             observability_options=observability_options,
             metadata=metadata,
         ), MetricsCapture():
-            method = functools.partial(
-                api.batch_write,
-                request=request,
-                metadata=metadata,
-            )
+            attempt = AtomicCounter(0)
+            nth_request = database._next_nth_request
+
+            def wrapped_method(*args, **kwargs):
+                return functools.partial(
+                    api.batch_write,
+                    request=request,
+                    metadata=database.metadata_with_request_id(
+                        nth_request,
+                        attempt.increment(),
+                        metadata,
+                    ),
+                )(*args, **kwargs)
+
             response = _retry(
-                method,
+                wrapped_method,
                 allowed_exceptions={
                     InternalServerError: _check_rst_stream_error,
                 },

google/cloud/spanner_v1/pool.py

@@ -15,15 +15,22 @@
 """Pools managing shared Session objects."""
 
 import datetime
+import functools
 import queue
 import time
 
+from google.api_core.exceptions import InternalServerError
+from google.api_core.exceptions import ServiceUnavailable
 from google.cloud.exceptions import NotFound
 from google.cloud.spanner_v1 import BatchCreateSessionsRequest
 from google.cloud.spanner_v1 import Session
 from google.cloud.spanner_v1._helpers import (
+    _check_rst_stream_error,
+    _check_unavailable,
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
+    _retry,
+    AtomicCounter,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import (
     add_span_event,
@@ -254,11 +261,25 @@ def bind(self, database):
                     f"Creating {request.session_count} sessions",
                     span_event_attributes,
                 )
-                resp = api.batch_create_sessions(
-                    request=request,
-                    metadata=database.metadata_with_request_id(
-                        database._next_nth_request, 1, metadata
-                    ),
+                attempt = AtomicCounter(0)
+                nth_request = database._next_nth_request
+
+                def wrapped_method(*args, **kwargs):
+                    method = functools.partial(
+                        api.batch_create_sessions,
+                        request=request,
+                        metadata=database.metadata_with_request_id(
+                            nth_request, attempt.increment(), metadata
+                        ),
+                    )
+                    return method(*args, **kwargs)
+
+                resp = _retry(
+                    wrapped_method,
+                    allowed_exceptions={
+                        InternalServerError: _check_rst_stream_error,
+                        ServiceUnavailable: _check_unavailable,
+                    },
                 )
 
                 add_span_event(
@@ -561,11 +582,23 @@ def bind(self, database):
         ) as span, MetricsCapture():
             returned_session_count = 0
             while returned_session_count < self.size:
-                resp = api.batch_create_sessions(
-                    request=request,
-                    metadata=database.metadata_with_request_id(
-                        database._next_nth_request, 1, metadata
-                    ),
+                attempt = AtomicCounter(0)
+                nth_request = database._next_nth_request
+
+                def wrapped_method(*args, **kwargs):
+                    return api.batch_create_sessions(
+                        request=request,
+                        metadata=database.metadata_with_request_id(
+                            database._next_nth_request, attempt.increment(), metadata
+                        ),
+                    )
+
+                resp = _retry(
+                    wrapped_method,
+                    allowed_exceptions={
+                        InternalServerError: _check_rst_stream_error,
+                        ServiceUnavailable: _check_unavailable,
+                    },
                 )
 
                 add_span_event(

google/cloud/spanner_v1/testing/interceptors.py

@@ -71,9 +71,6 @@ def reset(self):
 
 
 class XGoogRequestIDHeaderInterceptor(ClientInterceptor):
-    # TODO:(@odeke-em): delete this guard when PR #1367 is merged.
-    X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED = True
-
     def __init__(self):
         self._unary_req_segments = []
         self._stream_req_segments = []
@@ -87,24 +84,23 @@ def intercept(self, method, request_or_iterator, call_details):
                 x_goog_request_id = value
                 break
 
-        if self.X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED and not x_goog_request_id:
+        if not x_goog_request_id:
             raise Exception(
                 f"Missing {X_GOOG_REQUEST_ID} header in {call_details.method}"
             )
 
         response_or_iterator = method(request_or_iterator, call_details)
         streaming = getattr(response_or_iterator, "__iter__", None) is not None
 
-        if self.X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED:
-            with self.__lock:
-                if streaming:
-                    self._stream_req_segments.append(
-                        (call_details.method, parse_request_id(x_goog_request_id))
-                    )
-                else:
-                    self._unary_req_segments.append(
-                        (call_details.method, parse_request_id(x_goog_request_id))
-                    )
+        with self.__lock:
+            if streaming:
+                self._stream_req_segments.append(
+                    (call_details.method, parse_request_id(x_goog_request_id))
+                )
+            else:
+                self._unary_req_segments.append(
+                    (call_details.method, parse_request_id(x_goog_request_id))
+                )
 
         return response_or_iterator