From 8a8616123c69a0dc03204c9d6362a57f4e56402c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 3 Sep 2026 14:23:45 -0700 Subject: [PATCH 1/5] support custom on_error and predicates in shim retries --- .../bigtable/data/_async/_mutate_rows.py | 14 ++- .../bigtable/data/_metrics/tracked_retry.py | 29 ++++-- .../data/_sync_autogen/_mutate_rows.py | 13 ++- .../google/cloud/bigtable/table.py | 70 ++++++++++---- .../unit/data/_metrics/test_tracked_retry.py | 67 +++++++++++++ .../tests/unit/v2_client/test_table.py | 95 ++++++++++++++++++- 6 files changed, 255 insertions(+), 33 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py index 0007447a5505..d7db1aa7a140 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/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( diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py index 749ea04d2f08..5f84bddded73 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py @@ -46,6 +46,7 @@ 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 @@ -72,11 +73,16 @@ 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, + user_on_error: Optional[Callable[[Exception], None]] = None, ) -> ExceptionFactoryType: """ Used as input to api_core.Retry classes, to track when terminal errors are encountered @@ -108,7 +114,10 @@ def wrapper( ): # record ending attempt for timeout failures attempt_exc = exc_list[-1] - _track_retryable_error(operation)(attempt_exc) + if user_on_error is not None: + _track_retryable_error(operation, user_on_error)(attempt_exc) + else: + _track_retryable_error(operation)(attempt_exc) operation.end_with_status(source_exc) return source_exc, cause_exc @@ -122,15 +131,23 @@ 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 """ 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) + if user_on_error is not None: + on_error_fn = _track_retryable_error(operation, user_on_error) + terminal_fn = _track_terminal_error( + operation, in_exception_factory, user_on_error + ) + else: + on_error_fn = _track_retryable_error(operation) + terminal_fn = _track_terminal_error(operation, in_exception_factory) return retry_fn( sleep_generator=operation.backoff_generator, - on_error=_track_retryable_error(operation), - exception_factory=_track_terminal_error(operation, in_exception_factory), + on_error=on_error_fn, + exception_factory=terminal_fn, **kwargs, ) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index 8bb4e49e22eb..44794f8a9113 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/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 diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 097fdb832b43..26077cfeca04 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/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, @@ -85,6 +91,11 @@ class _BigtableRetryableError(Exception): """Retry-able error expected by the default retry strategy.""" +def _never_retry(exc: Exception) -> bool: + """Predicate that never retries any error.""" + return False + + DEFAULT_RETRY = Retry( predicate=if_exception_type(_BigtableRetryableError), initial=1.0, @@ -748,21 +759,25 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): retryable_errors = RETRYABLE_MUTATION_ERRORS - # 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: - 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: + if retry is None or getattr(retry, "deadline", None) == 0: operation_timeout = TABLE_DEFAULT.MUTATE_ROWS retryable_errors = [] + shim_predicate = _never_retry else: - operation_timeout = retry.deadline + operation_timeout = ( + retry.deadline + if retry.deadline is not None + else TABLE_DEFAULT.MUTATE_ROWS + ) + if ( + getattr(retry, "_predicate", None) is not None + and retry._predicate is not DEFAULT_RETRY._predicate + ): + shim_predicate = retry._predicate + else: + shim_predicate = None + + shim_on_error = getattr(retry, "_on_error", None) if retry is not None else None attempt_timeout = timeout mutation_entries = [] @@ -777,17 +792,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: # diff --git a/packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py b/packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py index 55d09c7829c5..9c2fb04888a2 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py +++ b/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): @@ -135,6 +151,30 @@ def test_timeout_active_attempt(self): operation.end_attempt_with_status.assert_called_once_with(last_exc) operation.end_with_status.assert_called_once() + def test_timeout_with_user_on_error(self): + """should call user_on_error if timeout occurs during active attempt.""" + from google.cloud.bigtable.data._metrics import OperationState + from google.cloud.bigtable.data._metrics.tracked_retry import ( + _track_terminal_error, + ) + + operation = mock.Mock() + operation.state = OperationState.ACTIVE_ATTEMPT + factory = mock.Mock() + factory.return_value = (RuntimeError("timeout"), None) + user_on_error = mock.Mock() + + wrapper = _track_terminal_error(operation, factory, user_on_error=user_on_error) + + last_exc = RuntimeError("last attempt error") + exc_list = [last_exc] + + wrapper(exc_list, RetryFailureReason.TIMEOUT, 1.0) + + operation.end_attempt_with_status.assert_called_once_with(last_exc) + user_on_error.assert_called_once_with(last_exc) + operation.end_with_status.assert_called_once() + def test_rpc_error_metadata(self): """should extract and add metadata from GoogleAPICallError in terminal errors.""" operation = mock.Mock() @@ -208,6 +248,33 @@ 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 and _track_terminal_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, user_on_error + ) + @pytest.mark.parametrize( "fn_name,type_verifier", [ diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py index a163f0a2a341..87ed5ab03b4d 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py @@ -731,6 +731,11 @@ def _table_mutate_rows_helper( from google.api_core import exceptions as api_exceptions from google.rpc import status_pb2 + from google.cloud.bigtable.data._helpers import ( + TABLE_DEFAULT, + _get_retryable_errors, + _get_timeouts, + ) from google.cloud.bigtable.data.exceptions import ( FailedMutationEntryError, MutationsExceptionGroup, @@ -765,6 +770,19 @@ def _table_mutate_rows_helper( table = _make_table(TABLE_ID, instance, **ctor_kwargs) + expected_operation_timeout, expected_attempt_timeout = _get_timeouts( + expected_operation_timeout, + ( + expected_attempt_timeout + if expected_attempt_timeout is not None + else TABLE_DEFAULT.MUTATE_ROWS + ), + table._table_impl, + ) + expected_retryable_errors = _get_retryable_errors( + expected_retryable_errors, table._table_impl + ) + call_kwargs = {} if retry is not _DEFAULT_SENTINEL: @@ -773,12 +791,16 @@ def _table_mutate_rows_helper( if timeout is not None: call_kwargs["timeout"] = timeout - with mock.patch.object(table._table_impl, "bulk_mutate_rows") as mutate_rows_mock: + with mock.patch( + "google.cloud.bigtable.table._MutateRowsOperation" + ) as mutate_rows_mock: + op_instance = mock.Mock() + mutate_rows_mock.return_value = op_instance # First entry = success # Second entry = api error # Third entry = non-api error # Fourth entry = retryexceptiongroup - mutate_rows_mock.side_effect = MutationsExceptionGroup( + op_instance.start.side_effect = MutationsExceptionGroup( excs=[ FailedMutationEntryError( failed_idx=1, @@ -834,14 +856,19 @@ def _table_mutate_rows_helper( # Check all call args other than mutation_entries mutate_rows_mock.assert_called_once_with( + table._table_impl.client._gapic_client, + table._table_impl, mock.ANY, operation_timeout=expected_operation_timeout, attempt_timeout=expected_attempt_timeout, - retryable_errors=expected_retryable_errors, + metric=mock.ANY, + retryable_exceptions=expected_retryable_errors, + shim_predicate=mock.ANY, + shim_on_error=mock.ANY, ) # Check that mutation entries are in order - mutation_entries = mutate_rows_mock.call_args.args[0] + mutation_entries = mutate_rows_mock.call_args.args[2] mutation_entry_keys = [row.row_key for row in mutation_entries] assert mutation_entry_keys == [ ROW_KEY, @@ -962,6 +989,66 @@ def test_table_mutate_rows_w_mutation_timeout_and_timeout_arg(): ) +def test_table_mutate_rows_w_retry_on_error(): + from google.cloud.bigtable.row import DirectRow + + on_error_calls = [] + + def on_error(exc): + on_error_calls.append(exc) + + retry = mock.Mock( + deadline=120.0, + _on_error=on_error, + ) + + credentials = _make_credentials() + client = _make_client(project="project-id", credentials=credentials, admin=True) + instance = client.instance(instance_id=INSTANCE_ID) + table = _make_table(TABLE_ID, instance) + row = mock.Mock(spec=DirectRow) + row.table = table + row.row_key = b"row-key" + row._get_mutations.return_value = [mock.MagicMock()] + + with mock.patch("google.cloud.bigtable.table._MutateRowsOperation") as mock_op_cls: + op_instance = mock.Mock() + mock_op_cls.return_value = op_instance + table.mutate_rows([row], retry=retry) + mock_op_cls.assert_called_once() + passed_on_error = mock_op_cls.call_args.kwargs["shim_on_error"] + assert passed_on_error is on_error + + +def test_table_mutate_rows_w_custom_predicate(): + from google.cloud.bigtable.row import DirectRow + + def custom_predicate(exc): + return True + + retry = mock.Mock( + deadline=120.0, + _predicate=custom_predicate, + ) + + credentials = _make_credentials() + client = _make_client(project="project-id", credentials=credentials, admin=True) + instance = client.instance(instance_id=INSTANCE_ID) + table = _make_table(TABLE_ID, instance) + row = mock.Mock(spec=DirectRow) + row.table = table + row.row_key = b"row-key" + row._get_mutations.return_value = [mock.MagicMock()] + + with mock.patch("google.cloud.bigtable.table._MutateRowsOperation") as mock_op_cls: + op_instance = mock.Mock() + mock_op_cls.return_value = op_instance + table.mutate_rows([row], retry=retry) + mock_op_cls.assert_called_once() + passed_predicate = mock_op_cls.call_args.kwargs["shim_predicate"] + assert passed_predicate is custom_predicate + + def test_table_read_rows(): from google.cloud._testing import _Monkey From 78d179384e828cb8fda8fe0df91784dafedb39be Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 3 Sep 2026 14:28:56 -0700 Subject: [PATCH 2/5] removed _never_retry --- .../google/cloud/bigtable/table.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 26077cfeca04..46583729ab25 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -91,11 +91,6 @@ class _BigtableRetryableError(Exception): """Retry-able error expected by the default retry strategy.""" -def _never_retry(exc: Exception) -> bool: - """Predicate that never retries any error.""" - return False - - DEFAULT_RETRY = Retry( predicate=if_exception_type(_BigtableRetryableError), initial=1.0, @@ -758,11 +753,15 @@ 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. + # 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. if retry is None or getattr(retry, "deadline", None) == 0: operation_timeout = TABLE_DEFAULT.MUTATE_ROWS retryable_errors = [] - shim_predicate = _never_retry else: operation_timeout = ( retry.deadline @@ -774,8 +773,6 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): and retry._predicate is not DEFAULT_RETRY._predicate ): shim_predicate = retry._predicate - else: - shim_predicate = None shim_on_error = getattr(retry, "_on_error", None) if retry is not None else None From eb83df9002dde810ef608518d11bf5c1c8057281 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 3 Sep 2026 14:32:40 -0700 Subject: [PATCH 3/5] update docstrings --- .../cloud/bigtable/data/_metrics/tracked_retry.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py index 5f84bddded73..34ef41c87331 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py @@ -52,6 +52,10 @@ def _track_retryable_error( 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: @@ -88,6 +92,11 @@ def _track_terminal_error( 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. + user_on_error: Optional callback to invoke if operation fails due to timeout. """ def wrapper( @@ -133,6 +142,11 @@ def tracked_retry( """ 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) user_on_error = kwargs.pop("on_error", None) From da6f70096a7b70b79a3dd55f65040a380a7947ec Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 3 Sep 2026 14:37:18 -0700 Subject: [PATCH 4/5] don't use callback for terminal error --- .../bigtable/data/_metrics/tracked_retry.py | 19 ++--------- .../unit/data/_metrics/test_tracked_retry.py | 32 ++----------------- 2 files changed, 6 insertions(+), 45 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py index 34ef41c87331..ca0ef1138e02 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_metrics/tracked_retry.py @@ -86,7 +86,6 @@ def wrapper(exc: Exception) -> None: def _track_terminal_error( operation: ActiveOperationMetric, exception_factory: ExceptionFactoryType, - user_on_error: Optional[Callable[[Exception], None]] = None, ) -> ExceptionFactoryType: """ Used as input to api_core.Retry classes, to track when terminal errors are encountered @@ -96,7 +95,6 @@ def _track_terminal_error( Args: operation: Active operation metric tracking the retry loop. exception_factory: Callback used to build the terminal exception. - user_on_error: Optional callback to invoke if operation fails due to timeout. """ def wrapper( @@ -123,10 +121,7 @@ def wrapper( ): # record ending attempt for timeout failures attempt_exc = exc_list[-1] - if user_on_error is not None: - _track_retryable_error(operation, user_on_error)(attempt_exc) - else: - _track_retryable_error(operation)(attempt_exc) + _track_retryable_error(operation)(attempt_exc) operation.end_with_status(source_exc) return source_exc, cause_exc @@ -151,17 +146,9 @@ def tracked_retry( in_exception_factory = kwargs.pop("exception_factory", _retry_exception_factory) user_on_error = kwargs.pop("on_error", None) kwargs.pop("sleep_generator", None) - if user_on_error is not None: - on_error_fn = _track_retryable_error(operation, user_on_error) - terminal_fn = _track_terminal_error( - operation, in_exception_factory, user_on_error - ) - else: - on_error_fn = _track_retryable_error(operation) - terminal_fn = _track_terminal_error(operation, in_exception_factory) return retry_fn( sleep_generator=operation.backoff_generator, - on_error=on_error_fn, - exception_factory=terminal_fn, + on_error=_track_retryable_error(operation, user_on_error), + exception_factory=_track_terminal_error(operation, in_exception_factory), **kwargs, ) diff --git a/packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py b/packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py index 9c2fb04888a2..e72c82769062 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_metrics/test_tracked_retry.py @@ -151,30 +151,6 @@ def test_timeout_active_attempt(self): operation.end_attempt_with_status.assert_called_once_with(last_exc) operation.end_with_status.assert_called_once() - def test_timeout_with_user_on_error(self): - """should call user_on_error if timeout occurs during active attempt.""" - from google.cloud.bigtable.data._metrics import OperationState - from google.cloud.bigtable.data._metrics.tracked_retry import ( - _track_terminal_error, - ) - - operation = mock.Mock() - operation.state = OperationState.ACTIVE_ATTEMPT - factory = mock.Mock() - factory.return_value = (RuntimeError("timeout"), None) - user_on_error = mock.Mock() - - wrapper = _track_terminal_error(operation, factory, user_on_error=user_on_error) - - last_exc = RuntimeError("last attempt error") - exc_list = [last_exc] - - wrapper(exc_list, RetryFailureReason.TIMEOUT, 1.0) - - operation.end_attempt_with_status.assert_called_once_with(last_exc) - user_on_error.assert_called_once_with(last_exc) - operation.end_with_status.assert_called_once() - def test_rpc_error_metadata(self): """should extract and add metadata from GoogleAPICallError in terminal errors.""" operation = mock.Mock() @@ -238,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( @@ -249,7 +225,7 @@ def test_tracked_retry_wraps_components(self): ) def test_tracked_retry_with_user_on_error(self): - """should pass user_on_error to _track_retryable_error and _track_terminal_error.""" + """should pass user_on_error to _track_retryable_error.""" from google.cloud.bigtable.data._metrics import tracked_retry module = sys.modules[tracked_retry.__module__] @@ -271,9 +247,7 @@ def test_tracked_retry_with_user_on_error(self): ) mock_track_retry.assert_called_once_with(operation, user_on_error) - mock_track_terminal.assert_called_once_with( - operation, custom_factory, user_on_error - ) + mock_track_terminal.assert_called_once_with(operation, custom_factory) @pytest.mark.parametrize( "fn_name,type_verifier", From 44f82fe0aec83c64520b3e7d48ac9e52d21372d9 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 3 Sep 2026 14:47:36 -0700 Subject: [PATCH 5/5] cleaned up logic --- .../google/cloud/bigtable/table.py | 26 +++++++++++-------- 1 file changed, 15 insertions(+), 11 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 46583729ab25..ac7220c32f7f 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -755,24 +755,28 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): retryable_errors = RETRYABLE_MUTATION_ERRORS shim_predicate = None + if retry is None: + operation_timeout = TABLE_DEFAULT.MUTATE_ROWS + retryable_errors = [] # 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. - if retry is None or getattr(retry, "deadline", None) == 0: + elif getattr(retry, "deadline", None) == 0: operation_timeout = TABLE_DEFAULT.MUTATE_ROWS retryable_errors = [] else: - operation_timeout = ( - retry.deadline - if retry.deadline is not None - else TABLE_DEFAULT.MUTATE_ROWS - ) - if ( - getattr(retry, "_predicate", None) is not None - and retry._predicate is not DEFAULT_RETRY._predicate - ): - shim_predicate = retry._predicate + 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