Initial ExecuteStreamingSql request in snapshot should have the header

Author: odeke-em

google/cloud/spanner_v1/_helpers.py

@@ -757,7 +757,6 @@ def patched_getattribute(obj, key, *args, **kwargs):
             return orig_getattribute(obj, key, *args, **kwargs)
 
         attr = orig_getattribute(obj, key, *args, **kwargs)
-        print("args", args, "attr.dir", dir(attr))
 
         # 0. If we already patched it, we can return immediately.
         if getattr(attr, "_patched", None) is not None:
@@ -772,22 +771,18 @@ def patched_getattribute(obj, key, *args, **kwargs):
         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 = (
@@ -802,7 +797,6 @@ def wrapped_attr(*args, **kwargs):
             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

google/cloud/spanner_v1/database.py

@@ -792,15 +792,19 @@ def execute_pdml():
                         query_options=query_options,
                         request_options=request_options,
                     )
-                    method = functools.partial(
-                        api.execute_streaming_sql,
-                        metadata=self.metadata_with_request_id(
-                            partial_nth_request, partial_attempt.increment(), metadata
-                        ),
-                    )
+
+                    def wrapped_method(*args, **kwargs):
+                        print("\033[34mwrapped_method\033[00m")
+                        method = functools.partial(
+                            api.execute_streaming_sql,
+                            metadata=self.metadata_with_request_id(
+                                partial_nth_request, partial_attempt.increment(), metadata
+                            ),
+                        )
+                        return method(*args, **kwargs)
 
                     iterator = _restart_on_unavailable(
-                        method=method,
+                        method=wrapped_method,
                         trace_name="CloudSpanner.ExecuteStreamingSql",
                         request=request,
                         metadata=metadata,

google/cloud/spanner_v1/snapshot.py

@@ -586,6 +586,10 @@ def wrapped_restart(*args, **kwargs):
             )
             return restart(*args, **kwargs)
 
+        # The initial request should contain the request-id.
+        augmented_metadata=database.metadata_with_request_id(
+            nth_request, attempt.increment(), metadata
+        )
         trace_attributes = {"db.statement": sql}
         observability_options = getattr(database, "observability_options", None)
 
@@ -595,7 +599,7 @@ def wrapped_restart(*args, **kwargs):
                 return self._get_streamed_result_set(
                     wrapped_restart,
                     request,
-                    metadata,
+                    augmented_metadata,
                     trace_attributes,
                     column_info,
                     observability_options,
@@ -605,7 +609,7 @@ def wrapped_restart(*args, **kwargs):
             return self._get_streamed_result_set(
                 wrapped_restart,
                 request,
-                metadata,
+                augmented_metadata,
                 trace_attributes,
                 column_info,
                 observability_options,