added routing cookie to streaming rpcs

Author: daniel-sanche

google/cloud/bigtable/data/_async/metrics_interceptor.py

@@ -16,6 +16,7 @@
 from typing import Sequence
 
 import time
+import contextvars
 from functools import wraps
 
 from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
@@ -95,6 +96,11 @@ class AsyncBigtableMetricsInterceptor(
     An async gRPC interceptor to add client metadata and print server metadata.
     """
 
+    _routing_cookie_context: ClassVar[
+        contextvars.ContextVar[dict[str, str|bytes]]
+    ] = contextvars.ContextVar("routing_cooke_context")
+
+
     @CrossSync.convert
     @_with_active_operation
     async def intercept_unary_unary(
@@ -130,6 +136,8 @@ async def intercept_unary_stream(
           - SampleRowKeys
         """
         try:
+            for k,v in self._routing_cookie_context.get({}).items():
+                client_call_details.metadata.add(k, v)
             return self._streaming_generator_wrapper(
                 operation, await continuation(client_call_details, request)
             )

google/cloud/bigtable/data/_metrics/tracked_retry.py

@@ -30,6 +30,7 @@
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
 from google.cloud.bigtable.data._metrics import ActiveOperationMetric
 from google.cloud.bigtable.data._metrics import OperationState
+from google.cloud.bigtable.data._async.metrics_interceptor import AsyncBigtableMetricsInterceptor
 
 
 T = TypeVar("T")
@@ -58,7 +59,12 @@ def wrapper(exc: Exception) -> None:
                 metadata = list(rpc_error.trailing_metadata()) + list(
                     rpc_error.initial_metadata()
                 )
-                operation.add_response_metadata({k: v for k, v in metadata})
+                metadata_dict = {k: v for k, v in metadata}
+                operation.add_response_metadata(metadata_dict)
+                # check for routing cookie:
+                cookie_headers = {k:v for k,v in metadata_dict.items() if k.startswith("x-goog-cbt-cookie")}
+                if cookie_headers:
+                    AsyncBigtableMetricsInterceptor._routing_cookie_context.set(cookie_headers)
         except Exception:
             # ignore errors in metadata collection
             pass