Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mutianf wants to merge 1 commit into
base: main
Choose a base branch
from
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -117,6 +117,9 @@ def __init__(
self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries]
self.remaining_indices = list(range(len(self.mutations)))
self.errors: dict[int, list[Exception]] = {}
# track which original entries received a response entry from the server,
# so we can detect entries the server never acknowledged (see start())
self._acknowledged_indices: set[int] = set()
# set up metrics
self._operation_metric = metric

Expand All @@ -138,6 +141,24 @@ async def start(self):
for idx in incomplete_indices:
self._handle_entry_error(idx, exc)
finally:
# response completeness check: the server must return one
# response entry for every request entry over the life of the
# operation. If fewer distinct entries were acknowledged than
# were sent, the stream closed without reporting an outcome for
# some mutations; we cannot assume they were applied, so fail
# them instead of silently treating them as successful.
if len(self._acknowledged_indices) != len(self.mutations):
for idx in range(len(self.mutations)):
if idx not in self._acknowledged_indices and (
idx not in self.errors
):
Comment on lines +150 to +154

@gemini-code-assist gemini-code-assist Bot Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Bug: Unacknowledged entries in retried attempts will not be correctly failed with ClientError

There is a subtle bug here when mutations are retried across multiple attempts:

  1. Scenario:

    • A mutation entry at index i fails with a retryable error (e.g., UNAVAILABLE) in Attempt 1.
    • This adds i to self._acknowledged_indices and the retryable error to self.errors[i].
    • Since it's retryable, it is retried in Attempt 2.
    • In Attempt 2, the stream finishes successfully but silently drops/ignores index i (i.e., no response entry is returned for it, and no stream-level exception is raised).

  2. The Issue:

    • In the finally block, len(self._acknowledged_indices) != len(self.mutations) is checked. Since i was acknowledged in Attempt 1, it is already in self._acknowledged_indices. If all other entries were acknowledged, this check will evaluate to False and the completeness check won't run.
    • Even if it runs, idx not in self._acknowledged_indices will be False and idx not in self.errors will be False (due to the error from Attempt 1).
    • Thus, the entry will not be marked with ClientError. Instead, the operation will fail with the obsolete retryable error from Attempt 1, which might cause the client to incorrectly retry the entire batch.

  3. Recommended Solution:
    At the start of each attempt (e.g., at the beginning of _run_attempt), we should clear the state of the entries being retried so that we only consider acknowledgments and errors from the current attempt:

    self._acknowledged_indices.difference_update(self.remaining_indices)
for idx in self.remaining_indices:
    self.errors.pop(idx, None)

