diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index 6f5e48244d25..189563378925 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -157,6 +157,12 @@ def mutate(self, row): """Add a row to the batch. If the current batch meets one of the size limits, the batch is sent asynchronously. + The row's mutations are copied into the batch right away and the row + itself is left untouched, so passing the same + :class:`~google.cloud.bigtable.row.DirectRow` twice, or committing it + afterwards, sends its mutations again. Use a fresh row per call, or + clear it yourself. + For example: .. literalinclude:: snippets_table.py diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row.py index c0131a80e345..1925c73e2ee5 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row.py @@ -14,7 +14,7 @@ """User-friendly container for Google Cloud Bigtable Row.""" -from google.api_core.exceptions import GoogleAPICallError +from google.api_core.exceptions import DeadlineExceeded, GoogleAPICallError from google.cloud._helpers import ( _datetime_from_microseconds, # type: ignore _microseconds_from_datetime, # type: ignore @@ -24,6 +24,8 @@ from google.cloud.bigtable.data import mutations from google.cloud.bigtable.data import read_modify_write_rules as rmw_rules +from google.cloud.bigtable.data._helpers import _get_status +from google.cloud.bigtable.data.exceptions import RetryExceptionGroup MAX_MUTATIONS = 100000 """The maximum number of mutations that a row can accumulate.""" @@ -475,18 +477,39 @@ def commit(self): try: self._table._table_impl.mutate_row(self.row_key, self._get_mutations()) - return status_pb2.Status(code=code_pb2.OK) + status = status_pb2.Status(code=code_pb2.OK) except GoogleAPICallError as e: - # If the RPC call returns an error, extract the error into a status object, if possible. - return status_pb2.Status( - code=e.grpc_status_code.value[0] - if e.grpc_status_code is not None - else code_pb2.UNKNOWN, - message=e.message, - details=getattr(e, "details", []), - ) - finally: - self.clear() + # On timeout the data client raises a synthetic DeadlineExceeded + # chained with the errors seen on the wire, which would mask the + # real failure, so report the terminal one instead. Only that case + # is unwrapped: otherwise the raised error already is the terminal + # one and the group holds only earlier attempts. The synthetic + # error is the one with no response, as server errors carry one. + cause = e.__cause__ + if ( + isinstance(e, DeadlineExceeded) + and e.response is None + and isinstance(cause, RetryExceptionGroup) + and cause.exceptions + ): + status = _get_status(cause.exceptions[-1]) + else: + # If the RPC call returns an error, extract the error into a status object, if possible. + status = status_pb2.Status( + code=( + e.grpc_status_code.value[0] + if e.grpc_status_code is not None + else code_pb2.UNKNOWN + ), + message=e.message, + details=getattr(e, "details", []), + ) + + # Clear once a status is produced, success or failure. An error that + # cannot be turned into a status propagates instead, and the mutations + # stay on the row so the caller can retry them. + self.clear() + return status def clear(self): """Removes all currently accumulated mutations on the current row. diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 36823c887f24..d86b7acf2ee9 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -783,6 +783,9 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): operation_timeout = retry.deadline attempt_timeout = timeout + # ``rows`` may be any iterable, so keep a reference to each one to + # clear the succeeded rows once the statuses are known. + row_list = [] mutation_entries = [] for row in rows: if not isinstance(row, DirectRow): @@ -794,6 +797,7 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): "Row %s is a part of %s table. Current table: %s" % (row.row_key, row.table.name, self.name) ) + row_list.append(row) mutation_entries.append(RowMutationEntry(row.row_key, row._get_mutations())) try: @@ -813,6 +817,12 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): for _ in range(len(mutation_entries)) ] + # Drop the mutations that were applied. Rows that failed keep theirs so + # that they can be applied again separately. + for row, status in zip(row_list, return_statuses): + if status.code == code_pb2.OK: + row.clear() + return return_statuses def sample_row_keys(self): diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py index 1f26db0a430e..b0100db89f07 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py @@ -479,6 +479,97 @@ def test_direct_row_commit_with_unknown_exception(): assert app_profile_id == call_args.app_profile_id[0] +def _make_row_for_commit(): + """Build a DirectRow holding a single pending mutation.""" + credentials = _make_credentials() + client = _make_client(project="project-id", credentials=credentials, admin=True) + table = _Table("projects/more-stuff", client=client, app_profile_id=None) + row = _make_direct_row(b"row_key", table) + row.set_cell("column_family_id", b"column", b"bytes-value") + return row + + +def _commit_with_operation_error(exc): + """Commit a row whose underlying ``mutate_row`` fails with ``exc``.""" + row = _make_row_for_commit() + + with mock.patch.object(row._table._table_impl, "mutate_row") as mutate_row_mock: + mutate_row_mock.side_effect = exc + return row.commit() + + +def test_direct_row_commit_unwraps_retry_exception_group(): + """On timeout the synthetic DeadlineExceeded is unwrapped, so that a single + row commit reports the same wire error as the equivalent bulk call.""" + from google.api_core.exceptions import ServiceUnavailable + from google.api_core.retry import RetryFailureReason + from google.rpc import code_pb2 + + from google.cloud.bigtable.data._helpers import _retry_exception_factory + + source_exc, _ = _retry_exception_factory( + [ServiceUnavailable("first failure"), ServiceUnavailable("final failure")], + RetryFailureReason.TIMEOUT, + 60.0, + ) + + result = _commit_with_operation_error(source_exc) + + assert result.code == code_pb2.Code.UNAVAILABLE + assert result.message == "final failure" + + +def test_direct_row_commit_reports_terminal_error_when_not_timed_out(): + """When retries stop for any other reason the raised exception already is + the terminal error, so the chained earlier attempts are not unwrapped.""" + from google.api_core.exceptions import PermissionDenied, ServiceUnavailable + from google.api_core.retry import RetryFailureReason + from google.rpc import code_pb2 + + from google.cloud.bigtable.data._helpers import _retry_exception_factory + + source_exc, _ = _retry_exception_factory( + [ServiceUnavailable("earlier attempt"), PermissionDenied("terminal failure")], + RetryFailureReason.NON_RETRYABLE_ERROR, + 60.0, + ) + + result = _commit_with_operation_error(source_exc) + + assert result.code == code_pb2.Code.PERMISSION_DENIED + assert result.message == "terminal failure" + + +def test_direct_row_commit_keeps_deadline_exceeded_returned_by_server(): + """A DeadlineExceeded returned by the server is the terminal error and is + reported as-is, unlike the locally built one.""" + from google.api_core.exceptions import DeadlineExceeded, ServiceUnavailable + from google.rpc import code_pb2 + + from google.cloud.bigtable.data.exceptions import RetryExceptionGroup + + exc = DeadlineExceeded("server deadline exceeded", response=mock.Mock()) + exc.__cause__ = RetryExceptionGroup([ServiceUnavailable("earlier attempt")]) + + result = _commit_with_operation_error(exc) + + assert result.code == code_pb2.Code.DEADLINE_EXCEEDED + assert result.message == "server deadline exceeded" + + +def test_direct_row_commit_keeps_mutations_when_it_raises(): + """An error that cannot be reported as a status escapes to the caller, who + is left with no outcome to act on, so the mutations are kept.""" + row = _make_row_for_commit() + + with mock.patch.object(row._table._table_impl, "mutate_row") as mutate_row_mock: + mutate_row_mock.side_effect = RuntimeError("failed with unspecified exception") + with pytest.raises(RuntimeError): + row.commit() + + assert len(row._get_mutations()) == 1 + + def test_direct_row_commit_with_invalid_argument(): from google.cloud.bigtable_v2.services.bigtable import BigtableClient 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 af1ed4004928..c786170500a9 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 @@ -780,6 +780,91 @@ def test_table_mutate_rows_w_retry_none(): ) +def _make_mutate_rows_table_and_rows(row_keys): + """Build a table and real ``DirectRow`` instances holding mutations. + + ``_MockRow`` fakes out ``_get_mutations``, so it cannot observe clearing. + """ + 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 in row_keys: + row = DirectRow(row_key=row_key) + row.set_cell(FAMILY_NAME, QUALIFIER, VALUE, timestamp=None) + rows.append(row) + + return table, rows + + +def test_table_mutate_rows_clears_succeeded_rows_only(): + """Succeeded rows are cleaned up, failed rows keep their mutations so that + they can be applied again separately.""" + from google.api_core import exceptions as api_exceptions + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + from google.cloud.bigtable.data.mutations import RowMutationEntry + + table, rows = _make_mutate_rows_table_and_rows([ROW_KEY, ROW_KEY_1]) + + with mock.patch.object(table._table_impl, "bulk_mutate_rows") as bulk_mock: + bulk_mock.side_effect = MutationsExceptionGroup( + excs=[ + FailedMutationEntryError( + failed_idx=1, + failed_mutation_entry=RowMutationEntry( + ROW_KEY_1, [mock.MagicMock()] + ), + cause=api_exceptions.InternalServerError("Failure"), + ) + ], + total_entries=2, + ) + statuses = table.mutate_rows(rows) + + assert [status.code for status in statuses] == [SUCCESS, STATUS_INTERNAL] + assert rows[0]._get_mutations() == [] + assert len(rows[1]._get_mutations()) == 1 + + +def test_table_mutate_rows_clears_all_rows_on_success(): + table, rows = _make_mutate_rows_table_and_rows([ROW_KEY, ROW_KEY_1]) + + with mock.patch.object(table._table_impl, "bulk_mutate_rows"): + statuses = table.mutate_rows(rows) + + assert [status.code for status in statuses] == [SUCCESS, SUCCESS] + assert all(row._get_mutations() == [] for row in rows) + + +def test_table_mutate_rows_keeps_mutations_when_call_raises(): + """An error that is not reported per entry leaves every row untouched.""" + table, rows = _make_mutate_rows_table_and_rows([ROW_KEY, ROW_KEY_1]) + + with mock.patch.object(table._table_impl, "bulk_mutate_rows") as bulk_mock: + bulk_mock.side_effect = DeadlineExceeded("deadline exceeded") + with pytest.raises(DeadlineExceeded): + table.mutate_rows(rows) + + assert all(len(row._get_mutations()) == 1 for row in rows) + + +def test_table_mutate_rows_w_iterator(): + """``rows`` may be any iterable, not just a sequence.""" + table, rows = _make_mutate_rows_table_and_rows([ROW_KEY, ROW_KEY_1]) + + with mock.patch.object(table._table_impl, "bulk_mutate_rows"): + statuses = table.mutate_rows(iter(rows)) + + assert [status.code for status in statuses] == [SUCCESS, SUCCESS] + assert all(row._get_mutations() == [] for row in rows) + + def test_table_mutate_rows_wrong_row_type(): import pytest