Correctly handle wrapping by class for api objects

odeke-em · odeke-em · commit df8f81faa8e0 · 2025-03-22T17:37:48.000+02:00
diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py
@@ -737,27 +737,73 @@ class InterceptingHeaderInjector:
     def __init__(self, original_callable: Callable):
         self._original_callable = original_callable
 
-    def __call__(self, *args, **kwargs):
-        metadata = kwargs.get("metadata", [])
-        # Find all the headers that match the x-goog-spanner-request-id
-        # header an on each retry increment the value.
-        all_metadata = []
-        for key, value in metadata:
-            if key is REQ_ID_HEADER_KEY:
-                # Otherwise now increment the count for the attempt number.
-                splits = value.split(".")
-                attempt_plus_one = int(splits[-1]) + 1
-                splits[-1] = str(attempt_plus_one)
-                value_before = value
-                value = ".".join(splits)
-                print("incrementing value on retry from", value_before, "to", value)
-
-            all_metadata.append(
-                (
-                    key,
-                    value,
-                )
-            )
 
-        kwargs["metadata"] = all_metadata
-        return self._original_callable(*args, **kwargs)
+patched = {}
+
+
+def inject_retry_header_control(api):
+    # For each method, add an _attempt value that'll then be
+    # retrieved for each retry.
+    # 1. Patch the __getattribute__ method to match items in our manifest.
+    target = type(api)
+    hex_id = hex(id(target))
+    if patched.get(hex_id, None) is not None:
+        return
+
+    orig_getattribute = getattr(target, "__getattribute__")
+
+    def patched_getattribute(*args, **kwargs):
+        attr = orig_getattribute(*args, **kwargs)
+
+        # 0. If we already patched it, we can return immediately.
+        if getattr(attr, "_patched", None) is not None:
+            return attr
+
+        # 1. Skip over non-methods.
+        if not callable(attr):
+            return attr
+
+        # 2. Skip modifying private and mangled methods.
+        mangled_or_private = attr.__name__.startswith("_")
+        if mangled_or_private:
+            return attr
+
+        print("\033[35mattr", attr, "hex_id", hex(id(attr)), "\033[00m")
+
+        # 3. Wrap the callable attribute and then capture its metadata keyed argument.
+        def wrapped_attr(*args, **kwargs):
+            metadata = kwargs.get("metadata", [])
+            if not metadata:
+                # Increment the reinvocation count.
+                print("not metatadata", attr.__name__)
+                wrapped_attr._attempt += 1
+                return attr(*args, **kwargs)
+
+            # 4. Find all the headers that match the target header key.
+            all_metadata = []
+            for key, value in metadata:
+                if key is REQ_ID_HEADER_KEY:
+                    print("key", key, "value", value, "attempt", wrapped_attr._attempt)
+                    # 5. Increment the original_attempt with that of our re-invocation count.
+                    splits = value.split(".")
+                    hdr_attempt_plus_reinvocation = (
+                        int(splits[-1]) + wrapped_attr._attempt
+                    )
+                    splits[-1] = str(hdr_attempt_plus_reinvocation)
+                    value = ".".join(splits)
+
+                all_metadata.append((key, value))
+
+            # Increment the reinvocation count.
+            wrapped_attr._attempt += 1
+
+            kwargs["metadata"] = all_metadata
+            print("\033[34mwrap_callable", hex(id(attr)), attr.__name__, "\033[00m")
+            return attr(*args, **kwargs)
+
+        wrapped_attr._attempt = 0
+        wrapped_attr._patched = True
+        return wrapped_attr
+
+    setattr(target, "__getattribute__", patched_getattribute)
+    patched[hex_id] = True
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
@@ -55,7 +55,7 @@
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
     _metadata_with_request_id,
-    InterceptingHeaderInjector,
+    inject_retry_header_control,
 )
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.batch import MutationGroups
@@ -434,42 +434,14 @@ def logger(self):
     @property
     def spanner_api(self):
         """Helper for session-related API calls."""
-        api = self._generate_spanner_api()
+        api = self.__generate_spanner_api()
         if not api:
             return api
 
