diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py index 0019e34518fe..1b71d7ec4e4a 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -16,15 +16,19 @@ import atexit import concurrent.futures +import logging import time import warnings from collections import deque -from typing import TYPE_CHECKING, Sequence, cast +from typing import TYPE_CHECKING, Any, Callable, Sequence, cast + +from google.rpc import code_pb2, status_pb2 from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( TABLE_DEFAULT, _get_retryable_errors, + _get_statuses_from_mutations_exception_group, _get_timeouts, ) from google.cloud.bigtable.data._metrics import ActiveOperationMetric, OperationType @@ -54,6 +58,7 @@ # used to make more readable default values _MB_SIZE = 1024 * 1024 +_LOGGER = logging.getLogger(__name__) @CrossSync.convert_class(sync_name="_FlowControl", add_mapping_for_name="_FlowControl") @@ -294,6 +299,9 @@ def __init__( self._newest_exceptions: deque[Exception] = deque( maxlen=self._exception_list_limit ) + self._user_batch_completed_callback: ( + Callable[[list[status_pb2.Status]], Any] | None + ) = None # clean up on program exit atexit.register(self._on_exit) @@ -410,6 +418,7 @@ async def _execute_mutate_rows( list of FailedMutationEntryError objects for mutations that failed. FailedMutationEntryError objects will not contain index information """ + statuses = [status_pb2.Status(code=code_pb2.UNKNOWN) for _ in range(len(batch))] try: operation = CrossSync._MutateRowsOperation( self._target.client._gapic_client, @@ -422,13 +431,26 @@ async def _execute_mutate_rows( ) await operation.start() except MutationsExceptionGroup as e: + statuses = _get_statuses_from_mutations_exception_group(e, len(batch)) + # strip index information from exceptions, since it is not useful in a batch context for subexc in e.exceptions: subexc.index = None return list(e.exceptions) + else: + statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(len(batch))] finally: # mark batch as complete in flow control await self._flow_control.remove_from_flow(batch) + + # Call batch done callback with list of statuses. + if self._user_batch_completed_callback: + try: + self._user_batch_completed_callback(statuses) + except Exception as exc: + _LOGGER.warning( + f"Exception raised in user batch completion callback: {exc}" + ) return [] def _add_exceptions(self, excs: list[Exception]): diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py index ceecf0eecaa1..97a6c06fe984 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py @@ -24,6 +24,7 @@ TYPE_CHECKING, Callable, List, + Optional, Sequence, Tuple, Union, @@ -32,8 +33,12 @@ from google.api_core import exceptions as core_exceptions from google.api_core import retry as retries from google.api_core.retry import RetryFailureReason, exponential_sleep_generator +from google.rpc import code_pb2, status_pb2 -from google.cloud.bigtable.data.exceptions import RetryExceptionGroup +from google.cloud.bigtable.data.exceptions import ( + MutationsExceptionGroup, + RetryExceptionGroup, +) from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery if TYPE_CHECKING: @@ -237,6 +242,66 @@ def _align_timeouts(operation: float, attempt: float | None) -> tuple[float, flo return operation, final_attempt +def _get_statuses_from_mutations_exception_group( + exc_group: MutationsExceptionGroup, batch_size: int +) -> list[status_pb2.Status]: + """ + Helper function that populates a list of Status objects with exception information from + the exception group. + + Args: + exc_group: The exception group from a mutate rows operation + batch_size: How many RowMutationGroups were provided to the batch + Returns: + list[status_pb2.Status]: A list of Status proto objects + """ + # We exception handle as follows: + # + # 1. Each exception in the error group is a FailedMutationEntryError, and its + # cause is either a singular exception or a RetryExceptionGroup consisting of + # multiple exceptions. + # + # 2. In the case of a singular exception, if the error does not have a gRPC status + # code, we return a status code of UNKNOWN. + # + # 3. In the case of a RetryExceptionGroup, we use terminal exception in the exception + # group and process that. + statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(batch_size)] + for error in exc_group.exceptions: + if isinstance(error.index, int) and 0 <= error.index < len(statuses): + cause = error.__cause__ + if isinstance(cause, RetryExceptionGroup): + statuses[error.index] = _get_status(cause.exceptions[-1]) + else: + statuses[error.index] = _get_status(cause) + return statuses + + +def _get_status(exc: Optional[Exception]) -> status_pb2.Status: + """ + Helper function that returns a Status object corresponding to the given exception. + + Args: + exc: An exception to be converted into a Status. + Returns: + status_pb2.Status: A Status proto object. + """ + if ( + isinstance(exc, core_exceptions.GoogleAPICallError) + and exc.grpc_status_code is not None + ): + return status_pb2.Status( # type: ignore[unreachable] + code=exc.grpc_status_code.value[0], + message=exc.message, + details=exc.details, + ) + + return status_pb2.Status( + code=code_pb2.UNKNOWN, + message=str(exc) if exc else "An unknown error has occurred", + ) + + def _validate_timeouts( operation_timeout: float, attempt_timeout: float | None, allow_none: bool = False ): diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py index eb1f0055f5c9..9a3b57d37512 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py @@ -19,15 +19,19 @@ import atexit import concurrent.futures +import logging import time import warnings from collections import deque -from typing import TYPE_CHECKING, Sequence, cast +from typing import TYPE_CHECKING, Any, Callable, Sequence, cast + +from google.rpc import code_pb2, status_pb2 from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( TABLE_DEFAULT, _get_retryable_errors, + _get_statuses_from_mutations_exception_group, _get_timeouts, ) from google.cloud.bigtable.data._metrics import ActiveOperationMetric, OperationType @@ -47,6 +51,7 @@ ) from google.cloud.bigtable.data.mutations import RowMutationEntry _MB_SIZE = 1024 * 1024 +_LOGGER = logging.getLogger(__name__) @CrossSync._Sync_Impl.add_mapping_decorator("_FlowControl") @@ -259,6 +264,9 @@ def __init__( self._newest_exceptions: deque[Exception] = deque( maxlen=self._exception_list_limit ) + self._user_batch_completed_callback: ( + Callable[[list[status_pb2.Status]], Any] | None + ) = None atexit.register(self._on_exit) def _timer_routine(self, interval: float | None) -> None: @@ -355,6 +363,7 @@ def _execute_mutate_rows( list[FailedMutationEntryError]: list of FailedMutationEntryError objects for mutations that failed. FailedMutationEntryError objects will not contain index information""" + statuses = [status_pb2.Status(code=code_pb2.UNKNOWN) for _ in range(len(batch))] try: operation = CrossSync._Sync_Impl._MutateRowsOperation( self._target.client._gapic_client, @@ -367,11 +376,21 @@ def _execute_mutate_rows( ) operation.start() except MutationsExceptionGroup as e: + statuses = _get_statuses_from_mutations_exception_group(e, len(batch)) for subexc in e.exceptions: subexc.index = None return list(e.exceptions) + else: + statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(len(batch))] finally: self._flow_control.remove_from_flow(batch) + if self._user_batch_completed_callback: + try: + self._user_batch_completed_callback(statuses) + except Exception as exc: + _LOGGER.warning( + f"Exception raised in user batch completion callback: {exc}" + ) return [] def _add_exceptions(self, excs: list[Exception]): diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py index 4ef489731166..c55fc2ac72f5 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py @@ -106,9 +106,7 @@ def __init__( self.rows: dict[bytes, PartialRowData] = {} @classmethod - def _from_generator( - cls, generator: Generator[Row, Any, Any] - ) -> PartialRowsData: + def _from_generator(cls, generator: Generator[Row, Any, Any]) -> PartialRowsData: """Internal constructor for Table.read_rows.""" return cls(generator=generator) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 5107d31b80c5..e0b2e27eb1f9 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -20,7 +20,6 @@ from google.api_core.exceptions import ( Aborted, DeadlineExceeded, - GoogleAPICallError, InternalServerError, NotFound, ServiceUnavailable, @@ -38,11 +37,11 @@ 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.exceptions import ( - MutationsExceptionGroup, - RetryExceptionGroup, +from google.cloud.bigtable.data._helpers import ( + TABLE_DEFAULT, + _get_statuses_from_mutations_exception_group, ) +from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup from google.cloud.bigtable.data.mutations import RowMutationEntry from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.encryption_info import EncryptionInfo @@ -783,9 +782,10 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): mutation_entries = [ RowMutationEntry(row.row_key, row._get_mutations()) for row in rows ] - return_statuses = [status_pb2.Status(code=code_pb2.Code.OK)] * len( - mutation_entries - ) # By default, return status OKs for everything + return_statuses = [ + status_pb2.Status(code=code_pb2.UNKNOWN) + for _ in range(len(mutation_entries)) + ] try: self._table_impl.bulk_mutate_rows( @@ -795,41 +795,16 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): retryable_errors=retryable_errors, ) except MutationsExceptionGroup as mut_exc_group: - # We exception handle as follows: - # - # 1. Each exception in the error group is a FailedMutationEntryError, and its - # cause is either a singular exception or a RetryExceptionGroup consisting of - # multiple exceptions. - # - # 2. In the case of a singular exception, if the error does not have a gRPC status - # code, we return a status code of UNKNOWN. - # - # 3. In the case of a RetryExceptionGroup, we use terminal exception in the exception - # group and process that. - for error in mut_exc_group.exceptions: - cause = error.__cause__ - if isinstance(cause, RetryExceptionGroup): - return_statuses[error.index] = self._get_status( - cause.exceptions[-1] - ) - else: - return_statuses[error.index] = self._get_status(cause) - - return return_statuses - - @staticmethod - def _get_status(error): - if isinstance(error, GoogleAPICallError) and error.grpc_status_code is not None: - return status_pb2.Status( - code=error.grpc_status_code.value[0], - message=error.message, - details=error.details, + return_statuses = _get_statuses_from_mutations_exception_group( + mut_exc_group, len(mutation_entries) ) + else: + return_statuses = [ + status_pb2.Status(code=code_pb2.OK) + for _ in range(len(mutation_entries)) + ] - return status_pb2.Status( - code=code_pb2.Code.UNKNOWN, - message=str(error), - ) + return return_statuses def sample_row_keys(self): """Read a sample of row keys in the table. diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py index 4ad84f01dda5..f88cf7719f90 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py @@ -512,6 +512,42 @@ async def test_mutations_batcher_timer_flush(self, client, target, temp_rows): # ensure cell is updated assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value + @pytest.mark.usefixtures("client") + @pytest.mark.usefixtures("target") + @CrossSync.Retry( + predicate=retry.if_exception_type(ClientError), initial=1, maximum=5 + ) + @CrossSync.pytest + async def test_mutations_batcher_completed_callback( + self, client, target, temp_rows + ): + """ + test batcher with batch completed callback. It should be called when the batcher flushes. + """ + import mock + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.mutations import RowMutationEntry + + callback = mock.Mock() + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await self._create_row_and_mutation( + target, temp_rows, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + flush_interval = 0.1 + async with target.mutations_batcher(flush_interval=flush_interval) as batcher: + batcher._user_batch_completed_callback = callback + await batcher.append(bulk_mutation) + await CrossSync.yield_to_event_loop() + assert len(batcher._staged_entries) == 1 + await CrossSync.sleep(flush_interval + 0.1) + assert len(batcher._staged_entries) == 0 + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + # ensure cell is updated + assert (await self._retrieve_cell_value(target, row_key)) == new_value + @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @CrossSync.Retry( diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py index 91844eb9bfce..23f069153eba 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py @@ -395,6 +395,35 @@ def test_mutations_batcher_timer_flush(self, client, target, temp_rows): assert len(batcher._staged_entries) == 0 assert temp_rows.retrieve_cell_value(target, row_key) == new_value + @pytest.mark.usefixtures("client") + @pytest.mark.usefixtures("target") + @CrossSync._Sync_Impl.Retry( + predicate=retry.if_exception_type(ClientError), initial=1, maximum=5 + ) + def test_mutations_batcher_completed_callback(self, client, target, temp_rows): + """test batcher with batch completed callback. It should be called when the batcher flushes.""" + import mock + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.mutations import RowMutationEntry + + callback = mock.Mock() + new_value = uuid.uuid4().hex.encode() + row_key, mutation = self._create_row_and_mutation( + target, temp_rows, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + flush_interval = 0.1 + with target.mutations_batcher(flush_interval=flush_interval) as batcher: + batcher._user_batch_completed_callback = callback + batcher.append(bulk_mutation) + CrossSync._Sync_Impl.yield_to_event_loop() + assert len(batcher._staged_entries) == 1 + CrossSync._Sync_Impl.sleep(flush_interval + 0.1) + assert len(batcher._staged_entries) == 0 + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + assert self._retrieve_cell_value(target, row_key) == new_value + @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @CrossSync._Sync_Impl.Retry( diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py index 2757e8a95cdb..5e5c82e10661 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py @@ -973,8 +973,85 @@ async def test__execute_mutate_rows_returns_errors(self): table.default_mutate_rows_attempt_timeout = 13 table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: + batch = [self._make_mutation(), self._make_mutation()] + result = await instance._execute_mutate_rows(batch, mock.Mock()) + assert len(result) == 2 + assert result[0] == err1 + assert result[1] == err2 + # indices should be set to None + assert result[0].index is None + assert result[1].index is None + + @CrossSync.pytest + async def test__execute_mutate_rows_batch_completed_callback(self): + from google.rpc import code_pb2, status_pb2 + + with mock.patch.object(CrossSync, "_MutateRowsOperation") as mutate_rows: + mutate_rows.return_value = CrossSync.Mock() + start_operation = mutate_rows().start + table = mock.Mock() + table.table_name = "test-table" + table.app_profile_id = "test-app-profile" + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + async with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback batch = [self._make_mutation()] result = await instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + assert start_operation.call_count == 1 + args, kwargs = mutate_rows.call_args + assert args[0] == table.client._gapic_client + assert args[1] == table + assert args[2] == batch + assert kwargs["operation_timeout"] == 17 + assert kwargs["attempt_timeout"] == 13 + assert result == [] + + @CrossSync.pytest + async def test__execute_mutate_rows_batch_completed_callback_errors(self): + from google.api_core import exceptions + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + with mock.patch.object(CrossSync._MutateRowsOperation, "start") as mutate_rows: + err1 = FailedMutationEntryError( + 1, mock.Mock(), exceptions.DataLoss("test error") + ) + err2 = FailedMutationEntryError( + 2, mock.Mock(), exceptions.DataLoss("test error") + ) + mutate_rows.side_effect = MutationsExceptionGroup([err1, err2], 10) + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + async with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [ + self._make_mutation(), + self._make_mutation(), + self._make_mutation(), + ] + result = await instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with( + [ + status_pb2.Status(code=code_pb2.OK), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + ] + ) assert len(result) == 2 assert result[0] == err1 assert result[1] == err2 @@ -982,6 +1059,22 @@ async def test__execute_mutate_rows_returns_errors(self): assert result[0].index is None assert result[1].index is None + @CrossSync.pytest + async def test__execute_mutate_rows_batch_completed_callback_exception(self): + with mock.patch.object(CrossSync, "_MutateRowsOperation") as mutate_rows: + mutate_rows.return_value = CrossSync.Mock() + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock(side_effect=RuntimeError("callback failed")) + async with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [self._make_mutation()] + result = await instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once() + assert result == [] + @CrossSync.pytest async def test__raise_exceptions(self): """Raise exceptions and reset error state""" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py index ad50533fe57c..e658c6ba26a4 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py @@ -854,14 +854,109 @@ def test__execute_mutate_rows_returns_errors(self): table.default_mutate_rows_attempt_timeout = 13 table.default_mutate_rows_retryable_errors = () with self._make_one(table) as instance: + batch = [self._make_mutation(), self._make_mutation()] + result = instance._execute_mutate_rows(batch, mock.Mock()) + assert len(result) == 2 + assert result[0] == err1 + assert result[1] == err2 + assert result[0].index is None + assert result[1].index is None + + def test__execute_mutate_rows_batch_completed_callback(self): + from google.rpc import code_pb2, status_pb2 + + with mock.patch.object( + CrossSync._Sync_Impl, "_MutateRowsOperation" + ) as mutate_rows: + mutate_rows.return_value = CrossSync._Sync_Impl.Mock() + start_operation = mutate_rows().start + table = mock.Mock() + table.table_name = "test-table" + table.app_profile_id = "test-app-profile" + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback batch = [self._make_mutation()] result = instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + assert start_operation.call_count == 1 + args, kwargs = mutate_rows.call_args + assert args[0] == table.client._gapic_client + assert args[1] == table + assert args[2] == batch + assert kwargs["operation_timeout"] == 17 + assert kwargs["attempt_timeout"] == 13 + assert result == [] + + def test__execute_mutate_rows_batch_completed_callback_errors(self): + from google.api_core import exceptions + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + with mock.patch.object( + CrossSync._Sync_Impl._MutateRowsOperation, "start" + ) as mutate_rows: + err1 = FailedMutationEntryError( + 1, mock.Mock(), exceptions.DataLoss("test error") + ) + err2 = FailedMutationEntryError( + 2, mock.Mock(), exceptions.DataLoss("test error") + ) + mutate_rows.side_effect = MutationsExceptionGroup([err1, err2], 10) + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [ + self._make_mutation(), + self._make_mutation(), + self._make_mutation(), + ] + result = instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with( + [ + status_pb2.Status(code=code_pb2.OK), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + ] + ) assert len(result) == 2 assert result[0] == err1 assert result[1] == err2 assert result[0].index is None assert result[1].index is None + def test__execute_mutate_rows_batch_completed_callback_exception(self): + with mock.patch.object( + CrossSync._Sync_Impl, "_MutateRowsOperation" + ) as mutate_rows: + mutate_rows.return_value = CrossSync._Sync_Impl.Mock() + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock(side_effect=RuntimeError("callback failed")) + with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [self._make_mutation()] + result = instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once() + assert result == [] + def test__raise_exceptions(self): """Raise exceptions and reset error state""" from google.cloud.bigtable.data import exceptions diff --git a/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py b/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py index ef65b5f30625..3e44ec7c3d04 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py @@ -16,8 +16,10 @@ import mock import pytest from google.api_core import exceptions as core_exceptions +from google.rpc import code_pb2, status_pb2 import google.cloud.bigtable.data._helpers as _helpers +import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import TABLE_DEFAULT @@ -264,6 +266,92 @@ def test_rst_stream_aware_predicate( assert predicate(exception) is expected_is_retryable +class TestGetStatusesFromMutationsExceptionGroup: + @pytest.mark.parametrize( + "failed_idx,cause_exc,expected_status", + [ + ( + 0, + core_exceptions.DeadlineExceeded( + "Operation timed out after 40 seconds" + ), + status_pb2.Status( + code=code_pb2.DEADLINE_EXCEEDED, + message="Operation timed out after 40 seconds", + ), + ), + ( + 0, + RuntimeError("Something happened"), + status_pb2.Status(code=code_pb2.UNKNOWN, message="Something happened"), + ), + ( + 0, + bt_exceptions.RetryExceptionGroup( + excs=[ + core_exceptions.ServiceUnavailable("Service Unavailable"), + core_exceptions.ServiceUnavailable("Service Unavailable"), + core_exceptions.DeadlineExceeded( + "Operation timed out after 40 seconds" + ), + ] + ), + status_pb2.Status( + code=code_pb2.DEADLINE_EXCEEDED, + message="Operation timed out after 40 seconds", + ), + ), + ( + 0, + bt_exceptions.RetryExceptionGroup( + excs=[ + core_exceptions.ServiceUnavailable("Service Unavailable"), + core_exceptions.ServiceUnavailable("Service Unavailable"), + RuntimeError("Something happened"), + ] + ), + status_pb2.Status(code=code_pb2.UNKNOWN, message="Something happened"), + ), + ( + 0, + None, + status_pb2.Status( + code=code_pb2.UNKNOWN, message="An unknown error has occurred" + ), + ), + ( + 100, + RuntimeError("Something happened"), + status_pb2.Status(code=code_pb2.OK), + ), + ( + None, + RuntimeError("Something happened"), + status_pb2.Status(code=code_pb2.OK), + ), + ], + ) + def test_get_statuses_from_mutations_exception_group( + self, failed_idx, cause_exc, expected_status + ): + mutation_exception_group = bt_exceptions.MutationsExceptionGroup( + excs=[ + bt_exceptions.FailedMutationEntryError( + failed_idx=failed_idx, + failed_mutation_entry=mock.Mock(), + cause=cause_exc, + ) + ], + total_entries=1, + message="Mutations failed.", + ) + + statuses = _helpers._get_statuses_from_mutations_exception_group( + mutation_exception_group, 1 + ) + assert statuses[0] == expected_status + + class TestGetRetryableErrors: @pytest.mark.parametrize( "input_codes,input_table,expected", diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py index 2b07b8191185..bc7f8202d9b3 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py @@ -553,6 +553,7 @@ def _make_cell_pb(value): def test_partial_rows_data_legacy_constructor_fallback(): import warnings + from google.cloud.bigtable.row_data import PartialRowsData read_method = object() @@ -575,6 +576,7 @@ def test_partial_rows_data_legacy_constructor_fallback(): def test_partial_rows_data_keyword_generator(): import warnings + from google.cloud.bigtable.row_data import PartialRowsData generator = _make_generator([]) @@ -589,6 +591,7 @@ def test_partial_rows_data_keyword_generator(): def test_partial_rows_data_deprecated_properties(): import warnings + from google.cloud.bigtable.row_data import PartialRowsData generator = _make_generator([]) @@ -612,8 +615,3 @@ def test_partial_rows_data_deprecated_properties(): assert f"The `{attr}` attribute on `PartialRowsData` is deprecated" in str( warned[0].message ) - - - - -