added operation logic to retry factory

Author: daniel-sanche

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

@@ -18,6 +18,7 @@
 from typing import Sequence, TYPE_CHECKING
 
 import time
+from functools import partial
 
 from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB
 from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB
@@ -121,7 +122,7 @@ def start_operation(self) -> CrossSync.Iterable[Row]:
             self._predicate,
             self._operation_metric.backoff_generator,
             self.operation_timeout,
-            exception_factory=_retry_exception_factory,
+            exception_factory=partial(_retry_exception_factory, operation=self._operation_metric),
         )
 
     def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
@@ -340,8 +341,6 @@ async def merge_rows(
                 except CrossSync.StopIteration:
                     raise InvalidChunk("premature end of stream")
         except Exception as generic_exception:
-            if not self._predicate(generic_exception):
-                self._operation_metric.end_attempt_with_status(generic_exception)
             raise generic_exception
         else:
             self._operation_metric.end_with_success()

google/cloud/bigtable/data/_helpers.py

@@ -90,6 +90,7 @@ 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
@@ -117,6 +118,16 @@ 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:
+        try:
+            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())
+                operation.add_response_metadata({k:v for k,v in metadata})
+        except Exception:
+            # ignore errors in metadata collection
+            pass
+        operation.end_with_status(source_exc)
     return source_exc, cause_exc