-        # Now wrap each method's __call__ method with our wrapped one.
-        # This is how to deal with the fact that there are no proper gRPC
-        # interceptors for Python hence the remedy is to replace callables
-        # with our custom wrapper.
-        attrs = dir(api)
-        for attr_name in attrs:
-            mangled = attr_name.startswith("__")
-            if mangled:
-                continue
-
-            non_public = attr_name.startswith("_")
-            if non_public:
-                continue
-
-            attr = getattr(api, attr_name)
-            callable_attr = callable(attr)
-            if callable_attr is None:
-                continue
-
-            # We should only be looking at bound methods to SpannerClient
-            # as those are the RPC invoking methods that need to be wrapped
-
-            is_method = type(attr).__name__ == "method"
-            if not is_method:
-                continue
-
-            print("attr_name", attr_name, "callable_attr", attr)
-            setattr(api, attr_name, InterceptingHeaderInjector(attr))
-
+        inject_retry_header_control(api)
         return api
 
-    def _generate_spanner_api(self):
+    def __generate_spanner_api(self):
         """Helper for session-related API calls."""
         if self._spanner_api is None:
             client_info = self._instance._client._client_info
@@ -804,7 +776,9 @@ def execute_pdml():
                         session=session.name,
                         options=txn_options,
                         metadata=self.metadata_with_request_id(
-                            begin_txn_nth_request, begin_txn_attempt.increment(), metadata
+                            begin_txn_nth_request,
+                            begin_txn_attempt.increment(),
+                            metadata,
                         ),
                     )
 
diff --git a/google/cloud/spanner_v1/pool.py b/google/cloud/spanner_v1/pool.py
@@ -268,7 +268,8 @@ def create_sessions(attempt):
                         metadata=all_metadata,
                     )
 
-                resp = retry_on_unavailable(create_sessions)
+                resp = retry_on_unavailable(create_sessions, "fixedpool")
+                # print("resp.FixedPool", resp)
 
                 add_span_event(
                     span,
@@ -581,7 +582,8 @@ def create_sessions(attempt):
                         metadata=all_metadata,
                     )
 
