From 4df3d78dc86ce19475e7fcef06926e40e15ff7c6 Mon Sep 17 00:00:00 2001 From: Kevin Zheng <147537668+gkevinzheng@users.noreply.github.com> Date: Tue, 24 Mar 2026 10:01:09 -0400 Subject: [PATCH 1/5] feat: Mutations Batcher shim (#1309) **Changes Made:** - Replaced mutations batcher implementation with one based off of the data client. - Reworked unit tests. - Added additional system tests. --- .../google/cloud/bigtable/batcher.py | 220 ++++--------- .../bigtable/data/_async/_mutate_rows.py | 18 ++ .../cloud/bigtable/data/_async/client.py | 9 + .../bigtable/data/_async/mutations_batcher.py | 60 +++- .../google/cloud/bigtable/data/_helpers.py | 24 +- .../data/_sync_autogen/_mutate_rows.py | 15 + .../data/_sync_autogen/mutations_batcher.py | 57 +++- .../google/cloud/bigtable/data/exceptions.py | 23 +- .../tests/system/v2_client/test_data_api.py | 105 +++++++ .../data/_async/test_mutations_batcher.py | 10 +- .../_sync_autogen/test_mutations_batcher.py | 10 +- .../tests/unit/v2_client/test_batcher.py | 289 +++++++++++++----- .../tests/unit/v2_client/test_table.py | 4 +- 13 files changed, 557 insertions(+), 287 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index e69b46382eb0..c82e4bef1dcb 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -13,12 +13,21 @@ # limitations under the License. """User friendly container for Google Cloud Bigtable MutationBatcher.""" +<<<<<<< ours import atexit import concurrent.futures import queue import threading from dataclasses import dataclass +======= +import queue +import atexit + + +from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup +from google.cloud.bigtable.data.mutations import RowMutationEntry +>>>>>>> theirs from google.api_core.exceptions import from_grpc_status @@ -40,131 +49,6 @@ 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 @@ -226,10 +110,8 @@ 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 +<<<<<<< ours self._executor = concurrent.futures.ThreadPoolExecutor() atexit.register(self.close) # ``flush_interval`` is retained for backwards compatibility but is no @@ -244,15 +126,42 @@ def __init__( ) 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, + } +>>>>>>> theirs self._user_batch_completed_callback = batch_completed_callback + self._init_batcher() + atexit.register(self.close) + self._exceptions = queue.Queue() @property def flush_count(self): - return self._rows.flush_count + return self._flush_count @property def max_row_bytes(self): - return self._rows.max_mutation_bytes + return self._max_row_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. + self._exceptions.put(error.__cause__) def __enter__(self): """Starting the MutationsBatcher as a context manager""" @@ -276,10 +185,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,6 +218,7 @@ def flush(self): :dedent: 4 :raises: +<<<<<<< ours * :exc:`.batcherMutationsBatchError` if there's any error in the mutations. """ rows_to_flush = [] @@ -399,37 +306,12 @@ def _row_fits_in_batch(self, row, batch_info): :rtype: bool :returns: True if the row can fit in the current batch. +======= + * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations. +>>>>>>> theirs """ - 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,8 +322,9 @@ 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. """ +<<<<<<< ours try: self.flush() except MutationsBatchError as exc: @@ -456,7 +339,10 @@ def close(self): # Record it like any other batch failure and continue. self.exceptions.put(exc) self._executor.shutdown(wait=True) +======= + self._close_batcher() +>>>>>>> theirs 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..412975fd113b 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 @@ -128,6 +128,7 @@ async def start(self): Raises: MutationsExceptionGroup: if any mutations failed """ +<<<<<<< ours with self._operation_metric: try: # trigger mutate_rows @@ -156,6 +157,23 @@ async def start(self): if all_errors: raise bt_exceptions.MutationsExceptionGroup( all_errors, len(self.mutations) +======= + try: + # trigger mutate_rows + await self._operation() + except Exception as exc: + # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations + incomplete_indices = self.remaining_indices.copy() + for idx in incomplete_indices: + self._handle_entry_error(idx, exc) + finally: + # raise exception detailing incomplete mutations + all_errors: list[bt_exceptions.FailedMutationEntryError] = [] + for idx, exc_list in self.errors.items(): + if len(exc_list) == 0: + raise core_exceptions.ClientError( + f"Mutation {idx} failed with no associated errors" +>>>>>>> theirs ) @CrossSync.convert diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 048171474648..05cf7f3c3fc3 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -24,7 +24,11 @@ import warnings from functools import partial from typing import ( +<<<<<<< ours TYPE_CHECKING, +======= + cast, +>>>>>>> theirs Any, AsyncIterable, Callable, @@ -146,7 +150,12 @@ ) if TYPE_CHECKING: +<<<<<<< ours from google.cloud.bigtable.data._helpers import RowKeySamples, ShardedQuery +======= + from google.cloud.bigtable.data._helpers import RowKeySamples + from google.cloud.bigtable.data._helpers import ShardedQuery +>>>>>>> theirs if CrossSync.is_async: from google.cloud.bigtable.data._async.mutations_batcher import ( 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..0c6589247dec 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,13 +293,20 @@ 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 ) +<<<<<<< ours self._user_batch_completed_callback: ( Callable[[list[status_pb2.Status]], Any] | None ) = None +======= + # only used by the shim right now. + self._user_batch_completed_callback: Optional[ + Callable[[list[status_pb2.Status]], None] + ] = None +>>>>>>> theirs # clean up on program exit atexit.register(self._on_exit) @@ -383,17 +390,26 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]): new_entries list of RowMutationEntry objects to flush """ # flush new entries +<<<<<<< ours in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = [] async for batch, metric in self._flow_control.add_to_flow_with_metrics( new_entries, self._target.client._metrics ): +======= + in_process_requests: list[ + tuple[ + CrossSync.Future[list[FailedMutationEntryError]], list[RowMutationEntry] + ] + ] = [] + async for batch in self._flow_control.add_to_flow(new_entries): +>>>>>>> theirs batch_task = CrossSync.create_task( self._execute_mutate_rows, batch, 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 +462,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 +562,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 +595,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..56450b2ae793 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py @@ -17,6 +17,11 @@ from __future__ import annotations +<<<<<<< ours +======= +from typing import cast, Callable, Sequence, List, Optional, Tuple, TYPE_CHECKING, Union +import time +>>>>>>> theirs import enum import time from collections import namedtuple @@ -286,14 +291,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..6112f005aa48 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 @@ -104,6 +104,7 @@ def start(self): Raises: MutationsExceptionGroup: if any mutations failed""" +<<<<<<< ours with self._operation_metric: try: self._operation() @@ -129,6 +130,20 @@ def start(self): if all_errors: raise bt_exceptions.MutationsExceptionGroup( all_errors, len(self.mutations) +======= + try: + self._operation() + except Exception as exc: + incomplete_indices = self.remaining_indices.copy() + for idx in incomplete_indices: + self._handle_entry_error(idx, exc) + finally: + all_errors: list[bt_exceptions.FailedMutationEntryError] = [] + for idx, exc_list in self.errors.items(): + if len(exc_list) == 0: + raise core_exceptions.ClientError( + f"Mutation {idx} failed with no associated errors" +>>>>>>> theirs ) def _run_attempt(self): 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..9bdde379a5d2 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 @@ -16,7 +16,11 @@ # This file is automatically generated by CrossSync. Do not edit manually. from __future__ import annotations +<<<<<<< ours +======= +from typing import Callable, Optional, Sequence, TYPE_CHECKING, cast +>>>>>>> theirs import atexit import concurrent.futures import time @@ -258,13 +262,19 @@ 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 ) +<<<<<<< ours self._user_batch_completed_callback: ( Callable[[list[status_pb2.Status]], Any] | None ) = None +======= + self._user_batch_completed_callback: Optional[ + Callable[[list[status_pb2.Status]], None] + ] = None +>>>>>>> theirs atexit.register(self._on_exit) def _timer_routine(self, interval: float | None) -> None: @@ -332,7 +342,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 +356,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 +399,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 +478,50 @@ 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: +<<<<<<< ours 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.""" +>>>>>>> theirs 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..d22443a994f9 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,108 @@ 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.cloud.bigtable.batcher import MutationsBatcher, MutationsBatchError + from google.cloud.bigtable_v2 import MutateRowsResponse + from google.rpc import code_pb2, status_pb2 + + 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.cloud.bigtable.batcher import MutationsBatcher + from google.rpc import status_pb2, code_pb2 + + 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..69a434da2086 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 @@ -23,55 +23,103 @@ ) 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(): + from google.cloud.bigtable.client import Client + from google.cloud.bigtable.table import Table + + import google.cloud.bigtable.data._sync_autogen.mutations_batcher + + 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_ELEMENTS + from google.cloud.bigtable.batcher import MAX_OUTSTANDING_BYTES + + 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 +128,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 +171,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 +186,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,10 +206,12 @@ 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 + operation_mock.assert_called_once() +<<<<<<< ours @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( @@ -348,3 +387,115 @@ def mutate_rows(self, rows): self.mutation_calls += 1 return [Status(code=0) for _ in rows] +======= +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() + + time.sleep(0.4) + operation_mock.assert_not_called() + + # Test could be flaky, so giving the thread some extra buffer time + time.sleep(0.25) + operation_mock.assert_called_once() + + mutation_batcher.close() + + +def test_mutations_batcher_response_with_error_codes(_setup_batcher): + from google.api_core import exceptions + from google.cloud.bigtable.data.exceptions import FailedMutationEntryError + from google.cloud.bigtable.data.exceptions import 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) + ] + 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 + from google.cloud.bigtable.data.exceptions import 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) + + operation_mock.return_value.start.side_effect = error1 + + 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() + + 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() + + 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 _AtexitMock: + def __init__(self): + self._functions = set() + + def register(self, func): + self._functions.add(func) + + def unregister(self, func): + self._functions.remove(func) +>>>>>>> theirs 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(): From 175b1f8c2eaba36cd7b769b15dd0dc59f4eddc76 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 21 Aug 2026 15:28:54 -0700 Subject: [PATCH 2/5] fix merge conflicts --- .../google/cloud/bigtable/batcher.py | 138 +----------- .../bigtable/data/_async/_mutate_rows.py | 20 +- .../cloud/bigtable/data/_async/client.py | 9 - .../bigtable/data/_async/mutations_batcher.py | 19 +- .../google/cloud/bigtable/data/_helpers.py | 6 +- .../data/_sync_autogen/_mutate_rows.py | 17 +- .../data/_sync_autogen/mutations_batcher.py | 20 +- .../tests/system/v2_client/test_data_api.py | 6 +- .../tests/unit/v2_client/test_batcher.py | 204 ++---------------- 9 files changed, 31 insertions(+), 408 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index c82e4bef1dcb..bcde32003150 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -12,24 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. + """User friendly container for Google Cloud Bigtable MutationBatcher.""" -<<<<<<< ours import atexit -import concurrent.futures -import queue -import threading -from dataclasses import dataclass -======= import queue -import atexit - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup from google.cloud.bigtable.data.mutations import RowMutationEntry ->>>>>>> theirs - -from google.api_core.exceptions import from_grpc_status FLUSH_COUNT = 100 # after this many elements, send out the batch @@ -111,22 +101,6 @@ def __init__( batch_completed_callback=None, ): self.table = table -<<<<<<< ours - 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, @@ -134,11 +108,10 @@ def __init__( "flow_control_max_mutation_count": MAX_OUTSTANDING_ELEMENTS, "flow_control_max_bytes": MAX_OUTSTANDING_BYTES, } ->>>>>>> theirs self._user_batch_completed_callback = batch_completed_callback self._init_batcher() atexit.register(self.close) - self._exceptions = queue.Queue() + self._exceptions: queue.Queue = queue.Queue() @property def flush_count(self): @@ -218,97 +191,7 @@ def flush(self): :dedent: 4 :raises: -<<<<<<< ours - * :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. - """ - 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. -======= * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations. ->>>>>>> theirs """ self._close_batcher() self._init_batcher() @@ -324,24 +207,7 @@ def close(self): :raises: * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations. """ -<<<<<<< ours - 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() ->>>>>>> theirs atexit.unregister(self.close) if self._exceptions.qsize() > 0: exc = list(self._exceptions.queue) 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 412975fd113b..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 @@ -128,7 +128,6 @@ async def start(self): Raises: MutationsExceptionGroup: if any mutations failed """ -<<<<<<< ours with self._operation_metric: try: # trigger mutate_rows @@ -140,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( @@ -157,23 +156,6 @@ async def start(self): if all_errors: raise bt_exceptions.MutationsExceptionGroup( all_errors, len(self.mutations) -======= - try: - # trigger mutate_rows - await self._operation() - except Exception as exc: - # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations - incomplete_indices = self.remaining_indices.copy() - for idx in incomplete_indices: - self._handle_entry_error(idx, exc) - finally: - # raise exception detailing incomplete mutations - all_errors: list[bt_exceptions.FailedMutationEntryError] = [] - for idx, exc_list in self.errors.items(): - if len(exc_list) == 0: - raise core_exceptions.ClientError( - f"Mutation {idx} failed with no associated errors" ->>>>>>> theirs ) @CrossSync.convert diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 05cf7f3c3fc3..048171474648 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -24,11 +24,7 @@ import warnings from functools import partial from typing import ( -<<<<<<< ours TYPE_CHECKING, -======= - cast, ->>>>>>> theirs Any, AsyncIterable, Callable, @@ -150,12 +146,7 @@ ) if TYPE_CHECKING: -<<<<<<< ours from google.cloud.bigtable.data._helpers import RowKeySamples, ShardedQuery -======= - from google.cloud.bigtable.data._helpers import RowKeySamples - from google.cloud.bigtable.data._helpers import ShardedQuery ->>>>>>> theirs if CrossSync.is_async: from google.cloud.bigtable.data._async.mutations_batcher import ( 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 0c6589247dec..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 @@ -297,16 +297,10 @@ def __init__( self._newest_exceptions: deque[FailedMutationEntryError] = deque( maxlen=self._exception_list_limit ) -<<<<<<< ours + # only used by the shim right now. self._user_batch_completed_callback: ( Callable[[list[status_pb2.Status]], Any] | None ) = None -======= - # only used by the shim right now. - self._user_batch_completed_callback: Optional[ - Callable[[list[status_pb2.Status]], None] - ] = None ->>>>>>> theirs # clean up on program exit atexit.register(self._on_exit) @@ -390,19 +384,14 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]): new_entries list of RowMutationEntry objects to flush """ # flush new entries -<<<<<<< ours - in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = [] - async for batch, metric in self._flow_control.add_to_flow_with_metrics( - new_entries, self._target.client._metrics - ): -======= in_process_requests: list[ tuple[ CrossSync.Future[list[FailedMutationEntryError]], list[RowMutationEntry] ] ] = [] - async for batch in self._flow_control.add_to_flow(new_entries): ->>>>>>> theirs + async for batch, metric in self._flow_control.add_to_flow_with_metrics( + new_entries, self._target.client._metrics + ): batch_task = CrossSync.create_task( self._execute_mutate_rows, batch, 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 56450b2ae793..eeaae90c915b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py @@ -17,11 +17,6 @@ from __future__ import annotations -<<<<<<< ours -======= -from typing import cast, Callable, Sequence, List, Optional, Tuple, TYPE_CHECKING, Union -import time ->>>>>>> theirs import enum import time from collections import namedtuple @@ -33,6 +28,7 @@ Sequence, Tuple, Union, + cast, ) from google.api_core import exceptions as core_exceptions 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 6112f005aa48..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 @@ -104,7 +104,6 @@ def start(self): Raises: MutationsExceptionGroup: if any mutations failed""" -<<<<<<< ours with self._operation_metric: try: self._operation() @@ -113,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( @@ -130,20 +129,6 @@ def start(self): if all_errors: raise bt_exceptions.MutationsExceptionGroup( all_errors, len(self.mutations) -======= - try: - self._operation() - except Exception as exc: - incomplete_indices = self.remaining_indices.copy() - for idx in incomplete_indices: - self._handle_entry_error(idx, exc) - finally: - all_errors: list[bt_exceptions.FailedMutationEntryError] = [] - for idx, exc_list in self.errors.items(): - if len(exc_list) == 0: - raise core_exceptions.ClientError( - f"Mutation {idx} failed with no associated errors" ->>>>>>> theirs ) def _run_attempt(self): 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 9bdde379a5d2..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 @@ -16,11 +16,7 @@ # This file is automatically generated by CrossSync. Do not edit manually. from __future__ import annotations -<<<<<<< ours -======= -from typing import Callable, Optional, Sequence, TYPE_CHECKING, cast ->>>>>>> theirs import atexit import concurrent.futures import time @@ -266,15 +262,9 @@ def __init__( self._newest_exceptions: deque[FailedMutationEntryError] = deque( maxlen=self._exception_list_limit ) -<<<<<<< ours self._user_batch_completed_callback: ( Callable[[list[status_pb2.Status]], Any] | None ) = None -======= - self._user_batch_completed_callback: Optional[ - Callable[[list[status_pb2.Status]], None] - ] = None ->>>>>>> theirs atexit.register(self._on_exit) def _timer_routine(self, interval: float | None) -> None: @@ -482,7 +472,7 @@ def _wait_for_batch_results( 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. @@ -491,18 +481,10 @@ def _wait_for_batch_results( *tasks: Tuples of futures representing _execute_mutate_rows or _flush_internal tasks, and their associated batches Returns: -<<<<<<< ours - 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.""" ->>>>>>> theirs if not tasks: return [] exceptions: list[FailedMutationEntryError] = [] 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 d22443a994f9..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 @@ -1250,9 +1250,10 @@ def callback(results): 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 - from google.rpc import code_pb2, status_pb2 num_sent = 5 @@ -1317,8 +1318,9 @@ def test_mutations_batcher_exceptions(data_table, rows_to_delete): 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 - from google.rpc import status_pb2, code_pb2 num_batches = 5 batch_size = 4 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 69a434da2086..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,13 +13,14 @@ # 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 @@ -33,11 +34,10 @@ @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 - import google.cloud.bigtable.data._sync_autogen.mutations_batcher - client = Client(project=PROJECT, credentials=_make_credentials()) instance = client.instance(INSTANCE_ID) @@ -58,8 +58,10 @@ def _atexit_mock(): def test_mutations_batcher_constructor(_setup_batcher, _atexit_mock): - from google.cloud.bigtable.batcher import MAX_OUTSTANDING_ELEMENTS - from google.cloud.bigtable.batcher import MAX_OUTSTANDING_BYTES + from google.cloud.bigtable.batcher import ( + MAX_OUTSTANDING_BYTES, + MAX_OUTSTANDING_ELEMENTS, + ) flush_count = 5 flush_interval = 0.1 @@ -211,183 +213,6 @@ def test_mutations_batcher_context_manager_flushed_when_closed(_setup_batcher): operation_mock.assert_called_once() -<<<<<<< ours -@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 - - 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_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 - - 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) - - 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.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_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 - - table = _Table(TABLE_NAME) - with MutationsBatcher(table=table) as mutation_batcher: - batch_info = _BatchInfo(rows_count=2, mutations_count=2, mutations_size=0) - - 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" - ) - 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() - - -@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) - - 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() - - # 3 batches submitted. 2 batches of 2 items, and the last one a single item batch. - assert mocked_executor_submit.call_count == 3 - - -class _Instance(object): - def __init__(self, client=None): - self._client = client - - -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 - - self.mutation_calls += 1 - - return [Status(code=0) for _ in rows] -======= def test_mutations_batcher_flush_interval(_setup_batcher): table, operation_mock = _setup_batcher flush_interval = 0.5 @@ -409,8 +234,11 @@ def test_mutations_batcher_flush_interval(_setup_batcher): def test_mutations_batcher_response_with_error_codes(_setup_batcher): from google.api_core import exceptions - from google.cloud.bigtable.data.exceptions import FailedMutationEntryError - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) table, operation_mock = _setup_batcher @@ -447,8 +275,11 @@ def test_mutations_batcher_response_with_error_codes(_setup_batcher): 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 - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) table, operation_mock = _setup_batcher @@ -498,4 +329,3 @@ def register(self, func): def unregister(self, func): self._functions.remove(func) ->>>>>>> theirs From 785f1388696cfaad333eea3b1f94c7fc2c16e9f6 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 21 Aug 2026 15:30:52 -0700 Subject: [PATCH 3/5] fixed property references --- .../google-cloud-bigtable/google/cloud/bigtable/batcher.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index bcde32003150..7d8964df9840 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -115,11 +115,11 @@ def __init__( @property def flush_count(self): - return self._flush_count + return self._batcher._flush_limit_count @property def max_row_bytes(self): - return self._max_row_bytes + return self._batcher._flush_limit_bytes def _init_batcher(self): self._batcher = self.table._table_impl.mutations_batcher(**self._batcher_kwargs) From ac62f4c123d9bb77502badb85a74ae96920620a3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 21 Aug 2026 15:38:33 -0700 Subject: [PATCH 4/5] un-deprecate flush_interval --- .../google/cloud/bigtable/batcher.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index 7d8964df9840..6824390e282c 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -42,17 +42,17 @@ def __init__(self, message, exc): 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()` @@ -82,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 From dbfc051e2a50750be2653a6c821ba83aac303809 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 21 Aug 2026 15:42:36 -0700 Subject: [PATCH 5/5] added guard against None causes --- .../google/cloud/bigtable/batcher.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index 6824390e282c..c3268cc09d7b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -131,9 +131,12 @@ def _close_batcher(self): 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. - self._exceptions.put(error.__cause__) + # 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"""