From e30bafb053706f3dbcfc5ef05e2920b0473bcff0 Mon Sep 17 00:00:00 2001 From: Kevin Zheng <147537668+gkevinzheng@users.noreply.github.com> Date: Wed, 18 Feb 2026 11:09:59 -0500 Subject: [PATCH 1/2] feat: Reworked MutateRows to use the data client (#1290) --- .../google/cloud/bigtable/table.py | 122 ++++++++++- .../tests/system/v2_client/test_data_api.py | 68 +++--- .../tests/unit/v2_client/test_table.py | 195 ++++++++++++++---- 3 files changed, 315 insertions(+), 70 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 590ce08aeec8..32bc6a91ad54 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -17,6 +17,7 @@ import warnings from typing import Set +<<<<<<< HEAD from google.api_core import timeout from google.api_core.exceptions import ( Aborted, @@ -26,12 +27,21 @@ RetryError, ServiceUnavailable, ) +======= +from google.api_core.exceptions import GoogleAPICallError +from google.api_core.exceptions import Aborted +from google.api_core.exceptions import DeadlineExceeded +from google.api_core.exceptions import NotFound +from google.api_core.exceptions import ServiceUnavailable +from google.api_core.exceptions import InternalServerError +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) from google.api_core.gapic_v1.method import DEFAULT from google.api_core.retry import Retry, if_exception_type from google.cloud._helpers import _to_bytes # type: ignore from google.cloud.bigtable import enums from google.cloud.bigtable.backup import Backup +<<<<<<< HEAD from google.cloud.bigtable.batcher import ( FLUSH_COUNT, MAX_MUTATION_SIZE, @@ -53,6 +63,35 @@ ) from google.cloud.bigtable_admin_v2.types import table as admin_messages_v2_pb2 from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 +======= +from google.cloud.bigtable.column_family import _gc_rule_from_pb +from google.cloud.bigtable.column_family import ColumnFamily +from google.cloud.bigtable.data._helpers import TABLE_DEFAULT +from google.cloud.bigtable.data.exceptions import ( + RetryExceptionGroup, + MutationsExceptionGroup, +) +from google.cloud.bigtable.data.mutations import RowMutationEntry +from google.cloud.bigtable.batcher import MutationsBatcher +from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_MUTATION_SIZE +from google.cloud.bigtable.encryption_info import EncryptionInfo +from google.cloud.bigtable.policy import Policy +from google.cloud.bigtable.row import AppendRow +from google.cloud.bigtable.row import ConditionalRow +from google.cloud.bigtable.row import DirectRow +from google.cloud.bigtable.row_data import PartialRowsData +from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS +from google.cloud.bigtable.row_set import RowSet +from google.cloud.bigtable.row_set import RowRange +from google.cloud.bigtable import enums +from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 +from google.cloud.bigtable.admin import BigtableTableAdminClient +from google.cloud.bigtable.admin.types import table as admin_messages_v2_pb2 +from google.cloud.bigtable.admin.types import ( + bigtable_table_admin as table_admin_messages_v2_pb2, +) +from google.rpc import code_pb2, status_pb2 +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) # Maximum number of mutations in bulk (MutateRowsRequest message): # (https://cloud.google.com/bigtable/docs/reference/data/rpc/ @@ -715,6 +754,9 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): specify a ``retry`` strategy of "do-nothing", a deadline of ``0.0`` can be specified. + If a deadline of ``None`` is specified, the deadline defaults to + a table-default of 600 seconds (10 minutes). + :type rows: list :param rows: List or other iterable of :class:`.DirectRow` instances. @@ -732,18 +774,78 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): :returns: A list of response statuses (`google.rpc.status_pb2.Status`) corresponding to success or failure of each row mutation sent. These will be in the same order as the `rows`. + + :raise: ValueError: If a row entry has no mutations, or too many mutations """ if timeout is DEFAULT: timeout = self.mutation_timeout - retryable_mutate_rows = _RetryableMutateRowsWorker( - self._instance._client, - self.name, - rows, - app_profile_id=self._app_profile_id, - timeout=timeout, + 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.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: + operation_timeout = TABLE_DEFAULT.MUTATE_ROWS + retryable_errors = [] + else: + operation_timeout = retry.deadline + + attempt_timeout = timeout + 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 + + try: + self._table_impl.bulk_mutate_rows( + mutation_entries, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + 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 status_pb2.Status( + code=code_pb2.Code.UNKNOWN, + message=str(error), ) - return retryable_mutate_rows(retry=retry) def sample_row_keys(self): """Read a sample of row keys in the table. @@ -1078,6 +1180,7 @@ def restore(self, new_table_id, cluster_id=None, backup_id=None, backup_name=Non ) +<<<<<<< HEAD class _RetryableMutateRowsWorker(object): """A callable worker that can retry to mutate rows with transient errors. @@ -1205,6 +1308,8 @@ def _do_mutate_retryable_rows(self): return self.responses_statuses +======= +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) class ClusterState(object): """Representation of a Cluster State. @@ -1351,6 +1456,7 @@ def _create_row_request( row_set._update_message_request(message) return message +<<<<<<< HEAD def _compile_mutation_entries(table_name, rows): @@ -1419,3 +1525,5 @@ def _check_row_type(row): raise TypeError( "Bulk processing can not be applied for conditional or append mutations." ) +======= +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) diff --git a/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py index 92d07153a13b..0bfdd3c0b0ea 100644 --- a/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py +++ b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py @@ -275,6 +275,7 @@ def test_table_mutate_rows(data_table, rows_to_delete): assert row2_data.cells[COLUMN_FAMILY_ID1][COL_NAME1][0].value == CELL_VAL4 +<<<<<<< HEAD def _add_test_error_handler(retry): """Overwrites the current on_error function to assert that backoff values are within expected bounds.""" import time @@ -308,6 +309,12 @@ def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): import mock from google.api_core import retry as retries from google.api_core.exceptions import InvalidArgument +======= +def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): + import mock + from google.cloud.bigtable_v2 import MutateRowsResponse + from google.cloud.bigtable.table import DEFAULT_RETRY +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) from google.rpc.code_pb2 import Code from google.rpc.status_pb2 import Status @@ -364,10 +371,7 @@ def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): rows_to_delete.append(row_2) row_2.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - # Testing the default retry - default_retry_copy = copy.copy(DEFAULT_RETRY) - _add_test_error_handler(default_retry_copy) - statuses = data_table.mutate_rows([row, row_2], retry=default_retry_copy) + statuses = data_table.mutate_rows([row, row_2]) assert statuses[0].code == Code.OK assert statuses[1].code == Code.OK @@ -387,28 +391,35 @@ def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): rows_to_delete.append(row_2) row_2.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - # Testing the default retry - default_retry_copy = copy.copy(DEFAULT_RETRY) - _add_test_error_handler(default_retry_copy) - statuses = data_table.mutate_rows([row, row_2], retry=default_retry_copy) + statuses = data_table.mutate_rows([row, row_2]) assert statuses[0].code == Code.OK - assert statuses[1].code == Code.INTERNAL + assert statuses[1].code == Code.DEADLINE_EXCEEDED - # Because of the way the retriable mutate worker class works, unusual things can happen - # when passing in custom retry predicates. - row = data_table.direct_row(ROW_KEY) - rows_to_delete.append(row) + # Retries with deadline 0 should do nothing. + with mock.patch.object( + data_table._instance._client.table_data_client, "mutate_rows" + ) as mutate_mock: + mutate_mock.side_effect = [ + initial_error_response, + followup_error_response, + followup_error_response, + final_success_response, + ] - row_2 = data_table.direct_row(ROW_KEY_ALT) - rows_to_delete.append(row_2) + row = data_table.direct_row(ROW_KEY) + rows_to_delete.append(row) + row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - retry = DEFAULT_RETRY.with_predicate( - retries.if_exception_type(_BigtableRetryableError, InvalidArgument) - ) - _add_test_error_handler(retry) - statuses = data_table.mutate_rows([row, row_2], retry=retry) - assert statuses[0] is None - assert statuses[1] is None + row_2 = data_table.direct_row(ROW_KEY_ALT) + rows_to_delete.append(row_2) + row_2.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) + + do_nothing_retry = DEFAULT_RETRY.with_deadline(0.0) + + statuses = data_table.mutate_rows([row, row_2], retry=do_nothing_retry) + assert statuses[0].code == Code.OK + assert statuses[1].code == Code.INTERNAL + mutate_mock.assert_called_once() def _populate_table( @@ -496,26 +507,29 @@ def test_table_mutate_rows_integers(data_table, rows_to_delete): def test_table_mutate_rows_input_errors(data_table, rows_to_delete): +<<<<<<< HEAD from google.api_core.exceptions import InvalidArgument from google.cloud.bigtable.table import _MAX_BULK_MUTATIONS, TooManyMutationsError +======= + from google.cloud.bigtable.table import _MAX_BULK_MUTATIONS +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) row = data_table.direct_row(ROW_KEY) rows_to_delete.append(row) - # Mutate row with 0 mutations gives an API error from the service, not - # from the client library. - with pytest.raises(InvalidArgument): + # Mutate row with 0 mutations gives a ValueError from the client library. + with pytest.raises(ValueError): data_table.mutate_rows([row]) row.clear() - # Mutate row with >100k mutations gives a TooManyMutationsError from the + # Mutate row with >100k mutations gives a ValueError from the # client library. for _ in range(0, _MAX_BULK_MUTATIONS + 1): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - with pytest.raises(TooManyMutationsError): + with pytest.raises(ValueError): data_table.mutate_rows([row]) 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 25cc35f7ef7d..bfa08b33f9a4 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 @@ -48,6 +48,7 @@ RETRYABLES = (RETRYABLE_1, RETRYABLE_2, RETRYABLE_3) NON_RETRYABLE = StatusCode.CANCELLED.value[0] STATUS_INTERNAL = StatusCode.INTERNAL.value[0] +<<<<<<< HEAD @mock.patch("google.cloud.bigtable.table._MAX_BULK_MUTATIONS", new=3) @@ -147,6 +148,9 @@ def test__check_row_type_w_right_row_type(): row = DirectRow(row_key=b"row_key", table="table") assert not _check_row_type(row) +======= +STATUS_UNKNOWN = StatusCode.UNKNOWN.value[0] +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) def _make_client(*args, **kwargs): @@ -813,11 +817,27 @@ def test_table_read_row_still_partial(): def _table_mutate_rows_helper( - mutation_timeout=None, app_profile_id=None, retry=None, timeout=None + mutation_timeout=None, + app_profile_id=None, + retry=None, + timeout=None, + expected_operation_timeout=None, + expected_attempt_timeout=None, + expected_retryable_errors=None, ): +<<<<<<< HEAD from google.rpc.status_pb2 import Status +======= + from google.api_core import exceptions as api_exceptions + from google.rpc import status_pb2 +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) from google.cloud.bigtable.table import DEFAULT_RETRY + from google.cloud.bigtable.table import RETRYABLE_MUTATION_ERRORS + from google.cloud.bigtable.data.exceptions import FailedMutationEntryError + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.exceptions import RetryExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -830,15 +850,20 @@ def _table_mutate_rows_helper( if app_profile_id is not None: ctor_kwargs["app_profile_id"] = app_profile_id - table = _make_table(TABLE_ID, instance, **ctor_kwargs) + if expected_operation_timeout is None: + expected_operation_timeout = DEFAULT_RETRY.deadline - rows = [mock.MagicMock(), mock.MagicMock()] - response = [Status(code=0), Status(code=1)] - instance_mock = mock.Mock(return_value=response) - klass_mock = mock.patch( - "google.cloud.bigtable.table._RetryableMutateRowsWorker", - new=mock.MagicMock(return_value=instance_mock), - ) + if expected_retryable_errors is None: + expected_retryable_errors = RETRYABLE_MUTATION_ERRORS + + rows = [ + _MockRow(ROW_KEY), + _MockRow(ROW_KEY_1), + _MockRow(ROW_KEY_2), + _MockRow(ROW_KEY_3), + ] + + table = _make_table(TABLE_ID, instance, **ctor_kwargs) call_kwargs = {} @@ -846,29 +871,84 @@ def _table_mutate_rows_helper( call_kwargs["retry"] = retry if timeout is not None: - expected_timeout = call_kwargs["timeout"] = timeout - else: - expected_timeout = mutation_timeout + call_kwargs["timeout"] = timeout + + with mock.patch.object(table._table_impl, "bulk_mutate_rows") as mutate_rows_mock: + # First entry = success + # Second entry = api error + # Third entry = non-api error + # Fourth entry = retryexceptiongroup + mutate_rows_mock.side_effect = MutationsExceptionGroup( + excs=[ + FailedMutationEntryError( + failed_idx=1, + failed_mutation_entry=RowMutationEntry( + ROW_KEY_1, [mock.MagicMock()] + ), + cause=api_exceptions.InternalServerError("Failure"), + ), + FailedMutationEntryError( + failed_idx=2, + failed_mutation_entry=RowMutationEntry( + ROW_KEY_2, [mock.MagicMock()] + ), + cause=ValueError("Invalid argument"), + ), + FailedMutationEntryError( + failed_idx=3, + failed_mutation_entry=RowMutationEntry( + ROW_KEY_3, [mock.MagicMock()] + ), + cause=RetryExceptionGroup( + [ + api_exceptions.InternalServerError("First failure"), + OSError("Out of memory"), + api_exceptions.InternalServerError("Final failure"), + ] + ), + ), + ], + total_entries=4, + ) - with klass_mock: statuses = table.mutate_rows(rows, **call_kwargs) - result = [status.code for status in statuses] - expected_result = [0, 1] - assert result == expected_result - - klass_mock.new.assert_called_once_with( - client, - TABLE_NAME, - rows, - app_profile_id=app_profile_id, - timeout=expected_timeout, - ) + assert statuses == [ + status_pb2.Status( + code=SUCCESS, + message="", + ), + status_pb2.Status( + code=STATUS_INTERNAL, + message="Failure", + ), + status_pb2.Status( + code=STATUS_UNKNOWN, + message="Invalid argument", + ), + status_pb2.Status( + code=STATUS_INTERNAL, + message="Final failure", + ), + ] - if retry is not None: - instance_mock.assert_called_once_with(retry=retry) - else: - instance_mock.assert_called_once_with(retry=DEFAULT_RETRY) + # Check all call args other than mutation_entries + mutate_rows_mock.assert_called_once_with( + mock.ANY, + operation_timeout=expected_operation_timeout, + attempt_timeout=expected_attempt_timeout, + retryable_errors=expected_retryable_errors, + ) + + # Check that mutation entries are in order + mutation_entries = mutate_rows_mock.call_args.args[0] + mutation_entry_keys = [row.row_key for row in mutation_entries] + assert mutation_entry_keys == [ + ROW_KEY, + ROW_KEY_1, + ROW_KEY_2, + ROW_KEY_3, + ] def test_table_mutate_rows_w_default_mutation_timeout_app_profile_id(): @@ -876,8 +956,10 @@ def test_table_mutate_rows_w_default_mutation_timeout_app_profile_id(): def test_table_mutate_rows_w_mutation_timeout(): - mutation_timeout = 123 - _table_mutate_rows_helper(mutation_timeout=mutation_timeout) + mutation_timeout = 50 + _table_mutate_rows_helper( + mutation_timeout=mutation_timeout, expected_attempt_timeout=mutation_timeout + ) def test_table_mutate_rows_w_app_profile_id(): @@ -886,19 +968,49 @@ def test_table_mutate_rows_w_app_profile_id(): def test_table_mutate_rows_w_retry(): + deadline = 456.0 retry = mock.Mock() - _table_mutate_rows_helper(retry=retry) + retry.deadline = deadline + _table_mutate_rows_helper(retry=retry, expected_operation_timeout=deadline) + + +def test_table_mutate_rows_w_zero_deadline_retry(): + from google.cloud.bigtable.data._helpers import TABLE_DEFAULT + + deadline = 0.0 + retry = mock.Mock() + retry.deadline = deadline + _table_mutate_rows_helper( + retry=retry, + expected_operation_timeout=TABLE_DEFAULT.MUTATE_ROWS, + expected_retryable_errors=[], + ) + + +def test_table_mutate_rows_w_none_deadline_retry(): + from google.cloud.bigtable.data._helpers import TABLE_DEFAULT + + deadline = None + retry = mock.Mock() + retry.deadline = deadline + _table_mutate_rows_helper( + retry=retry, expected_operation_timeout=TABLE_DEFAULT.MUTATE_ROWS + ) def test_table_mutate_rows_w_timeout_arg(): - timeout = 123 - _table_mutate_rows_helper(timeout=timeout) + timeout = 40 + _table_mutate_rows_helper(timeout=timeout, expected_attempt_timeout=timeout) def test_table_mutate_rows_w_mutation_timeout_and_timeout_arg(): - mutation_timeout = 123 - timeout = 456 - _table_mutate_rows_helper(mutation_timeout=mutation_timeout, timeout=timeout) + mutation_timeout = 50 + timeout = 100 + _table_mutate_rows_helper( + mutation_timeout=mutation_timeout, + timeout=timeout, + expected_attempt_timeout=timeout, + ) def test_table_read_rows(): @@ -1558,6 +1670,7 @@ def test_table_restore_table_w_backup_name(): _table_restore_helper(backup_name=BACKUP_NAME) +<<<<<<< HEAD def _make_worker(*args, **kwargs): from google.cloud.bigtable.table import _RetryableMutateRowsWorker @@ -2060,6 +2173,8 @@ def test_rmrw_do_mutate_retryable_rows_mismatch_num_responses(): _do_mutate_retryable_rows_helper(row_cells, responses) +======= +>>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) def test__create_row_request_table_name_only(): from google.cloud.bigtable.table import _create_row_request @@ -2281,6 +2396,14 @@ def _ReadRowsResponsePB(*args, **kw): return messages_v2_pb2.ReadRowsResponse(*args, **kw) +class _MockRow(object): + def __init__(self, row_key): + self.row_key = row_key + + def _get_mutations(self): + return [mock.MagicMock()] + + class _MockReadRowsIterator(object): def __init__(self, *values): self.iter_values = iter(values) From 734f5f0dcd803139a6072c46e28eb229effa0ea3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 21 Aug 2026 13:40:23 -0700 Subject: [PATCH 2/2] fix merge conflicts --- .../google/cloud/bigtable/table.py | 251 +------ .../tests/system/v2_client/test_data_api.py | 46 +- .../tests/unit/v2_client/test_table.py | 624 +----------------- 3 files changed, 16 insertions(+), 905 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 32bc6a91ad54..16222408d877 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -17,44 +17,39 @@ import warnings from typing import Set -<<<<<<< HEAD -from google.api_core import timeout from google.api_core.exceptions import ( Aborted, DeadlineExceeded, + GoogleAPICallError, InternalServerError, NotFound, - RetryError, ServiceUnavailable, ) -======= -from google.api_core.exceptions import GoogleAPICallError -from google.api_core.exceptions import Aborted -from google.api_core.exceptions import DeadlineExceeded -from google.api_core.exceptions import NotFound -from google.api_core.exceptions import ServiceUnavailable -from google.api_core.exceptions import InternalServerError ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) from google.api_core.gapic_v1.method import DEFAULT from google.api_core.retry import Retry, if_exception_type from google.cloud._helpers import _to_bytes # type: ignore +from google.rpc import code_pb2, status_pb2 from google.cloud.bigtable import enums from google.cloud.bigtable.backup import Backup -<<<<<<< HEAD from google.cloud.bigtable.batcher import ( FLUSH_COUNT, MAX_MUTATION_SIZE, 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.mutations import RowMutationEntry from google.cloud.bigtable.encryption_info import EncryptionInfo from google.cloud.bigtable.policy import Policy from google.cloud.bigtable.row import AppendRow, ConditionalRow, DirectRow from google.cloud.bigtable.row_data import ( DEFAULT_RETRY_READ_ROWS, PartialRowsData, - _retriable_internal_server_error, ) from google.cloud.bigtable.row_set import RowRange, RowSet from google.cloud.bigtable_admin_v2 import BaseBigtableTableAdminClient @@ -63,35 +58,6 @@ ) from google.cloud.bigtable_admin_v2.types import table as admin_messages_v2_pb2 from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 -======= -from google.cloud.bigtable.column_family import _gc_rule_from_pb -from google.cloud.bigtable.column_family import ColumnFamily -from google.cloud.bigtable.data._helpers import TABLE_DEFAULT -from google.cloud.bigtable.data.exceptions import ( - RetryExceptionGroup, - MutationsExceptionGroup, -) -from google.cloud.bigtable.data.mutations import RowMutationEntry -from google.cloud.bigtable.batcher import MutationsBatcher -from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_MUTATION_SIZE -from google.cloud.bigtable.encryption_info import EncryptionInfo -from google.cloud.bigtable.policy import Policy -from google.cloud.bigtable.row import AppendRow -from google.cloud.bigtable.row import ConditionalRow -from google.cloud.bigtable.row import DirectRow -from google.cloud.bigtable.row_data import PartialRowsData -from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS -from google.cloud.bigtable.row_set import RowSet -from google.cloud.bigtable.row_set import RowRange -from google.cloud.bigtable import enums -from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 -from google.cloud.bigtable.admin import BigtableTableAdminClient -from google.cloud.bigtable.admin.types import table as admin_messages_v2_pb2 -from google.cloud.bigtable.admin.types import ( - bigtable_table_admin as table_admin_messages_v2_pb2, -) -from google.rpc import code_pb2, status_pb2 ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) # Maximum number of mutations in bulk (MutateRowsRequest message): # (https://cloud.google.com/bigtable/docs/reference/data/rpc/ @@ -1180,136 +1146,6 @@ def restore(self, new_table_id, cluster_id=None, backup_id=None, backup_name=Non ) -<<<<<<< HEAD -class _RetryableMutateRowsWorker(object): - """A callable worker that can retry to mutate rows with transient errors. - - This class is a callable that can retry mutating rows that result in - transient errors. After all rows are successful or none of the rows - are retryable, any subsequent call on this callable will be a no-op. - """ - - def __init__(self, client, table_name, rows, app_profile_id=None, timeout=None): - self.client = client - self.table_name = table_name - self.rows = rows - self.app_profile_id = app_profile_id - self.responses_statuses = [None] * len(self.rows) - self.timeout = timeout - - def __call__(self, retry=DEFAULT_RETRY): - """Attempt to mutate all rows and retry rows with transient errors. - - Will retry the rows with transient errors until all rows succeed or - ``deadline`` specified in the `retry` is reached. - - :rtype: list - :returns: A list of response statuses (`google.rpc.status_pb2.Status`) - corresponding to success or failure of each row mutation - sent. These will be in the same order as the ``rows``. - """ - mutate_rows = self._do_mutate_retryable_rows - if retry: - mutate_rows = retry(self._do_mutate_retryable_rows) - - try: - mutate_rows() - except (_BigtableRetryableError, RetryError): - # - _BigtableRetryableError raised when no retry strategy is used - # and a retryable error on a mutation occurred. - # - RetryError raised when retry deadline is reached. - # In both cases, just return current `responses_statuses`. - pass - - return self.responses_statuses - - @staticmethod - def _is_retryable(status): - return status is None or status.code in RETRYABLE_CODES - - def _do_mutate_retryable_rows(self): - """Mutate all the rows that are eligible for retry. - - A row is eligible for retry if it has not been tried or if it resulted - in a transient error in a previous call. - - :rtype: list - :return: The responses statuses, which is a list of - :class:`~google.rpc.status_pb2.Status`. - :raises: One of the following: - - * :exc:`~.table._BigtableRetryableError` if any - row returned a transient error. - * :exc:`RuntimeError` if the number of responses doesn't - match the number of rows that were retried - """ - retryable_rows = [] - index_into_all_rows = [] - for index, status in enumerate(self.responses_statuses): - if self._is_retryable(status): - retryable_rows.append(self.rows[index]) - index_into_all_rows.append(index) - - if not retryable_rows: - # All mutations are either successful or non-retryable now. - return self.responses_statuses - - entries = _compile_mutation_entries(self.table_name, retryable_rows) - data_client = self.client.table_data_client - - kwargs = {} - if self.timeout is not None: - kwargs["timeout"] = timeout.ExponentialTimeout(deadline=self.timeout) - - try: - responses = data_client.mutate_rows( - table_name=self.table_name, - entries=entries, - app_profile_id=self.app_profile_id, - retry=None, - **kwargs, - ) - except RETRYABLE_MUTATION_ERRORS as exc: - # If an exception, considered retryable by `RETRYABLE_MUTATION_ERRORS`, is - # returned from the initial call, consider - # it to be retryable. Wrap as a Bigtable Retryable Error. - # For InternalServerError, it is only retriable if the message is related to RST Stream messages - if _retriable_internal_server_error(exc) or not isinstance( - exc, InternalServerError - ): - raise _BigtableRetryableError - else: - # re-raise the original exception - raise - - num_responses = 0 - num_retryable_responses = 0 - for response in responses: - for entry in response.entries: - num_responses += 1 - index = index_into_all_rows[entry.index] - self.responses_statuses[index] = entry.status - if self._is_retryable(entry.status): - num_retryable_responses += 1 - if entry.status.code == 0: - self.rows[index].clear() - - if len(retryable_rows) != num_responses: - raise RuntimeError( - "Unexpected number of responses", - num_responses, - "Expected", - len(retryable_rows), - ) - - if num_retryable_responses: - raise _BigtableRetryableError - - return self.responses_statuses - - -======= ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) class ClusterState(object): """Representation of a Cluster State. @@ -1456,74 +1292,3 @@ def _create_row_request( row_set._update_message_request(message) return message -<<<<<<< HEAD - - -def _compile_mutation_entries(table_name, rows): - """Create list of mutation entries - - :type table_name: str - :param table_name: The name of the table to write to. - - :type rows: list - :param rows: List or other iterable of :class:`.DirectRow` instances. - - :rtype: List[:class:`data_messages_v2_pb2.MutateRowsRequest.Entry`] - :returns: entries corresponding to the inputs. - :raises: :exc:`~.table.TooManyMutationsError` if the number of mutations is - greater than the max ({}) - """.format(_MAX_BULK_MUTATIONS) - entries = [] - mutations_count = 0 - entry_klass = data_messages_v2_pb2.MutateRowsRequest.Entry - - for row in rows: - _check_row_table_name(table_name, row) - _check_row_type(row) - mutations = row._get_mutation_pbs() - entries.append(entry_klass(row_key=row.row_key, mutations=mutations)) - mutations_count += len(mutations) - - if mutations_count > _MAX_BULK_MUTATIONS: - raise TooManyMutationsError( - "Maximum number of mutations is %s" % (_MAX_BULK_MUTATIONS,) - ) - return entries - - -def _check_row_table_name(table_name, row): - """Checks that a row belongs to a table. - - :type table_name: str - :param table_name: The name of the table. - - :type row: :class:`~google.cloud.bigtable.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.row.Row` - subclasses. - - :raises: :exc:`~.table.TableMismatchError` if the row does not belong to - the table. - """ - if row.table is not None and row.table.name != table_name: - raise TableMismatchError( - "Row %s is a part of %s table. Current table: %s" - % (row.row_key, row.table.name, table_name) - ) - - -def _check_row_type(row): - """Checks that a row is an instance of :class:`.DirectRow`. - - :type row: :class:`~google.cloud.bigtable.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.row.Row` - subclasses. - - :raises: :class:`TypeError ` if the row is not an - instance of DirectRow. - """ - if not isinstance(row, DirectRow): - raise TypeError( - "Bulk processing can not be applied for conditional or append mutations." - ) -======= ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) diff --git a/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py index 0bfdd3c0b0ea..dda1d8bfc15b 100644 --- a/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py +++ b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py @@ -275,50 +275,12 @@ def test_table_mutate_rows(data_table, rows_to_delete): assert row2_data.cells[COLUMN_FAMILY_ID1][COL_NAME1][0].value == CELL_VAL4 -<<<<<<< HEAD -def _add_test_error_handler(retry): - """Overwrites the current on_error function to assert that backoff values are within expected bounds.""" - import time - - curr_time = time.monotonic() - times_triggered = 0 - - # Assert that the retry handler works properly. - def test_error_handler(exc): - nonlocal curr_time, times_triggered - next_time = time.monotonic() - if times_triggered >= 1: - gap = next_time - curr_time - - # Exponential backoff = uniform randomness from 0 to max_gap - max_gap = min( - retry._initial * retry._multiplier**times_triggered, - retry._maximum, - ) - # Allow a small tolerance margin (1.0s) for OS sleep scheduling latency - assert gap <= max_gap + 1.0 - times_triggered += 1 - curr_time = next_time - - retry._on_error = test_error_handler - - -def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): - import copy - - import mock - from google.api_core import retry as retries - from google.api_core.exceptions import InvalidArgument -======= def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): import mock - from google.cloud.bigtable_v2 import MutateRowsResponse - from google.cloud.bigtable.table import DEFAULT_RETRY ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) from google.rpc.code_pb2 import Code from google.rpc.status_pb2 import Status - from google.cloud.bigtable.table import DEFAULT_RETRY, _BigtableRetryableError + from google.cloud.bigtable.table import DEFAULT_RETRY from google.cloud.bigtable_v2 import MutateRowsResponse # Simulate a server error on row 2, and a normal response on row 1, followed by a bunch of error @@ -507,13 +469,7 @@ def test_table_mutate_rows_integers(data_table, rows_to_delete): def test_table_mutate_rows_input_errors(data_table, rows_to_delete): -<<<<<<< HEAD - from google.api_core.exceptions import InvalidArgument - - from google.cloud.bigtable.table import _MAX_BULK_MUTATIONS, TooManyMutationsError -======= from google.cloud.bigtable.table import _MAX_BULK_MUTATIONS ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) row = data_table.direct_row(ROW_KEY) rows_to_delete.append(row) 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 bfa08b33f9a4..5959061ba5d6 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 @@ -48,109 +48,7 @@ RETRYABLES = (RETRYABLE_1, RETRYABLE_2, RETRYABLE_3) NON_RETRYABLE = StatusCode.CANCELLED.value[0] STATUS_INTERNAL = StatusCode.INTERNAL.value[0] -<<<<<<< HEAD - - -@mock.patch("google.cloud.bigtable.table._MAX_BULK_MUTATIONS", new=3) -def test__compile_mutation_entries_w_too_many_mutations(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import ( - TooManyMutationsError, - _compile_mutation_entries, - ) - - table = mock.Mock(name="table", spec=["name"]) - table.name = "table" - rows = [ - DirectRow(row_key=b"row_key", table=table), - DirectRow(row_key=b"row_key_2", table=table), - ] - rows[0].set_cell("cf1", b"c1", 1) - rows[0].set_cell("cf1", b"c1", 2) - rows[1].set_cell("cf1", b"c1", 3) - rows[1].set_cell("cf1", b"c1", 4) - - with pytest.raises(TooManyMutationsError): - _compile_mutation_entries("table", rows) - - -def test__compile_mutation_entries_normal(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _compile_mutation_entries - from google.cloud.bigtable_v2.types import MutateRowsRequest, data - - table = mock.Mock(spec=["name"]) - table.name = "table" - rows = [ - DirectRow(row_key=b"row_key", table=table), - DirectRow(row_key=b"row_key_2"), - ] - rows[0].set_cell("cf1", b"c1", b"1") - rows[1].set_cell("cf1", b"c1", b"2") - - result = _compile_mutation_entries("table", rows) - - entry_1 = MutateRowsRequest.Entry() - entry_1.row_key = b"row_key" - mutations_1 = data.Mutation() - mutations_1.set_cell.family_name = "cf1" - mutations_1.set_cell.column_qualifier = b"c1" - mutations_1.set_cell.timestamp_micros = -1 - mutations_1.set_cell.value = b"1" - entry_1.mutations.append(mutations_1) - - entry_2 = MutateRowsRequest.Entry() - entry_2.row_key = b"row_key_2" - mutations_2 = data.Mutation() - mutations_2.set_cell.family_name = "cf1" - mutations_2.set_cell.column_qualifier = b"c1" - mutations_2.set_cell.timestamp_micros = -1 - mutations_2.set_cell.value = b"2" - entry_2.mutations.append(mutations_2) - assert result == [entry_1, entry_2] - - -def test__check_row_table_name_w_wrong_table_name(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import TableMismatchError, _check_row_table_name - - table = mock.Mock(name="table", spec=["name"]) - table.name = "table" - row = DirectRow(row_key=b"row_key", table=table) - - with pytest.raises(TableMismatchError): - _check_row_table_name("other_table", row) - - -def test__check_row_table_name_w_right_table_name(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _check_row_table_name - - table = mock.Mock(name="table", spec=["name"]) - table.name = "table" - row = DirectRow(row_key=b"row_key", table=table) - - assert not _check_row_table_name("table", row) - - -def test__check_row_type_w_wrong_row_type(): - from google.cloud.bigtable.row import ConditionalRow - from google.cloud.bigtable.table import _check_row_type - - row = ConditionalRow(row_key=b"row_key", table="table", filter_=None) - with pytest.raises(TypeError): - _check_row_type(row) - - -def test__check_row_type_w_right_row_type(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _check_row_type - - row = DirectRow(row_key=b"row_key", table="table") - assert not _check_row_type(row) -======= STATUS_UNKNOWN = StatusCode.UNKNOWN.value[0] ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) def _make_client(*args, **kwargs): @@ -825,19 +723,16 @@ def _table_mutate_rows_helper( expected_attempt_timeout=None, expected_retryable_errors=None, ): -<<<<<<< HEAD - from google.rpc.status_pb2 import Status - -======= from google.api_core import exceptions as api_exceptions from google.rpc import status_pb2 ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) - from google.cloud.bigtable.table import DEFAULT_RETRY - from google.cloud.bigtable.table import RETRYABLE_MUTATION_ERRORS - from google.cloud.bigtable.data.exceptions import FailedMutationEntryError - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.exceptions import RetryExceptionGroup + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + RetryExceptionGroup, + ) from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.table import DEFAULT_RETRY, RETRYABLE_MUTATION_ERRORS credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -1670,511 +1565,6 @@ def test_table_restore_table_w_backup_name(): _table_restore_helper(backup_name=BACKUP_NAME) -<<<<<<< HEAD -def _make_worker(*args, **kwargs): - from google.cloud.bigtable.table import _RetryableMutateRowsWorker - - return _RetryableMutateRowsWorker(*args, **kwargs) - - -def _make_responses_statuses(codes): - from google.rpc.status_pb2 import Status - - response = [Status(code=code) for code in codes] - return response - - -def _make_responses(codes): - from google.rpc.status_pb2 import Status - - from google.cloud.bigtable_v2.types.bigtable import MutateRowsResponse - - entries = [ - MutateRowsResponse.Entry(index=i, status=Status(code=codes[i])) - for i in range(len(codes)) - ] - return MutateRowsResponse(entries=entries) - - -def test_rmrw_callable_empty_rows(): - 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) - gapic_api = _make_gapic_api(client) - gapic_api.mutate_rows.return_value = [] - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - - worker = _make_worker(client, table.name, []) - statuses = worker() - - assert len(statuses) == 0 - - -def test_rmrw_callable_no_retry_strategy(): - from google.cloud.bigtable.row import DirectRow - - # Setup: - # - Mutate 3 rows. - # Action: - # - Attempt to mutate the rows w/o any retry strategy. - # Expectation: - # - Since no retry, should return statuses as they come back. - # - Even if there are retryable errors, no retry attempt is made. - # - State of responses_statuses should be - # [success, retryable, non-retryable] - 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_1 = DirectRow(row_key=b"row_key", table=table) - row_1.set_cell("cf", b"col", b"value1") - row_2 = DirectRow(row_key=b"row_key_2", table=table) - row_2.set_cell("cf", b"col", b"value2") - row_3 = DirectRow(row_key=b"row_key_3", table=table) - row_3.set_cell("cf", b"col", b"value3") - - response_codes = [SUCCESS, RETRYABLE_1, NON_RETRYABLE] - response = _make_responses(response_codes) - - gapic_api = _make_gapic_api(client) - gapic_api.mutate_rows.return_value = [response] - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - worker = _make_worker(client, table.name, [row_1, row_2, row_3]) - - statuses = worker(retry=None) - - result = [status.code for status in statuses] - assert result == response_codes - - gapic_api.mutate_rows.assert_called_once() - - -def test_rmrw_callable_retry(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import DEFAULT_RETRY - - # Setup: - # - Mutate 3 rows. - # Action: - # - Initial attempt will mutate all 3 rows. - # Expectation: - # - First attempt will result in one retryable error. - # - Second attempt will result in success for the retry-ed row. - # - Check MutateRows is called twice. - # - State of responses_statuses should be - # [success, success, non-retryable] - - 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_1 = DirectRow(row_key=b"row_key", table=table) - row_1.set_cell("cf", b"col", b"value1") - row_2 = DirectRow(row_key=b"row_key_2", table=table) - row_2.set_cell("cf", b"col", b"value2") - row_3 = DirectRow(row_key=b"row_key_3", table=table) - row_3.set_cell("cf", b"col", b"value3") - - response_1 = _make_responses([SUCCESS, RETRYABLE_1, NON_RETRYABLE]) - response_2 = _make_responses([SUCCESS]) - gapic_api = _make_gapic_api(client) - gapic_api.mutate_rows.side_effect = [[response_1], [response_2]] - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - worker = _make_worker(client, table.name, [row_1, row_2, row_3]) - retry = DEFAULT_RETRY.with_delay(initial=0.1) - - statuses = worker(retry=retry) - - result = [status.code for status in statuses] - - assert result == [SUCCESS, SUCCESS, NON_RETRYABLE] - - assert client._table_data_client._gapic_client.mutate_rows.call_count == 2 - - -def _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=None, - expected_result=None, - raising_retry=False, - retryable_error=False, - timeout=None, - mutate_rows_side_effect=None, -): - from google.api_core.exceptions import ServiceUnavailable - - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _BigtableRetryableError - from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 - - # Setup: - # - Mutate 2 rows. - # Action: - # - Initial attempt will mutate all 2 rows. - # Expectation: - # - Expect [success, non-retryable] - - 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) - - rows = [] - for row_key, cell_data in row_cells: - row = DirectRow(row_key=row_key, table=table) - row.set_cell(*cell_data) - rows.append(row) - - response = _make_responses(responses) - - gapic_api = _make_gapic_api(client) - if retryable_error: - if mutate_rows_side_effect is not None: - gapic_api.mutate_rows.side_effect = mutate_rows_side_effect - else: - gapic_api.mutate_rows.side_effect = ServiceUnavailable("testing") - else: - if mutate_rows_side_effect is not None: - gapic_api.mutate_rows.side_effect = mutate_rows_side_effect - gapic_api.mutate_rows.return_value = [response] - - worker = _make_worker(client, table.name, rows=rows) - - if prior_statuses is not None: - assert len(prior_statuses) == len(rows) - worker.responses_statuses = _make_responses_statuses(prior_statuses) - - expected_entries = [] - for row, prior_status in zip(rows, worker.responses_statuses): - if prior_status is None or prior_status.code in RETRYABLES: - entry = data_messages_v2_pb2.MutateRowsRequest.Entry( - row_key=row.row_key, - mutations=row._get_mutation_pbs().copy(), # row clears on success - ) - expected_entries.append(entry) - - expected_kwargs = {} - if timeout is not None: - worker.timeout = timeout - expected_kwargs["timeout"] = mock.ANY - - if retryable_error or raising_retry: - with pytest.raises(_BigtableRetryableError): - worker._do_mutate_retryable_rows() - statuses = worker.responses_statuses - else: - statuses = worker._do_mutate_retryable_rows() - - if not retryable_error: - result = [status.code for status in statuses] - - if expected_result is None: - expected_result = responses - - assert result == expected_result - - if len(responses) == 0 and not retryable_error: - gapic_api.mutate_rows.assert_not_called() - else: - gapic_api.mutate_rows.assert_called_once_with( - table_name=table.name, - entries=expected_entries, - app_profile_id=None, - retry=None, - **expected_kwargs, - ) - if timeout is not None: - called = gapic_api.mutate_rows.mock_calls[0] - assert called.kwargs["timeout"]._deadline == timeout - - -def test_rmrw_do_mutate_retryable_rows_empty_rows(): - # - # Setup: - # - No mutated rows. - # Action: - # - No API call made. - # Expectation: - # - No change. - # - row_cells = [] - responses = [] - - _do_mutate_retryable_rows_helper(row_cells, responses) - - -def test_rmrw_do_mutate_retryable_rows_w_timeout(): - # - # Setup: - # - Mutate 2 rows. - # Action: - # - Initial attempt will mutate all 2 rows. - # Expectation: - # - No retryable error codes, so don't expect a raise. - # - State of responses_statuses should be [success, non-retryable]. - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = [SUCCESS, NON_RETRYABLE] - - timeout = 5 # seconds - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - timeout=timeout, - ) - - -def test_rmrw_do_mutate_retryable_rows_w_retryable_error(): - # - # Setup: - # - Mutate 2 rows. - # Action: - # - Initial attempt will mutate all 2 rows. - # Expectation: - # - No retryable error codes, so don't expect a raise. - # - State of responses_statuses should be [success, non-retryable]. - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = () - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - retryable_error=True, - ) - - -def test_rmrw_do_mutate_retryable_rows_w_retryable_error_internal_rst_stream_error(): - # Mutate two rows - # Raise internal server error with RST STREAM error messages - # There should be no error raised and that the request is retried - from google.api_core.exceptions import InternalServerError - - from google.cloud.bigtable.row_data import RETRYABLE_INTERNAL_ERROR_MESSAGES - - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - responses = () - - for retryable_internal_error_message in RETRYABLE_INTERNAL_ERROR_MESSAGES: - for message in [ - retryable_internal_error_message, - retryable_internal_error_message.upper(), - ]: - _do_mutate_retryable_rows_helper( - row_cells, - responses, - retryable_error=True, - mutate_rows_side_effect=InternalServerError(message), - ) - - -def test_rmrw_do_mutate_rows_w_retryable_error_internal_not_retryable(): - # Mutate two rows - # Raise internal server error but not RST STREAM error messages - # mutate_rows should raise Internal Server Error - from google.api_core.exceptions import InternalServerError - - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - responses = () - - with pytest.raises(InternalServerError): - _do_mutate_retryable_rows_helper( - row_cells, - responses, - mutate_rows_side_effect=InternalServerError("Error not retryable."), - ) - - -def test_rmrw_do_mutate_retryable_rows_retry(): - # - # Setup: - # - Mutate 3 rows. - # Action: - # - Initial attempt will mutate all 3 rows. - # Expectation: - # - Second row returns retryable error code, so expect a raise. - # - State of responses_statuses should be - # [success, retryable, non-retryable] - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - (b"row_key_3", ("cf", b"col", b"value3")), - ] - - responses = [SUCCESS, RETRYABLE_1, NON_RETRYABLE] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - raising_retry=True, - ) - - -def test_rmrw_do_mutate_retryable_rows_second_retry(): - # - # Setup: - # - Mutate 4 rows. - # - First try results: - # [success, retryable, non-retryable, retryable] - # Action: - # - Second try should re-attempt the 'retryable' rows. - # Expectation: - # - After second try: - # [success, success, non-retryable, retryable] - # - One of the rows tried second time returns retryable error code, - # so expect a raise. - # - Exception contains response whose index should be '3' even though - # only two rows were retried. - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - (b"row_key_3", ("cf", b"col", b"value3")), - (b"row_key_4", ("cf", b"col", b"value4")), - ] - - responses = [SUCCESS, RETRYABLE_1] - - prior_statuses = [ - SUCCESS, - RETRYABLE_1, - NON_RETRYABLE, - RETRYABLE_2, - ] - - expected_result = [ - SUCCESS, - SUCCESS, - NON_RETRYABLE, - RETRYABLE_1, - ] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=prior_statuses, - expected_result=expected_result, - raising_retry=True, - ) - - -def test_rmrw_do_mutate_retryable_rows_second_try(): - # - # Setup: - # - Mutate 4 rows. - # - First try results: - # [success, retryable, non-retryable, retryable] - # Action: - # - Second try should re-attempt the 'retryable' rows. - # Expectation: - # - After second try: - # [success, non-retryable, non-retryable, success] - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - (b"row_key_3", ("cf", b"col", b"value3")), - (b"row_key_4", ("cf", b"col", b"value4")), - ] - - responses = [NON_RETRYABLE, SUCCESS] - - prior_statuses = [ - SUCCESS, - RETRYABLE_1, - NON_RETRYABLE, - RETRYABLE_2, - ] - - expected_result = [ - SUCCESS, - NON_RETRYABLE, - NON_RETRYABLE, - SUCCESS, - ] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=prior_statuses, - expected_result=expected_result, - ) - - -def test_rmrw_do_mutate_retryable_rows_second_try_no_retryable(): - # - # Setup: - # - Mutate 2 rows. - # - First try results: [success, non-retryable] - # Action: - # - Second try has no row to retry. - # Expectation: - # - After second try: [success, non-retryable] - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = [] # no calls will be made - - prior_statuses = [ - SUCCESS, - NON_RETRYABLE, - ] - - expected_result = [ - SUCCESS, - NON_RETRYABLE, - ] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=prior_statuses, - expected_result=expected_result, - ) - - -def test_rmrw_do_mutate_retryable_rows_mismatch_num_responses(): - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = [SUCCESS] - - with pytest.raises(RuntimeError): - _do_mutate_retryable_rows_helper(row_cells, responses) - - -======= ->>>>>>> 8e1634b4aa9 (feat: Reworked MutateRows to use the data client (#1290)) def test__create_row_request_table_name_only(): from google.cloud.bigtable.table import _create_row_request