address gemini comments

Author: daniel-sanche

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

@@ -16,6 +16,7 @@
 
 import atexit
 import concurrent.futures
+import logging
 import time
 import warnings
 from collections import deque
@@ -57,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")
@@ -416,7 +418,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.Code.UNKNOWN)] * len(batch)
+        statuses = [status_pb2.Status(code=code_pb2.UNKNOWN) for _ in range(len(batch))]
         try:
             operation = CrossSync._MutateRowsOperation(
                 self._target.client._gapic_client,
@@ -436,14 +438,19 @@ async def _execute_mutate_rows(
                 subexc.index = None
             return list(e.exceptions)
         else:
-            statuses = [status_pb2.Status(code=code_pb2.Code.OK)] * len(batch)
+            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:
-                self._user_batch_completed_callback(statuses)
+                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]):

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

@@ -266,7 +266,7 @@ def _get_statuses_from_mutations_exception_group(
     #
     # 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)] * batch_size
+    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__
@@ -297,7 +297,7 @@ def _get_status(exc: Optional[Exception]) -> status_pb2.Status:
         )
 
     return status_pb2.Status(
-        code=code_pb2.Code.UNKNOWN,
+        code=code_pb2.UNKNOWN,
         message=str(exc) if exc else "An unknown error has occurred",
     )
 

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

@@ -19,6 +19,7 @@
 
 import atexit
 import concurrent.futures
+import logging
 import time
 import warnings
 from collections import deque
@@ -50,6 +51,7 @@
     )
     from google.cloud.bigtable.data.mutations import RowMutationEntry
 _MB_SIZE = 1024 * 1024
+_LOGGER = logging.getLogger(__name__)
 
 
 @CrossSync._Sync_Impl.add_mapping_decorator("_FlowControl")
@@ -361,7 +363,7 @@ def _execute_mutate_rows(
             list[FailedMutationEntryError]:
                 list of FailedMutationEntryError objects for mutations that failed.
                 FailedMutationEntryError objects will not contain index information"""
-        statuses = [status_pb2.Status(code=code_pb2.Code.UNKNOWN)] * len(batch)
+        statuses = [status_pb2.Status(code=code_pb2.UNKNOWN) for _ in range(len(batch))]
         try:
             operation = CrossSync._Sync_Impl._MutateRowsOperation(
                 self._target.client._gapic_client,
@@ -379,11 +381,16 @@ def _execute_mutate_rows(
                 subexc.index = None
             return list(e.exceptions)
         else:
-            statuses = [status_pb2.Status(code=code_pb2.Code.OK)] * len(batch)
+            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:
-                self._user_batch_completed_callback(statuses)
+                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]):

packages/google-cloud-bigtable/google/cloud/bigtable/table.py

@@ -782,9 +782,10 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
         mutation_entries = [
             RowMutationEntry(row.row_key, row._get_mutations()) for row in rows
         ]
-        return_statuses = [status_pb2.Status(code=code_pb2.Code.UNKNOWN)] * len(
-            mutation_entries
-        )
+        return_statuses = [
+            status_pb2.Status(code=code_pb2.UNKNOWN)
+            for _ in range(len(mutation_entries))
+        ]
 
         try:
             self._table_impl.bulk_mutate_rows(
@@ -798,9 +799,10 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
                 mut_exc_group, len(mutation_entries)
             )
         else:
-            return_statuses = [status_pb2.Status(code=code_pb2.Code.OK)] * len(
-                mutation_entries
-            )
+            return_statuses = [
+                status_pb2.Status(code=code_pb2.OK)
+                for _ in range(len(mutation_entries))
+            ]
 
         return return_statuses
 

packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py

@@ -1059,6 +1059,22 @@ async def test__execute_mutate_rows_batch_completed_callback_errors(self):
                 assert result[0].index is None
                 assert result[1].index is None
 
+    @CrossSync.pytest
+    async def test__execute_mutate_rows_batch_completed_callback_exception(self):
+        with mock.patch.object(CrossSync, "_MutateRowsOperation") as mutate_rows:
+            mutate_rows.return_value = CrossSync.Mock()
+            table = mock.Mock()
+            table.default_mutate_rows_operation_timeout = 17
+            table.default_mutate_rows_attempt_timeout = 13
+            table.default_mutate_rows_retryable_errors = ()
+            callback = mock.Mock(side_effect=RuntimeError("callback failed"))
+            async with self._make_one(table) as instance:
+                instance._user_batch_completed_callback = callback
+                batch = [self._make_mutation()]
+                result = await instance._execute_mutate_rows(batch, mock.Mock())
+                callback.assert_called_once()
+                assert result == []
+
     @CrossSync.pytest
     async def test__raise_exceptions(self):
         """Raise exceptions and reset error state"""

packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py

@@ -940,6 +940,23 @@ def test__execute_mutate_rows_batch_completed_callback_errors(self):
                 assert result[0].index is None
                 assert result[1].index is None
 
+    def test__execute_mutate_rows_batch_completed_callback_exception(self):
+        with mock.patch.object(
+            CrossSync._Sync_Impl, "_MutateRowsOperation"
+        ) as mutate_rows:
+            mutate_rows.return_value = CrossSync._Sync_Impl.Mock()
+            table = mock.Mock()
+            table.default_mutate_rows_operation_timeout = 17
+            table.default_mutate_rows_attempt_timeout = 13
+            table.default_mutate_rows_retryable_errors = ()
+            callback = mock.Mock(side_effect=RuntimeError("callback failed"))
+            with self._make_one(table) as instance:
+                instance._user_batch_completed_callback = callback
+                batch = [self._make_mutation()]
+                result = instance._execute_mutate_rows(batch, mock.Mock())
+                callback.assert_called_once()
+                assert result == []
+
     def test__raise_exceptions(self):
         """Raise exceptions and reset error state"""
         from google.cloud.bigtable.data import exceptions