refacotred tracked exception factory

Author: daniel-sanche

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

@@ -23,7 +23,7 @@
 import google.cloud.bigtable_v2.types.bigtable as types_pb
 import google.cloud.bigtable.data.exceptions as bt_exceptions
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._helpers import _retry_exception_factory
+from google.cloud.bigtable.data._helpers import _tracked_exception_factory
 
 # mutate_rows requests are limited to this number of mutations
 from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT
@@ -108,9 +108,7 @@ def __init__(
             self.is_retryable,
             metric.backoff_generator,
             operation_timeout,
-            exception_factory=partial(
-                _retry_exception_factory, operation=metric
-            ),
+            exception_factory=_tracked_exception_factory(self._operation_metric),
         )
         # initialize state
         self.timeout_generator = _attempt_timeout_generator(

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

@@ -31,7 +31,7 @@
 from google.cloud.bigtable.data.exceptions import _RowSetComplete
 from google.cloud.bigtable.data.exceptions import _ResetRow
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._helpers import _retry_exception_factory
+from google.cloud.bigtable.data._helpers import _tracked_exception_factory
 
 from google.api_core import retry as retries
 
@@ -122,7 +122,7 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
             self._predicate,
             self._operation_metric.backoff_generator,
             self.operation_timeout,
-            exception_factory=partial(_retry_exception_factory, operation=self._operation_metric),
+            exception_factory=_tracked_exception_factory(self._operation_metric),
         )
 
     def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:

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

@@ -73,6 +73,7 @@
 from google.cloud.bigtable.data._helpers import _WarmedInstanceKey
 from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
+from google.cloud.bigtable.data._helpers import _tracked_exception_factory
 from google.cloud.bigtable.data._helpers import _validate_timeouts
 from google.cloud.bigtable.data._helpers import _get_error_type
 from google.cloud.bigtable.data._helpers import _get_retryable_errors
@@ -1391,7 +1392,7 @@ async def execute_rpc():
                 predicate,
                 operation_metric.backoff_generator,
                 operation_timeout,
-                exception_factory=partial(_retry_exception_factory, operation=operation_metric)
+                exception_factory=_tracked_exception_factory(operation_metric),
             )
 
     @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"})

google/cloud/bigtable/data/_helpers.py

@@ -16,7 +16,7 @@
 """
 from __future__ import annotations
 
-from typing import Sequence, List, Tuple, TYPE_CHECKING, Union
+from typing import Callable, Sequence, List, Tuple, TYPE_CHECKING, Union
 import time
 import enum
 from collections import namedtuple
@@ -31,6 +31,7 @@
     import grpc
     from google.cloud.bigtable.data._async.client import _DataApiTargetAsync
     from google.cloud.bigtable.data._sync_autogen.client import _DataApiTarget
+    from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
 
 """
 Helper functions used in various places in the library.
@@ -90,7 +91,6 @@ def _retry_exception_factory(
     exc_list: list[Exception],
     reason: RetryFailureReason,
     timeout_val: float | None,
-    operation: "ActiveOperationMetric" | None=None,
 ) -> tuple[Exception, Exception | None]:
     """
     Build retry error based on exceptions encountered during operation
@@ -118,8 +118,35 @@ def _retry_exception_factory(
     # use the retry exception group as the cause of the exception
     cause_exc: Exception | None = RetryExceptionGroup(exc_list) if exc_list else None
     source_exc.__cause__ = cause_exc
-    if operation:
+    return source_exc, cause_exc
+
+
+def _tracked_exception_factory(
+    operation: "ActiveOperationMetric",
+) -> Callable[[list[Exception], RetryFailureReason, float | None], tuple[Exception, Exception | None]]:
+    """
+    wraps and extends _retry_exception_factory to add client-side metrics tracking.
+
+    When the rpc raises a terminal error, record any discovered metadata and finalize
+    the associated operation
+
+    Used by streaming rpcs, which can't always be perfectly captured by context managers
+    for operation termination.
+
+    Args:
+        exc_list: list of exceptions encountered during operation
+        is_timeout: whether the operation failed due to timeout
+        timeout_val: the operation timeout value in seconds, for constructing
+            the error message
+        operation: the operation to finalize when an exception is built
+    Returns:
+        tuple[Exception, Exception|None]:
+            tuple of the exception to raise, and a cause exception if applicabl
+    """
+    def wrapper(exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None) -> tuple[Exception, Exception | None]:
+        source_exc, cause_exc = _retry_exception_factory(exc_list, reason, timeout_val)
         try:
+            # record metadata from failed rpc
             if isinstance(source_exc, core_exceptions.GoogleAPICallError) and source_exc.errors:
                 rpc_error = source_exc.errors[-1]
                 metadata = list(rpc_error.trailing_metadata()) + list(rpc_error.initial_metadata())
@@ -128,7 +155,8 @@ def _retry_exception_factory(
             # ignore errors in metadata collection
             pass
         operation.end_with_status(source_exc)
-    return source_exc, cause_exc
+        return source_exc, cause_exc
+    return wrapper
 
 
 def _get_timeouts(