fix(bigtable): data client should acknowledge all mutations in batch

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

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
 

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

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(

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

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(