feat(bigtable): support custom retry functions in mutate_rows shim

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

@@ -14,7 +14,7 @@
 #
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Callable, Sequence
 
 from google.api_core import exceptions as core_exceptions
 from google.api_core import retry as retries
@@ -73,6 +73,8 @@ class _MutateRowsOperationAsync:
             If not specified, the request will run until operation_timeout is reached.
         metric: the metric object representing the active operation
         retryable_exceptions: a list of exceptions that should be retried
+        shim_predicate: optional predicate callback, used only to support the legacy client shim.
+        shim_on_error: optional error callback, used only to support the legacy client shim.
     """
 
     @CrossSync.convert
@@ -85,6 +87,8 @@ def __init__(
         attempt_timeout: float | None,
         metric: ActiveOperationMetric,
         retryable_exceptions: Sequence[type[Exception]] = (),
+        shim_predicate: Callable[[Exception], bool] | None = None,
+        shim_on_error: Callable[[Exception], None] | None = None,
     ):
         # check that mutations are within limits
         total_mutations = sum(len(entry.mutations) for entry in mutation_entries)
@@ -97,18 +101,24 @@ def __init__(
         self._target = target
         self._gapic_fn = gapic_client.mutate_rows
         # create predicate for determining which errors are retryable
-        self.is_retryable = retries.if_exception_type(
+        base_predicate = retries.if_exception_type(
             # RPC level errors
             *retryable_exceptions,
             # Entry level errors
             bt_exceptions._MutateRowsIncomplete,
         )
+        if shim_predicate is not None:
+            self.is_retryable = lambda exc: shim_predicate(exc) and base_predicate(exc)
+        else:
+            self.is_retryable = base_predicate
+
         self._operation = lambda: tracked_retry(
             retry_fn=CrossSync.retry_target,
             operation=metric,
             target=self._run_attempt,
             predicate=self.is_retryable,
             timeout=operation_timeout,
+            on_error=shim_on_error,
         )
         # initialize state
         self.timeout_generator = _attempt_timeout_generator(

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

@@ -46,11 +46,16 @@
 
 def _track_retryable_error(
     operation: ActiveOperationMetric,
+    user_on_error: Optional[Callable[[Exception], None]] = None,
 ) -> Callable[[Exception], None]:
     """
     Used as input to api_core.Retry classes, to track when retryable errors are encountered
 
     Should be passed as on_error callback
+
+    Args:
+        operation: Active operation metric tracking the retry loop.
+        user_on_error: Optional callback to invoke when an error is encountered.
     """
 
     def wrapper(exc: Exception) -> None:
@@ -72,16 +77,24 @@ def wrapper(exc: Exception) -> None:
         else:
             operation.end_attempt_with_status(exc)
 
+        if user_on_error is not None:
+            user_on_error(exc)
+
     return wrapper
 
 
 def _track_terminal_error(
-    operation: ActiveOperationMetric, exception_factory: ExceptionFactoryType
+    operation: ActiveOperationMetric,
+    exception_factory: ExceptionFactoryType,
 ) -> ExceptionFactoryType:
     """
     Used as input to api_core.Retry classes, to track when terminal errors are encountered
 
     Should be used as a wrapper over an exception_factory callback
+
+    Args:
+        operation: Active operation metric tracking the retry loop.
+        exception_factory: Callback used to build the terminal exception.
     """
 
     def wrapper(
@@ -122,15 +135,20 @@ def tracked_retry(
     **kwargs,
 ) -> T:
     """
-    Wrapper for retry_rarget or retry_target_stream, which injects methods to
+    Wrapper for retry_target or retry_target_stream, which injects methods to
     track the lifecycle of the retry using the provided ActiveOperationMetric
