google-cloud-bigtable 2.46.0__tar.gz → 2.48.0__tar.gz

{google_cloud_bigtable-2.46.0/google_cloud_bigtable.egg-info → google_cloud_bigtable-2.48.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-cloud-bigtable
-Version: 2.46.0
+Version: 2.48.0
 Summary: Google Cloud Bigtable API client library
 Home-page: https://github.com/googleapis/google-cloud-python/tree/main/packages/google-cloud-bigtable
 Author: Google LLC

google_cloud_bigtable-2.48.0/google/cloud/bigtable/batcher.py

@@ -0,0 +1,224 @@
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+"""User friendly container for Google Cloud Bigtable MutationBatcher."""
+
+import atexit
+import queue
+
+from google.api_core import exceptions as core_exceptions
+
+from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
+from google.cloud.bigtable.data.mutations import RowMutationEntry
+
+FLUSH_COUNT = 100  # after this many elements, send out the batch
+
+MAX_MUTATION_SIZE = 20 * 1024 * 1024  # 20MB # after this many bytes, send out the batch
+
+MAX_OUTSTANDING_BYTES = 100 * 1024 * 1024  # 100MB # max inflight byte size.
+
+MAX_OUTSTANDING_ELEMENTS = 100000  # max inflight mutations.
+
+
+class MutationsBatchError(Exception):
+    """Error in the batch request"""
+
+    def __init__(self, message, exc):
+        self.exc = exc
+        self.message = message
+        super().__init__(self.message)
+
+
+class MutationsBatcher(object):
+    """A MutationsBatcher is used in batch cases where the number of mutations
+    is large or unknown. It will store :class:`DirectRow` in memory until one of the
+    size limits is reached, the ``flush_interval`` timer fires, or an explicit call
+    to :func:`flush()` is performed. When a flush event occurs, the :class:`DirectRow`
+    in memory will be sent to Cloud Bigtable. Batching mutations is more efficient than
+    sending individual requests.
+
+    This class is not suited for usage in systems where each mutation
+    must be guaranteed to be sent, since calling :func:`mutate()` may only
+    result in an in-memory change. Rows are sent to the service when a size
+    limit is reached, when the ``flush_interval`` timer fires, when :func:`flush()`
+    is called explicitly, or when the batcher is closed (:func:`close()` is also
+    registered to run at interpreter exit). As a result, if the process
+    terminates abruptly -- e.g. a crash, ``SIGKILL``, or ``os._exit`` where the
+    ``atexit`` handler never runs -- any :class:`DirectRow` still buffered in
+    memory is silently dropped and never sent, even after :func:`mutate()`
+    returned.
+
+    Note on thread safety: The same :class:`MutationBatcher` cannot be shared by multiple end-user threads.
+
+    .. warning::
+
+       If using ``MutationsBatcher``, please ensure you are using
+       ``google-cloud-bigtable >= 2.42.0``, or consider migrating to the
+       batcher provided by the synchronous/asynchronous data client
+       (:mod:`google.cloud.bigtable.data`).
+
+    :type table: class
+    :param table: class:`~google.cloud.bigtable.table.Table`.
+
+    :type flush_count: int
+    :param flush_count: (Optional) Max number of rows to flush. If it
+        reaches the max number of rows it calls finish_batch() to mutate the
+        current row batch. Default is FLUSH_COUNT (1000 rows).
+
+    :type max_row_bytes: int
+    :param max_row_bytes: (Optional) Max number of row mutations size to
+        flush. If it reaches the max number of row mutations size it calls
+        finish_batch() to mutate the current row batch. Default is MAX_ROW_BYTES
+        (5 MB).
+
+    :type flush_interval: float
+    :param flush_interval: (Optional) Automatically flush every flush_interval seconds.
+        If None or <= 0, no time-based flushing is performed. Default is 1 second.
+
+    :type batch_completed_callback: Callable[list:[`~google.rpc.status_pb2.Status`]] = None
+    :param batch_completed_callback: (Optional) A callable for handling responses
+        after the current batch is sent. The callable function expect a list of grpc
+        Status.
+    """
+
+    def __init__(
+        self,
+        table,
+        flush_count=FLUSH_COUNT,
+        max_row_bytes=MAX_MUTATION_SIZE,
+        flush_interval=1,
+        batch_completed_callback=None,
+    ):
+        self.table = table
+        self._batcher_kwargs = {
+            "flush_interval": flush_interval,
+            "flush_limit_mutation_count": flush_count,
+            "flush_limit_bytes": max_row_bytes,
+            "flow_control_max_mutation_count": MAX_OUTSTANDING_ELEMENTS,
+            "flow_control_max_bytes": MAX_OUTSTANDING_BYTES,
+        }
+        self._user_batch_completed_callback = batch_completed_callback
+        self._init_batcher()
+        atexit.register(self.close)
+        self._exceptions: queue.Queue = queue.Queue()
+
+    @property
+    def flush_count(self):
+        return self._batcher._flush_limit_count
+
+    @property
+    def max_row_bytes(self):
+        return self._batcher._flush_limit_bytes
+
+    def _init_batcher(self):
+        self._batcher = self.table._table_impl.mutations_batcher(**self._batcher_kwargs)
+        self._batcher._user_batch_completed_callback = (
+            self._user_batch_completed_callback
+        )
+
+    def _close_batcher(self):
+        try:
+            self._batcher.close()
+        except MutationsExceptionGroup as exc_group:
+            for error in exc_group.exceptions:
+                # Unpack the root cause of the FailedMutationEntryError and
+                # return that error to the user. In standard execution paths,
+                # FailedMutationEntryError always has an Exception cause;
+                # defensively wrap in Unknown if __cause__ is None.
+                if error.__cause__ is not None:
+                    self._exceptions.put(error.__cause__)
+                else:
+                    unknown_err = core_exceptions.Unknown(str(error))
+                    unknown_err.__cause__ = error
+                    self._exceptions.put(unknown_err)
+        except Exception as exc:
+            self._exceptions.put(exc)
+
+    def __enter__(self):
+        """Starting the MutationsBatcher as a context manager"""
+        return self
+
+    def mutate(self, row):
+        """Add a row to the batch. If the current batch meets one of the size
+        limits, the batch is sent asynchronously.
+
+        For example:
+
+        .. literalinclude:: snippets_table.py
+            :start-after: [START bigtable_api_batcher_mutate]
+            :end-before: [END bigtable_api_batcher_mutate]
+            :dedent: 4
+
+        :type row: class
+        :param row: :class:`~google.cloud.bigtable.row.DirectRow`.
+
+        :raises: One of the following:
+            * :exc:`~.table._BigtableRetryableError` if any row returned a transient error.
+            * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried
+        """
+        self._batcher.append(RowMutationEntry(row.row_key, row._get_mutations()))
+
+    def mutate_rows(self, rows):
+        """Add multiple rows to the batch. If the current batch meets one of the size
+        limits, the batch is sent asynchronously.
+
+        For example:
+
+        .. literalinclude:: snippets_table.py
+            :start-after: [START bigtable_api_batcher_mutate_rows]
+            :end-before: [END bigtable_api_batcher_mutate_rows]
+            :dedent: 4
+
+        :type rows: list:[`~google.cloud.bigtable.row.DirectRow`]
+        :param rows: list:[`~google.cloud.bigtable.row.DirectRow`].
+
+        :raises: One of the following:
+            * :exc:`~.table._BigtableRetryableError` if any row returned a transient error.
+            * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried
+        """
+        for row in rows:
+            self.mutate(row)
+
+    def flush(self):
+        """Sends the current batch to Cloud Bigtable synchronously.
+        For example:
+
+        .. literalinclude:: snippets_table.py
+            :start-after: [START bigtable_api_batcher_flush]
+            :end-before: [END bigtable_api_batcher_flush]
+            :dedent: 4
+
+        :raises:
+            * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations.
+        """
+        self._close_batcher()
+        self._init_batcher()
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        """Clean up resources. Flush and shutdown the ThreadPoolExecutor."""
+        self.close()
+
+    def close(self):
+        """Clean up resources. Flush and shutdown the ThreadPoolExecutor.
+        Any errors will be raised.
+
+        :raises:
+            * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations.
+        """
+        self._close_batcher()
+        atexit.unregister(self.close)
+        if self._exceptions.qsize() > 0:
+            exc = list(self._exceptions.queue)
+            raise MutationsBatchError("Errors in batch mutations.", exc=exc)

{google_cloud_bigtable-2.46.0 → google_cloud_bigtable-2.48.0}/google/cloud/bigtable/data/_async/_mutate_rows.py

@@ -139,7 +139,7 @@ class _MutateRowsOperationAsync:
                     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(

{google_cloud_bigtable-2.46.0 → google_cloud_bigtable-2.48.0}/google/cloud/bigtable/data/_async/mutations_batcher.py

@@ -16,15 +16,20 @@ from __future__ import annotations
 
 import atexit
 import concurrent.futures
+import inspect
+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 +59,7 @@ __CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen.mutations_batc
 
 # 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 +296,14 @@ class MutationsBatcherAsync:
         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 +387,11 @@ class MutationsBatcherAsync:
             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 +401,7 @@ class MutationsBatcherAsync:
                 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 +424,7 @@ class MutationsBatcherAsync:
                 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 +437,41 @@ class MutationsBatcherAsync:
             )
             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.
+            # This is an internal callback used by the legacy synchronous shim,
+            # though async callbacks are awaited when running in async mode.
+            if self._user_batch_completed_callback:
+                try:
+                    result = self._user_batch_completed_callback(statuses)
+                    if CrossSync.is_async:
+                        if inspect.isawaitable(result):
+                            await result
+                    else:
+                        if inspect.isawaitable(result):
+                            if inspect.iscoroutine(result):
+                                result.close()
+                            raise TypeError(
+                                "_user_batch_completed_callback must be a synchronous callable"
+                            )
+                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 +571,28 @@ class MutationsBatcherAsync:
     @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
@@ -563,5 +605,13 @@ class MutationsBatcherAsync:
                         exc.index = None
                     exceptions.extend(exc_list)
             except Exception as e:
-                exceptions.append(e)
+                exceptions.extend(
+                    [
+                        FailedMutationEntryError(
+                            failed_idx=None, failed_mutation_entry=entry, cause=e
+                        )
+                        for entry in batch
+                    ]
+                )
+
         return exceptions

{google_cloud_bigtable-2.46.0 → google_cloud_bigtable-2.48.0}/google/cloud/bigtable/data/_helpers.py

@@ -24,6 +24,7 @@ from typing import (
     TYPE_CHECKING,
     Callable,
     List,
+    Optional,
     Sequence,
     Tuple,
     Union,
@@ -32,8 +33,12 @@ from typing import (
 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:
@@ -238,6 +243,66 @@ 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)
+        and exc.grpc_status_code is not None
+    ):
+        return status_pb2.Status(  # type: ignore[unreachable]
+            code=exc.grpc_status_code.value[0],
+            message=exc.message,
+            details=exc.details,
+        )
+
+    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
 ):