-                resp = retry_on_unavailable(create_sessions)
+                resp = retry_on_unavailable(create_sessions, "pingpool")
+                print("resp.PingingPool", resp)
 
                 add_span_event(
                     span,
@@ -822,20 +824,23 @@ def __exit__(self, *ignored):
         self._pool.put(self._session)
 
 
-def retry_on_unavailable(fn, max=6):
+def retry_on_unavailable(fn, kind, max=6):
     """
     Retries `fn` to a maximum of `max` times on encountering UNAVAILABLE exceptions,
     each time passing in the iteration's ordinal number to signal
     the nth attempt. It retries with exponential backoff with jitter.
     """
     last_exc = None
     for i in range(max):
+        print("retry_on_unavailable", kind, i)
         try:
             return fn(i + 1)
         except ServiceUnavailable as exc:
+            print("exc", exc)
             last_exc = exc
             time.sleep(i**2 + random.random())
-        except:
+        except Exception as e:
+            print("got exception", e)
             raise
 
     raise last_exc
diff --git a/google/cloud/spanner_v1/services/spanner/transports/grpc.py b/google/cloud/spanner_v1/services/spanner/transports/grpc.py
@@ -413,7 +413,9 @@ def batch_create_sessions(
                 request_serializer=spanner.BatchCreateSessionsRequest.serialize,
                 response_deserializer=spanner.BatchCreateSessionsResponse.deserialize,
             )
-        return self._stubs["batch_create_sessions"]
+        fn = self._stubs["batch_create_sessions"]
+        print("\033[32minvoking batch_create_sessionhex_id", hex(id(fn)), "\033[00m")
+        return fn
 
     @property
     def get_session(self) -> Callable[[spanner.GetSessionRequest], spanner.Session]:
diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py
@@ -332,9 +332,7 @@ def read(
         )
 
         nth_request = getattr(database, "_next_nth_request", 0)
-        all_metadata = database.metadata_with_request_id(
-            nth_request, 1, metadata
-        )
+        all_metadata = database.metadata_with_request_id(nth_request, 1, metadata)
 
         restart = functools.partial(
             api.streaming_read,
@@ -574,15 +572,19 @@ def execute_sql(
         if not isinstance(nth_request, int):
             raise Exception(f"failed to get an integer back: {nth_request}")
 
-        restart = functools.partial(
-            api.execute_streaming_sql,
-            request=request,
-            metadata=database.metadata_with_request_id(
-                nth_request, 1, metadata
-            ),
-            retry=retry,
-            timeout=timeout,
-        )
+        attempt = AtomicCounter(0)
+
+        def wrapped_restart(*args, **kwargs):
+            restart = functools.partial(
+                api.execute_streaming_sql,
+                request=request,
+                metadata=database.metadata_with_request_id(
+                    nth_request, attempt.increment(), metadata
+                ),
+                retry=retry,
+                timeout=timeout,
+            )
+            return restart(*args, **kwargs)
 
         trace_attributes = {"db.statement": sql}
         observability_options = getattr(database, "observability_options", None)
@@ -591,7 +593,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,
                     metadata,
                     trace_attributes,
@@ -601,7 +603,7 @@ def execute_sql(
                 )
         else:
             return self._get_streamed_result_set(
-                restart,
+                wrapped_restart,
                 request,
                 metadata,
                 trace_attributes,
@@ -733,9 +735,7 @@ def partition_read(
             metadata=metadata,
         ), MetricsCapture():
             nth_request = getattr(database, "_next_nth_request", 0)
-            all_metadata = database.metadata_with_request_id(
-                nth_request, 1, metadata
-            )
+            all_metadata = database.metadata_with_request_id(nth_request, 1, metadata)
             method = functools.partial(
                 api.partition_read,
                 request=request,
@@ -842,9 +842,7 @@ def partition_query(
             metadata=metadata,
         ), MetricsCapture():
             nth_request = getattr(database, "_next_nth_request", 0)
-            all_metadata = database.metadata_with_request_id(
-                nth_request, 1, metadata
-            )
+            all_metadata = database.metadata_with_request_id(nth_request, 1, metadata)
             method = functools.partial(
                 api.partition_query,
                 request=request,
@@ -994,9 +992,7 @@ def begin(self):
             metadata=metadata,
         ), MetricsCapture():
             nth_request = getattr(database, "_next_nth_request", 0)
-            all_metadata = database.metadata_with_request_id(
-                nth_request, 1, metadata
-            )
+            all_metadata = database.metadata_with_request_id(nth_request, 1, metadata)
             method = functools.partial(
                 api.begin_transaction,
                 session=self._session.name,
diff --git a/google/cloud/spanner_v1/testing/database_test.py b/google/cloud/spanner_v1/testing/database_test.py
@@ -27,6 +27,7 @@
     MethodAbortInterceptor,
     XGoogRequestIDHeaderInterceptor,
 )
+from google.cloud.spanner_v1._helpers import inject_retry_header_control
 
 
 class TestDatabase(Database):
@@ -70,6 +71,14 @@ def __init__(
 
     @property
     def spanner_api(self):
+        api = self.__generate_spanner_api()
+        if not api:
+            return api
+
+        inject_retry_header_control(api)
+        return api
+
+    def __generate_spanner_api(self):
         """Helper for session-related API calls."""
         if self._spanner_api is None:
             client = self._instance._client
diff --git a/google/cloud/spanner_v1/testing/mock_spanner.py b/google/cloud/spanner_v1/testing/mock_spanner.py
@@ -53,6 +53,7 @@ def pop_error(self, context):
         name = inspect.currentframe().f_back.f_code.co_name
         error: _Status | None = self.errors.pop(name, None)
         if error:
+            print("context.abort_with_status", error)
             context.abort_with_status(error)
 
     def get_result_as_partial_result_sets(
diff --git a/tests/mockserver_tests/test_request_id_header.py b/tests/mockserver_tests/test_request_id_header.py
@@ -309,6 +309,7 @@ def test_unary_retryable_error(self):
             )
         ]
 
+        print("got_unaries", got_unary_segments)
         assert got_unary_segments == want_unary_segments
         assert got_stream_segments == want_stream_segments
 

Original file line number	Diff line number	Diff line change
`@@ -309,6 +309,7 @@ def test_unary_retryable_error(self):`
`309`	`309`	`)`
`310`	`310`	`]`
`311`	`311`
	`312`	`+ print("got_unaries", got_unary_segments)`
`312`	`313`	`assert got_unary_segments == want_unary_segments`
`313`	`314`	`assert got_stream_segments == want_stream_segments`
`314`	`315`