diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index e69b46382eb0..c3268cc09d7b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -12,15 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. + """User friendly container for Google Cloud Bigtable MutationBatcher.""" import atexit -import concurrent.futures import queue -import threading -from dataclasses import dataclass -from google.api_core.exceptions import from_grpc_status +from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup +from google.cloud.bigtable.data.mutations import RowMutationEntry FLUSH_COUNT = 100 # after this many elements, send out the batch @@ -40,145 +39,20 @@ def __init__(self, message, exc): super().__init__(self.message) -class _MutationsBatchQueue(object): - """Private Threadsafe Queue to hold rows for batching.""" - - def __init__(self, max_mutation_bytes=MAX_MUTATION_SIZE, flush_count=FLUSH_COUNT): - """Specify the queue constraints""" - self._queue = queue.Queue() - self.total_mutation_count = 0 - self.total_size = 0 - self.max_mutation_bytes = max_mutation_bytes - self.flush_count = flush_count - - def get(self): - """ - Retrieve an item from the queue. Recalculate queue size. - - If the queue is empty, return None. - """ - try: - row = self._queue.get_nowait() - mutation_size = row.get_mutations_size() - self.total_mutation_count -= len(row._get_mutations()) - self.total_size -= mutation_size - return row - except queue.Empty: - return None - - def put(self, item): - """Insert an item to the queue. Recalculate queue size.""" - - mutation_count = len(item._get_mutations()) - - self._queue.put(item) - - self.total_size += item.get_mutations_size() - self.total_mutation_count += mutation_count - - def full(self): - """Check if the queue is full.""" - if ( - self.total_mutation_count >= self.flush_count - or self.total_size >= self.max_mutation_bytes - ): - return True - return False - - -@dataclass -class _BatchInfo: - """Keeping track of size of a batch""" - - mutations_count: int = 0 - rows_count: int = 0 - mutations_size: int = 0 - - -class _FlowControl(object): - def __init__( - self, - max_mutations=MAX_OUTSTANDING_ELEMENTS, - max_mutation_bytes=MAX_OUTSTANDING_BYTES, - ): - """Control the inflight requests. Keep track of the mutations, row bytes and row counts. - As requests to backend are being made, adjust the number of mutations being processed. - - If threshold is reached, block the flow. - Reopen the flow as requests are finished. - """ - self.max_mutations = max_mutations - self.max_mutation_bytes = max_mutation_bytes - self.inflight_mutations = 0 - self.inflight_size = 0 - self.event = threading.Event() - self.event.set() - self._lock = threading.Lock() - - def is_blocked(self): - """Returns True if: - - - inflight mutations >= max_mutations, or - - inflight bytes size >= max_mutation_bytes, or - """ - - return ( - self.inflight_mutations >= self.max_mutations - or self.inflight_size >= self.max_mutation_bytes - ) - - def control_flow(self, batch_info): - """ - Calculate the resources used by this batch - """ - - with self._lock: - self.inflight_mutations += batch_info.mutations_count - self.inflight_size += batch_info.mutations_size - self.set_flow_control_status() - - def wait(self): - """ - Wait until flow control pushback has been released. - It awakens as soon as `event` is set. - """ - self.event.wait() - - def set_flow_control_status(self): - """Check the inflight mutations and size. - - If values exceed the allowed threshold, block the event. - """ - if self.is_blocked(): - self.event.clear() # sleep - else: - self.event.set() # awaken the threads - - def release(self, batch_info): - """ - Release the resources. - Decrement the row size to allow enqueued mutations to be run. - """ - with self._lock: - self.inflight_mutations -= batch_info.mutations_count - self.inflight_size -= batch_info.mutations_size - self.set_flow_control_status() - - class MutationsBatcher(object): """A MutationsBatcher is used in batch cases where the number of mutations is large or unknown. It will store :class:`DirectRow` in memory until one of the - size limits is reached, or an explicit call to :func:`flush()` is performed. When - a flush event occurs, the :class:`DirectRow` in memory will be sent to Cloud - Bigtable. Batching mutations is more efficient than sending individual - request. + size limits is reached, the ``flush_interval`` timer fires, or an explicit call + to :func:`flush()` is performed. When a flush event occurs, the :class:`DirectRow` + in memory will be sent to Cloud Bigtable. Batching mutations is more efficient than + sending individual requests. This class is not suited for usage in systems where each mutation must be guaranteed to be sent, since calling :func:`mutate()` may only - result in an in-memory change. Rows are only sent to the service when a size - limit is reached, when :func:`flush()` is called explicitly, or when the - batcher is closed (:func:`close()` is also registered to run at interpreter - exit). There is no time-based background flush. As a result, if the process + result in an in-memory change. Rows are sent to the service when a size + limit is reached, when the ``flush_interval`` timer fires, when :func:`flush()` + is called explicitly, or when the batcher is closed (:func:`close()` is also + registered to run at interpreter exit). As a result, if the process terminates abruptly -- e.g. a crash, ``SIGKILL``, or ``os._exit`` where the ``atexit`` handler never runs -- any :class:`DirectRow` still buffered in memory is silently dropped and never sent, even after :func:`mutate()` @@ -208,9 +82,8 @@ class MutationsBatcher(object): (5 MB). :type flush_interval: float - :param flush_interval: (Deprecated) No longer used. Retained only for - backwards compatibility. There is no time-based background flush; see the - class docstring for when rows are sent. + :param flush_interval: (Optional) Automatically flush every flush_interval seconds. + If None or <= 0, no time-based flushing is performed. Default is 1 second. :type batch_completed_callback: Callable[list:[`~google.rpc.status_pb2.Status`]] = None :param batch_completed_callback: (Optional) A callable for handling responses @@ -226,33 +99,44 @@ def __init__( flush_interval=1, batch_completed_callback=None, ): - self._rows = _MutationsBatchQueue( - max_mutation_bytes=max_row_bytes, flush_count=flush_count - ) self.table = table - self._executor = concurrent.futures.ThreadPoolExecutor() - atexit.register(self.close) - # ``flush_interval`` is retained for backwards compatibility but is no - # longer used: the previous background ``threading.Timer`` was one-shot - # (never re-armed), so it fired at most once and could silently drop the - # rows it dequeued if that single flush raised on the timer thread. - # Flushing now happens only on size thresholds, explicit ``flush()``, - # or ``close()`` (also registered via ``atexit``). - self.flow_control = _FlowControl( - max_mutations=MAX_OUTSTANDING_ELEMENTS, - max_mutation_bytes=MAX_OUTSTANDING_BYTES, - ) - self.futures_mapping = {} - self.exceptions = queue.Queue() + 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) + self._exceptions: queue.Queue = queue.Queue() @property def flush_count(self): - return self._rows.flush_count + return self._batcher._flush_limit_count @property def max_row_bytes(self): - return self._rows.max_mutation_bytes + return self._batcher._flush_limit_bytes + + def _init_batcher(self): + self._batcher = self.table._table_impl.mutations_batcher(**self._batcher_kwargs) + self._batcher._user_batch_completed_callback = ( + self._user_batch_completed_callback + ) + + def _close_batcher(self): + try: + self._batcher.close() + except MutationsExceptionGroup as exc_group: + for error in exc_group.exceptions: + # Unpack the root cause of the FailedMutationEntryError and + # return that error to the user. In standard execution paths, + # FailedMutationEntryError always has an Exception cause; + # defensively fall back to error itself if __cause__ is None. + cause = error.__cause__ if error.__cause__ is not None else error + self._exceptions.put(cause) def __enter__(self): """Starting the MutationsBatcher as a context manager""" @@ -276,10 +160,7 @@ def mutate(self, row): * :exc:`~.table._BigtableRetryableError` if any row returned a transient error. * :exc:`RuntimeError` if the number of responses doesn't match the number of rows that were retried """ - self._rows.put(row) - - if self._rows.full(): - self._flush_async() + self._batcher.append(RowMutationEntry(row.row_key, row._get_mutations())) def mutate_rows(self, rows): """Add multiple rows to the batch. If the current batch meets one of the size @@ -312,124 +193,10 @@ def flush(self): :dedent: 4 :raises: - * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. - """ - rows_to_flush = [] - row = self._rows.get() - while row is not None: - rows_to_flush.append(row) - row = self._rows.get() - response = self._flush_rows(rows_to_flush) - return response - - def _flush_async(self): - """Sends the current batch to Cloud Bigtable asynchronously. - - :raises: - * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. - """ - next_row = self._rows.get() - while next_row is not None: - # start a new batch - rows_to_flush = [next_row] - batch_info = _BatchInfo( - mutations_count=len(next_row._get_mutations()), - rows_count=1, - mutations_size=next_row.get_mutations_size(), - ) - # fill up batch with rows - next_row = self._rows.get() - while next_row is not None and self._row_fits_in_batch( - next_row, batch_info - ): - rows_to_flush.append(next_row) - batch_info.mutations_count += len(next_row._get_mutations()) - batch_info.rows_count += 1 - batch_info.mutations_size += next_row.get_mutations_size() - next_row = self._rows.get() - # send batch over network - # wait for resources to become available - self.flow_control.wait() - # once unblocked, submit the batch - # event flag will be set by control_flow to block subsequent thread, but not blocking this one - self.flow_control.control_flow(batch_info) - future = self._executor.submit(self._flush_rows, rows_to_flush) - # schedule release of resources from flow control - self.futures_mapping[future] = batch_info - future.add_done_callback(self._batch_completed_callback) - - def _batch_completed_callback(self, future): - """Callback for when the mutation has finished to clean up the current batch - and release items from the flow controller. - Raise exceptions if there's any. - Release the resources locked by the flow control and allow enqueued tasks to be run. + * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations. """ - processed_rows = self.futures_mapping[future] - self.flow_control.release(processed_rows) - del self.futures_mapping[future] - # Surface any exception raised inside the async flush. Without this, an - # exception raised by ``_flush_rows`` (e.g. a non-retryable RPC error, a - # retry deadline, or a response-count mismatch) would be stored on the - # future and silently discarded, so the failed mutations would never be - # reported to the user -- effectively silent data loss. Per-row errors - # from a successful RPC are already recorded in ``self.exceptions`` by - # ``_flush_rows``; here the whole batch failed with a single exception, - # so record it once per row in the batch to keep the reported error - # count aligned with the number of affected mutations. - # - # A cancelled future is "done", so this callback still runs for it, but - # ``future.exception()`` would raise ``CancelledError``. Nothing here - # cancels futures today, but guard against it so the callback stays - # correct if cancellation is ever introduced. - if future.cancelled(): - return - exc = future.exception() - if exc is not None: - for _ in range(processed_rows.rows_count): - self.exceptions.put(exc) - - def _row_fits_in_batch(self, row, batch_info): - """Checks if a row can fit in the current batch. - - :type row: class - :param row: :class:`~google.cloud.bigtable.row.DirectRow`. - - :type batch_info: :class:`_BatchInfo` - :param batch_info: Information about the current batch. - - :rtype: bool - :returns: True if the row can fit in the current batch. - """ - new_rows_count = batch_info.rows_count + 1 - new_mutations_count = batch_info.mutations_count + len(row._get_mutations()) - new_mutations_size = batch_info.mutations_size + row.get_mutations_size() - return ( - new_rows_count <= self.flush_count - and new_mutations_size <= self.max_row_bytes - and new_mutations_count <= self.flow_control.max_mutations - and new_mutations_size <= self.flow_control.max_mutation_bytes - ) - - def _flush_rows(self, rows_to_flush): - """Mutate the specified rows. - - :raises: - * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. - """ - responses = [] - if len(rows_to_flush) > 0: - response = self.table.mutate_rows(rows_to_flush) - - if self._user_batch_completed_callback: - self._user_batch_completed_callback(response) - - for result in response: - if result.code != 0: - exc = from_grpc_status(result.code, result.message) - self.exceptions.put(exc) - responses.append(result) - - return responses + self._close_batcher() + self._init_batcher() def __exit__(self, exc_type, exc_value, exc_traceback): """Clean up resources. Flush and shutdown the ThreadPoolExecutor.""" @@ -440,23 +207,10 @@ 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. """ - try: - self.flush() - except MutationsBatchError as exc: - for e in exc.exc: - self.exceptions.put(e) - except Exception as exc: - # A failure in this final synchronous flush must not abort cleanup. - # If it propagated here it would skip the executor shutdown (leaving - # in-flight async flushes un-awaited) and skip draining - # self.exceptions, masking every error already captured from - # earlier async flushes -- silently discarding those failures. - # Record it like any other batch failure and continue. - self.exceptions.put(exc) - self._executor.shutdown(wait=True) + self._close_batcher() atexit.unregister(self.close) - if self.exceptions.qsize() > 0: - exc = list(self.exceptions.queue) + if self._exceptions.qsize() > 0: + exc = list(self._exceptions.queue) raise MutationsBatchError("Errors in batch mutations.", exc=exc) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py index 0007447a5505..2c6bc47ecd2d 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -139,7 +139,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( diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py index e8d56a5f9008..584c99891b6b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -293,10 +293,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: ( Callable[[list[status_pb2.Status]], Any] | None ) = None @@ -383,7 +384,11 @@ 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_requests: list[ + tuple[ + CrossSync.Future[list[FailedMutationEntryError]], list[RowMutationEntry] + ] + ] = [] async for batch, metric in self._flow_control.add_to_flow_with_metrics( new_entries, self._target.client._metrics ): @@ -393,7 +398,7 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]): metric, sync_executor=self._sync_rpc_executor, ) - in_process_requests.append(batch_task) + 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) # update exception data to reflect any new errors @@ -446,7 +451,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, @@ -546,26 +551,28 @@ def _on_exit(self): @staticmethod @CrossSync.convert async def _wait_for_batch_results( - *tasks: CrossSync.Future[list[FailedMutationEntryError]] - | CrossSync.Future[None], - ) -> list[Exception]: + *tasks: tuple[ + CrossSync.Future[list[FailedMutationEntryError]] | CrossSync.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[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 tasks: if CrossSync.is_async: # futures don't need to be awaited in sync mode await task @@ -577,6 +584,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 diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py index 2e7281d269c6..eeaae90c915b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py @@ -28,6 +28,7 @@ Sequence, Tuple, Union, + cast, ) from google.api_core import exceptions as core_exceptions @@ -286,14 +287,17 @@ def _get_status(exc: Optional[Exception]) -> status_pb2.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, + if isinstance(exc, core_exceptions.GoogleAPICallError): + status_code = cast(Optional["grpc.StatusCode"], exc.grpc_status_code) + if status_code is not None: + return status_pb2.Status( + code=status_code.value[0], + message=exc.message, + details=exc.details, + ) + return status_pb2.Status( + code=code_pb2.Code.UNKNOWN, + message="An unknown error has occurred", ) return status_pb2.Status( diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index 8bb4e49e22eb..54edd06b24ab 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -112,7 +112,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( diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py index e3483a70f43b..a7d02144e828 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py @@ -258,8 +258,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: ( @@ -332,7 +332,10 @@ 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], + ] ] = [] for batch, metric in self._flow_control.add_to_flow_with_metrics( new_entries, self._target.client._metrics @@ -343,7 +346,7 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]): metric, sync_executor=self._sync_rpc_executor, ) - in_process_requests.append(batch_task) + 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) @@ -386,7 +389,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. @@ -465,30 +468,42 @@ 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: tuple[ + CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]] + | 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[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 tasks: 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 diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py index 475ea1e9742a..840e44161200 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py @@ -141,19 +141,18 @@ 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 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 +170,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 +191,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 +214,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: diff --git a/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py index 779cbf26bb67..897d7a174d17 100644 --- a/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py +++ b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py @@ -1245,3 +1245,110 @@ def callback(results): time.sleep(0.01) # ensure all mutations were sent assert len(all_results) == num_sent + + +def test_mutations_batcher_exceptions(data_table, rows_to_delete): + """Test the mutations batcher exception handling""" + import mock + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.batcher import MutationsBatcher, MutationsBatchError + from google.cloud.bigtable_v2 import MutateRowsResponse + + num_sent = 5 + + error_response = [ + MutateRowsResponse( + entries=[ + MutateRowsResponse.Entry( + index=i, + status=status_pb2.Status( + code=code_pb2.INTERNAL, + message="Test error", + ), + ) + for i in range(num_sent) + ] + ) + ] + + # Simulate only failures + with pytest.raises(MutationsBatchError): + with mock.patch.object( + data_table._instance._client.table_data_client, "mutate_rows" + ) as mutate_mock: + mutate_mock.side_effect = [error_response] * 100000 + with MutationsBatcher( + data_table, + flush_count=10, + flush_interval=1, + ) as batcher: + for i in range(num_sent): + row = data_table.direct_row("row{}".format(i)) + row.set_cell( + COLUMN_FAMILY_ID1, COL_NAME1, "val{}".format(i).encode("utf-8") + ) + rows_to_delete.append(row) + batcher.mutate(row) + batcher.flush() + + # Test that exceptions are only raised on close. + with mock.patch.object( + data_table._instance._client.table_data_client, "mutate_rows" + ) as mutate_mock: + mutate_mock.side_effect = [error_response] * 100000 + batcher = MutationsBatcher( + data_table, + flush_count=10, + flush_interval=1, + ) + for i in range(num_sent): + row = data_table.direct_row("row{}".format(i)) + row.set_cell( + COLUMN_FAMILY_ID1, COL_NAME1, "val{}".format(i).encode("utf-8") + ) + rows_to_delete.append(row) + batcher.mutate(row) + batcher.flush() + + with pytest.raises(MutationsBatchError): + batcher.close() + + +def test_mutations_batcher_manual_flush(data_table, rows_to_delete): + """Test the mutations batcher manual flush""" + import mock + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.batcher import MutationsBatcher + + num_batches = 5 + batch_size = 4 + callback = mock.MagicMock() + + with MutationsBatcher( + data_table, + flush_count=500, + flush_interval=5, + batch_completed_callback=callback, + ) as batcher: + for i in range(num_batches): + for j in range(batch_size): + num = i * batch_size + j + row = data_table.direct_row(f"row{num}".encode("utf-8")) + row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, f"val{num}".encode("utf-8")) + rows_to_delete.append(row) + batcher.mutate(row) + batcher.flush() + callback.assert_called_with( + [status_pb2.Status(code=code_pb2.OK)] * batch_size + ) + + # ensure all mutations were sent + rows = data_table.read_rows() + rows.consume_all() + for row_num in range(0, num_batches * batch_size): + row = rows.rows[f"row{row_num}".encode("utf-8")] + assert row.cells[COLUMN_FAMILY_ID1][COL_NAME1][ + 0 + ].value == f"val{row_num}".encode("utf-8") diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py index 17ba411d3bc0..fe6eac3b4cd4 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py @@ -644,7 +644,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( + *[(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 sent_batch = op_mock.call_args[0][0] @@ -744,7 +747,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( + *[(job, mock.MagicMock()) for job in jobs] + ) duration = time.monotonic() - start_time assert len(instance._oldest_exceptions) == 0 assert len(instance._newest_exceptions) == 0 diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py index d8375984f006..d45e266f8ee7 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py @@ -561,7 +561,10 @@ 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( + *[(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 @@ -651,7 +654,10 @@ 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( + *[(job, mock.MagicMock()) for job in jobs] + ) duration = time.monotonic() - start_time assert len(instance._oldest_exceptions) == 0 assert len(instance._newest_exceptions) == 0 diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py index 758ca226a12d..ed20ea49983e 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py @@ -13,65 +13,115 @@ # limitations under the License. +import time + import mock import pytest from google.cloud.bigtable.batcher import ( MutationsBatcher, MutationsBatchError, - _FlowControl, ) from google.cloud.bigtable.row import DirectRow +from ._testing import _make_credentials + +PROJECT = "PROJECT" +INSTANCE_ID = "instance-id" TABLE_ID = "table-id" TABLE_NAME = "/tables/" + TABLE_ID -def test_mutation_batcher_constructor(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table) as mutation_batcher: - assert table is mutation_batcher.table - - -def test_mutation_batcher_w_user_callback(): - table = _Table(TABLE_NAME) - - def callback_fn(response): - callback_fn.count = len(response) +@pytest.fixture +def _setup_batcher(): + import google.cloud.bigtable.data._sync_autogen.mutations_batcher + from google.cloud.bigtable.client import Client + from google.cloud.bigtable.table import Table + + client = Client(project=PROJECT, credentials=_make_credentials()) + instance = client.instance(INSTANCE_ID) + + with mock.patch.object( + google.cloud.bigtable.data._sync_autogen.mutations_batcher.CrossSync._Sync_Impl, + "_MutateRowsOperation", + ) as operation_mock: + yield Table(TABLE_ID, instance=instance), operation_mock + + +@pytest.fixture +def _atexit_mock(): + atexit_mock = _AtexitMock() + with mock.patch.multiple( + "atexit", register=atexit_mock.register, unregister=atexit_mock.unregister + ): + yield atexit_mock + + +def test_mutations_batcher_constructor(_setup_batcher, _atexit_mock): + from google.cloud.bigtable.batcher import ( + MAX_OUTSTANDING_BYTES, + MAX_OUTSTANDING_ELEMENTS, + ) + + flush_count = 5 + flush_interval = 0.1 + max_row_bytes = 10000 + table, _ = _setup_batcher + with mock.patch.object( + table._table_impl, "mutations_batcher" + ) as batcher_impl_constructor: + with MutationsBatcher( + table, + flush_count=flush_count, + flush_interval=flush_interval, + max_row_bytes=max_row_bytes, + ) as mutation_batcher: + assert table is mutation_batcher.table + batcher_impl_constructor.assert_called_once_with( + 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 + + +def test_mutations_batcher_w_user_callback(_setup_batcher): + table, _ = _setup_batcher + + callback_fn = mock.Mock() + batch_size = 4 with MutationsBatcher( - table, flush_count=1, batch_completed_callback=callback_fn + table, flush_count=batch_size, batch_completed_callback=callback_fn ) as mutation_batcher: - rows = [ - DirectRow(row_key=b"row_key"), - DirectRow(row_key=b"row_key_2"), - DirectRow(row_key=b"row_key_3"), - DirectRow(row_key=b"row_key_4"), - ] + rows = [DirectRow(row_key=f"row_key_{i}".encode()) for i in range(batch_size)] + for row in rows: + row.delete() mutation_batcher.mutate_rows(rows) - assert callback_fn.count == 4 + assert len(callback_fn.call_args[0][0]) == batch_size -def test_mutation_batcher_mutate_row(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: - rows = [ - DirectRow(row_key=b"row_key"), - DirectRow(row_key=b"row_key_2"), - DirectRow(row_key=b"row_key_3"), - DirectRow(row_key=b"row_key_4"), - ] +def test_mutations_batcher_mutate_row(_setup_batcher): + table, operation_mock = _setup_batcher + batch_size = 4 + + with MutationsBatcher(table, flush_count=batch_size) as mutation_batcher: + rows = [DirectRow(row_key=f"row_key_{i}".encode()) for i in range(batch_size)] + for row in rows: + row.delete() mutation_batcher.mutate_rows(rows) - assert table.mutation_calls == 1 + operation_mock.assert_called_once() -def test_mutation_batcher_mutate(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: +def test_mutations_batcher_mutate(_setup_batcher): + table, operation_mock = _setup_batcher + with MutationsBatcher(table=table, flush_count=1) as mutation_batcher: row = DirectRow(row_key=b"row_key") row.set_cell("cf1", b"c1", 1) row.set_cell("cf1", b"c2", 2) @@ -80,47 +130,36 @@ def test_mutation_batcher_mutate(): mutation_batcher.mutate(row) - assert table.mutation_calls == 1 + operation_mock.assert_called_once() -def test_mutation_batcher_flush_w_no_rows(): - table = _Table(TABLE_NAME) +def test_mutations_batcher_manual_flush(_setup_batcher, _atexit_mock): + table, operation_mock = _setup_batcher with MutationsBatcher(table=table) as mutation_batcher: - mutation_batcher.flush() - - assert table.mutation_calls == 0 + original_batcher_impl = mutation_batcher._batcher + assert original_batcher_impl._on_exit in _atexit_mock._functions + row = DirectRow(row_key=b"row_key") + row.set_cell("cf1", b"c1", 1) + mutation_batcher.mutate(row) -def test_mutation_batcher_mutate_w_max_flush_count(): - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table, flush_count=3) as mutation_batcher: - row_1 = DirectRow(row_key=b"row_key_1") - row_2 = DirectRow(row_key=b"row_key_2") - row_3 = DirectRow(row_key=b"row_key_3") - - mutation_batcher.mutate(row_1) - mutation_batcher.mutate(row_2) - mutation_batcher.mutate(row_3) + mutation_batcher.flush() - assert table.mutation_calls == 1 + operation_mock.assert_called_once() + assert mutation_batcher._batcher != original_batcher_impl + assert original_batcher_impl._on_exit not in _atexit_mock._functions -@mock.patch("google.cloud.bigtable.batcher.MAX_OUTSTANDING_ELEMENTS", new=3) -def test_mutation_batcher_mutate_w_max_mutations(): - table = _Table(TABLE_NAME) +def test_mutations_batcher_flush_w_no_rows(_setup_batcher): + table, operation_mock = _setup_batcher with MutationsBatcher(table=table) as mutation_batcher: - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", 1) - row.set_cell("cf1", b"c2", 2) - row.set_cell("cf1", b"c3", 3) - - mutation_batcher.mutate(row) + mutation_batcher.flush() - assert table.mutation_calls == 1 + operation_mock.assert_not_called() -def test_mutation_batcher_mutate_w_max_row_bytes(): - table = _Table(TABLE_NAME) +def test_mutations_batcher_mutate_w_max_row_bytes(_setup_batcher): + table, operation_mock = _setup_batcher with MutationsBatcher( table=table, max_row_bytes=3 * 1024 * 1024 ) as mutation_batcher: @@ -134,11 +173,11 @@ def test_mutation_batcher_mutate_w_max_row_bytes(): mutation_batcher.mutate(row) - assert table.mutation_calls == 1 + operation_mock.assert_called_once() -def test_mutations_batcher_flushed_when_closed(): - table = _Table(TABLE_NAME) +def test_mutations_batcher_flushed_when_closed(_setup_batcher): + table, operation_mock = _setup_batcher mutation_batcher = MutationsBatcher(table=table, max_row_bytes=3 * 1024 * 1024) number_of_bytes = 1 * 1024 * 1024 @@ -149,15 +188,15 @@ def test_mutations_batcher_flushed_when_closed(): row.set_cell("cf1", b"c2", max_value) mutation_batcher.mutate(row) - assert table.mutation_calls == 0 + operation_mock.assert_not_called() mutation_batcher.close() - assert table.mutation_calls == 1 + operation_mock.assert_called_once() -def test_mutations_batcher_context_manager_flushed_when_closed(): - table = _Table(TABLE_NAME) +def test_mutations_batcher_context_manager_flushed_when_closed(_setup_batcher): + table, operation_mock = _setup_batcher with MutationsBatcher( table=table, max_row_bytes=3 * 1024 * 1024 ) as mutation_batcher: @@ -169,182 +208,124 @@ def test_mutations_batcher_context_manager_flushed_when_closed(): row.set_cell("cf1", b"c2", max_value) mutation_batcher.mutate(row) + operation_mock.assert_not_called() - assert table.mutation_calls == 1 - - -@mock.patch("google.cloud.bigtable.batcher.threading.Timer") -@mock.patch("google.cloud.bigtable.batcher.MutationsBatcher.flush") -def test_mutations_batcher_flush_interval_does_not_start_timer( - mocked_flush, mocked_timer -): - # ``flush_interval`` is accepted for backwards compatibility but no longer - # starts a background timer. Constructing the batcher must not create a - # timer or trigger a flush. - table = _Table(TABLE_NAME) - MutationsBatcher(table=table, flush_interval=0.5) - - mocked_timer.assert_not_called() - mocked_flush.assert_not_called() - - -def test_mutations_batcher_response_with_error_codes(): - from google.rpc.status_pb2 import Status - - mocked_response = [Status(code=1), Status(code=5)] - - with mock.patch("tests.unit.v2_client.test_batcher._Table") as mocked_table: - table = mocked_table.return_value - mutation_batcher = MutationsBatcher(table=table) - - row1 = DirectRow(row_key=b"row_key") - row2 = DirectRow(row_key=b"row_key") - table.mutate_rows.return_value = mocked_response + operation_mock.assert_called_once() - mutation_batcher.mutate_rows([row1, row2]) - with pytest.raises(MutationsBatchError) as exc: - mutation_batcher.close() - assert exc.value.message == "Errors in batch mutations." - assert len(exc.value.exc) == 2 - - assert exc.value.exc[0].message == mocked_response[0].message - assert exc.value.exc[1].message == mocked_response[1].message +def test_mutations_batcher_flush_interval(_setup_batcher): + table, operation_mock = _setup_batcher + flush_interval = 0.5 + mutation_batcher = MutationsBatcher(table=table, flush_interval=flush_interval) + row = DirectRow(row_key=b"row_key") + row.set_cell("cf1", b"c1", b"1") + mutation_batcher.mutate(row) + operation_mock.assert_not_called() -def test_mutations_batcher_asynchronous_flush_exception_is_surfaced(): - """An exception raised by the underlying ``mutate_rows`` call (e.g. a - non-retryable RPC error or a response-count mismatch) is raised inside the - async flush task. It must be captured and re-raised at ``close()`` rather - than being silently swallowed by the executor -- otherwise the failed - mutations are never reported to the user (silent data loss).""" - from google.api_core.exceptions import PermissionDenied + time.sleep(0.4) + operation_mock.assert_not_called() - with mock.patch("tests.unit.v2_client.test_batcher._Table") as mocked_table: - table = mocked_table.return_value - # flush_count=2 forces the batch to flush asynchronously (through the - # executor) as soon as the second row is added - mutation_batcher = MutationsBatcher(table=table, flush_count=2) + # Test could be flaky, so giving the thread some extra buffer time + time.sleep(0.25) + operation_mock.assert_called_once() - row1 = DirectRow(row_key=b"row_key") - row1.set_cell("cf1", b"c1", b"1") - row2 = DirectRow(row_key=b"row_key") - row2.set_cell("cf1", b"c1", b"2") - table.mutate_rows.side_effect = PermissionDenied("denied") + mutation_batcher.close() - mutation_batcher.mutate_rows([row1, row2]) - with pytest.raises(MutationsBatchError) as exc: - mutation_batcher.close() - assert exc.value.message == "Errors in batch mutations." - # the whole batch (both rows) failed, so both are reported -- the error - # count stays aligned with the number of affected mutations - assert len(exc.value.exc) == 2 - assert all(isinstance(e, PermissionDenied) for e in exc.value.exc) +def test_mutations_batcher_response_with_error_codes(_setup_batcher): + from google.api_core import exceptions -def test_batch_completed_callback_ignores_cancelled_future(): - """A cancelled future is still "done", so the completion callback runs for - it, but ``future.exception()`` would raise ``CancelledError``. The callback - must short-circuit on a cancelled future instead of letting that propagate.""" - from google.cloud.bigtable.batcher import _BatchInfo + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: - batch_info = _BatchInfo(rows_count=2, mutations_count=2, mutations_size=0) + table, operation_mock = _setup_batcher - cancelled_future = mock.Mock() - cancelled_future.cancelled.return_value = True - cancelled_future.exception.side_effect = AssertionError( - "exception() must not be called on a cancelled future" + causes = [ + exceptions.InternalServerError("Something happened"), + exceptions.DataLoss("Data loss"), + ] + excs = [ + FailedMutationEntryError( + failed_idx=i, failed_mutation_entry=mock.Mock(), cause=cause ) - mutation_batcher.futures_mapping[cancelled_future] = batch_info - - # Should not raise, should not record any exceptions - mutation_batcher._batch_completed_callback(cancelled_future) - - assert cancelled_future not in mutation_batcher.futures_mapping - assert mutation_batcher.exceptions.qsize() == 0 - - -def test_mutations_batcher_close_surfaces_errors_when_final_flush_raises(): - """If the final flush in ``close()`` raises, ``close()`` must still shut - down the executor and surface every accumulated error -- including ones - already captured from async flushes -- instead of letting the flush - exception mask them and abort cleanup (silent data loss).""" - from google.api_core.exceptions import PermissionDenied, ServiceUnavailable - - with mock.patch("tests.unit.v2_client.test_batcher._Table") as mocked_table: - table = mocked_table.return_value - mutation_batcher = MutationsBatcher(table=table) - - # Simulate an error already captured earlier (e.g. from an async flush). - prior_error = ServiceUnavailable("earlier async failure") - mutation_batcher.exceptions.put(prior_error) - - # The row stays queued (below flush_count), so it is only flushed by - # close(); make that final flush raise. - row = DirectRow(row_key=b"row_key") - row.set_cell("cf1", b"c1", b"1") - mutation_batcher.mutate(row) - table.mutate_rows.side_effect = PermissionDenied("denied") - - with pytest.raises(MutationsBatchError) as exc: - mutation_batcher.close() - - # both the pre-existing and the flush-time errors are reported - assert prior_error in exc.value.exc - assert any(isinstance(e, PermissionDenied) for e in exc.value.exc) - # cleanup still ran despite the flush raising - assert mutation_batcher._executor._shutdown is True - - -def test_flow_control_event_is_set_when_not_blocked(): - flow_control = _FlowControl() - - flow_control.set_flow_control_status() - assert flow_control.event.is_set() - - -def test_flow_control_event_is_not_set_when_blocked(): - flow_control = _FlowControl() - - flow_control.inflight_mutations = flow_control.max_mutations - flow_control.inflight_size = flow_control.max_mutation_bytes - - flow_control.set_flow_control_status() - assert not flow_control.event.is_set() - + for i, cause in enumerate(causes) + ] + error = MutationsExceptionGroup(excs=excs, total_entries=len(excs)) + + operation_mock.return_value.start.side_effect = error + + mutations_batcher = MutationsBatcher(table=table) + row1 = DirectRow(row_key=b"row_key") + row1.set_cell("cf1", b"c1", b"1") + row2 = DirectRow(row_key=b"row_key_2") + row2.set_cell("cf1", b"c1", b"1") + mutations_batcher.mutate_rows([row1, row2]) + mutations_batcher.flush() + + with pytest.raises(MutationsBatchError) as raised_error: + mutations_batcher.close() + assert raised_error.value.message == "Errors in batch mutations." + assert len(raised_error.value.exc) == 2 + + assert raised_error.value.exc[0].message == causes[0].message + assert raised_error.value.exc[1].message == causes[1].message + + +def test_mutations_batcher_response_with_error_codes_multiple_flushes(_setup_batcher): + from google.api_core import exceptions + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + table, operation_mock = _setup_batcher + + causes = [ + exceptions.InternalServerError("Something happened"), + exceptions.DataLoss("Data loss"), + ] + excs = [ + FailedMutationEntryError( + failed_idx=i, failed_mutation_entry=mock.Mock(), cause=cause + ) + for i, cause in enumerate(causes) + ] + error1 = MutationsExceptionGroup(excs=excs[0:1], total_entries=1) + error2 = MutationsExceptionGroup(excs=excs[1:2], total_entries=1) -@mock.patch("concurrent.futures.ThreadPoolExecutor.submit") -def test_flush_async_batch_count(mocked_executor_submit): - table = _Table(TABLE_NAME) - mutation_batcher = MutationsBatcher(table=table, flush_count=2) + operation_mock.return_value.start.side_effect = error1 - number_of_bytes = 1 * 1024 * 1024 - max_value = b"1" * number_of_bytes - for index in range(5): - row = DirectRow(row_key=f"row_key_{index}") - row.set_cell("cf1", b"c1", max_value) - mutation_batcher.mutate(row) - mutation_batcher._flush_async() + mutations_batcher = MutationsBatcher(table=table) + row1 = DirectRow(row_key=b"row_key") + row1.set_cell("cf1", b"c1", b"1") + mutations_batcher.mutate(row1) + mutations_batcher.flush() - # 3 batches submitted. 2 batches of 2 items, and the last one a single item batch. - assert mocked_executor_submit.call_count == 3 + operation_mock.return_value.start.side_effect = error2 + row2 = DirectRow(row_key=b"row_key_2") + row2.set_cell("cf1", b"c1", b"1") + mutations_batcher.mutate(row2) + mutations_batcher.flush() -class _Instance(object): - def __init__(self, client=None): - self._client = client + with pytest.raises(MutationsBatchError) as raised_error: + mutations_batcher.close() + assert raised_error.value.message == "Errors in batch mutations." + assert len(raised_error.value.exc) == 2 + assert raised_error.value.exc[0].message == causes[0].message + assert raised_error.value.exc[1].message == causes[1].message -class _Table(object): - def __init__(self, name, client=None): - self.name = name - self._instance = _Instance(client) - self.mutation_calls = 0 - def mutate_rows(self, rows): - from google.rpc.status_pb2 import Status +class _AtexitMock: + def __init__(self): + self._functions = set() - self.mutation_calls += 1 + def register(self, func): + self._functions.add(func) - return [Status(code=0) for _ in rows] + def unregister(self, func): + self._functions.remove(func) diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py index 96d018dffe84..8977db1b6989 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py @@ -980,8 +980,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():