chore(bigtable): Added an internal batch completed callback to the data client mutations batcher

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

@@ -16,15 +16,19 @@
 
 import atexit
 import concurrent.futures
+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 +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")
@@ -294,6 +299,9 @@ def __init__(
         self._newest_exceptions: deque[Exception] = deque(
             maxlen=self._exception_list_limit
         )
+        self._user_batch_completed_callback: (
+            Callable[[list[status_pb2.Status]], Any] | None
+        ) = None
         # clean up on program exit
         atexit.register(self._on_exit)
 
@@ -410,6 +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.UNKNOWN) for _ in range(len(batch))]
         try:
             operation = CrossSync._MutateRowsOperation(
                 self._target.client._gapic_client,
@@ -422,13 +431,26 @@ async def _execute_mutate_rows(
             )
             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.
+            if self._user_batch_completed_callback:
+                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

@@ -24,6 +24,7 @@
     TYPE_CHECKING,
     Callable,
     List,
+    Optional,
     Sequence,
     Tuple,
     Union,
@@ -32,8 +33,12 @@
 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
 ):

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

@@ -19,15 +19,19 @@
 
 import atexit
 import concurrent.futures
+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 +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")
@@ -259,6 +264,9 @@ def __init__(
         self._newest_exceptions: deque[Exception] = 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:
@@ -355,6 +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.UNKNOWN) for _ in range(len(batch))]
         try:
             operation = CrossSync._Sync_Impl._MutateRowsOperation(
                 self._target.client._gapic_client,
@@ -367,11 +376,21 @@ def _execute_mutate_rows(
             )
             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:
+                    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

@@ -20,7 +20,6 @@
 from google.api_core.exceptions import (
     Aborted,
     DeadlineExceeded,
-    GoogleAPICallError,
     InternalServerError,
     NotFound,
     ServiceUnavailable,
@@ -38,11 +37,11 @@
     MutationsBatcher,
 )
 from google.cloud.bigtable.column_family import ColumnFamily, _gc_rule_from_pb
-from google.cloud.bigtable.data._helpers import TABLE_DEFAULT
-from google.cloud.bigtable.data.exceptions import (
-    MutationsExceptionGroup,
-    RetryExceptionGroup,
+from google.cloud.bigtable.data._helpers import (
+    TABLE_DEFAULT,
+    _get_statuses_from_mutations_exception_group,
 )
+from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
 from google.cloud.bigtable.data.mutations import RowMutationEntry
 from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
 from google.cloud.bigtable.encryption_info import EncryptionInfo
@@ -796,9 +795,6 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
                     % (row.row_key, row.table.name, self.name)
                 )
             mutation_entries.append(RowMutationEntry(row.row_key, row._get_mutations()))
-        return_statuses = [
-            status_pb2.Status(code=code_pb2.OK) for _ in range(len(mutation_entries))
-        ]  # By default, return status OKs for everything
 
         try:
             self._table_impl.bulk_mutate_rows(
@@ -808,41 +804,16 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
                 retryable_errors=retryable_errors,
             )
         except MutationsExceptionGroup as mut_exc_group:
-            # 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.
-            for error in mut_exc_group.exceptions:
-                cause = error.__cause__
-                if isinstance(cause, RetryExceptionGroup):
-                    return_statuses[error.index] = self._get_status(
-                        cause.exceptions[-1]
-                    )
-                else:
-                    return_statuses[error.index] = self._get_status(cause)
-
-        return return_statuses
-
-    @staticmethod
-    def _get_status(error):
-        if isinstance(error, GoogleAPICallError) and error.grpc_status_code is not None:
-            return status_pb2.Status(
-                code=error.grpc_status_code.value[0],
-                message=error.message,
-                details=error.details,
+            return_statuses = _get_statuses_from_mutations_exception_group(
+                mut_exc_group, len(mutation_entries)
             )
+        else:
+            return_statuses = [
+                status_pb2.Status(code=code_pb2.OK)
+                for _ in range(len(mutation_entries))
+            ]
 
-        return status_pb2.Status(
-            code=code_pb2.UNKNOWN,
-            message=str(error),
-        )
+        return return_statuses
 
     def sample_row_keys(self):
         """Read a sample of row keys in the table.

packages/google-cloud-bigtable/tests/system/data/test_system_async.py

@@ -512,6 +512,42 @@ async def test_mutations_batcher_timer_flush(self, client, target, temp_rows):
             # ensure cell is updated
             assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value
 
+    @pytest.mark.usefixtures("client")
+    @pytest.mark.usefixtures("target")
+    @CrossSync.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    @CrossSync.pytest
+    async def test_mutations_batcher_completed_callback(
+        self, client, target, temp_rows
+    ):
+        """
+        test batcher with batch completed callback. It should be called when the batcher flushes.
+        """
+        import mock
+        from google.rpc import code_pb2, status_pb2
+
+        from google.cloud.bigtable.data.mutations import RowMutationEntry
+
+        callback = mock.Mock()
+
+        new_value = uuid.uuid4().hex.encode()
+        row_key, mutation = await temp_rows.create_row_and_mutation(
+            target, new_value=new_value
+        )
+        bulk_mutation = RowMutationEntry(row_key, [mutation])
+        flush_interval = 0.1
+        async with target.mutations_batcher(flush_interval=flush_interval) as batcher:
+            batcher._user_batch_completed_callback = callback
+            await batcher.append(bulk_mutation)
+            await CrossSync.yield_to_event_loop()
+            assert len(batcher._staged_entries) == 1
+            await CrossSync.sleep(flush_interval + 0.1)
+            assert len(batcher._staged_entries) == 0
+            callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)])
+            # ensure cell is updated
+            assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value
+
     @pytest.mark.usefixtures("client")
     @pytest.mark.usefixtures("target")
     @CrossSync.Retry(

packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py

@@ -395,6 +395,35 @@ def test_mutations_batcher_timer_flush(self, client, target, temp_rows):
             assert len(batcher._staged_entries) == 0
             assert temp_rows.retrieve_cell_value(target, row_key) == new_value
 
+    @pytest.mark.usefixtures("client")
+    @pytest.mark.usefixtures("target")
+    @CrossSync._Sync_Impl.Retry(
+        predicate=retry.if_exception_type(ClientError), initial=1, maximum=5
+    )
+    def test_mutations_batcher_completed_callback(self, client, target, temp_rows):
+        """test batcher with batch completed callback. It should be called when the batcher flushes."""
+        import mock
+        from google.rpc import code_pb2, status_pb2
+
+        from google.cloud.bigtable.data.mutations import RowMutationEntry
+
+        callback = mock.Mock()
+        new_value = uuid.uuid4().hex.encode()
+        row_key, mutation = temp_rows.create_row_and_mutation(
+            target, new_value=new_value
+        )
+        bulk_mutation = RowMutationEntry(row_key, [mutation])
+        flush_interval = 0.1
+        with target.mutations_batcher(flush_interval=flush_interval) as batcher:
+            batcher._user_batch_completed_callback = callback
+            batcher.append(bulk_mutation)
+            CrossSync._Sync_Impl.yield_to_event_loop()
+            assert len(batcher._staged_entries) == 1
+            CrossSync._Sync_Impl.sleep(flush_interval + 0.1)
+            assert len(batcher._staged_entries) == 0
+            callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)])
+            assert temp_rows.retrieve_cell_value(target, row_key) == new_value
+
     @pytest.mark.usefixtures("client")
     @pytest.mark.usefixtures("target")
     @CrossSync._Sync_Impl.Retry(