Monkey patch updates

Author: odeke-em

google/cloud/spanner_v1/__init__.py

@@ -75,6 +75,7 @@
 from google.cloud.spanner_v1.pool import FixedSizePool
 from google.cloud.spanner_v1.pool import PingingPool
 from google.cloud.spanner_v1.pool import TransactionPingingPool
+from google.cloud.spanner_v1._helpers import monkey_patch
 
 
 COMMIT_TIMESTAMP = "spanner.commit_timestamp()"
@@ -83,6 +84,8 @@
 ``(allow_commit_timestamp=true)`` in the schema.
 """
 
+monkey_patch(Transaction)
+
 
 __all__ = (
     # google.cloud.spanner_v1

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
@@ -739,9 +740,100 @@ def __init__(self, original_callable: Callable):
 
 
 patched = {}
+patched_mu = threading.Lock()
 
 
 def inject_retry_header_control(api):
+    return
+    monkey_patch(type(api))
+
+memoize_map = dict()
+
+def monkey_patch(obj):
+    return
+
+    """
+    klass = obj
+    attrs = dir(klass)
+    for attr_key in attrs:
+        if attr_key.startswith('_'):
+            continue
+
+        attr_value = getattr(obj, attr_key)
+        if not callable(attr_value):
+            continue
+
+        signature = inspect.signature(attr_value)
+        print(attr_key, signature.parameters)
+
+        call = attr_value
+        # Our goal is to replace the runtime pass through.
+        def wrapped(*args, **kwargs):
+            print(attr_key, 'called')
+            return call(*args, **kwargs)
+
+        setattr(klass, attr_key, wrapped)
+
+    return
+    """
+
+    orig_get_attr = getattr(obj, "__getattribute__")
+    def patched_getattribute(obj, key, *args, **kwargs):
+        if key.startswith('_'):
+            return orig_get_attr(obj, key, *args, **kwargs)
+
+        orig_value = orig_get_attr(obj, key, *args, **kwargs)
+        if not callable(orig_value):
+            return orig_value
+
+        map_key = hex(id(key)) + hex(id(obj))
+        memoized = memoize_map.get(map_key, None)
+        if memoized:
+            print("memoized_hit", key, '\033[35m', inspect.getsource(orig_value), '\033[00m')
+            return memoized
+
+        signature = inspect.signature(orig_value)
+        if signature.parameters.get('metadata', None) is None:
+            return orig_value
+
+        print(key, '\033[34m', map_key, '\033[00m', signature, signature.parameters.get('metadata', None))
+        counters = dict(attempt=0)
+        def patched_method(*aargs, **kkwargs):
+            counters['attempt'] += 1
+            metadata = kkwargs.get('metadata', None)
+            if not metadata:
+                return orig_value(*aargs, **kkwargs)
+
+            # 4. Find all the headers that match the target header key.
+            all_metadata = []
+            for mkey, value in metadata:
+                if mkey is REQ_ID_HEADER_KEY:
+                    attempt = counters['attempt']
+                    if attempt > 1:
+                        # 5. Increment the original_attempt with that of our re-invocation count.
+                        splits = value.split(".")
+                        print('\033[34mkey', mkey, '\033[00m', splits)
+                        hdr_attempt_plus_reinvocation = (
+                            int(splits[-1]) + attempt
+                        )
+                        splits[-1] = str(hdr_attempt_plus_reinvocation)
+                        value = ".".join(splits)
+
+                all_metadata.append((mkey, value))
+
+            kwargs["metadata"] = all_metadata
+            return orig_value(*aargs, **kkwargs)
+
+        memoize_map[map_key] = patched_method
+        return patched_method
+
+    setattr(obj, '__getattribute__', patched_getattribute)
+
+
+def foo(api):
+    global patched
+    global patched_mu
+    
     # 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.
@@ -753,55 +845,66 @@ def inject_retry_header_control(api):
     orig_getattribute = getattr(target, "__getattribute__")
 
     def patched_getattribute(obj, key, *args, **kwargs):
+        # 1. Skip modifying private and mangled methods.
         if key.startswith("_"):
             return orig_getattribute(obj, key, *args, **kwargs)
 
         attr = orig_getattribute(obj, key, *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.
+        # 2. Skip over non-methods.
         if not callable(attr):
+            patched_mu.release()
             return attr
 
-        # 2. Skip modifying private and mangled methods.
-        mangled_or_private = attr.__name__.startswith("_")
-        if mangled_or_private:
-            return attr
-
+        patched_key = hex(id(key)) + hex(id(obj))
+        patched_mu.acquire()
+        already_patched = patched.get(patched_key, None)
+        
+        other_attempts = dict(attempts=0)
         # 3. Wrap the callable attribute and then capture its metadata keyed argument.
         def wrapped_attr(*args, **kwargs):
+            print("\033[31m", key, "attempt", other_attempts['attempts'], "\033[00m")
+            other_attempts['attempts'] += 1
+
             metadata = kwargs.get("metadata", [])
             if not metadata:
                 # Increment the reinvocation count.
                 wrapped_attr._attempt += 1
                 return attr(*args, **kwargs)
 
+            print("\033[35mwrapped_attr", key, args, kwargs, 'attempt', wrapped_attr._attempt, "\033[00m")
+
             # 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:
-                    # 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
+            for mkey, value in metadata:
+                if mkey is REQ_ID_HEADER_KEY:
+                    if wrapped_attr._attempt > 0:
+                        # 5. Increment the original_attempt with that of our re-invocation count.
+                        splits = value.split(".")
+                        print('\033[34mkey', mkey, '\033[00m', splits)
+                        hdr_attempt_plus_reinvocation = (
+                            int(splits[-1]) + wrapped_attr._attempt
+                        )
+                        splits[-1] = str(hdr_attempt_plus_reinvocation)
+                        value = ".".join(splits)
+
+                all_metadata.append((mkey, value))
 
             kwargs["metadata"] = all_metadata
+            wrapped_attr._attempt += 1
+            print(key, "\033[36mreplaced_all_metadata", all_metadata, "\033[00m")
             return attr(*args, **kwargs)
 
-        wrapped_attr._attempt = 0
-        wrapped_attr._patched = True
+        if already_patched:
+            print("patched_key \033[32m", patched_key, key, "\033[00m", already_patched)
+            setattr(attr, 'patched', True)
+            # Increment the reinvocation count.
+            patched_mu.release()
+            return already_patched
+
+        patched[patched_key] = wrapped_attr
+        setattr(wrapped_attr, '_attempt', 0)
+        patched_mu.release()
         return wrapped_attr
 
     setattr(target, "__getattribute__", patched_getattribute)
-    patched[hex_id] = True

google/cloud/spanner_v1/database.py

@@ -55,7 +55,7 @@
     _metadata_with_prefix,
     _metadata_with_leader_aware_routing,
     _metadata_with_request_id,
-    inject_retry_header_control,
+    monkey_patch,
 )
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.batch import MutationGroups
@@ -438,7 +438,7 @@ def spanner_api(self):
         if not api:
             return api
 
-        inject_retry_header_control(api)
+        monkey_patch(api)
         return api
 
     def __generate_spanner_api(self):
@@ -813,6 +813,7 @@ def execute_pdml():
     def _next_nth_request(self):
         if self._instance and self._instance._client:
             return self._instance._client._next_nth_request
+        raise Exception("returning 1 for next_nth_request")
         return 1
 
     @property

google/cloud/spanner_v1/pool.py

@@ -249,6 +249,7 @@ def bind(self, database):
             attempt = 1
             returned_session_count = 0
             while not self._sessions.full():
+                print("fixedPool.batchCreateSessions")
                 request.session_count = requested_session_count - self._sessions.qsize()
                 add_span_event(
                     span,
@@ -562,6 +563,7 @@ def bind(self, database):
         ) as span, MetricsCapture():
             returned_session_count = 0
             while returned_session_count < self.size:
+                print("pingingPool.batchCreateSessions")
                 resp = api.batch_create_sessions(
                     request=request,
                     metadata=database.metadata_with_request_id(

google/cloud/spanner_v1/snapshot.py

@@ -591,21 +591,24 @@ def execute_sql(
             directed_read_options=directed_read_options,
         )
 
-        restart = functools.partial(
-            api.execute_streaming_sql,
-            request=request,
-            metadata=metadata,
-            retry=retry,
-            timeout=timeout,
-        )
+        def wrapped_restart(*args, **kwargs):
+            restart = functools.partial(
+                api.execute_streaming_sql,
+                request=request,
+                metadata=kwargs.get('metadata', metadata),
+                retry=retry,
+                timeout=timeout,
+            )
+            return restart(*args, **kwargs)
+
         trace_attributes = {"db.statement": sql}
         observability_options = getattr(database, "observability_options", None)
 
         if self._transaction_id is None:
             # 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,
@@ -615,7 +618,7 @@ def execute_sql(
                 )
         else:
             return self._get_streamed_result_set(
-                restart,
+                wrapped_restart,
                 request,
                 metadata,
                 trace_attributes,

tests/mockserver_tests/test_request_id_header.py

@@ -310,6 +310,7 @@ def test_unary_retryable_error(self):
         ]
 
         print("got_unaries", got_unary_segments)
+        print("got_stream", got_stream_segments)
         assert got_unary_segments == want_unary_segments
         assert got_stream_segments == want_stream_segments