{google_cloud_bigtable-2.46.0 → google_cloud_bigtable-2.48.0}/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py

@@ -112,7 +112,7 @@ class _MutateRowsOperation:
                 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(

{google_cloud_bigtable-2.46.0 → google_cloud_bigtable-2.48.0}/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py

@@ -19,15 +19,20 @@ from __future__ import annotations
 
 import atexit
 import concurrent.futures
+import inspect
+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
@@ -47,6 +52,7 @@ if TYPE_CHECKING:
     )
     from google.cloud.bigtable.data.mutations import RowMutationEntry
 _MB_SIZE = 1024 * 1024
+_LOGGER = logging.getLogger(__name__)
 
 
 @CrossSync._Sync_Impl.add_mapping_decorator("_FlowControl")
@@ -255,10 +261,13 @@ class MutationsBatcher:
         self._entries_processed_since_last_raise: int = 0
         self._exceptions_since_last_raise: int = 0
         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
         )
+        self._user_batch_completed_callback: (
+            Callable[[list[status_pb2.Status]], Any] | None
+        ) = None
         atexit.register(self._on_exit)
 
     def _timer_routine(self, interval: float | None) -> None:
@@ -326,7 +335,10 @@ class MutationsBatcher:
         Args:
             new_entries list of RowMutationEntry objects to flush"""
         in_process_requests: list[
-            CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
+            tuple[
+                CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]],
+                list[RowMutationEntry],
+            ]
         ] = []
         for batch, metric in self._flow_control.add_to_flow_with_metrics(
             new_entries, self._target.client._metrics
@@ -337,7 +349,7 @@ class MutationsBatcher:
                 metric,
                 sync_executor=self._sync_rpc_executor,
             )
-            in_process_requests.append(batch_task)
+            in_process_requests.append((batch_task, batch))
         found_exceptions = self._wait_for_batch_results(*in_process_requests)
         self._entries_processed_since_last_raise += len(new_entries)
         self._add_exceptions(found_exceptions)
@@ -355,6 +367,7 @@ class MutationsBatcher:
             list[FailedMutationEntryError]:
                 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._Sync_Impl._MutateRowsOperation(
                 self._target.client._gapic_client,
@@ -367,14 +380,30 @@ class MutationsBatcher:
             )
             operation.start()
         except MutationsExceptionGroup as e:
+            statuses = _get_statuses_from_mutations_exception_group(e, len(batch))
             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:
             self._flow_control.remove_from_flow(batch)
+            if self._user_batch_completed_callback:
+                try:
+                    result = self._user_batch_completed_callback(statuses)
+                    if inspect.isawaitable(result):
+                        if inspect.iscoroutine(result):
+                            result.close()
+                        raise TypeError(
+                            "_user_batch_completed_callback must be a synchronous callable"
+                        )
+                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,
         and discard any in between.
@@ -453,24 +482,27 @@ class MutationsBatcher:
 
     @staticmethod
     def _wait_for_batch_results(
-        *tasks: CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
-        | CrossSync._Sync_Impl.Future[None],
-    ) -> list[Exception]:
+        *tasks: tuple[
+            CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
+            | CrossSync._Sync_Impl.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:
             try:
                 exc_list = task.result()
                 if exc_list:
@@ -478,5 +510,12 @@ class MutationsBatcher:
                         exc.index = None
                     exceptions.extend(exc_list)
             except Exception as e:
-                exceptions.append(e)
+                exceptions.extend(
+                    [
+                        FailedMutationEntryError(
+                            failed_idx=None, failed_mutation_entry=entry, cause=e
+                        )
+                        for entry in batch
+                    ]
+                )
         return exceptions