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..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 @@ -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, ) 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..ac7220c32f7f 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, @@ -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: # 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..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 @@ -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", [ 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