+
+    Args:
+        retry_fn: The retry function to invoke (retry_target or retry_target_stream).
+        operation: Active operation metric tracking the retry loop.
+        **kwargs: Keyword arguments passed to retry_fn (predicate, timeout, on_error, etc).
     """
     in_exception_factory = kwargs.pop("exception_factory", _retry_exception_factory)
-    kwargs.pop("on_error", None)
+    user_on_error = kwargs.pop("on_error", None)
     kwargs.pop("sleep_generator", None)
     return retry_fn(
         sleep_generator=operation.backoff_generator,
-        on_error=_track_retryable_error(operation),
+        on_error=_track_retryable_error(operation, user_on_error),
         exception_factory=_track_terminal_error(operation, in_exception_factory),
         **kwargs,
     )

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py

@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Callable, Sequence
 
 from google.api_core import exceptions as core_exceptions
 from google.api_core import retry as retries
@@ -62,6 +62,8 @@ class _MutateRowsOperation:
             If not specified, the request will run until operation_timeout is reached.
         metric: the metric object representing the active operation
         retryable_exceptions: a list of exceptions that should be retried
+        shim_predicate: optional predicate callback, used only to support the legacy client shim.
+        shim_on_error: optional error callback, used only to support the legacy client shim.
     """
 
     def __init__(
@@ -73,6 +75,8 @@ def __init__(
         attempt_timeout: float | None,
         metric: ActiveOperationMetric,
         retryable_exceptions: Sequence[type[Exception]] = (),
+        shim_predicate: Callable[[Exception], bool] | None = None,
+        shim_on_error: Callable[[Exception], None] | None = None,
     ):
         total_mutations = sum((len(entry.mutations) for entry in mutation_entries))
         if total_mutations > _MUTATE_ROWS_REQUEST_MUTATION_LIMIT:
@@ -81,15 +85,20 @@ def __init__(
             )
         self._target = target
         self._gapic_fn = gapic_client.mutate_rows
-        self.is_retryable = retries.if_exception_type(
+        base_predicate = retries.if_exception_type(
             *retryable_exceptions, bt_exceptions._MutateRowsIncomplete
         )
+        if shim_predicate is not None:
+            self.is_retryable = lambda exc: shim_predicate(exc) and base_predicate(exc)
+        else:
+            self.is_retryable = base_predicate
         self._operation = lambda: tracked_retry(
             retry_fn=CrossSync._Sync_Impl.retry_target,
             operation=metric,
             target=self._run_attempt,
             predicate=self.is_retryable,
             timeout=operation_timeout,
+            on_error=shim_on_error,
         )
         self.timeout_generator = _attempt_timeout_generator(
             attempt_timeout, operation_timeout

packages/google-cloud-bigtable/google/cloud/bigtable/table.py

@@ -38,7 +38,13 @@
     MutationsBatcher,
 )
 from google.cloud.bigtable.column_family import ColumnFamily, _gc_rule_from_pb
-from google.cloud.bigtable.data._helpers import TABLE_DEFAULT
+from google.cloud.bigtable.data._helpers import (
+    TABLE_DEFAULT,
+    _get_retryable_errors,
+    _get_timeouts,
+)
+from google.cloud.bigtable.data._metrics import OperationType
+from google.cloud.bigtable.data._sync_autogen._mutate_rows import _MutateRowsOperation
 from google.cloud.bigtable.data.exceptions import (
     MutationsExceptionGroup,
     RetryExceptionGroup,
@@ -747,23 +753,33 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
             timeout = self.mutation_timeout
 
         retryable_errors = RETRYABLE_MUTATION_ERRORS
+        shim_predicate = None
 
-        # The data client cannot take in zero or null values for deadline, so we set it to
-        # the default if that is the case.
         if retry is None:
             operation_timeout = TABLE_DEFAULT.MUTATE_ROWS
             retryable_errors = []
-        elif retry.deadline is None:
+        # The data client cannot take in zero or null values for deadline, so we set it to
+        # the default if that is the case.
+        elif getattr(retry, "deadline", None) is None:
             operation_timeout = TABLE_DEFAULT.MUTATE_ROWS
-
         # To adhere to the retry strategy of do-nothing being achievable with a deadline
         # of 0.0, we modify the retryable errors to be empty if such a deadline is passed.
-        elif retry.deadline == 0:
+        elif getattr(retry, "deadline", None) == 0:
             operation_timeout = TABLE_DEFAULT.MUTATE_ROWS
             retryable_errors = []
         else:
             operation_timeout = retry.deadline
 
+        if (
+            retry is not None
+            and retryable_errors
+            and getattr(retry, "_predicate", None) is not None
+            and retry._predicate is not DEFAULT_RETRY._predicate
+        ):
+            shim_predicate = retry._predicate
+
+        shim_on_error = getattr(retry, "_on_error", None) if retry is not None else None
+
         attempt_timeout = timeout
         mutation_entries = []
         for row in rows:
@@ -777,17 +793,34 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
                     % (row.row_key, row.table.name, self.name)
                 )
             mutation_entries.append(RowMutationEntry(row.row_key, row._get_mutations()))
+
         return_statuses = [
             status_pb2.Status(code=code_pb2.OK) for _ in range(len(mutation_entries))
         ]  # By default, return status OKs for everything
 
+        operation_timeout, attempt_timeout = _get_timeouts(
+            operation_timeout,
+            attempt_timeout
+            if attempt_timeout is not None
+            else TABLE_DEFAULT.MUTATE_ROWS,
+            self._table_impl,
+        )
+        retryable_excs = _get_retryable_errors(retryable_errors, self._table_impl)
+
+        operation = _MutateRowsOperation(
+            self._table_impl.client._gapic_client,
+            self._table_impl,
+            mutation_entries,
+            operation_timeout=operation_timeout,
+            attempt_timeout=attempt_timeout,
+            metric=self._table_impl._create_operation(OperationType.BULK_MUTATE_ROWS),
+            retryable_exceptions=retryable_excs,
+            shim_predicate=shim_predicate,
+            shim_on_error=shim_on_error,
+        )
+
         try:
-            self._table_impl.bulk_mutate_rows(
-                mutation_entries,
-                operation_timeout=operation_timeout,
-                attempt_timeout=attempt_timeout,
-                retryable_errors=retryable_errors,
-            )
+            operation.start()
         except MutationsExceptionGroup as mut_exc_group:
             # We exception handle as follows:
             #

packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py

@@ -86,6 +86,22 @@ def test_metadata_error_ignored(self):
 
         operation.end_attempt_with_status.assert_called_once_with(exc)
 
+    def test_user_on_error_called(self):
+        """should call user_on_error with exception if provided."""
+        from google.cloud.bigtable.data._metrics.tracked_retry import (
+            _track_retryable_error,
+        )
+
+        operation = mock.Mock()
+        user_on_error = mock.Mock()
+        wrapper = _track_retryable_error(operation, user_on_error=user_on_error)
+
+        exc = RuntimeError("test")
+        wrapper(exc)
+
+        operation.end_attempt_with_status.assert_called_once_with(exc)
+        user_on_error.assert_called_once_with(exc)
+
 
 class TestTrackTerminalError:
     def _call_fut(self, operation, factory):
@@ -198,7 +214,7 @@ def test_tracked_retry_wraps_components(self):
                     arg=1,
                 )
 
-                mock_track_retry.assert_called_once_with(operation)
+                mock_track_retry.assert_called_once_with(operation, None)
                 mock_track_terminal.assert_called_once_with(operation, custom_factory)
 
                 retry_fn.assert_called_once_with(
@@ -208,6 +224,31 @@ def test_tracked_retry_wraps_components(self):
                     arg=1,
                 )
 
+    def test_tracked_retry_with_user_on_error(self):
+        """should pass user_on_error to _track_retryable_error."""
+        from google.cloud.bigtable.data._metrics import tracked_retry
+
+        module = sys.modules[tracked_retry.__module__]
+
+        with mock.patch.object(module, "_track_retryable_error") as mock_track_retry:
+            with mock.patch.object(
+                module, "_track_terminal_error"
+            ) as mock_track_terminal:
+                operation = mock.Mock()
+                retry_fn = mock.Mock()
+                custom_factory = mock.Mock()
+                user_on_error = mock.Mock()
+
+                self._call_fut(
+                    retry_fn=retry_fn,
+                    operation=operation,
+                    exception_factory=custom_factory,
+                    on_error=user_on_error,
+                )
+
+                mock_track_retry.assert_called_once_with(operation, user_on_error)
+                mock_track_terminal.assert_called_once_with(operation, custom_factory)
+
     @pytest.mark.parametrize(
         "fn_name,type_verifier",
         [