fix(spanner): ensure request ID retry invariance, add ResourceExhausted handling, and expand mock server test suite

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database.py

@@ -19,7 +19,6 @@
 import functools
 import logging
 import re
-import threading
 from typing import Optional
 
 import google.auth.credentials
@@ -62,6 +61,7 @@
 from google.cloud.spanner_v1.keyset import KeySet
 from google.cloud.spanner_v1.merged_result_set import MergedResultSet
 from google.cloud.spanner_v1.request_id_header import (
+    REQ_CHANNEL_ID,
     REQ_ID_HEADER_KEY,
     X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR,
     _CachedPrefixDescriptor,
@@ -184,9 +184,7 @@ class Database(object):
 
     _spanner_api: SpannerClient = None
 
-    __transport_lock = threading.Lock()
-    __transports_to_channel_id = dict()
-    _channel_id_val = 0
+    _channel_id_val = REQ_CHANNEL_ID
     _req_id_prefix = _CachedPrefixDescriptor()
 
     def __init__(
@@ -235,7 +233,7 @@ def __init__(
             self._directed_read_options = None
             self.default_transaction_options = None
         self._proto_descriptors = proto_descriptors
-        self._channel_id = 0  # It'll be created when _spanner_api is created.
+        self._channel_id = REQ_CHANNEL_ID
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
@@ -561,14 +559,6 @@ def spanner_api(self):
                 client_options=client_options,
             )
 
-            with self.__transport_lock:
-                transport = self._spanner_api.transport
-                channel_id = self.__transports_to_channel_id.get(transport, None)
-                if channel_id is None:
-                    channel_id = len(self.__transports_to_channel_id) + 1
-                    self.__transports_to_channel_id[transport] = channel_id
-                self._channel_id = channel_id
-
         return self._spanner_api
 
     @property

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/snapshot.py

@@ -16,15 +16,22 @@
 
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.snapshot"
 import functools
+import time
 from typing import List, Optional, Union
 
 from google.api_core import gapic_v1
 from google.api_core.exceptions import (
     Aborted,
     InternalServerError,
     InvalidArgument,
+    ResourceExhausted,
     ServiceUnavailable,
 )
+from google.api_core.retry import (
+    RetryFailureReason,
+    build_retry_error,
+    exponential_sleep_generator,
+)
 from google.protobuf.struct_pb2 import Struct
 
 from google.cloud.aio._cross_sync import CrossSync
@@ -89,6 +96,7 @@ async def _restart_on_unavailable(
     observability_options=None,
     request_id_manager=None,
     resource_info=None,
+    retry=None,
 ):
     """Restart iteration after :exc:`.ServiceUnavailable`.
 
@@ -123,6 +131,39 @@ async def _restart_on_unavailable(
     current_request_id = None
     stream_finished = False
 
+    deadline = None
+    sleep_iter = None
+    retry_predicate = None
+    timeout = None
+    # GAPIC has no default retry for streaming reads/queries, so without an
+    # explicit ``Retry`` the stream is resumed on UNAVAILABLE as before.
+    if retry is not None and retry is not gapic_v1.method.DEFAULT:
+        # api_core's Retry exposes only `timeout` publicly; the backoff
+        # settings and predicate have no public accessors.
+        timeout = retry.timeout
+        retry_predicate = retry._predicate
+        if timeout is not None:
+            deadline = time.monotonic() + timeout
+        sleep_iter = exponential_sleep_generator(
+            retry._initial, retry._maximum, multiplier=retry._multiplier
+        )
+
+    async def sleep_or_raise(exc):
+        if sleep_iter is not None:
+            next_sleep = next(sleep_iter)
+            if deadline is not None and time.monotonic() + next_sleep > deadline:
+                final_exc, source_exc = build_retry_error(
+                    [exc], RetryFailureReason.TIMEOUT, timeout
+                )
+                final_exc = _augment_error_with_request_id(
+                    final_exc, current_request_id
+                )
+                source_exc = _augment_error_with_request_id(
+                    source_exc, current_request_id
+                )
+                raise final_exc from source_exc
+            await CrossSync.sleep(next_sleep)
+
     metrics_tracer = None
     factory = SpannerMetricsTracerFactory()
     if factory.enabled:
@@ -203,7 +244,15 @@ async def _restart_on_unavailable(
                         resume_token = item.resume_token
                         break
 
-            except ServiceUnavailable:
+            except (ServiceUnavailable, ResourceExhausted) as exc:
+                if retry_predicate is None:
+                    # No explicit Retry: resume on UNAVAILABLE only, as before.
+                    if isinstance(exc, ResourceExhausted):
+                        raise _augment_error_with_request_id(exc, current_request_id)
+                elif not retry_predicate(exc):
+                    raise _augment_error_with_request_id(exc, current_request_id)
+                await sleep_or_raise(exc)
+
                 del item_buffer[:]
                 request.resume_token = resume_token
                 if transaction is not None:
@@ -220,6 +269,8 @@ async def _restart_on_unavailable(
                 )
                 if not resumable_error:
                     raise _augment_error_with_request_id(exc, current_request_id)
+                await sleep_or_raise(exc)
+
                 del item_buffer[:]
                 request.resume_token = resume_token
                 if transaction is not None:
@@ -502,7 +553,7 @@ async def read(
             api.streaming_read,
             request=read_request,
             metadata=metadata,
-            retry=retry,
+            retry=None,
             timeout=timeout,
         )
 
@@ -517,6 +568,7 @@ async def read(
             },
             column_info=column_info,
             lazy_decode=lazy_decode,
+            retry=retry,
         )
 
     @CrossSync.convert
@@ -698,7 +750,7 @@ async def execute_sql(
             api.execute_streaming_sql,
             request=execute_sql_request,
             metadata=metadata,
-            retry=retry,
+            retry=None,
             timeout=timeout,
         )
 
@@ -709,11 +761,19 @@ async def execute_sql(
             trace_attributes={"db.statement": sql, "request_options": request_options},
             column_info=column_info,
             lazy_decode=lazy_decode,
+            retry=retry,
         )
 
     @CrossSync.convert
     async def _get_streamed_result_set(
-        self, method, request, metadata, trace_attributes, column_info, lazy_decode
+        self,
+        method,
+        request,
+        metadata,
+        trace_attributes,
+        column_info,
+        lazy_decode,
+        retry=None,
     ):
         """Returns the streamed result set for a read or execute SQL request."""
         session = self._session
@@ -740,6 +800,7 @@ async def _get_streamed_result_set(
                 observability_options=getattr(database, "observability_options", None),
                 request_id_manager=database,
                 resource_info=self._resource_info,
+                retry=retry,
             )
 
             if is_execute_sql_request:

packages/google-cloud-spanner/google/cloud/spanner_v1/database.py

@@ -21,7 +21,6 @@
 import functools
 import logging
 import re
-import threading
 from typing import Optional
 
 import google.auth.credentials
@@ -68,6 +67,7 @@
 )
 from google.cloud.spanner_v1.pool import BurstyPool
 from google.cloud.spanner_v1.request_id_header import (
+    REQ_CHANNEL_ID,
     REQ_ID_HEADER_KEY,
     X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR,
     _CachedPrefixDescriptor,
@@ -156,9 +156,7 @@ class Database(object):
                               statements in 'ddl_statements' above."""
 
     _spanner_api: SpannerClient = None
-    __transport_lock = threading.Lock()
-    __transports_to_channel_id = dict()
-    _channel_id_val = 0
+    _channel_id_val = REQ_CHANNEL_ID
     _req_id_prefix = _CachedPrefixDescriptor()
 
     def __init__(
@@ -207,7 +205,7 @@ def __init__(
             self._directed_read_options = None
             self.default_transaction_options = None
         self._proto_descriptors = proto_descriptors
-        self._channel_id = 0
+        self._channel_id = REQ_CHANNEL_ID
         if pool is None:
             pool = BurstyPool(database_role=database_role)
         self._pool = pool
@@ -486,13 +484,6 @@ def spanner_api(self):
                 client_info=client_info,
                 client_options=client_options,
             )
-            with self.__transport_lock:
-                transport = self._spanner_api.transport
-                channel_id = self.__transports_to_channel_id.get(transport, None)
-                if channel_id is None:
-                    channel_id = len(self.__transports_to_channel_id) + 1
-                    self.__transports_to_channel_id[transport] = channel_id
-                self._channel_id = channel_id
         return self._spanner_api
 
     @property

packages/google-cloud-spanner/google/cloud/spanner_v1/request_id_header.py

@@ -16,13 +16,30 @@
 
 REQ_ID_VERSION = 1  # The version of the x-goog-spanner-request-id spec.
 REQ_ID_HEADER_KEY = "x-goog-spanner-request-id"
+# Python clients always report channel 1.
+REQ_CHANNEL_ID = 1
 
 
 def generate_rand_uint64():
     return int.from_bytes(os.urandom(8), "big")
 
 
-REQ_RAND_PROCESS_ID = generate_rand_uint64()
+def _get_process_id():
+    """Return the process ID used in every request ID of this process.
+
+    Users can override it with the ``SPANNER_PROCESS_ID`` or
+    ``GOOGLE_CLOUD_SPANNER_PROCESS_ID`` environment variable. Otherwise it is a
+    64-bit random value formatted as 16 lower-case hexadecimal characters,
+    matching the Java and Go clients.
+    """
+    return (
+        os.environ.get("SPANNER_PROCESS_ID")
+        or os.environ.get("GOOGLE_CLOUD_SPANNER_PROCESS_ID")
+        or f"{generate_rand_uint64():016x}"
+    )
+
+
+REQ_RAND_PROCESS_ID = _get_process_id()
 _REQ_ID_PROCESS_PREFIX = f"{REQ_ID_VERSION}.{REQ_RAND_PROCESS_ID}."
 X_GOOG_SPANNER_REQUEST_ID_SPAN_ATTR = "x_goog_spanner_request_id"
 
@@ -78,14 +95,23 @@ def build_request_id(client_id, channel_id, nth_request, attempt):
 
 
 def parse_request_id(request_id_str):
-    version, rand_process_id, client_id, channel_id, nth_request, nth_attempt = map(
-        int, request_id_str.split(".")
-    )
+    """Parse ``<version>.<process>.<client>.<channel>.<request>.<attempt>``.
+
+    Returns a 6-tuple. The process ID is returned as its hexadecimal string;
+    the other five components are returned as ints.
+    """
+    splits = request_id_str.split(".")
+    if len(splits) != 6:
+        raise ValueError(
+            f"Request ID must have 6 dot-separated components, got {len(splits)}: "
+            f"{request_id_str!r}"
+        )
+    version, rand_process_id, client_id, channel_id, nth_request, nth_attempt = splits
     return (
-        version,
+        int(version),
         rand_process_id,
-        client_id,
-        channel_id,
-        nth_request,
-        nth_attempt,
+        int(client_id),
+        int(channel_id),
+        int(nth_request),
+        int(nth_attempt),
     )