Wrap Snapshot methods with x-goog-request-id metadata injector

Author: odeke-em

google/cloud/spanner_v1/database.py

@@ -700,7 +700,8 @@ def execute_partitioned_dml(
             )
 
         nth_request = getattr(self, "_next_nth_request", 0)
-        attempt = AtomicCounter(1)  # It'll be incremented inside _restart_on_unavailable
+        # Attempt will be incremented inside _restart_on_unavailable.
+        attempt = AtomicCounter(1)
 
         def execute_pdml():
             with SessionCheckout(self._pool) as session:

google/cloud/spanner_v1/snapshot.py

@@ -35,9 +35,11 @@
     _merge_query_options,
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
+    _metadata_with_request_id,
     _retry,
     _check_rst_stream_error,
     _SessionWrapper,
+    AtomicCounter,
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
 from google.cloud.spanner_v1.streamed import StreamedResultSet
@@ -320,13 +322,26 @@ def read(
             data_boost_enabled=data_boost_enabled,
             directed_read_options=directed_read_options,
         )
-        restart = functools.partial(
-            api.streaming_read,
-            request=request,
-            metadata=metadata,
-            retry=retry,
-            timeout=timeout,
-        )
+
+        nth_request = getattr(database, "_next_nth_request", 0)
+        attempt = AtomicCounter(0)
+
+        def wrapped_restart(*args, **kwargs):
+            attempt.increment()
+            channel_id = getattr(self._session, "_channel_id", 0)
+            client_id = getattr(database, "_nth_client_id", 0)
+            all_metadata = _metadata_with_request_id(
+                client_id, channel_id, nth_request, attempt.value, metadata
+            )
+
+            restart = functools.partial(
+                api.streaming_read,
+                request=request,
+                metadata=all_metadata,
+                retry=retry,
+                timeout=timeout,
+            )
+            return restart(*args, **kwargs)
 
         trace_attributes = {"table_id": table, "columns": columns}
         observability_options = getattr(database, "observability_options", None)
