addressed more feedback

Author: gkevinzheng

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

@@ -363,18 +363,18 @@ 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_batches: list[list[RowMutationEntry]] = []
+        in_process_requests: list[
+            tuple[
+                CrossSync.Future[list[FailedMutationEntryError]], 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)
+            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, in_process_batches
-        )
+        found_exceptions = await self._wait_for_batch_results(*in_process_requests)
         # update exception data to reflect any new errors
         self._entries_processed_since_last_raise += len(new_entries)
         self._add_exceptions(found_exceptions)
@@ -524,17 +524,18 @@ def _on_exit(self):
     @staticmethod
     @CrossSync.convert
     async def _wait_for_batch_results(
-        tasks: Sequence[
-            CrossSync.Future[list[FailedMutationEntryError]] | CrossSync.Future[None]
+        *tasks: tuple[
+            CrossSync.Future[list[FailedMutationEntryError]] | CrossSync.Future[None],
+            list[RowMutationEntry],
         ],
-        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
+            *tasks: Tuples of futures representing _execute_mutate_rows or
+                    _flush_internal tasks, and their associated batches
         Returns:
             list[FailedMutationEntryError]:
                 list of FailedMutationEntryError encountered by any of the tasks,
@@ -544,7 +545,7 @@ async def _wait_for_batch_results(
         if not tasks:
             return []
         exceptions: list[FailedMutationEntryError] = []
-        for task, batch in list(zip(tasks, batches)):
+        for task, batch in tasks:
             if CrossSync.is_async:
                 # futures don't need to be awaited in sync mode
                 await task

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

@@ -308,18 +308,17 @@ 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],
+            ]
         ] = []
-        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)
-            in_process_batches.append(batch)
-        found_exceptions = self._wait_for_batch_results(
-            in_process_requests, in_process_batches
-        )
+            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)
 
@@ -439,17 +438,18 @@ def _on_exit(self):
 
     @staticmethod
     def _wait_for_batch_results(
-        tasks: Sequence[
+        *tasks: tuple[
             CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
-            | CrossSync._Sync_Impl.Future[None]
-        ],
-        batches: Sequence[list[RowMutationEntry]],
+            | 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[FailedMutationEntryError]:
                 list of FailedMutationEntryError encountered by any of the tasks,
@@ -458,7 +458,7 @@ def _wait_for_batch_results(
         if not tasks:
             return []
         exceptions: list[FailedMutationEntryError] = []
-        for task, batch in list(zip(tasks, batches)):
+        for task, batch in tasks:
             try:
                 exc_list = task.result()
                 if exc_list:

google/cloud/bigtable/data/exceptions.py

@@ -141,14 +141,11 @@ 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

tests/unit/data/_async/test_mutations_batcher.py

@@ -642,7 +642,7 @@ async def mock_call(*args, **kwargs):
                 # let any flush jobs finish
                 jobs = instance._flush_jobs
                 await instance._wait_for_batch_results(
-                    jobs, [mock.MagicMock()] * len(jobs)
+                    *[(job, mock.MagicMock()) for job in jobs]
                 )
                 # should have only flushed once, with large mutation and first mutation in loop
                 assert op_mock.call_count == 1
@@ -745,7 +745,7 @@ async def mock_call(*args, **kwargs):
                 # allow flushes to complete
                 jobs = instance._flush_jobs
                 await instance._wait_for_batch_results(
-                    jobs, [mock.MagicMock()] * len(jobs)
+                    *[(job, mock.MagicMock()) for job in jobs]
                 )
                 duration = time.monotonic() - start_time
                 assert len(instance._oldest_exceptions) == 0

tests/unit/data/_sync_autogen/test_mutations_batcher.py

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