googleapis · daniel-sanche · Jan 21, 2026 · Aug 19, 2026 · Aug 21, 2026 · Jan 29, 2026
@@ -139,7 +139,7 @@ async def start(self):
                     self._handle_entry_error(idx, exc)
             finally:
                 # raise exception detailing incomplete mutations
-                all_errors: list[Exception] = []
+                all_errors: list[bt_exceptions.FailedMutationEntryError] = []
                 for idx, exc_list in self.errors.items():
                     if len(exc_list) == 0:
                         raise core_exceptions.ClientError(

@@ -18,12 +18,12 @@
 import time
 from typing import TYPE_CHECKING, Sequence
 
-from google.api_core import retry as retries
 from grpc import StatusCode
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable.data._helpers import (
     _attempt_timeout_generator,
+    _rst_stream_aware_predicate,
 )
 from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 from google.cloud.bigtable.data.exceptions import (
@@ -108,7 +108,7 @@ def __init__(
         else:
             self.request = query._to_pb(target)
         self.target = target
-        self._predicate = retries.if_exception_type(*retryable_exceptions)
+        self._predicate = _rst_stream_aware_predicate(*retryable_exceptions)
         self._last_yielded_row_key: bytes | None = None
         self._remaining_count: int | None = self.request.rows_limit or None
         self._operation_metric = metric

@@ -16,15 +16,19 @@
 
 import atexit
 import concurrent.futures
+import logging
 import time
 import warnings
 from collections import deque
-from typing import TYPE_CHECKING, Sequence, cast
+from typing import TYPE_CHECKING, Any, Callable, Sequence, cast
+
+from google.rpc import code_pb2, status_pb2
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from google.cloud.bigtable.data._helpers import (
     TABLE_DEFAULT,
     _get_retryable_errors,
+    _get_statuses_from_mutations_exception_group,
     _get_timeouts,
 )
 from google.cloud.bigtable.data._metrics import ActiveOperationMetric, OperationType
@@ -54,6 +58,7 @@
 
 # used to make more readable default values
 _MB_SIZE = 1024 * 1024
+_LOGGER = logging.getLogger(__name__)
 
 
 @CrossSync.convert_class(sync_name="_FlowControl", add_mapping_for_name="_FlowControl")
@@ -290,10 +295,14 @@ def __init__(
         self._exceptions_since_last_raise: int = 0
         # keep track of the first and last _exception_list_limit exceptions
         self._exception_list_limit: int = 10
-        self._oldest_exceptions: list[Exception] = []
-        self._newest_exceptions: deque[Exception] = deque(
+        self._oldest_exceptions: list[FailedMutationEntryError] = []
+        self._newest_exceptions: deque[FailedMutationEntryError] = deque(
             maxlen=self._exception_list_limit
         )
+        # only used by the shim right now.
+        self._user_batch_completed_callback: (
+            Callable[[list[status_pb2.Status]], Any] | None
+        ) = None
         # clean up on program exit
         atexit.register(self._on_exit)
 
@@ -377,7 +386,11 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
             new_entries list of RowMutationEntry objects to flush
         """
         # flush new entries
-        in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = []
+        in_process_requests: list[
+            tuple[
+                CrossSync.Future[list[FailedMutationEntryError]], list[RowMutationEntry]
+            ]
+        ] = []
         async for batch, metric in self._flow_control.add_to_flow_with_metrics(
             new_entries, self._target.client._metrics
         ):
@@ -387,7 +400,7 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
                 metric,
                 sync_executor=self._sync_rpc_executor,
             )
-            in_process_requests.append(batch_task)
+            in_process_requests.append((batch_task, batch))
         # wait for all inflight requests to complete
         found_exceptions = await self._wait_for_batch_results(*in_process_requests)
         # update exception data to reflect any new errors
@@ -410,6 +423,7 @@ async def _execute_mutate_rows(
                 list of FailedMutationEntryError objects for mutations that failed.
                 FailedMutationEntryError objects will not contain index information
         """
+        statuses = [status_pb2.Status(code=code_pb2.UNKNOWN) for _ in range(len(batch))]
         try:
             operation = CrossSync._MutateRowsOperation(
                 self._target.client._gapic_client,
@@ -422,16 +436,29 @@ async def _execute_mutate_rows(
             )
             await operation.start()
         except MutationsExceptionGroup as e:
+            statuses = _get_statuses_from_mutations_exception_group(e, len(batch))
+
             # strip index information from exceptions, since it is not useful in a batch context
             for subexc in e.exceptions:
                 subexc.index = None
             return list(e.exceptions)
+        else:
+            statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(len(batch))]
         finally:
             # mark batch as complete in flow control
             await self._flow_control.remove_from_flow(batch)
+
+            # Call batch done callback with list of statuses.
+            if self._user_batch_completed_callback:
+                try:
+                    self._user_batch_completed_callback(statuses)
+                except Exception as exc:
+                    _LOGGER.warning(
+                        f"Exception raised in user batch completion callback: {exc}"
+                    )
         return []
 
-    def _add_exceptions(self, excs: list[Exception]):
+    def _add_exceptions(self, excs: list[FailedMutationEntryError]):
         """
         Add new list of exceptions to internal store. To avoid unbounded memory,
         the batcher will store the first and last _exception_list_limit exceptions,
@@ -531,26 +558,28 @@ def _on_exit(self):
     @staticmethod
     @CrossSync.convert
     async def _wait_for_batch_results(
-        *tasks: CrossSync.Future[list[FailedMutationEntryError]]
-        | CrossSync.Future[None],
-    ) -> list[Exception]:
+        *tasks: tuple[
+            CrossSync.Future[list[FailedMutationEntryError]] | CrossSync.Future[None],
+            list[RowMutationEntry],
+        ],
+    ) -> list[FailedMutationEntryError]:
         """
         Takes in a list of futures representing _execute_mutate_rows tasks,
         waits for them to complete, and returns a list of errors encountered.
 
         Args:
-            *tasks: futures representing _execute_mutate_rows or _flush_internal tasks
+            *tasks: Tuples of futures representing _execute_mutate_rows or
+                    _flush_internal tasks, and their associated batches
         Returns:
-            list[Exception]:
-                list of Exceptions encountered by any of the tasks. Errors are expected
-                to be FailedMutationEntryError, representing a failed mutation operation.
-                If a task fails with a different exception, it will be included in the
-                output list. Successful tasks will not be represented in the output list.
+            list[FailedMutationEntryError]:
+                list of FailedMutationEntryError encountered by any of the tasks,
+                representing a failed mutation operation.
+                Successful tasks will not be represented in the output list.
         """
         if not tasks:
             return []
-        exceptions: list[Exception] = []
-        for task in tasks:
+        exceptions: list[FailedMutationEntryError] = []
+        for task, batch in tasks:
             if CrossSync.is_async:
                 # futures don't need to be awaited in sync mode
                 await task
@@ -562,6 +591,16 @@ async def _wait_for_batch_results(
                         # strip index information
                         exc.index = None
                     exceptions.extend(exc_list)
-            except Exception as e:
+            except FailedMutationEntryError as e:
                 exceptions.append(e)
+            except Exception as e:
+                exceptions.extend(
+                    [
+                        FailedMutationEntryError(
+                            failed_idx=None, failed_mutation_entry=entry, cause=e
+                        )
+                        for entry in batch
+                    ]
+                )
+
         return exceptions
@@ -20,12 +20,26 @@
 import enum
 import time
 from collections import namedtuple
-from typing import TYPE_CHECKING, List, Sequence, Tuple, Union
+from typing import (
+    TYPE_CHECKING,
+    Callable,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    cast,
+)
 
 from google.api_core import exceptions as core_exceptions
+from google.api_core import retry as retries
 from google.api_core.retry import RetryFailureReason, exponential_sleep_generator
+from google.rpc import code_pb2, status_pb2
 
-from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
+from google.cloud.bigtable.data.exceptions import (
+    MutationsExceptionGroup,
+    RetryExceptionGroup,
+)
 from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
 
 if TYPE_CHECKING:
@@ -50,6 +64,15 @@
 # used by every data client as a default project name for testing on Bigtable emulator.
 _DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator"
 
+# Internal error messages that can be retried during ReadRows. Internal error messages with this error
+# text should be treated as Unavailable error messages with the same error text, and will therefore be
+# treated as Unavailable errors rather than Internal errors.
+_RETRYABLE_INTERNAL_ERROR_MESSAGES = (
+    "rst_stream",
+    "rst stream",
+    "received unexpected eos on data frame from server",
+)
+
 # used to identify an active bigtable resource that needs to be warmed through PingAndWarm
 # each instance/app_profile_id pair needs to be individually tracked
 _WarmedInstanceKey = namedtuple(
@@ -126,6 +149,35 @@ def _retry_exception_factory(
     return source_exc, cause_exc
 
 
+def _rst_stream_aware_predicate(
+    *exception_types: type[Exception],
+) -> Callable[[Exception], bool]:
+    """A custom retry predicate.
+
+    This predicate treats Internal error messages with RST_STREAM errors as
+    ServiceUnavailable errors and will retry them if the Unavailable exception is retryable.
+
+    Args:
+        exception_types: Exception types to be retried during operation
+
+    Returns:
+        Callable[[Exception], bool]: A retry predicate that takes in an exception and
+            returns whether or not that exception is retryable
+    """
+    # predicate to check for retryable error types
+    if_exception_type = retries.if_exception_type(*exception_types)
+
+    # special case: treat InternalServerError with rst_stream error message as ServiceUnavailable
+    def rst_check(e):
+        return (
+            core_exceptions.ServiceUnavailable in exception_types
+            and isinstance(e, core_exceptions.InternalServerError)
+            and any(m in e.message.lower() for m in _RETRYABLE_INTERNAL_ERROR_MESSAGES)
+        )
+
+    return lambda e: if_exception_type(e) or rst_check(e)
+
+
 def _get_timeouts(
     operation: float | TABLE_DEFAULT,
     attempt: float | None | TABLE_DEFAULT,
@@ -191,6 +243,69 @@ def _align_timeouts(operation: float, attempt: float | None) -> tuple[float, flo
     return operation, final_attempt
 
 
+def _get_statuses_from_mutations_exception_group(
+    exc_group: MutationsExceptionGroup, batch_size: int
+) -> list[status_pb2.Status]:
+    """
+    Helper function that populates a list of Status objects with exception information from
+    the exception group.
+
+    Args:
+        exc_group: The exception group from a mutate rows operation
+        batch_size: How many RowMutationGroups were provided to the batch
+    Returns:
+        list[status_pb2.Status]: A list of Status proto objects
+    """
+    # We exception handle as follows:
+    #
+    # 1. Each exception in the error group is a FailedMutationEntryError, and its
+    #    cause is either a singular exception or a RetryExceptionGroup consisting of
+    #    multiple exceptions.
+    #
+    # 2. In the case of a singular exception, if the error does not have a gRPC status
+    #    code, we return a status code of UNKNOWN.
+    #
+    # 3. In the case of a RetryExceptionGroup, we use terminal exception in the exception
+    #    group and process that.
+    statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(batch_size)]
+    for error in exc_group.exceptions:
+        if isinstance(error.index, int) and 0 <= error.index < len(statuses):
+            cause = error.__cause__
+            if isinstance(cause, RetryExceptionGroup):
+                statuses[error.index] = _get_status(cause.exceptions[-1])
+            else:
+                statuses[error.index] = _get_status(cause)
+    return statuses
+
+
+def _get_status(exc: Optional[Exception]) -> status_pb2.Status:
+    """
+    Helper function that returns a Status object corresponding to the given exception.
+
+    Args:
+        exc: An exception to be converted into a Status.
+    Returns:
+        status_pb2.Status: A Status proto object.
+    """
+    if isinstance(exc, core_exceptions.GoogleAPICallError):
+        status_code = cast(Optional["grpc.StatusCode"], exc.grpc_status_code)
+        if status_code is not None:
+            return status_pb2.Status(
+                code=status_code.value[0],
+                message=exc.message,
+                details=exc.details,
+            )
+        return status_pb2.Status(
+            code=code_pb2.Code.UNKNOWN,
+            message="An unknown error has occurred",
+        )
+
+    return status_pb2.Status(
+        code=code_pb2.UNKNOWN,
+        message=str(exc) if exc else "An unknown error has occurred",
+    )
+
+
 def _validate_timeouts(
     operation_timeout: float, attempt_timeout: float | None, allow_none: bool = False
 ):

@@ -112,7 +112,7 @@ def start(self):
                 for idx in incomplete_indices:
                     self._handle_entry_error(idx, exc)
             finally:
-                all_errors: list[Exception] = []
+                all_errors: list[bt_exceptions.FailedMutationEntryError] = []
                 for idx, exc_list in self.errors.items():
                     if len(exc_list) == 0:
                         raise core_exceptions.ClientError(

@@ -21,11 +21,13 @@
 import time
 from typing import TYPE_CHECKING, Sequence
 
-from google.api_core import retry as retries
 from grpc import StatusCode
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
-from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
+from google.cloud.bigtable.data._helpers import (
+    _attempt_timeout_generator,
+    _rst_stream_aware_predicate,
+)
 from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry
 from google.cloud.bigtable.data.exceptions import (
     InvalidChunk,
@@ -98,7 +100,7 @@ def __init__(
         else:
             self.request = query._to_pb(target)
         self.target = target
-        self._predicate = retries.if_exception_type(*retryable_exceptions)
+        self._predicate = _rst_stream_aware_predicate(*retryable_exceptions)
         self._last_yielded_row_key: bytes | None = None
         self._remaining_count: int | None = self.request.rows_limit or None
         self._operation_metric = metric