addressed review feedback

Author: gkevinzheng

google/cloud/bigtable/batcher.py

@@ -88,9 +88,13 @@ def __init__(
         batch_completed_callback=None,
     ):
         self.table = table
-        self._flush_count = flush_count
-        self._max_row_bytes = max_row_bytes
-        self._flush_interval = flush_interval
+        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)
@@ -105,11 +109,7 @@ def max_row_bytes(self):
         return self._max_row_bytes
 
     def _init_batcher(self):
-        self._batcher = self.table._table_impl.mutations_batcher(
-            flush_interval=self._flush_interval,
-            flush_limit_mutation_count=self._flush_count,
-            flush_limit_bytes=self._max_row_bytes,
-        )
+        self._batcher = self.table._table_impl.mutations_batcher(**self._batcher_kwargs)
         self._batcher._user_batch_completed_callback = (
             self._user_batch_completed_callback
         )
@@ -119,8 +119,8 @@ def _close_batcher(self):
             self._batcher.close()
         except MutationsExceptionGroup as exc_group:
             for error in exc_group.exceptions:
-                # Return the cause of the FailedMutationEntryError to the user,
-                # as this might be more what they're expecting.
+                # Unpack the root cause of the FailedMutationEntryError
+                # and return that error to the user.
                 self._exceptions.put(error.__cause__)
 
     def __enter__(self):
@@ -178,7 +178,7 @@ def flush(self):
             :dedent: 4
 
         :raises:
-            * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
+            * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations.
         """
         self._close_batcher()
         self._init_batcher()
@@ -192,7 +192,7 @@ def close(self):
         Any errors will be raised.
 
         :raises:
-            * :exc:`.batcherMutationsBatchError` if there's any error in the mutations.
+            * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations.
         """
         self._close_batcher()
         atexit.unregister(self.close)

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

