feat(bigtable): Reroute Mutations Batcher to use data client

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py

@@ -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(

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py

@@ -296,10 +296,11 @@ 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
@@ -386,7 +387,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
         ):
@@ -396,7 +401,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
@@ -466,7 +471,7 @@ async def _execute_mutate_rows(
                     )
         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,
@@ -566,26 +571,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
@@ -598,5 +605,13 @@ async def _wait_for_batch_results(
                         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

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py

@@ -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(

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py

@@ -261,8 +261,8 @@ def __init__(
         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: (
@@ -335,7 +335,10 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]):
         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
@@ -346,7 +349,7 @@ 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))
         found_exceptions = self._wait_for_batch_results(*in_process_requests)
         self._entries_processed_since_last_raise += len(new_entries)
         self._add_exceptions(found_exceptions)
@@ -400,7 +403,7 @@ def _execute_mutate_rows(
                     )
         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.
@@ -479,30 +482,40 @@ def _on_exit(self):
 
     @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:
                     for exc in exc_list:
                         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

packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py

@@ -141,19 +141,18 @@ def __repr__(self):
         return f"{self.__class__.__name__}({message!r}, {self.exceptions!r})"
 
 
-# TODO: When working on mutations batcher, rework exception handling to guarantee that
-# MutationsExceptionGroup only stores FailedMutationEntryErrors.
 class MutationsExceptionGroup(_BigtableExceptionGroup):
     """
     Represents one or more exceptions that occur during a bulk mutation operation
 
-    Exceptions will typically be of type FailedMutationEntryError, but other exceptions may
-    be included if they are raised during the mutation operation
+    Exceptions will be of type FailedMutationEntryError.
     """
 
     @staticmethod
     def _format_message(
-        excs: list[Exception], total_entries: int, exc_count: int | None = None
+        excs: list[FailedMutationEntryError],
+        total_entries: int,
+        exc_count: int | None = None,
     ) -> str:
         """
         Format a message for the exception group
@@ -171,7 +170,10 @@ def _format_message(
         return f"{exc_count} failed {entry_str} from {total_entries} attempted."
 
     def __init__(
-        self, excs: list[Exception], total_entries: int, message: str | None = None
+        self,
+        excs: list[FailedMutationEntryError],
+        total_entries: int,
+        message: str | None = None,
     ):
         """
         Args:
@@ -189,7 +191,10 @@ def __init__(
         self.total_entries_attempted = total_entries
 
     def __new__(
-        cls, excs: list[Exception], total_entries: int, message: str | None = None
+        cls,
+        excs: list[FailedMutationEntryError],
+        total_entries: int,
+        message: str | None = None,
     ):
         """
         Args:
@@ -209,8 +214,8 @@ def __new__(
     @classmethod
     def from_truncated_lists(
         cls,
-        first_list: list[Exception],
-        last_list: list[Exception],
+        first_list: list[FailedMutationEntryError],
+        last_list: list[FailedMutationEntryError],
         total_excs: int,
         entry_count: int,
     ) -> MutationsExceptionGroup: