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

fail V3 mutate_rows entries the server never acknowledged

Author: mutianf

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

@@ -1591,6 +1591,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
@@ -3492,9 +3511,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

@@ -1307,6 +1307,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
@@ -2959,7 +2974,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(