refactored interceptor

Author: daniel-sanche

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

@@ -15,7 +15,6 @@
 
 import time
 from functools import wraps
-from grpc import RpcError
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
 )
@@ -66,6 +65,26 @@ def wrapper(self, continuation, client_call_details, request):
     return wrapper
 
 
+def _end_attempt(operation, exc, metadata):
+    """Helper to add metadata and exception to an operation"""
+    if metadata is not None:
+        operation.add_response_metadata(metadata)
+    if exc is not None:
+        # end attempt. If it succeeded, let higher levels decide when to end operation
+        operation.end_attempt_with_status(exc)
+
+
+@CrossSync.convert
+async def _get_metadata(source):
+    """Helper to extract metadata from a call or RpcError"""
+    try:
+        return (await source.trailing_metadata() or []) + (
+            await source.initial_metadata() or []
+        )
+    except Exception:
+        # ignore errors while fetching metadata
+        return None
+
 @CrossSync.convert_class(sync_name="BigtableMetricsInterceptor")
 class AsyncBigtableMetricsInterceptor(
     UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor, MetricsHandler
@@ -109,25 +128,14 @@ async def intercept_unary_unary(
         metadata = None
         try:
             call = await continuation(client_call_details, request)
-            metadata = (await call.trailing_metadata() or []) + (await call.initial_metadata() or [])
+            metadata = await _get_metadata(call)
             return call
-        except RpcError as rpc_error:
-            # attempt extracting metadata from error
-            try:
-                metadata = (await rpc_error.trailing_metadata() or []) + (await rpc_error.initial_metadata() or [])
-            except Exception:
-                pass
+        except Exception as rpc_error:
+            metadata = await _get_metadata(rpc_error)
             encountered_exc = rpc_error
             raise rpc_error
-        except Exception as e:
-            encountered_exc = e
-            raise
         finally:
-            if metadata is not None:
-                operation.add_response_metadata(metadata)
-            if encountered_exc is not None:
-                # end attempt. If it succeeded, let higher levels decide when to end operation
-                operation.end_attempt_with_status(encountered_exc)
+            _end_attempt(operation, encountered_exc, metadata)
 
     @CrossSync.convert
     @_with_operation_from_metadata
@@ -146,30 +154,17 @@ async def response_wrapper(call):
                         )
                         has_first_response = True
                     yield response
-
-
             except Exception as e:
+                # handle errors while processing stream
                 encountered_exc = e
                 raise
             finally:
                 if call is not None:
-                    metadata = (await call.trailing_metadata() or []) + (await call.initial_metadata() or [])
-                    operation.add_response_metadata(metadata)
-                    if encountered_exc is not None:
-                        # end attempt. If it succeeded, let higher levels decide when to end operation
-                        operation.end_attempt_with_status(encountered_exc)
+                    _end_attempt(operation, encountered_exc, await _get_metadata(call))
 
         try:
             return response_wrapper(await continuation(client_call_details, request))
-        except RpcError as rpc_error:
-            # attempt extracting metadata from error
-            try:
-                metadata = (await rpc_error.trailing_metadata() or []) + (await rpc_error.initial_metadata() or [])
-                operation.add_response_metadata(metadata)
-            except Exception:
-                pass
-            operation.end_attempt_with_status(rpc_error)
-            raise rpc_error
-        except Exception as e:
-            operation.end_attempt_with_status(e)
-            raise
+        except Exception as rpc_error:
+            # handle errors while intializing stream
+            _end_attempt(operation, rpc_error, await _get_metadata(rpc_error))
+            raise rpc_error

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

@@ -17,7 +17,6 @@
 from __future__ import annotations
 import time
 from functools import wraps
-from grpc import RpcError
 from google.cloud.bigtable.data._metrics.data_model import (
     OPERATION_INTERCEPTOR_METADATA_KEY,
 )
@@ -56,6 +55,22 @@ def wrapper(self, continuation, client_call_details, request):
     return wrapper
 
 
+def _end_attempt(operation, exc, metadata):
+    """Helper to add metadata and exception to an operation"""
+    if metadata is not None:
+        operation.add_response_metadata(metadata)
+    if exc is not None:
+        operation.end_attempt_with_status(exc)
+
+
+def _get_metadata(source):
+    """Helper to extract metadata from a call or RpcError"""
+    try:
+        return (source.trailing_metadata() or []) + (source.initial_metadata() or [])
+    except Exception:
+        return None
+
+
 class BigtableMetricsInterceptor(
     UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor, MetricsHandler
 ):
@@ -95,27 +110,14 @@ def intercept_unary_unary(
         metadata = None
         try:
             call = continuation(client_call_details, request)
-            metadata = (call.trailing_metadata() or []) + (
-                call.initial_metadata() or []
-            )
+            metadata = _get_metadata(call)
             return call
-        except RpcError as rpc_error:
-            try:
-                metadata = (rpc_error.trailing_metadata() or []) + (
-                    rpc_error.initial_metadata() or []
-                )
-            except Exception:
-                pass
+        except Exception as rpc_error:
+            metadata = _get_metadata(rpc_error)
             encountered_exc = rpc_error
             raise rpc_error
-        except Exception as e:
-            encountered_exc = e
-            raise
         finally:
-            if metadata is not None:
-                operation.add_response_metadata(metadata)
-            if encountered_exc is not None:
-                operation.end_attempt_with_status(encountered_exc)
+            _end_attempt(operation, encountered_exc, metadata)
 
     @_with_operation_from_metadata
     def intercept_unary_stream(
@@ -137,25 +139,10 @@ def response_wrapper(call):
                 raise
             finally:
                 if call is not None:
-                    metadata = (call.trailing_metadata() or []) + (
-                        call.initial_metadata() or []
-                    )
-                    operation.add_response_metadata(metadata)
-                    if encountered_exc is not None:
-                        operation.end_attempt_with_status(encountered_exc)
+                    _end_attempt(operation, encountered_exc, _get_metadata(call))
 
         try:
             return response_wrapper(continuation(client_call_details, request))
-        except RpcError as rpc_error:
-            try:
-                metadata = (rpc_error.trailing_metadata() or []) + (
-                    rpc_error.initial_metadata() or []
-                )
-                operation.add_response_metadata(metadata)
-            except Exception:
-                pass
-            operation.end_attempt_with_status(rpc_error)
+        except Exception as rpc_error:
+            _end_attempt(operation, rpc_error, _get_metadata(rpc_error))
             raise rpc_error
-        except Exception as e:
-            operation.end_attempt_with_status(e)
-            raise