diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py index 974e450d232b..c54f3c97bfb3 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -214,6 +214,18 @@ async def _run_attempt(self): self._handle_entry_error(idx, exc) # bubble up exception to be handled by retry wrapper raise + # Any entries that were sent but never received a response entry (a + # successfully-closed but incomplete stream) must not be treated as + # successful. Record a retryable error so idempotent entries are retried + # and non-idempotent entries surface as failures instead of being + # silently dropped. + for idx in active_request_indices.values(): + self._handle_entry_error( + idx, + bt_exceptions._MutateRowsIncomplete( + "no response entry received for mutation" + ), + ) # check if attempt succeeded, or needs to be retried if self.remaining_indices: # unfinished work; raise exception to trigger retry diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index 40e19dd85847..ac2c1d9edac3 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -173,6 +173,13 @@ def _run_attempt(self): for idx in active_request_indices.values(): self._handle_entry_error(idx, exc) raise + for idx in active_request_indices.values(): + self._handle_entry_error( + idx, + bt_exceptions._MutateRowsIncomplete( + "no response entry received for mutation" + ), + ) if self.remaining_indices: raise bt_exceptions._MutateRowsIncomplete diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py index 8ff6e42532b4..5a17b2582e00 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py @@ -53,8 +53,13 @@ def _make_mutation(self, count=1, size=1): return mutation @CrossSync.convert - async def _mock_stream(self, mutation_list, error_dict): + async def _mock_stream(self, mutation_list, error_dict, omit_indices=None): + omit_indices = omit_indices or set() for idx, entry in enumerate(mutation_list): + if idx in omit_indices: + # simulate a server that closes the stream OK without returning + # a response entry for this mutation + continue code = error_dict.get(idx, 0) yield MutateRowsResponse( entries=[ @@ -64,12 +69,12 @@ async def _mock_stream(self, mutation_list, error_dict): ] ) - def _make_mock_gapic(self, mutation_list, error_dict=None): + def _make_mock_gapic(self, mutation_list, error_dict=None, omit_indices=None): mock_fn = CrossSync.Mock() if error_dict is None: error_dict = {} mock_fn.side_effect = lambda *args, **kwargs: self._mock_stream( - mutation_list, error_dict + mutation_list, error_dict, omit_indices ) return mock_fn @@ -374,3 +379,47 @@ async def test_run_attempt_partial_success_non_retryable(self): assert len(instance.errors[1]) == 1 assert instance.errors[1][0].grpc_status_code == 300 assert 2 not in instance.errors + + @CrossSync.pytest + async def test_run_attempt_missing_entry_retryable(self): + """If the server closes the stream successfully but omits a response + entry, the unanswered mutation must not be treated as successful. It + should be recorded as a retryable _MutateRowsIncomplete error so + idempotent entries are retried.""" + from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete + + mutations = [ + self._make_mutation(), + self._make_mutation(), + self._make_mutation(), + ] + # server omits the response entry for index 1 + mock_gapic_fn = self._make_mock_gapic(mutations, omit_indices={1}) + instance = self._make_one(mutation_entries=mutations) + instance.is_retryable = lambda x: True + with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): + with pytest.raises(_MutateRowsIncomplete): + await instance._run_attempt() + assert instance.remaining_indices == [1] + assert 0 not in instance.errors + assert len(instance.errors[1]) == 1 + assert isinstance(instance.errors[1][0], _MutateRowsIncomplete) + assert 2 not in instance.errors + + @CrossSync.pytest + async def test_run_attempt_missing_entry_non_retryable(self): + """A missing response entry for a non-retryable mutation is surfaced as + a failure rather than being silently dropped.""" + from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete + + mutations = [self._make_mutation(), self._make_mutation()] + # server omits the response entry for index 0 + mock_gapic_fn = self._make_mock_gapic(mutations, omit_indices={0}) + instance = self._make_one(mutation_entries=mutations) + instance.is_retryable = lambda x: False + with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): + await instance._run_attempt() + assert instance.remaining_indices == [] + assert len(instance.errors[0]) == 1 + assert isinstance(instance.errors[0][0], _MutateRowsIncomplete) + assert 1 not in instance.errors diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index 62ab9b5f96f0..df081761665a 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -1505,6 +1505,25 @@ async def test_call_metadata(self, include_app_profile, fn_name, fn_args, gapic_ transport_mock = mock.MagicMock() rpc_mock = CrossSync.Mock() transport_mock._wrapped_methods.__getitem__.return_value = rpc_mock + if gapic_fn == "mutate_rows": + # An unacknowledged MutateRows entry is now treated as a retryable + # incomplete mutation, so an empty mock stream would retry until the + # operation timeout. Return a success entry for the single mutation + # so the operation completes after a single attempt. + from google.cloud.bigtable_v2.types import MutateRowsResponse + from google.rpc import status_pb2 + + @CrossSync.convert + async def mutate_rows_stream(*args, **kwargs): + yield MutateRowsResponse( + entries=[ + MutateRowsResponse.Entry( + index=0, status=status_pb2.Status(code=0) + ) + ] + ) + + rpc_mock.side_effect = mutate_rows_stream gapic_client = client._gapic_client if CrossSync.is_async: # inner BigtableClient is held as ._client for BigtableAsyncClient @@ -3349,9 +3368,10 @@ async def test_bulk_mutate_error_recovery(self): async with self._make_client(project="project") as client: table = client.get_table("instance", "table") with mock.patch.object(client._gapic_client, "mutate_rows") as mock_gapic: - # fail with a retryable error, then a non-retryable one + # first entry fails with a retryable error (other two succeed), + # then the retried entry succeeds mock_gapic.side_effect = [ - self._mock_response([DeadlineExceeded("mock")]), + self._mock_response([DeadlineExceeded("mock"), None, None]), self._mock_response([None]), ] mutation = mutations.SetCell( diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py index 2fe86a41fef0..e3979f6152fa 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py @@ -52,8 +52,11 @@ def _make_mutation(self, count=1, size=1): mutation.size = lambda: size return mutation - def _mock_stream(self, mutation_list, error_dict): + def _mock_stream(self, mutation_list, error_dict, omit_indices=None): + omit_indices = omit_indices or set() for idx, entry in enumerate(mutation_list): + if idx in omit_indices: + continue code = error_dict.get(idx, 0) yield MutateRowsResponse( entries=[ @@ -63,12 +66,12 @@ def _mock_stream(self, mutation_list, error_dict): ] ) - def _make_mock_gapic(self, mutation_list, error_dict=None): + def _make_mock_gapic(self, mutation_list, error_dict=None, omit_indices=None): mock_fn = CrossSync._Sync_Impl.Mock() if error_dict is None: error_dict = {} mock_fn.side_effect = lambda *args, **kwargs: self._mock_stream( - mutation_list, error_dict + mutation_list, error_dict, omit_indices ) return mock_fn @@ -323,3 +326,43 @@ def test_run_attempt_partial_success_non_retryable(self): assert len(instance.errors[1]) == 1 assert instance.errors[1][0].grpc_status_code == 300 assert 2 not in instance.errors + + def test_run_attempt_missing_entry_retryable(self): + """If the server closes the stream successfully but omits a response + entry, the unanswered mutation must not be treated as successful. It + should be recorded as a retryable _MutateRowsIncomplete error so + idempotent entries are retried.""" + from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete + + mutations = [ + self._make_mutation(), + self._make_mutation(), + self._make_mutation(), + ] + mock_gapic_fn = self._make_mock_gapic(mutations, omit_indices={1}) + instance = self._make_one(mutation_entries=mutations) + instance.is_retryable = lambda x: True + with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): + with pytest.raises(_MutateRowsIncomplete): + instance._run_attempt() + assert instance.remaining_indices == [1] + assert 0 not in instance.errors + assert len(instance.errors[1]) == 1 + assert isinstance(instance.errors[1][0], _MutateRowsIncomplete) + assert 2 not in instance.errors + + def test_run_attempt_missing_entry_non_retryable(self): + """A missing response entry for a non-retryable mutation is surfaced as + a failure rather than being silently dropped.""" + from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete + + mutations = [self._make_mutation(), self._make_mutation()] + mock_gapic_fn = self._make_mock_gapic(mutations, omit_indices={0}) + instance = self._make_one(mutation_entries=mutations) + instance.is_retryable = lambda x: False + with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn): + instance._run_attempt() + assert instance.remaining_indices == [] + assert len(instance.errors[0]) == 1 + assert isinstance(instance.errors[0][0], _MutateRowsIncomplete) + assert 1 not in instance.errors diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index eb954beb5113..0c3c3b49b7b3 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -1222,6 +1222,21 @@ def test_call_metadata(self, include_app_profile, fn_name, fn_args, gapic_fn): transport_mock = mock.MagicMock() rpc_mock = CrossSync._Sync_Impl.Mock() transport_mock._wrapped_methods.__getitem__.return_value = rpc_mock + if gapic_fn == "mutate_rows": + from google.rpc import status_pb2 + + from google.cloud.bigtable_v2.types import MutateRowsResponse + + def mutate_rows_stream(*args, **kwargs): + yield MutateRowsResponse( + entries=[ + MutateRowsResponse.Entry( + index=0, status=status_pb2.Status(code=0) + ) + ] + ) + + rpc_mock.side_effect = mutate_rows_stream gapic_client = client._gapic_client gapic_client._transport = transport_mock gapic_client._is_universe_domain_valid = True @@ -2817,7 +2832,7 @@ def test_bulk_mutate_error_recovery(self): table = client.get_table("instance", "table") with mock.patch.object(client._gapic_client, "mutate_rows") as mock_gapic: mock_gapic.side_effect = [ - self._mock_response([DeadlineExceeded("mock")]), + self._mock_response([DeadlineExceeded("mock"), None, None]), self._mock_response([None]), ] mutation = mutations.SetCell(