chore(x-goog-request-id): commit testing scaffold

This change commits the scaffolding for which testing
will be used. This is a carve out of PRs #1264 and #1364,
meant to make those changes lighter and much easier to
review then merge.

Updates #1261

Author: odeke-em

google/cloud/spanner_v1/request_id_header.py

@@ -37,6 +37,6 @@ def generate_rand_uint64():
 
 def with_request_id(client_id, channel_id, nth_request, attempt, other_metadata=[]):
     req_id = f"{REQ_ID_VERSION}.{REQ_RAND_PROCESS_ID}.{client_id}.{channel_id}.{nth_request}.{attempt}"
-    all_metadata = other_metadata.copy()
+    all_metadata = (other_metadata or []).copy()
     all_metadata.append((REQ_ID_HEADER_KEY, req_id))
     return all_metadata

google/cloud/spanner_v1/testing/database_test.py

@@ -25,6 +25,7 @@
 from google.cloud.spanner_v1.testing.interceptors import (
     MethodCountInterceptor,
     MethodAbortInterceptor,
+    XGoogRequestIDHeaderInterceptor,
 )
 
 
@@ -34,6 +35,8 @@ class TestDatabase(Database):
     currently, and we don't want to make changes in the Database class for
     testing purpose as this is a hack to use interceptors in tests."""
 
+    _interceptors = []
+
     def __init__(
         self,
         database_id,
@@ -74,6 +77,8 @@ def spanner_api(self):
             client_options = client._client_options
             if self._instance.emulator_host is not None:
                 channel = grpc.insecure_channel(self._instance.emulator_host)
+                self._x_goog_request_id_interceptor = XGoogRequestIDHeaderInterceptor()
+                self._interceptors.append(self._x_goog_request_id_interceptor)
                 channel = grpc.intercept_channel(channel, *self._interceptors)
                 transport = SpannerGrpcTransport(channel=channel)
                 self._spanner_api = SpannerClient(
@@ -110,3 +115,7 @@ def _create_spanner_client_for_tests(self, client_options, credentials):
             client_options=client_options,
             transport=transport,
         )
+
+    def reset(self):
+        if self._x_goog_request_id_interceptor:
+            self._x_goog_request_id_interceptor.reset()

google/cloud/spanner_v1/testing/interceptors.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 from collections import defaultdict
+import threading
+
 from grpc_interceptor import ClientInterceptor
 from google.api_core.exceptions import Aborted
 
@@ -63,3 +65,69 @@ def reset(self):
         self._method_to_abort = None
         self._count = 0
         self._connection = None
+
+
+X_GOOG_REQUEST_ID = "x-goog-spanner-request-id"
+# TODO:(@odeke-em): delete this guard when PR #1367 is merged.
+__X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED = False
+
+
+class XGoogRequestIDHeaderInterceptor(ClientInterceptor):
+    def __init__(self):
+        self._unary_req_segments = []
+        self._stream_req_segments = []
+        self.__lock = threading.Lock()
+
+    def intercept(self, method, request_or_iterator, call_details):
+        metadata = call_details.metadata
+        x_goog_request_id = None
+        for key, value in metadata:
+            if key == X_GOOG_REQUEST_ID:
+                x_goog_request_id = value
+                break
+
+        if __X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED and 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
+        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
+
+    @property
+    def unary_request_ids(self):
+        return self._unary_req_segments
+
+    @property
+    def stream_request_ids(self):
+        return self._stream_req_segments
+
+    def reset(self):
+        self._stream_req_segments.clear()
+        self._unary_req_segments.clear()
+
+
+def parse_request_id(request_id_str):
+    splits = request_id_str.split(".")
+    version, rand_process_id, client_id, channel_id, nth_request, nth_attempt = list(
+        map(lambda v: int(v), splits)
+    )
+    return (
+        version,
+        rand_process_id,
+        client_id,
+        channel_id,
+        nth_request,
+        nth_attempt,
+    )

google/cloud/spanner_v1/testing/mock_spanner.py

@@ -22,8 +22,6 @@
 from google.cloud.spanner_v1 import (
     TransactionOptions,
     ResultSetMetadata,
-    ExecuteSqlRequest,
-    ExecuteBatchDmlRequest,
 )
 from google.cloud.spanner_v1.testing.mock_database_admin import DatabaseAdminServicer
 import google.cloud.spanner_v1.testing.spanner_database_admin_pb2_grpc as database_admin_grpc
@@ -107,6 +105,7 @@ def CreateSession(self, request, context):
 
     def BatchCreateSessions(self, request, context):
         self._requests.append(request)
+        self.mock_spanner.pop_error(context)
         sessions = []
         for i in range(request.session_count):
             sessions.append(
@@ -186,9 +185,7 @@ def BeginTransaction(self, request, context):
         self._requests.append(request)
         return self.__create_transaction(request.session, request.options)
 
-    def __maybe_create_transaction(
-        self, request: ExecuteSqlRequest | ExecuteBatchDmlRequest
-    ):
+    def __maybe_create_transaction(self, request):
         started_transaction = None
         if not request.transaction.begin == TransactionOptions():
             started_transaction = self.__create_transaction(

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
@@ -78,6 +79,27 @@ def unavailable_status() -> _Status:
     return status
 
 
+# Creates an UNAVAILABLE status with the smallest possible retry delay.
+def unavailable_status() -> _Status:
+    error = status_pb2.Status(
+        code=code_pb2.UNAVAILABLE,
+        message="Service unavailable.",
+    )
+    retry_info = RetryInfo(retry_delay=Duration(seconds=0, nanos=1))
+    status = _Status(
+        code=code_to_grpc_status_code(error.code),
+        details=error.message,
+        trailing_metadata=(
+            ("grpc-status-details-bin", error.SerializeToString()),
+            (
+                "google.rpc.retryinfo-bin",
+                retry_info.SerializeToString(),
+            ),
+        ),
+    )
+    return status
+
+
 def add_error(method: str, error: status_pb2.Status):
     MockServerTestBase.spanner_service.mock_spanner.add_error(method, error)
 
@@ -153,6 +175,7 @@ def setup_class(cls):
     def teardown_class(cls):
         if MockServerTestBase.server is not None:
             MockServerTestBase.server.stop(grace=None)
+            Client.NTH_CLIENT.reset()
             MockServerTestBase.server = None
 
     def setup_method(self, *args, **kwargs):
@@ -186,6 +209,8 @@ def instance(self) -> Instance:
     def database(self) -> Database:
         if self._database is None:
             self._database = self.instance.database(
-                "test-database", pool=FixedSizePool(size=10)
+                "test-database",
+                pool=FixedSizePool(size=10),
+                enable_interceptors_in_tests=True,
             )
         return self._database

tests/unit/test_transaction.py

@@ -21,6 +21,11 @@
 from google.cloud.spanner_v1 import TypeCode
 from google.api_core.retry import Retry
 from google.api_core import gapic_v1
+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
 
 from tests._helpers import (
     HAS_OPENTELEMETRY_INSTALLED,
@@ -197,6 +202,10 @@ def test_begin_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}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
             ],
         )
 
@@ -301,6 +310,10 @@ def test_rollback_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}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
             ],
         )
 
@@ -492,6 +505,10 @@ def _commit_helper(
             [
                 ("google-cloud-resource-prefix", database.name),
                 ("x-goog-spanner-route-to-leader", "true"),
+                (
+                    "x-goog-spanner-request-id",
+                    f"1.{REQ_RAND_PROCESS_ID}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
             ],
         )
         self.assertEqual(actual_request_options, expected_request_options)
@@ -666,6 +683,10 @@ def _execute_update_helper(
             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}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
             ],
         )
 
@@ -859,6 +880,10 @@ def _batch_update_helper(
             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}.{_Client.NTH_CLIENT.value}.1.1.1",
+                ),
             ],
             retry=retry,
             timeout=timeout,
@@ -974,6 +999,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}.{_Client.NTH_CLIENT.value}.1.2.1",
+                ),
             ],
         )
 
@@ -1004,11 +1033,19 @@ def test_context_mgr_failure(self):
 
 
 class _Client(object):
+    NTH_CLIENT = AtomicCounter()
+
     def __init__(self):
         from google.cloud.spanner_v1 import ExecuteSqlRequest
 
         self._query_options = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
         self.directed_read_options = None
+        self._nth_client_id = _Client.NTH_CLIENT.increment()
+        self._nth_request = AtomicCounter()
+
+    @property
+    def _next_nth_request(self):
+        return self._nth_request.increment()
 
 
 class _Instance(object):
@@ -1024,6 +1061,27 @@ def __init__(self):
         self._directed_read_options = None
         self.default_transaction_options = DefaultTransactionOptions()
 
+    @property
+    def _next_nth_request(self):
+        return self._instance._client._next_nth_request
+
+    @property
+    def _nth_client_id(self):
+        return self._instance._client._nth_client_id
+
+    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 _Session(object):
     _transaction = None