feat: mutate rows batching

google/cloud/bigtable/_mutate_rows.py

@@ -31,6 +31,9 @@
     from google.cloud.bigtable.client import Table
     from google.cloud.bigtable.mutations import RowMutationEntry
 
+# mutate_rows requests are limited to this value
+MUTATE_ROWS_REQUEST_MUTATION_LIMIT = 100_000
+
 
 class _MutateRowsIncomplete(RuntimeError):
     """
@@ -68,6 +71,14 @@ def __init__(
           - per_request_timeout: the timeoutto use for each mutate_rows attempt, in seconds.
               If not specified, the request will run until operation_timeout is reached.
         """
+        # check that mutations are within limits
+        total_mutations = sum(len(entry.mutations) for entry in mutation_entries)
+        if total_mutations > MUTATE_ROWS_REQUEST_MUTATION_LIMIT:
+            raise ValueError(
+                "mutate_rows requests can contain at most "
+                f"{MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations across "
+                f"all entries. Found {total_mutations}."
+            )
         # create partial function to pass to trigger rpc call
         metadata = _make_metadata(table.table_name, table.app_profile_id)
         self._gapic_fn = functools.partial(
@@ -119,7 +130,7 @@ async def start(self):
                 self._handle_entry_error(idx, exc)
         finally:
             # raise exception detailing incomplete mutations
-            all_errors = []
+            all_errors: list[Exception] = []
             for idx, exc_list in self.errors.items():
                 if len(exc_list) == 0:
                     raise core_exceptions.ClientError(

google/cloud/bigtable/client.py

@@ -20,8 +20,6 @@
     Any,
     Optional,
     Set,
-    Callable,
-    Coroutine,
     TYPE_CHECKING,
 )
 
@@ -58,6 +56,8 @@
 from google.cloud.bigtable._mutate_rows import _MutateRowsOperation
 from google.cloud.bigtable._helpers import _make_metadata
 from google.cloud.bigtable._helpers import _convert_retry_deadline
+from google.cloud.bigtable.mutations_batcher import MutationsBatcher
+from google.cloud.bigtable.mutations_batcher import _MB_SIZE
 from google.cloud.bigtable._helpers import _attempt_timeout_generator
 
 from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule
@@ -67,7 +67,6 @@
 from google.cloud.bigtable.row_filters import RowFilterChain
 
 if TYPE_CHECKING:
-    from google.cloud.bigtable.mutations_batcher import MutationsBatcher
     from google.cloud.bigtable import RowKeySamples
     from google.cloud.bigtable import ShardedQuery
 
@@ -722,17 +721,48 @@ async def execute_rpc():
         )
         return await wrapped_fn()
 
-    def mutations_batcher(self, **kwargs) -> MutationsBatcher:
+    def mutations_batcher(
+        self,
+        *,
+        flush_interval: float | None = 5,
+        flush_limit_mutation_count: int | None = 1000,
+        flush_limit_bytes: int = 20 * _MB_SIZE,
+        flow_control_max_mutation_count: int = 100_000,
+        flow_control_max_bytes: int = 100 * _MB_SIZE,
+        batch_operation_timeout: float | None = None,
+        batch_per_request_timeout: float | None = None,
+    ) -> MutationsBatcher:
         """
         Returns a new mutations batcher instance.
 
         Can be used to iteratively add mutations that are flushed as a group,
         to avoid excess network calls
 
+        Args:
+          - flush_interval: Automatically flush every flush_interval seconds. If None,
+              a table default will be used
+          - flush_limit_mutation_count: Flush immediately after flush_limit_mutation_count
+              mutations are added across all entries. If None, this limit is ignored.
+          - flush_limit_bytes: Flush immediately after flush_limit_bytes bytes are added.
+          - flow_control_max_mutation_count: Maximum number of inflight mutations.
+          - flow_control_max_bytes: Maximum number of inflight bytes.
+          - batch_operation_timeout: timeout for each mutate_rows operation, in seconds. If None,
+              table default_operation_timeout will be used
+          - batch_per_request_timeout: timeout for each individual request, in seconds. If None,
+              table default_per_request_timeout will be used
         Returns:
             - a MutationsBatcher context manager that can batch requests
         """