@@ -131,7 +131,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(

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

@@ -272,10 +272,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: Optional[
             Callable[[list[status_pb2.Status]], None]
         ] = None
@@ -363,13 +364,17 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
         """
         # flush new entries
         in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = []
+        in_process_batches: list[list[RowMutationEntry]] = []
         async for batch in self._flow_control.add_to_flow(new_entries):
             batch_task = CrossSync.create_task(
                 self._execute_mutate_rows, batch, sync_executor=self._sync_rpc_executor
             )
             in_process_requests.append(batch_task)
+            in_process_batches.append(batch)
         # wait for all inflight requests to complete
-        found_exceptions = await self._wait_for_batch_results(*in_process_requests)
+        found_exceptions = await self._wait_for_batch_results(
+            in_process_requests, in_process_batches
+        )
         # update exception data to reflect any new errors
         self._entries_processed_since_last_raise += len(new_entries)
         self._add_exceptions(found_exceptions)
@@ -419,7 +424,7 @@ async def _execute_mutate_rows(
                 self._user_batch_completed_callback(statuses)
         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,
@@ -519,26 +524,27 @@ def _on_exit(self):
     @staticmethod
     @CrossSync.convert
     async def _wait_for_batch_results(
-        *tasks: CrossSync.Future[list[FailedMutationEntryError]]
-        | CrossSync.Future[None],
-    ) -> list[Exception]:
+        tasks: Sequence[
+            CrossSync.Future[list[FailedMutationEntryError]] | CrossSync.Future[None]
+        ],
+        batches: Sequence[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
         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 list(zip(tasks, batches)):
             if CrossSync.is_async:
                 # futures don't need to be awaited in sync mode
                 await task
@@ -550,6 +556,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

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

@@ -102,7 +102,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(

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

@@ -234,8 +234,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: Optional[
@@ -310,12 +310,16 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]):
         in_process_requests: list[
             CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
         ] = []
+        in_process_batches: list[list[RowMutationEntry]] = []
         for batch in self._flow_control.add_to_flow(new_entries):
             batch_task = CrossSync._Sync_Impl.create_task(
                 self._execute_mutate_rows, batch, sync_executor=self._sync_rpc_executor
             )
             in_process_requests.append(batch_task)
-        found_exceptions = self._wait_for_batch_results(*in_process_requests)
+            in_process_batches.append(batch)
+        found_exceptions = self._wait_for_batch_results(
+            in_process_requests, in_process_batches
+        )
         self._entries_processed_since_last_raise += len(new_entries)
         self._add_exceptions(found_exceptions)
 
@@ -356,7 +360,7 @@ def _execute_mutate_rows(
                 self._user_batch_completed_callback(statuses)
         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.
@@ -435,31 +439,41 @@ 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: Sequence[
+            CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
+            | CrossSync._Sync_Impl.Future[None]
+        ],
+        batches: Sequence[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
         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 list(zip(tasks, batches)):
             try:
                 exc_list = task.result()
                 if exc_list:
                     for exc in exc_list:
                         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

google/cloud/bigtable/data/exceptions.py

@@ -153,7 +153,9 @@ class MutationsExceptionGroup(_BigtableExceptionGroup):
 
     @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 +173,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 +194,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 +217,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:

tests/unit/data/_async/test_mutations_batcher.py

@@ -640,7 +640,10 @@ async def mock_call(*args, **kwargs):
                 for _ in range(num_entries):
                     await instance.append(self._make_mutation(size=1))
                 # let any flush jobs finish
-                await instance._wait_for_batch_results(*instance._flush_jobs)
+                jobs = instance._flush_jobs
+                await instance._wait_for_batch_results(
+                    jobs, [mock.MagicMock()] * len(jobs)
+                )
                 # should have only flushed once, with large mutation and first mutation in loop
                 assert op_mock.call_count == 1
                 sent_batch = op_mock.call_args[0][0]
@@ -740,7 +743,10 @@ async def mock_call(*args, **kwargs):
                     )
                     await CrossSync.sleep(0.01)
                 # allow flushes to complete
-                await instance._wait_for_batch_results(*instance._flush_jobs)
+                jobs = instance._flush_jobs
+                await instance._wait_for_batch_results(
+                    jobs, [mock.MagicMock()] * len(jobs)
+                )
                 duration = time.monotonic() - start_time
                 assert len(instance._oldest_exceptions) == 0
                 assert len(instance._newest_exceptions) == 0

tests/unit/data/_sync_autogen/test_mutations_batcher.py

@@ -556,7 +556,8 @@ def mock_call(*args, **kwargs):
                 num_entries = 10
                 for _ in range(num_entries):
                     instance.append(self._make_mutation(size=1))
-                instance._wait_for_batch_results(*instance._flush_jobs)
+                jobs = instance._flush_jobs
+                instance._wait_for_batch_results(jobs, [mock.MagicMock()] * len(jobs))
                 assert op_mock.call_count == 1
                 sent_batch = op_mock.call_args[0][0]
                 assert len(sent_batch) == 2
@@ -646,7 +647,8 @@ def mock_call(*args, **kwargs):
                         [self._make_mutation(count=1)]
                     )
                     CrossSync._Sync_Impl.sleep(0.01)
-                instance._wait_for_batch_results(*instance._flush_jobs)
+                jobs = instance._flush_jobs
+                instance._wait_for_batch_results(jobs, [mock.MagicMock()] * len(jobs))
                 duration = time.monotonic() - start_time
                 assert len(instance._oldest_exceptions) == 0
                 assert len(instance._newest_exceptions) == 0

tests/unit/v2_client/test_batcher.py

@@ -59,6 +59,9 @@ def _atexit_mock():
 
 
 def test_mutations_batcher_constructor(_setup_batcher, _atexit_mock):
+    from google.cloud.bigtable.batcher import MAX_OUTSTANDING_ELEMENTS
+    from google.cloud.bigtable.batcher import MAX_OUTSTANDING_BYTES
+
     flush_count = 5
     flush_interval = 0.1
     max_row_bytes = 10000
@@ -77,6 +80,8 @@ def test_mutations_batcher_constructor(_setup_batcher, _atexit_mock):
                 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,
             )
             assert mutation_batcher.close in _atexit_mock._functions
 

tests/unit/v2_client/test_table.py

@@ -976,8 +976,8 @@ def test_table_mutations_batcher_factory():
     )
 
     assert mutation_batcher.table.table_id == TABLE_ID
-    assert mutation_batcher.flush_count == flush_count
-    assert mutation_batcher.max_row_bytes == max_row_bytes
+    assert mutation_batcher._batcher_kwargs["flush_limit_mutation_count"] == flush_count
+    assert mutation_batcher._batcher_kwargs["flush_limit_bytes"] == max_row_bytes
 
 
 def test_table_get_iam_policy():