combined wrapped predicate with wrapped exc factory

Author: daniel-sanche

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

@@ -115,31 +115,17 @@ def start_operation(self) -> AsyncGenerator[Row, None]:
         self._operation_metrics.backoff_generator = sleep_generator
 
         # Metrics:
-        # track attempt failures using build_wrapped_predicate() for raised exceptions
-        # and _metric_wrapped_exception_factory for operation timeouts
+        # track attempt failures using build_wrapped_fn_handlers() for raised exceptions
+        # and operation timeouts
+        metric_predicate, metric_excs = self._operation_metrics.build_wrapped_fn_handlers(self._predicate)
         return retries.retry_target_stream_async(
             self._read_rows_attempt,
-            self._operation_metrics.build_wrapped_predicate(self._predicate),
+            metric_predicate,
             sleep_generator,
             self.operation_timeout,
-            exception_factory=self._metric_wrapped_exception_factory,
+            exception_factory=metric_excs,
         )
 
-    def _metric_wrapped_exception_factory(
-        self,
-        exc_list: list[Exception],
-        reason: retries.RetryFailureReason,
-        timeout_val: float | None,
-    ) -> tuple[Exception, Exception | None]:
-        """
-        Wrap the retry exception builder to alert the metrics class
-        when we are going to emit an operation timeout.
-        """
-        exc, source = _retry_exception_factory(exc_list, reason, timeout_val)
-        if reason != retries.RetryFailureReason.NON_RETRYABLE_ERROR:
-            self._operation_metrics.end_with_status(exc)
-        return exc, source
-
     def _read_rows_attempt(self) -> AsyncGenerator[Row, None]:
         """
         Attempt a single read_rows rpc call.

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

@@ -30,8 +30,10 @@
 from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
 from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup
 from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
+from google.cloud.bigtable.data._helpers import _retry_exception_factory
 from google.cloud.bigtable_v2.types.response_params import ResponseParams
 from google.protobuf.message import DecodeError
+from google.api_core.retry import RetryFailureReason
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
@@ -377,13 +379,17 @@ def end_with_success(self):
         """
         return self.end_with_status(StatusCode.OK)
 
-    def build_wrapped_predicate(
-        self, inner_predicate: Callable[[Exception], bool]
+    def build_wrapped_fn_handlers(
+        self,
+        inner_predicate: Callable[[Exception], bool],
     ) -> Callable[[Exception], bool]:
         """
-        Wrapps a predicate to include metrics tracking. Any call to the resulting predicate
-        is assumed to be an rpc failure, and will either mark the end of the active attempt
-        or the end of the operation.
+        One way to track metrics is by wrapping the `predicate` and `exception_factory`
+        arguments of `api_core.Retry`. This will notify us when an exception occurs so
+        we can track it.
+
+        This function retruns wrapped versions of the `predicate` and `exception_factory`
+        to be passed down when building the `Retry` object.
 
         Args:
           - predicate: The predicate to wrap.
@@ -397,7 +403,17 @@ def wrapped_predicate(exc: Exception) -> bool:
                 self.end_with_status(exc)
             return inner_result
 
-        return wrapped_predicate
+        def wrapped_exception_factory(
+            exc_list: list[Exception],
+            reason: RetryFailureReason,
+            timeout_val: float | None,
+        ) -> tuple[Exception, Exception | None]:
+            exc, source = _retry_exception_factory(exc_list, reason, timeout_val)
+            if reason != RetryFailureReason.NON_RETRYABLE_ERROR:
+                self.end_with_status(exc)
+            return exc, source
+
+        return wrapped_predicate, wrapped_exception_factory
 
     @staticmethod
     def _exc_to_status(exc: Exception) -> StatusCode:

tests/unit/data/_metrics/test_data_model.py

@@ -601,7 +601,7 @@ def test_end_on_empty_operation(self):
         assert final_op.final_status == StatusCode.OK
         assert final_op.completed_attempts == []
 
-    def test_build_wrapped_predicate(self):
+    def test_build_wrapped_fn_handlers_predicate(self):
         """
         predicate generated by object should terminate attempt or operation
         based on passed in predicate
@@ -610,23 +610,61 @@ def test_build_wrapped_predicate(self):
         cls = type(self._make_one(object()))
         # ensure predicate is called with the exception
         mock_predicate = mock.Mock()
-        cls.build_wrapped_predicate(mock.Mock(), mock_predicate)(input_exc)
+        pred, _ = cls.build_wrapped_fn_handlers(mock.Mock(), mock_predicate)
+        pred(input_exc)
         assert mock_predicate.call_count == 1
         assert mock_predicate.call_args[0][0] == input_exc
         assert len(mock_predicate.call_args[0]) == 1
         # if predicate is true, end the attempt
         mock_instance = mock.Mock()
-        cls.build_wrapped_predicate(mock_instance, lambda x: True)(input_exc)
+        pred, _ = cls.build_wrapped_fn_handlers(mock_instance, lambda x: True)
+        pred(input_exc)
         assert mock_instance.end_attempt_with_status.call_count == 1
         assert mock_instance.end_attempt_with_status.call_args[0][0] == input_exc
         assert len(mock_instance.end_attempt_with_status.call_args[0]) == 1
         # if predicate is false, end the operation
         mock_instance = mock.Mock()
-        cls.build_wrapped_predicate(mock_instance, lambda x: False)(input_exc)
+        pred, _ = cls.build_wrapped_fn_handlers(mock_instance, lambda x: False)
+        pred(input_exc)
         assert mock_instance.end_with_status.call_count == 1
         assert mock_instance.end_with_status.call_args[0][0] == input_exc
         assert len(mock_instance.end_with_status.call_args[0]) == 1
 
+    def test_build_wrapped_fn_handlers_exc_factory(self):
+        """
+        exception factory generated by object should terminate operation
+        on timeout
+        """
+        from google.api_core.retry import RetryFailureReason
+        from google.api_core.exceptions import DeadlineExceeded
+
+        cls = type(self._make_one(object()))
+        # ensure inner factory is called with the exception
+        _, factory = cls.build_wrapped_fn_handlers(mock.Mock(), None)
+        with mock.patch(
+            "google.cloud.bigtable.data._metrics.data_model._retry_exception_factory"
+        ) as mock_factory:
+            expected_return = (object(), object())
+            mock_factory.return_value = expected_return
+            args = ("a", "b", "c")
+            got_result = factory(*args)
+            assert expected_return == got_result
+            assert mock_factory.call_count == 1
+            assert mock_factory.call_args[0] == args
+
+        # if called with reason == TIMEOUT, end the operation
+        mock_instance = mock.Mock()
+        _, factory = cls.build_wrapped_fn_handlers(mock_instance, None)
+        factory([], RetryFailureReason.TIMEOUT, None)
+        assert mock_instance.end_with_status.call_count == 1
+        assert type(mock_instance.end_with_status.call_args[0][0]) == DeadlineExceeded
+
+        # if called with reason==NON_RETRYABLE_ERROR, do not
+        mock_instance = mock.Mock()
+        _, factory = cls.build_wrapped_fn_handlers(mock_instance, None)
+        factory([], RetryFailureReason.NON_RETRYABLE_ERROR, None)
+        assert mock_instance.end_with_status.call_count == 0
+
     def test__exc_to_status(self):
         """
         Should return grpc_status_code if grpc error, otherwise UNKNOWN