self.errors[idx] = [
core_exceptions.ClientError(
"No response entry received for mutation "
"entry; the server acknowledged fewer "
"entries than were sent"
)
]
# raise exception detailing incomplete mutations
all_errors: list[Exception] = []
for idx, exc_list in self.errors.items():
Expand Down Expand Up @@ -194,6 +215,8 @@ async def _run_attempt(self):
for result in result_list.entries:
# convert sub-request index to global index
orig_idx = active_request_indices[result.index]
# record that the server returned a response for this entry
self._acknowledged_indices.add(orig_idx)
entry_error = core_exceptions.from_grpc_status(
result.status.code,
result.status.message,
Expand Down
13 changes: 13 additions & 0 deletions packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,7 @@ def __init__(
self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries]
self.remaining_indices = list(range(len(self.mutations)))
self.errors: dict[int, list[Exception]] = {}
self._acknowledged_indices: set[int] = set()
self._operation_metric = metric

def start(self):
Expand All @@ -112,6 +113,17 @@ def start(self):
for idx in incomplete_indices:
self._handle_entry_error(idx, exc)
finally:
if len(self._acknowledged_indices) != len(self.mutations):
for idx in range(len(self.mutations)):
if (
idx not in self._acknowledged_indices
and idx not in self.errors
):
Comment on lines +116 to +121

@gemini-code-assist gemini-code-assist Bot Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Bug: Unacknowledged entries in retried attempts will not be correctly failed with ClientError

This is the same issue as identified in the async version (_async/_mutate_rows.py). When mutations are retried across multiple attempts, obsolete acknowledgments and errors from previous attempts are not cleared, preventing the completeness check from correctly identifying and failing silently dropped entries in the final attempt.

Recommended Solution:
At the start of each attempt (e.g., at the beginning of _run_attempt), clear the state of the entries being retried:

self._acknowledged_indices.difference_update(self.remaining_indices)
for idx in self.remaining_indices:
    self.errors.pop(idx, None)

self.errors[idx] = [
core_exceptions.ClientError(
"No response entry received for mutation entry; the server acknowledged fewer entries than were sent"
)
]
all_errors: list[Exception] = []
for idx, exc_list in self.errors.items():
if len(exc_list) == 0:
Expand Down Expand Up @@ -159,6 +171,7 @@ def _run_attempt(self):
for result_list in result_generator:
for result in result_list.entries:
orig_idx = active_request_indices[result.index]
self._acknowledged_indices.add(orig_idx)
entry_error = core_exceptions.from_grpc_status(
result.status.code,
result.status.message,
Expand Down
54 changes: 51 additions & 3 deletions packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 without returning a
# response entry for this mutation
continue
code = error_dict.get(idx, 0)
yield MutateRowsResponse(
entries=[
Expand All @@ -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

Expand Down Expand Up @@ -163,6 +168,9 @@ async def test_mutate_rows_operation(self):
instance = self._make_one(
client, table, entries, operation_timeout, operation_timeout, metric
)
# _run_attempt is mocked out, so simulate it acknowledging every
# entry to satisfy the response-completeness check in start()
instance._acknowledged_indices = set(range(len(entries)))
await instance.start()
assert attempt_mock.call_count == 1

Expand Down Expand Up @@ -265,9 +273,49 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type):
metric,
retryable_exceptions=(exc_type,),
)
# _run_attempt is mocked out, so simulate it acknowledging every
# entry to satisfy the response-completeness check in start()
instance._acknowledged_indices = set(range(len(entries)))
await instance.start()
assert attempt_mock.call_count == num_retries + 1

@CrossSync.pytest
async def test_mutate_rows_unacknowledged_entries_fail(self):
"""
If the server closes the stream without returning a response entry for
every request entry, the unacknowledged entries must be surfaced as
failures instead of being silently treated as successful.
"""
from google.cloud.bigtable.data.exceptions import (
FailedMutationEntryError,
MutationsExceptionGroup,
)

mutations = [
self._make_mutation(),
self._make_mutation(),
self._make_mutation(),
]
# server returns responses for indices 0 and 2, but omits index 1
mock_gapic_fn = self._make_mock_gapic(mutations, omit_indices={1})
instance = self._make_one(
mutation_entries=mutations,
)
with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn):
with pytest.raises(MutationsExceptionGroup) as exc_info:
await instance.start()
# only the omitted entry should be reported as failed
assert len(exc_info.value.exceptions) == 1
failed = exc_info.value.exceptions[0]
assert isinstance(failed, FailedMutationEntryError)
assert failed.index == 1
assert "fewer" in str(failed.__cause__)
# a single attempt is enough; omitted entries are not retried
assert mock_gapic_fn.call_count == 1
# acknowledged entries should not be reported as failures
assert 0 not in instance.errors
assert 2 not in instance.errors

@CrossSync.pytest
async def test_mutate_rows_incomplete_ignored(self):
"""
Expand Down
5 changes: 3 additions & 2 deletions packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3349,9 +3349,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, the others succeed;
# the retry then resolves the first entry
mock_gapic.side_effect = [
self._mock_response([DeadlineExceeded("mock")]),
self._mock_response([DeadlineExceeded("mock"), None, None]),
self._mock_response([None]),
]
mutation = mutations.SetCell(
Expand Down
39 changes: 36 additions & 3 deletions packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=[
Expand All @@ -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

Expand Down Expand Up @@ -145,6 +148,7 @@ def test_mutate_rows_operation(self):
instance = self._make_one(
client, table, entries, operation_timeout, operation_timeout, metric
)
instance._acknowledged_indices = set(range(len(entries)))
instance.start()
assert attempt_mock.call_count == 1

Expand Down Expand Up @@ -230,9 +234,38 @@ def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type):
metric,
retryable_exceptions=(exc_type,),
)
instance._acknowledged_indices = set(range(len(entries)))
instance.start()
assert attempt_mock.call_count == num_retries + 1

def test_mutate_rows_unacknowledged_entries_fail(self):
"""If the server closes the stream without returning a response entry for
every request entry, the unacknowledged entries must be surfaced as
failures instead of being silently treated as successful."""
from google.cloud.bigtable.data.exceptions import (
FailedMutationEntryError,
MutationsExceptionGroup,
)

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)
with mock.patch.object(instance, "_gapic_fn", mock_gapic_fn):
with pytest.raises(MutationsExceptionGroup) as exc_info:
instance.start()
assert len(exc_info.value.exceptions) == 1
failed = exc_info.value.exceptions[0]
assert isinstance(failed, FailedMutationEntryError)
assert failed.index == 1
assert "fewer" in str(failed.__cause__)
assert mock_gapic_fn.call_count == 1
assert 0 not in instance.errors
assert 2 not in instance.errors

def test_mutate_rows_incomplete_ignored(self):
"""MutateRowsIncomplete exceptions should not be added to error list"""
from google.api_core.exceptions import DeadlineExceeded
Expand Down
2 changes: 1 addition & 1 deletion packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2817,7 +2817,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(
Expand Down
Loading