@@ -335,7 +350,7 @@ def read(
             # lock is added to handle the inline begin for first rpc
             with self._lock:
                 iterator = _restart_on_unavailable(
-                    restart,
+                    wrapped_restart,
                     request,
                     "CloudSpanner.ReadOnlyTransaction",
                     self._session,
@@ -357,7 +372,7 @@ def read(
                     )
         else:
             iterator = _restart_on_unavailable(
-                restart,
+                wrapped_restart,
                 request,
                 "CloudSpanner.ReadOnlyTransaction",
                 self._session,
@@ -536,13 +551,27 @@ def execute_sql(
             data_boost_enabled=data_boost_enabled,
             directed_read_options=directed_read_options,
         )
-        restart = functools.partial(
-            api.execute_streaming_sql,
-            request=request,
-            metadata=metadata,
-            retry=retry,
-            timeout=timeout,
-        )
+
+        nth_request = getattr(database, "_next_nth_request", 0)
+        attempt = AtomicCounter(0)
+
+        def wrapped_restart(*args, **kwargs):
+            attempt.increment()
+            channel_id = getattr(self._session, "_channel_id", 0)
+            client_id = getattr(database, "_nth_client_id", 0)
+            all_metadata = _metadata_with_request_id(
+                client_id, channel_id, nth_request, attempt.value, metadata
+            )
+
+            restart = functools.partial(
+                api.execute_streaming_sql,
+                request=request,
+                metadata=all_metadata,
+                retry=retry,
+                timeout=timeout,
+            )
+
+            return restart(*args, **kwargs)
 
         trace_attributes = {"db.statement": sql}
         observability_options = getattr(database, "observability_options", None)
@@ -551,7 +580,7 @@ def execute_sql(
             # lock is added to handle the inline begin for first rpc
             with self._lock:
                 return self._get_streamed_result_set(
-                    restart,
+                    wrapped_restart,
                     request,
                     trace_attributes,
                     column_info,
@@ -560,7 +589,7 @@ def execute_sql(
                 )
         else:
             return self._get_streamed_result_set(
-                restart,
+                wrapped_restart,
                 request,
                 trace_attributes,
                 column_info,
@@ -683,15 +712,27 @@ def partition_read(
             trace_attributes,
             observability_options=getattr(database, "observability_options", None),
         ):
-            method = functools.partial(
-                api.partition_read,
-                request=request,
-                metadata=metadata,
-                retry=retry,
-                timeout=timeout,
-            )
+            nth_request = getattr(database, "_next_nth_request", 0)
+            attempt = AtomicCounter(0)
+
+            def wrapped_method(*args, **kwargs):
+                attempt.increment()
+                channel_id = getattr(self._session, "_channel_id", 0)
+                client_id = getattr(database, "_nth_client_id", 0)
+                all_metadata = _metadata_with_request_id(
+                    client_id, channel_id, nth_request, attempt.value, metadata
+                )
+                method = functools.partial(
+                    api.partition_read,
+                    request=request,
+                    metadata=all_metadata,
+                    retry=retry,
+                    timeout=timeout,
+                )
+                return method(*args, **kwargs)
+
             response = _retry(
-                method,
+                wrapped_method,
                 allowed_exceptions={InternalServerError: _check_rst_stream_error},
             )
 
@@ -786,15 +827,28 @@ def partition_query(
             trace_attributes,
             observability_options=getattr(database, "observability_options", None),
         ):
-            method = functools.partial(
-                api.partition_query,
-                request=request,
-                metadata=metadata,
-                retry=retry,
-                timeout=timeout,
-            )
+            nth_request = getattr(database, "_next_nth_request", 0)
+            attempt = AtomicCounter(0)
+
+            def wrapped_method(*args, **kwargs):
+                attempt.increment()
+                channel_id = getattr(self._session, "_channel_id", 0)
+                client_id = getattr(database, "_nth_client_id", 0)
+                all_metadata = _metadata_with_request_id(
+                    client_id, channel_id, nth_request, attempt.value, metadata
+                )
+
+                method = functools.partial(
+                    api.partition_query,
+                    request=request,
+                    metadata=all_metadata,
+                    retry=retry,
+                    timeout=timeout,
+                )
+                return method(*args, **kwargs)
+
             response = _retry(
-                method,
+                wrapped_method,
                 allowed_exceptions={InternalServerError: _check_rst_stream_error},
             )
 
@@ -932,14 +986,27 @@ def begin(self):
             self._session,
             observability_options=getattr(database, "observability_options", None),
         ):
-            method = functools.partial(
-                api.begin_transaction,
-                session=self._session.name,
-                options=txn_selector.begin,
-                metadata=metadata,
-            )
+            nth_request = getattr(database, "_next_nth_request", 0)
+            attempt = AtomicCounter(0)
+
+            def wrapped_method(*args, **kwargs):
+                attempt.increment()
+                channel_id = getattr(self._session, "_channel_id", 0)
+                client_id = getattr(database, "_nth_client_id", 0)
+                all_metadata = _metadata_with_request_id(
+                    client_id, channel_id, nth_request, attempt.value, metadata
+                )
+
+                method = functools.partial(
+                    api.begin_transaction,
+                    session=self._session.name,
+                    options=txn_selector.begin,
+                    metadata=all_metadata,
+                )
+                return method(*args, **kwargs)
+
             response = _retry(
-                method,
+                wrapped_method,
                 allowed_exceptions={InternalServerError: _check_rst_stream_error},
             )
         self._transaction_id = response.id

tests/unit/test_atomic_counter.py

@@ -18,6 +18,7 @@
 import unittest
 from google.cloud.spanner_v1._helpers import AtomicCounter
 
+
 class TestAtomicCounter(unittest.TestCase):
     def test_initialization(self):
         ac_default = AtomicCounter()
@@ -54,7 +55,6 @@ def test_plus_call(self):
         assert n == 201
         assert ac.value == 1
 
-
     def test_multiple_threads_incrementing(self):
         ac = AtomicCounter()
         n = 200
@@ -78,4 +78,4 @@ def do_work():
             assert th.is_alive() == False
 
         # Finally the result should be n*m
-        assert ac.value == n*m
+        assert ac.value == n * m