-        return MutationsBatcher(self, **kwargs)
+        return MutationsBatcher(
+            self,
+            flush_interval=flush_interval,
+            flush_limit_mutation_count=flush_limit_mutation_count,
+            flush_limit_bytes=flush_limit_bytes,
+            flow_control_max_mutation_count=flow_control_max_mutation_count,
+            flow_control_max_bytes=flow_control_max_bytes,
+            batch_operation_timeout=batch_operation_timeout,
+            batch_per_request_timeout=batch_per_request_timeout,
+        )
 
     async def mutate_row(
         self,
@@ -830,10 +860,6 @@ async def bulk_mutate_rows(
         *,
         operation_timeout: float | None = 60,
         per_request_timeout: float | None = None,
-        on_success: Callable[
-            [int, RowMutationEntry], None | Coroutine[None, None, None]
-        ]
-        | None = None,
     ):
         """
         Applies mutations for multiple rows in a single batched request.
@@ -859,9 +885,6 @@ async def bulk_mutate_rows(
                 in seconds. If it takes longer than this time to complete, the request
                 will be cancelled with a DeadlineExceeded exception, and a retry will
                 be attempted if within operation_timeout budget
-            - on_success: a callback function that will be called when each mutation
-                entry is confirmed to be applied successfully. Will be passed the
-                index and the entry itself.
         Raises:
             - MutationsExceptionGroup if one or more mutations fails
                 Contains details about any failed entries in .exceptions

google/cloud/bigtable/exceptions.py

@@ -85,19 +85,96 @@ def __str__(self):
 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
     """
 
     @staticmethod
-    def _format_message(excs: list[FailedMutationEntryError], total_entries: int):
-        entry_str = "entry" if total_entries == 1 else "entries"
-        plural_str = "" if len(excs) == 1 else "s"
-        return f"{len(excs)} sub-exception{plural_str} (from {total_entries} {entry_str} attempted)"
+    def _format_message(
+        excs: list[Exception], total_entries: int, exc_count: int | None = None
+    ) -> str:
+        """
+        Format a message for the exception group
+
+        Args:
+          - excs: the exceptions in the group
+          - total_entries: the total number of entries attempted, successful or not
+          - exc_count: the number of exceptions associated with the request
+             if None, this will be len(excs)
+        """
+        exc_count = exc_count if exc_count is not None else len(excs)
+        entry_str = "entry" if exc_count == 1 else "entries"
+        return f"{exc_count} failed {entry_str} from {total_entries} attempted."
+
+    def __init__(
+        self, excs: list[Exception], total_entries: int, message: str | None = None
+    ):
+        """
+        Args:
+          - excs: the exceptions in the group
+          - total_entries: the total number of entries attempted, successful or not
+          - message: the message for the exception group. If None, a default message
+              will be generated
+        """
+        message = (
+            message
+            if message is not None
+            else self._format_message(excs, total_entries)
+        )
+        super().__init__(message, excs)
+        self.total_entries_attempted = total_entries
 
-    def __init__(self, excs: list[FailedMutationEntryError], total_entries: int):
-        super().__init__(self._format_message(excs, total_entries), excs)
+    def __new__(
+        cls, excs: list[Exception], total_entries: int, message: str | None = None
+    ):
+        """
+        Args:
+          - excs: the exceptions in the group
+          - total_entries: the total number of entries attempted, successful or not
+          - message: the message for the exception group. If None, a default message
+        """
+        message = (
+            message if message is not None else cls._format_message(excs, total_entries)
+        )
+        instance = super().__new__(cls, message, excs)
+        instance.total_entries_attempted = total_entries
+        return instance
 
-    def __new__(cls, excs: list[FailedMutationEntryError], total_entries: int):
-        return super().__new__(cls, cls._format_message(excs, total_entries), excs)
+    @classmethod
+    def from_truncated_lists(
+        cls,
+        first_list: list[Exception],
+        last_list: list[Exception],
+        total_excs: int,
+        entry_count: int,
+    ) -> MutationsExceptionGroup:
+        """
+        Create a MutationsExceptionGroup from two lists of exceptions, representing
+        a larger set that has been truncated. The MutationsExceptionGroup will
+        contain the union of the two lists as sub-exceptions, and the error message
+        describe the number of exceptions that were truncated.
+
+        Args:
+          - first_list: the set of oldest exceptions to add to the ExceptionGroup
+          - last_list: the set of newest exceptions to add to the ExceptionGroup
+          - total_excs: the total number of exceptions associated with the request
+             Should be len(first_list) + len(last_list) + number of dropped exceptions
+             in the middle
+          - entry_count: the total number of entries attempted, successful or not
+        """
+        first_count, last_count = len(first_list), len(last_list)
+        if first_count + last_count >= total_excs:
+            # no exceptions were dropped
+            return cls(first_list + last_list, entry_count)
+        excs = first_list + last_list
+        truncation_count = total_excs - (first_count + last_count)
+        base_message = cls._format_message(excs, entry_count, total_excs)
+        first_message = f"first {first_count}" if first_count else ""
+        last_message = f"last {last_count}" if last_count else ""
+        conjunction = " and " if first_message and last_message else ""
+        message = f"{base_message} ({first_message}{conjunction}{last_message} attached as sub-exceptions; {truncation_count} truncated)"
+        return cls(excs, entry_count, message)
 
 
 class FailedMutationEntryError(Exception):
@@ -108,14 +185,17 @@ class FailedMutationEntryError(Exception):
 
     def __init__(
         self,
-        failed_idx: int,
+        failed_idx: int | None,
         failed_mutation_entry: "RowMutationEntry",
         cause: Exception,
     ):
         idempotent_msg = (
             "idempotent" if failed_mutation_entry.is_idempotent() else "non-idempotent"
         )
-        message = f"Failed {idempotent_msg} mutation entry at index {failed_idx} with cause: {cause!r}"
+        index_msg = f" at index {failed_idx} " if failed_idx is not None else " "
+        message = (
+            f"Failed {idempotent_msg} mutation entry{index_msg}with cause: {cause!r}"
+        )
         super().__init__(message)
         self.index = failed_idx
         self.entry = failed_mutation_entry

google/cloud/bigtable/mutations.py

@@ -17,6 +17,11 @@
 import time
 from dataclasses import dataclass
 from abc import ABC, abstractmethod
+from sys import getsizeof
+
+# mutation entries above this should be rejected
+from google.cloud.bigtable._mutate_rows import MUTATE_ROWS_REQUEST_MUTATION_LIMIT
+
 
 from google.cloud.bigtable.read_modify_write_rules import MAX_INCREMENT_VALUE
 
@@ -41,6 +46,12 @@ def is_idempotent(self) -> bool:
     def __str__(self) -> str:
         return str(self._to_dict())
 
+    def size(self) -> int:
+        """
+        Get the size of the mutation in bytes
+        """
+        return getsizeof(self._to_dict())
+
     @classmethod
     def _from_dict(cls, input_dict: dict[str, Any]) -> Mutation:
         instance: Mutation | None = None
@@ -195,6 +206,12 @@ def __init__(self, row_key: bytes | str, mutations: Mutation | list[Mutation]):
             row_key = row_key.encode("utf-8")
         if isinstance(mutations, Mutation):
             mutations = [mutations]
+        if len(mutations) == 0:
+            raise ValueError("mutations must not be empty")
+        elif len(mutations) > MUTATE_ROWS_REQUEST_MUTATION_LIMIT:
+            raise ValueError(
+                f"entries must have <= {MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations"
+            )
         self.row_key = row_key
         self.mutations = tuple(mutations)
 
@@ -208,6 +225,12 @@ def is_idempotent(self) -> bool:
         """Check if the mutation is idempotent"""
         return all(mutation.is_idempotent() for mutation in self.mutations)
 
+    def size(self) -> int:
+        """
+        Get the size of the mutation in bytes
+        """
+        return getsizeof(self._to_dict())
+
     @classmethod
     def _from_dict(cls, input_dict: dict[str, Any]) -> RowMutationEntry:
         return RowMutationEntry(