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/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py index 9244d9cf193e..3e6c2b533b1e 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py @@ -18,12 +18,12 @@ import time from typing import TYPE_CHECKING, Sequence -from google.api_core import retry as retries from grpc import StatusCode from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( _attempt_timeout_generator, + _rst_stream_aware_predicate, ) from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry from google.cloud.bigtable.data.exceptions import ( @@ -108,7 +108,7 @@ def __init__( else: self.request = query._to_pb(target) self.target = target - self._predicate = retries.if_exception_type(*retryable_exceptions) + self._predicate = _rst_stream_aware_predicate(*retryable_exceptions) self._last_yielded_row_key: bytes | None = None self._remaining_count: int | None = self.request.rows_limit or None self._operation_metric = metric 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 0019e34518fe..90d29838be7b 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 @@ -16,15 +16,19 @@ import atexit import concurrent.futures +import logging import time import warnings from collections import deque -from typing import TYPE_CHECKING, Sequence, cast +from typing import TYPE_CHECKING, Any, Callable, Sequence, cast + +from google.rpc import code_pb2, status_pb2 from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( TABLE_DEFAULT, _get_retryable_errors, + _get_statuses_from_mutations_exception_group, _get_timeouts, ) from google.cloud.bigtable.data._metrics import ActiveOperationMetric, OperationType @@ -54,6 +58,7 @@ # used to make more readable default values _MB_SIZE = 1024 * 1024 +_LOGGER = logging.getLogger(__name__) @CrossSync.convert_class(sync_name="_FlowControl", add_mapping_for_name="_FlowControl") @@ -290,10 +295,14 @@ 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 # clean up on program exit atexit.register(self._on_exit) @@ -377,7 +386,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 ): @@ -387,7 +400,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 @@ -410,6 +423,7 @@ async def _execute_mutate_rows( list of FailedMutationEntryError objects for mutations that failed. FailedMutationEntryError objects will not contain index information """ + statuses = [status_pb2.Status(code=code_pb2.UNKNOWN) for _ in range(len(batch))] try: operation = CrossSync._MutateRowsOperation( self._target.client._gapic_client, @@ -422,16 +436,29 @@ async def _execute_mutate_rows( ) await operation.start() except MutationsExceptionGroup as e: + statuses = _get_statuses_from_mutations_exception_group(e, len(batch)) + # strip index information from exceptions, since it is not useful in a batch context for subexc in e.exceptions: subexc.index = None return list(e.exceptions) + else: + statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(len(batch))] finally: # mark batch as complete in flow control await self._flow_control.remove_from_flow(batch) + + # Call batch done callback with list of statuses. + if self._user_batch_completed_callback: + try: + self._user_batch_completed_callback(statuses) + except Exception as exc: + _LOGGER.warning( + f"Exception raised in user batch completion callback: {exc}" + ) return [] - def _add_exceptions(self, excs: list[Exception]): + 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, @@ -531,26 +558,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 @@ -562,6 +591,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 c21b50064a6d..c5b17055b581 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py @@ -20,12 +20,26 @@ import enum import time from collections import namedtuple -from typing import TYPE_CHECKING, List, Sequence, Tuple, Union +from typing import ( + TYPE_CHECKING, + Callable, + List, + Optional, + Sequence, + Tuple, + Union, + cast, +) from google.api_core import exceptions as core_exceptions +from google.api_core import retry as retries from google.api_core.retry import RetryFailureReason, exponential_sleep_generator +from google.rpc import code_pb2, status_pb2 -from google.cloud.bigtable.data.exceptions import RetryExceptionGroup +from google.cloud.bigtable.data.exceptions import ( + MutationsExceptionGroup, + RetryExceptionGroup, +) from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery if TYPE_CHECKING: @@ -50,6 +64,15 @@ # used by every data client as a default project name for testing on Bigtable emulator. _DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator" +# Internal error messages that can be retried during ReadRows. Internal error messages with this error +# text should be treated as Unavailable error messages with the same error text, and will therefore be +# treated as Unavailable errors rather than Internal errors. +_RETRYABLE_INTERNAL_ERROR_MESSAGES = ( + "rst_stream", + "rst stream", + "received unexpected eos on data frame from server", +) + # used to identify an active bigtable resource that needs to be warmed through PingAndWarm # each instance/app_profile_id pair needs to be individually tracked _WarmedInstanceKey = namedtuple( @@ -126,6 +149,35 @@ def _retry_exception_factory( return source_exc, cause_exc +def _rst_stream_aware_predicate( + *exception_types: type[Exception], +) -> Callable[[Exception], bool]: + """A custom retry predicate. + + This predicate treats Internal error messages with RST_STREAM errors as + ServiceUnavailable errors and will retry them if the Unavailable exception is retryable. + + Args: + exception_types: Exception types to be retried during operation + + Returns: + Callable[[Exception], bool]: A retry predicate that takes in an exception and + returns whether or not that exception is retryable + """ + # predicate to check for retryable error types + if_exception_type = retries.if_exception_type(*exception_types) + + # special case: treat InternalServerError with rst_stream error message as ServiceUnavailable + def rst_check(e): + return ( + core_exceptions.ServiceUnavailable in exception_types + and isinstance(e, core_exceptions.InternalServerError) + and any(m in e.message.lower() for m in _RETRYABLE_INTERNAL_ERROR_MESSAGES) + ) + + return lambda e: if_exception_type(e) or rst_check(e) + + def _get_timeouts( operation: float | TABLE_DEFAULT, attempt: float | None | TABLE_DEFAULT, @@ -191,6 +243,69 @@ def _align_timeouts(operation: float, attempt: float | None) -> tuple[float, flo return operation, final_attempt +def _get_statuses_from_mutations_exception_group( + exc_group: MutationsExceptionGroup, batch_size: int +) -> list[status_pb2.Status]: + """ + Helper function that populates a list of Status objects with exception information from + the exception group. + + Args: + exc_group: The exception group from a mutate rows operation + batch_size: How many RowMutationGroups were provided to the batch + Returns: + list[status_pb2.Status]: A list of Status proto objects + """ + # We exception handle as follows: + # + # 1. Each exception in the error group is a FailedMutationEntryError, and its + # cause is either a singular exception or a RetryExceptionGroup consisting of + # multiple exceptions. + # + # 2. In the case of a singular exception, if the error does not have a gRPC status + # code, we return a status code of UNKNOWN. + # + # 3. In the case of a RetryExceptionGroup, we use terminal exception in the exception + # group and process that. + statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(batch_size)] + for error in exc_group.exceptions: + if isinstance(error.index, int) and 0 <= error.index < len(statuses): + cause = error.__cause__ + if isinstance(cause, RetryExceptionGroup): + statuses[error.index] = _get_status(cause.exceptions[-1]) + else: + statuses[error.index] = _get_status(cause) + return statuses + + +def _get_status(exc: Optional[Exception]) -> status_pb2.Status: + """ + Helper function that returns a Status object corresponding to the given exception. + + Args: + exc: An exception to be converted into a Status. + Returns: + status_pb2.Status: A Status proto object. + """ + if isinstance(exc, core_exceptions.GoogleAPICallError): + 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( + code=code_pb2.UNKNOWN, + message=str(exc) if exc else "An unknown error has occurred", + ) + + def _validate_timeouts( operation_timeout: float, attempt_timeout: float | None, allow_none: bool = False ): 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/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index 3c2f99283789..ce29226cac53 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -21,11 +21,13 @@ import time from typing import TYPE_CHECKING, Sequence -from google.api_core import retry as retries from grpc import StatusCode from google.cloud.bigtable.data._cross_sync import CrossSync -from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data._helpers import ( + _attempt_timeout_generator, + _rst_stream_aware_predicate, +) from google.cloud.bigtable.data._metrics.tracked_retry import tracked_retry from google.cloud.bigtable.data.exceptions import ( InvalidChunk, @@ -98,7 +100,7 @@ def __init__( else: self.request = query._to_pb(target) self.target = target - self._predicate = retries.if_exception_type(*retryable_exceptions) + self._predicate = _rst_stream_aware_predicate(*retryable_exceptions) self._last_yielded_row_key: bytes | None = None self._remaining_count: int | None = self.request.rows_limit or None self._operation_metric = metric 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 eb1f0055f5c9..e05bf493c7c8 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 @@ -19,15 +19,19 @@ import atexit import concurrent.futures +import logging import time import warnings from collections import deque -from typing import TYPE_CHECKING, Sequence, cast +from typing import TYPE_CHECKING, Any, Callable, Sequence, cast + +from google.rpc import code_pb2, status_pb2 from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( TABLE_DEFAULT, _get_retryable_errors, + _get_statuses_from_mutations_exception_group, _get_timeouts, ) from google.cloud.bigtable.data._metrics import ActiveOperationMetric, OperationType @@ -47,6 +51,7 @@ ) from google.cloud.bigtable.data.mutations import RowMutationEntry _MB_SIZE = 1024 * 1024 +_LOGGER = logging.getLogger(__name__) @CrossSync._Sync_Impl.add_mapping_decorator("_FlowControl") @@ -255,10 +260,13 @@ 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: ( + Callable[[list[status_pb2.Status]], Any] | None + ) = None atexit.register(self._on_exit) def _timer_routine(self, interval: float | None) -> None: @@ -326,7 +334,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 @@ -337,7 +348,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) @@ -355,6 +366,7 @@ def _execute_mutate_rows( list[FailedMutationEntryError]: list of FailedMutationEntryError objects for mutations that failed. FailedMutationEntryError objects will not contain index information""" + statuses = [status_pb2.Status(code=code_pb2.UNKNOWN) for _ in range(len(batch))] try: operation = CrossSync._Sync_Impl._MutateRowsOperation( self._target.client._gapic_client, @@ -367,14 +379,24 @@ def _execute_mutate_rows( ) operation.start() except MutationsExceptionGroup as e: + statuses = _get_statuses_from_mutations_exception_group(e, len(batch)) for subexc in e.exceptions: subexc.index = None return list(e.exceptions) + else: + statuses = [status_pb2.Status(code=code_pb2.OK) for _ in range(len(batch))] finally: self._flow_control.remove_from_flow(batch) + if self._user_batch_completed_callback: + try: + self._user_batch_completed_callback(statuses) + except Exception as exc: + _LOGGER.warning( + f"Exception raised in user batch completion callback: {exc}" + ) return [] - def _add_exceptions(self, excs: list[Exception]): + 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. @@ -453,30 +475,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 bfc1a9eadddd..840e44161200 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py @@ -145,13 +145,14 @@ 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 @@ -169,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: @@ -187,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: @@ -207,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/google/cloud/bigtable/helpers.py b/packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py index 6bc423f50439..87478d94b845 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/helpers.py @@ -28,3 +28,30 @@ def batched(iterable: Iterable[T], n) -> Generator[Tuple[T, ...], None, None]: while batch: yield batch batch = tuple(islice(it, n)) + + +class _MappableAttributesMixin: + """ + Mixin for classes that need some of their attribute names remapped. + + This is for taking some of the classes from the data client row filters + and row range classes that are 1:1 with their legacy client counterparts but with + some of their attributes renamed. To use in a class, override the base class with this mixin + class and define a map _attribute_map from legacy client attributes to data client + attributes. + + Attributes are remapped and redefined in __init__ as well as getattr/setattr. + """ + + def __init__(self, *args, **kwargs): + new_kwargs = {self._attribute_map.get(k, k): v for (k, v) in kwargs.items()} + super(_MappableAttributesMixin, self).__init__(*args, **new_kwargs) + + def __getattr__(self, name): + if name not in self._attribute_map: + raise AttributeError + return getattr(self, self._attribute_map[name]) + + def __setattr__(self, name, value): + attribute = self._attribute_map.get(name, name) + super(_MappableAttributesMixin, self).__setattr__(attribute, value) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row.py index d769f2d24644..f6d04b8d9146 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row.py @@ -14,17 +14,16 @@ """User-friendly container for Google Cloud Bigtable Row.""" -import struct - +from google.api_core.exceptions import GoogleAPICallError from google.cloud._helpers import ( _datetime_from_microseconds, # type: ignore _microseconds_from_datetime, # type: ignore _to_bytes, # type: ignore ) +from google.rpc import code_pb2, status_pb2 -from google.cloud.bigtable_v2.types import data as data_v2_pb2 - -_PACK_I64 = struct.Struct(">q").pack +from google.cloud.bigtable.data import mutations +from google.cloud.bigtable.data import read_modify_write_rules as rmw_rules MAX_MUTATIONS = 100000 """The maximum number of mutations that a row can accumulate.""" @@ -158,26 +157,21 @@ def _set_cell(self, column_family_id, column, value, timestamp=None, state=None) :param state: (Optional) The state that is passed along to :meth:`_get_mutations`. """ - column = _to_bytes(column) - if isinstance(value, int): - value = _PACK_I64(value) - value = _to_bytes(value) if timestamp is None: - # Use -1 for current Bigtable server time. - timestamp_micros = -1 + # Use current Bigtable server time. + timestamp_micros = mutations._SERVER_SIDE_TIMESTAMP else: timestamp_micros = _microseconds_from_datetime(timestamp) # Truncate to millisecond granularity. timestamp_micros -= timestamp_micros % 1000 - mutation_val = data_v2_pb2.Mutation.SetCell( - family_name=column_family_id, - column_qualifier=column, + mutation = mutations.SetCell( + family=column_family_id, + qualifier=column, + new_value=value, timestamp_micros=timestamp_micros, - value=value, ) - mutation_pb = data_v2_pb2.Mutation(set_cell=mutation_val) - self._get_mutations(state).append(mutation_pb) + self._get_mutations(state).append(mutation) def _delete(self, state=None): """Helper for :meth:`delete` @@ -192,9 +186,7 @@ def _delete(self, state=None): :param state: (Optional) The state that is passed along to :meth:`_get_mutations`. """ - mutation_val = data_v2_pb2.Mutation.DeleteFromRow() - mutation_pb = data_v2_pb2.Mutation(delete_from_row=mutation_val) - self._get_mutations(state).append(mutation_pb) + self._get_mutations(state).append(mutations.DeleteAllFromRow()) def _delete_cells(self, column_family_id, columns, time_range=None, state=None): """Helper for :meth:`delete_cell` and :meth:`delete_cells`. @@ -221,33 +213,30 @@ def _delete_cells(self, column_family_id, columns, time_range=None, state=None): :param state: (Optional) The state that is passed along to :meth:`_get_mutations`. """ - mutations_list = self._get_mutations(state) if columns is self.ALL_COLUMNS: - mutation_val = data_v2_pb2.Mutation.DeleteFromFamily( - family_name=column_family_id + self._get_mutations(state).append( + mutations.DeleteAllFromFamily(family_to_delete=column_family_id) ) - mutation_pb = data_v2_pb2.Mutation(delete_from_family=mutation_val) - mutations_list.append(mutation_pb) else: - delete_kwargs = {} - if time_range is not None: - delete_kwargs["time_range"] = time_range._to_pb() + timestamps = time_range._to_dict() if time_range else {} + start_timestamp_micros = timestamps.get("start_timestamp_micros") + end_timestamp_micros = timestamps.get("end_timestamp_micros") to_append = [] for column in columns: column = _to_bytes(column) - # time_range will never change if present, but the rest of - # delete_kwargs will - delete_kwargs.update( - family_name=column_family_id, column_qualifier=column + to_append.append( + mutations.DeleteRangeFromColumn( + family=column_family_id, + qualifier=column, + start_timestamp_micros=start_timestamp_micros, + end_timestamp_micros=end_timestamp_micros, + ) ) - mutation_val = data_v2_pb2.Mutation.DeleteFromColumn(**delete_kwargs) - mutation_pb = data_v2_pb2.Mutation(delete_from_column=mutation_val) - to_append.append(mutation_pb) # We don't add the mutations until all columns have been # processed without error. - mutations_list.extend(to_append) + self._get_mutations(state).extend(to_append) class DirectRow(_SetDeleteRow): @@ -285,7 +274,7 @@ class DirectRow(_SetDeleteRow): def __init__(self, row_key, table=None): super(DirectRow, self).__init__(row_key, table) - self._pb_mutations = [] + self._mutations = [] def _get_mutations(self, state=None): # pylint: disable=unused-argument """Gets the list of mutations for a given state. @@ -300,7 +289,12 @@ def _get_mutations(self, state=None): # pylint: disable=unused-argument :rtype: list :returns: The list to add new mutations to (for the current state). """ - return self._pb_mutations + return self._mutations + + def _get_mutation_pbs(self): + """Gets the list of mutation protos.""" + + return [mut._to_pb() for mut in self._get_mutations()] def get_mutations_size(self): """Gets the total mutations size for current row @@ -314,7 +308,7 @@ def get_mutations_size(self): """ mutation_size = 0 - for mutation in self._get_mutations(): + for mutation in self._get_mutation_pbs(): mutation_size += mutation._pb.ByteSize() return mutation_size @@ -449,7 +443,8 @@ def delete_cells(self, column_family_id, columns, time_range=None): def commit(self): """Makes a ``MutateRow`` API request. - If no mutations have been created in the row, no request is made. + If no mutations have been created in the row, no request is made and a + ValueError is raised instead. Mutations are applied atomically and in order, meaning that earlier mutations can be masked / negated by later ones. Cells already present @@ -468,14 +463,22 @@ def commit(self): :rtype: :class:`~google.rpc.status_pb2.Status` :returns: A response status (`google.rpc.status_pb2.Status`) representing success or failure of the row committed. - :raises: :exc:`~.table.TooManyMutationsError` if the number of - mutations is greater than 100,000. + :raises: ValueError: if no mutations have been created in the row """ - response = self._table.mutate_rows([self]) - - self.clear() - - return response[0] + try: + self._table._table_impl.mutate_row(self.row_key, self._get_mutations()) + return status_pb2.Status(code=code_pb2.OK) + except GoogleAPICallError as e: + # If the RPC call returns an error, extract the error into a status object, if possible. + return status_pb2.Status( + code=e.grpc_status_code.value[0] + if e.grpc_status_code is not None + else code_pb2.UNKNOWN, + message=e.message, + details=e.details, + ) + finally: + self.clear() def clear(self): """Removes all currently accumulated mutations on the current row. @@ -487,7 +490,7 @@ def clear(self): :end-before: [END bigtable_api_row_clear] :dedent: 4 """ - del self._pb_mutations[:] + del self._mutations[:] class ConditionalRow(_SetDeleteRow): @@ -598,17 +601,15 @@ def commit(self): % (MAX_MUTATIONS, num_true_mutations, num_false_mutations) ) - data_client = self._table._instance._client.table_data_client - resp = data_client.check_and_mutate_row( - table_name=self._table.name, + table = self._table._table_impl + resp = table.check_and_mutate_row( row_key=self._row_key, - predicate_filter=self._filter._to_pb(), - app_profile_id=self._table._app_profile_id, - true_mutations=true_mutations, - false_mutations=false_mutations, + predicate=self._filter, + true_case_mutations=true_mutations, + false_case_mutations=false_mutations, ) self.clear() - return resp.predicate_matched + return resp # pylint: disable=arguments-differ def set_cell(self, column_family_id, column, value, timestamp=None, state=True): @@ -798,7 +799,7 @@ class AppendRow(Row): def __init__(self, row_key, table): super(AppendRow, self).__init__(row_key, table) - self._rule_pb_list = [] + self._rule_list = [] def clear(self): """Removes all currently accumulated modifications on current row. @@ -810,7 +811,7 @@ def clear(self): :end-before: [END bigtable_api_row_clear] :dedent: 4 """ - del self._rule_pb_list[:] + del self._rule_list[:] def append_cell_value(self, column_family_id, column, value): """Appends a value to an existing cell. @@ -843,12 +844,11 @@ def append_cell_value(self, column_family_id, column, value): the targeted cell is unset, it will be treated as containing the empty string. """ - column = _to_bytes(column) - value = _to_bytes(value) - rule_pb = data_v2_pb2.ReadModifyWriteRule( - family_name=column_family_id, column_qualifier=column, append_value=value + self._rule_list.append( + rmw_rules.AppendValueRule( + family=column_family_id, qualifier=column, append_value=value + ) ) - self._rule_pb_list.append(rule_pb) def increment_cell_value(self, column_family_id, column, int_value): """Increments a value in an existing cell. @@ -887,13 +887,11 @@ def increment_cell_value(self, column_family_id, column, int_value): big-endian signed integer), or the entire request will fail. """ - column = _to_bytes(column) - rule_pb = data_v2_pb2.ReadModifyWriteRule( - family_name=column_family_id, - column_qualifier=column, - increment_amount=int_value, + self._rule_list.append( + rmw_rules.IncrementRule( + family=column_family_id, qualifier=column, increment_amount=int_value + ) ) - self._rule_pb_list.append(rule_pb) def commit(self): """Makes a ``ReadModifyWriteRow`` API request. @@ -926,7 +924,7 @@ def commit(self): :raises: :class:`ValueError ` if the number of mutations exceeds the :data:`MAX_MUTATIONS`. """ - num_mutations = len(self._rule_pb_list) + num_mutations = len(self._rule_list) if num_mutations == 0: return {} if num_mutations > MAX_MUTATIONS: @@ -935,12 +933,10 @@ def commit(self): "allowable %d." % (num_mutations, MAX_MUTATIONS) ) - data_client = self._table._instance._client.table_data_client - row_response = data_client.read_modify_write_row( - table_name=self._table.name, + table = self._table._table_impl + row_response = table.read_modify_write_row( row_key=self._row_key, - rules=self._rule_pb_list, - app_profile_id=self._table._app_profile_id, + rules=self._rule_list, ) # Reset modifications after commit-ing request. @@ -984,47 +980,13 @@ def _parse_rmw_row_response(row_response): } """ result = {} - for column_family in row_response.row.families: - column_family_id, curr_family = _parse_family_pb(column_family) - result[column_family_id] = curr_family + for cell in row_response.cells: + column_family = result.setdefault(cell.family, {}) + column = column_family.setdefault(cell.qualifier, []) + column.append((cell.value, _datetime_from_microseconds(cell.timestamp_micros))) return result -def _parse_family_pb(family_pb): - """Parses a Family protobuf into a dictionary. - - :type family_pb: :class:`._generated.data_pb2.Family` - :param family_pb: A protobuf - - :rtype: tuple - :returns: A string and dictionary. The string is the name of the - column family and the dictionary has column names (within the - family) as keys and cell lists as values. Each cell is - represented with a two-tuple with the value (in bytes) and the - timestamp for the cell. For example: - - .. code:: python - - { - b'col-name1': [ - (b'cell-val', datetime.datetime(...)), - (b'cell-val-newer', datetime.datetime(...)), - ], - b'col-name2': [ - (b'altcol-cell-val', datetime.datetime(...)), - ], - } - """ - result = {} - for column in family_pb.columns: - result[column.qualifier] = cells = [] - for cell in column.cells: - val_pair = (cell.value, _datetime_from_microseconds(cell.timestamp_micros)) - cells.append(val_pair) - - return family_pb.name, result - - class PartialRowData(object): """Representation of partial row in a Google Cloud Bigtable Table. @@ -1039,6 +1001,16 @@ def __init__(self, row_key): self._row_key = row_key self._cells = {} + @classmethod + def _from_data_client_row(cls, row): + partial_row_data = cls(row.row_key) + for column_family in row._index: + columns = {} + for column, items in row._index[column_family].items(): + columns[column] = [Cell._from_data_client_cell(item) for item in items] + partial_row_data._cells[column_family] = columns + return partial_row_data + def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented @@ -1242,6 +1214,14 @@ def from_pb(cls, cell_pb): else: return cls(cell_pb.value, cell_pb.timestamp_micros) + @classmethod + def _from_data_client_cell(cls, cell): + return cls( + cell.value, + cell.timestamp_micros, + cell.labels, + ) + @property def timestamp(self): return _datetime_from_microseconds(self.timestamp_micros) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py index f00b9d568bc4..c55fc2ac72f5 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_data.py @@ -14,17 +14,16 @@ """Container for Google Cloud Bigtable Cells and Streaming Row Contents.""" -import copy -import warnings +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Generator -import grpc # type: ignore from google.api_core import exceptions, retry -from google.cloud._helpers import _to_bytes # type: ignore from google.cloud.bigtable.row import Cell, InvalidChunk, PartialRowData -from google.cloud.bigtable.row_merger import _RowMerger, _State -from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 -from google.cloud.bigtable_v2.types import data as data_v2_pb2 + +if TYPE_CHECKING: + from google.cloud.bigtable.data.row import Row # Some classes need to be re-exported here to keep backwards # compatibility. Those classes were moved to row_merger, but we dont want to @@ -57,36 +56,7 @@ class InvalidRetryRequest(RuntimeError): """Exception raised when retry request is invalid.""" -RETRYABLE_INTERNAL_ERROR_MESSAGES = ( - "rst_stream", - "rst stream", - "received unexpected eos on data frame from server", -) -"""Internal error messages that can be retried during read row and mutation.""" - - -def _retriable_internal_server_error(exc): - """ - Return True if the internal server error is retriable. - """ - return isinstance(exc, exceptions.InternalServerError) and any( - retryable_message in exc.message.lower() - for retryable_message in RETRYABLE_INTERNAL_ERROR_MESSAGES - ) - - -def _retry_read_rows_exception(exc): - """Return True if the exception is retriable for read row requests.""" - if isinstance(exc, grpc.RpcError): - exc = exceptions.from_grpc_error(exc) - - return _retriable_internal_server_error(exc) or isinstance( - exc, (exceptions.ServiceUnavailable, exceptions.DeadlineExceeded) - ) - - DEFAULT_RETRY_READ_ROWS = retry.Retry( - predicate=_retry_read_rows_exception, initial=1.0, maximum=15.0, multiplier=2.0, @@ -102,97 +72,119 @@ def _retry_read_rows_exception(exc): class PartialRowsData(object): """Convenience wrapper for consuming a ``ReadRows`` streaming response. + This class will be returned by the ``read_rows`` method, and should not + be constructed manually. + :type read_method: :class:`client._table_data_client.read_rows` :param read_method: ``ReadRows`` method. - :type request: :class:`data_messages_v2_pb2.ReadRowsRequest` - :param request: The ``ReadRowsRequest`` message used to create a - ReadRowsResponse iterator. If the iterator fails, a new - iterator is created, allowing the scan to continue from - the point just beyond the last successfully read row, - identified by self.last_scanned_row_key. The retry happens - inside of the Retry class, using a predicate for the - expected exceptions during iteration. - - :type retry: :class:`~google.api_core.retry.Retry` - :param retry: (Optional) Retry delay and deadline arguments. To override, - the default value :attr:`DEFAULT_RETRY_READ_ROWS` can be - used and modified with the - :meth:`~google.api_core.retry.Retry.with_delay` method - or the - :meth:`~google.api_core.retry.Retry.with_deadline` method. + :type generator: :class:`Iterable[Row]` + :param generator: The `Row` iterator from :meth:`Table.read_rows`. """ - NEW_ROW = "New row" # No cells yet complete for row - ROW_IN_PROGRESS = "Row in progress" # Some cells complete for row - CELL_IN_PROGRESS = "Cell in progress" # Incomplete cell for row - - STATE_NEW_ROW = 1 - STATE_ROW_IN_PROGRESS = 2 - STATE_CELL_IN_PROGRESS = 3 - - read_states = { - STATE_NEW_ROW: NEW_ROW, - STATE_ROW_IN_PROGRESS: ROW_IN_PROGRESS, - STATE_CELL_IN_PROGRESS: CELL_IN_PROGRESS, - } - - def __init__(self, read_method, request, retry=DEFAULT_RETRY_READ_ROWS): - # Counter for rows returned to the user - self._counter = 0 - self._row_merger = _RowMerger() - - # May be cached from previous response - self.last_scanned_row_key = None - self.read_method = read_method - self.request = request - self.retry = retry - - # The `timeout` parameter must be somewhat greater than the value - # contained in `self.retry`, in order to avoid race-like condition and - # allow registering the first deadline error before invoking the retry. - # Otherwise there is a risk of entering an infinite loop that resets - # the timeout counter just before it being triggered. The increment - # by 1 second here is customary but should not be much less than that. - self.response_iterator = read_method( - request, timeout=self.retry._deadline + 1, retry=self.retry + def __init__( + self, + *args: Any, + generator: Generator[Row, Any, Any] | None = None, + **kwargs: Any, + ) -> None: + self._generator: Generator[Row, Any, Any] + if generator is not None and not args and not kwargs: + self._generator = generator + else: + import warnings + + warnings.warn( + "Direct instantiation of `PartialRowsData` is deprecated " + "and will be removed in a future release. `PartialRowsData` should only be obtained " + "by calling `Table.read_rows()`. An empty stream has been initialized.", + DeprecationWarning, + stacklevel=2, + ) + self._generator = (x for x in ()) + + self.rows: dict[bytes, PartialRowData] = {} + + @classmethod + def _from_generator(cls, generator: Generator[Row, Any, Any]) -> PartialRowsData: + """Internal constructor for Table.read_rows.""" + return cls(generator=generator) + + @property + def state(self): + import warnings + + warnings.warn( + "The `state` attribute on `PartialRowsData` is deprecated " + "and will return None.", + DeprecationWarning, + stacklevel=2, ) + return None - self.rows = {} + @property + def last_scanned_row_key(self): + import warnings - # Flag to stop iteration, for any reason not related to self.retry() - self._cancelled = False + warnings.warn( + "The `last_scanned_row_key` attribute on `PartialRowsData` is deprecated " + "and will return None.", + DeprecationWarning, + stacklevel=2, + ) + return None @property - def state(self): # pragma: NO COVER - """ - DEPRECATED: this property is deprecated and will be removed in the - future. - """ + def read_method(self): + import warnings + warnings.warn( - "`PartialRowsData#state()` is deprecated and will be removed in the future", + "The `read_method` attribute on `PartialRowsData` is deprecated " + "and will return None.", DeprecationWarning, stacklevel=2, ) + return None - # Best effort: try to map internal RowMerger states to old strings for - # backwards compatibility - internal_state = self._row_merger.state - if internal_state == _State.ROW_START: - return self.NEW_ROW - # note: _State.CELL_START, _State.CELL_COMPLETE are transient states - # and will not be visible in between chunks - elif internal_state == _State.CELL_IN_PROGRESS: - return self.CELL_IN_PROGRESS - elif internal_state == _State.ROW_COMPLETE: - return self.NEW_ROW - else: - raise RuntimeError("unexpected internal state: " + self._) + @property + def request(self): + import warnings + + warnings.warn( + "The `request` attribute on `PartialRowsData` is deprecated " + "and will return None.", + DeprecationWarning, + stacklevel=2, + ) + return None + + @property + def retry(self): + import warnings + + warnings.warn( + "The `retry` attribute on `PartialRowsData` is deprecated " + "and will return None.", + DeprecationWarning, + stacklevel=2, + ) + return None + + @property + def response_iterator(self): + import warnings + + warnings.warn( + "The `response_iterator` attribute on `PartialRowsData` is deprecated " + "and will return None.", + DeprecationWarning, + stacklevel=2, + ) + return None def cancel(self): """Cancels the iterator, closing the stream.""" - self._cancelled = True - self.response_iterator.cancel() + self._generator.close() def consume_all(self, max_loops=None): """Consume the streamed responses until there are no more. @@ -202,176 +194,27 @@ def consume_all(self, max_loops=None): class as a generator instead. :type max_loops: int - :param max_loops: (Optional) Maximum number of times to try to consume - an additional ``ReadRowsResponse``. You can use this - to avoid long wait times. + :param max_loops: (Deprecated). Maximum number of times to try to consume + an additional ``ReadRowsResponse``. This parameter is + deprecated and is only kept for backwards compatibility. """ for row in self: self.rows[row.row_key] = row - def _create_retry_request(self): - """Helper for :meth:`__iter__`.""" - req_manager = _ReadRowsRequestManager( - self.request, self.last_scanned_row_key, self._counter - ) - return req_manager.build_updated_request() - - def _on_error(self, exc): - """Helper for :meth:`__iter__`.""" - # restart the read scan from AFTER the last successfully read row - retry_request = self.request - if self.last_scanned_row_key: - retry_request = self._create_retry_request() - - self._row_merger = _RowMerger(self._row_merger.last_seen_row_key) - self.response_iterator = self.read_method(retry_request, retry=self.retry) - - def _read_next(self): - """Helper for :meth:`__iter__`.""" - return next(self.response_iterator) - - def _read_next_response(self): - """Helper for :meth:`__iter__`.""" - resp_protoplus = self.retry(self._read_next, on_error=self._on_error)() - # unwrap the underlying protobuf, there is a significant amount of - # overhead that protoplus imposes for very little gain. The protos - # are not user visible, so we just use the raw protos for merging. - return data_messages_v2_pb2.ReadRowsResponse.pb(resp_protoplus) - def __iter__(self): - """Consume the ``ReadRowsResponse`` s from the stream. - Read the rows and yield each to the reader - - Parse the response and its chunks into a new/existing row in - :attr:`_rows`. Rows are returned in order by row key. + """Consume the ``Row`` s from the stream. + Convert them to ``PartialRowData`` and yield each to the reader. """ - while not self._cancelled: - try: - response = self._read_next_response() - except StopIteration: - self._row_merger.finalize() - break - except InvalidRetryRequest: - self._cancelled = True - break - - for row in self._row_merger.process_chunks(response): - self.last_scanned_row_key = self._row_merger.last_seen_row_key - self._counter += 1 - - yield row - - if self._cancelled: - break - # The last response might not have generated any rows, but it - # could've updated last_scanned_row_key - self.last_scanned_row_key = self._row_merger.last_seen_row_key - - -class _ReadRowsRequestManager(object): - """Update the ReadRowsRequest message in case of failures by - filtering the already read keys. - - :type message: class:`data_messages_v2_pb2.ReadRowsRequest` - :param message: Original ReadRowsRequest containing all of the parameters - of API call - - :type last_scanned_key: bytes - :param last_scanned_key: last successfully scanned key - - :type rows_read_so_far: int - :param rows_read_so_far: total no of rows successfully read so far. - this will be used for updating rows_limit - - """ - - def __init__(self, message, last_scanned_key, rows_read_so_far): - self.message = message - self.last_scanned_key = last_scanned_key - self.rows_read_so_far = rows_read_so_far - - def build_updated_request(self): - """Updates the given message request as per last scanned key""" - - resume_request = data_messages_v2_pb2.ReadRowsRequest() - data_messages_v2_pb2.ReadRowsRequest.copy_from(resume_request, self.message) - - if self.message.rows_limit != 0: - row_limit_remaining = self.message.rows_limit - self.rows_read_so_far - if row_limit_remaining > 0: - resume_request.rows_limit = row_limit_remaining - else: - raise InvalidRetryRequest - - # if neither RowSet.row_keys nor RowSet.row_ranges currently exist, - # add row_range that starts with last_scanned_key as start_key_open - # to request only rows that have not been returned yet - if "rows" not in self.message: - row_range = data_v2_pb2.RowRange(start_key_open=self.last_scanned_key) - resume_request.rows = data_v2_pb2.RowSet(row_ranges=[row_range]) - else: - row_keys = self._filter_rows_keys() - row_ranges = self._filter_row_ranges() - - if len(row_keys) == 0 and len(row_ranges) == 0: - # Avoid sending empty row_keys and row_ranges - # if that was not the intention - raise InvalidRetryRequest - - resume_request.rows = data_v2_pb2.RowSet( - row_keys=row_keys, row_ranges=row_ranges - ) - return resume_request - - def _filter_rows_keys(self): - """Helper for :meth:`build_updated_request`""" - return [ - row_key - for row_key in self.message.rows.row_keys - if row_key > self.last_scanned_key - ] - - def _filter_row_ranges(self): - """Helper for :meth:`build_updated_request`""" - new_row_ranges = [] - - for row_range in self.message.rows.row_ranges: - # if current end_key (open or closed) is set, return its value, - # if not, set to empty string (''). - # NOTE: Empty string in end_key means "end of table" - end_key = self._end_key_set(row_range) - # if end_key is already read, skip to the next row_range - if end_key and self._key_already_read(end_key): - continue - - # if current start_key (open or closed) is set, return its value, - # if not, then set to empty string ('') - # NOTE: Empty string in start_key means "beginning of table" - start_key = self._start_key_set(row_range) - - # if start_key was already read or doesn't exist, - # create a row_range with last_scanned_key as start_key_open - # to be passed to retry request - retry_row_range = row_range - if self._key_already_read(start_key): - retry_row_range = copy.deepcopy(row_range) - retry_row_range.start_key_closed = _to_bytes("") - retry_row_range.start_key_open = self.last_scanned_key - - new_row_ranges.append(retry_row_range) - - return new_row_ranges - - def _key_already_read(self, key): - """Helper for :meth:`_filter_row_ranges`""" - return key <= self.last_scanned_key - - @staticmethod - def _start_key_set(row_range): - """Helper for :meth:`_filter_row_ranges`""" - return row_range.start_key_open or row_range.start_key_closed - - @staticmethod - def _end_key_set(row_range): - """Helper for :meth:`_filter_row_ranges`""" - return row_range.end_key_open or row_range.end_key_closed + try: + for row in self._generator: + yield PartialRowData._from_data_client_row(row) + + # Any exception from the generator should cancel the iterator. A + # timeout, defined by catching a DeadlineExceeded, should be reraised + # as a RetryError instead. + except exceptions.DeadlineExceeded as e: + self.cancel() + raise exceptions.RetryError(e.message, e.__cause__) + except Exception as e: + self.cancel() + raise e diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py index 449543bc132c..433ac7c04ef5 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py @@ -50,6 +50,7 @@ from google.cloud.bigtable.data.row_filters import ( TimestampRangeFilter as BaseTimestampRangeFilter, ) +from google.cloud.bigtable.helpers import _MappableAttributesMixin _PACK_I64 = struct.Struct(">q").pack @@ -70,35 +71,6 @@ def _deprecated_to_pb(self): # Provide to_pb alias with DeprecationWarning for backwards compatibility with legacy client code RowFilter.to_pb = _deprecated_to_pb # type: ignore[attr-defined] TimestampRange.to_pb = _deprecated_to_pb # type: ignore[attr-defined] - - -class _MappableAttributesMixin: - """ - Mixin for classes that need some of their attribute names remapped. - - This is for taking some of the classes from the data client row filters - that are 1:1 with their legacy client counterparts but with some of their - attributes renamed. To use in a class, override the base class with this mixin - class and define a map _attribute_map from legacy client attributes to data client - attributes. - - Attributes are remapped and redefined in __init__ as well as getattr/setattr. - """ - - def __init__(self, *args, **kwargs): - new_kwargs = {self._attribute_map.get(k, k): v for (k, v) in kwargs.items()} - super(_MappableAttributesMixin, self).__init__(*args, **new_kwargs) - - def __getattr__(self, name): - if name not in self._attribute_map: - raise AttributeError - return getattr(self, self._attribute_map[name]) - - def __setattr__(self, name, value): - attribute = self._attribute_map.get(name, name) - super(_MappableAttributesMixin, self).__setattr__(attribute, value) - - # The classes defined below are to provide constructors and members # that have an interface that does not match the one used by the data # client, for backwards compatibility purposes. diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row_merger.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_merger.py deleted file mode 100644 index 6fe9cdca505a..000000000000 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_merger.py +++ /dev/null @@ -1,251 +0,0 @@ -from collections import OrderedDict -from enum import Enum - -from google.cloud.bigtable.row import Cell, InvalidChunk, PartialRowData - -_MISSING_COLUMN_FAMILY = "Column family {} is not among the cells stored in this row." -_MISSING_COLUMN = ( - "Column {} is not among the cells stored in this row in the column family {}." -) -_MISSING_INDEX = ( - "Index {!r} is not valid for the cells stored in this row for column {} " - "in the column family {}. There are {} such cells." -) - - -class _State(Enum): - ROW_START = "ROW_START" - CELL_START = "CELL_START" - CELL_IN_PROGRESS = "CELL_IN_PROGRESS" - CELL_COMPLETE = "CELL_COMPLETE" - ROW_COMPLETE = "ROW_COMPLETE" - - -class _PartialRow(object): - __slots__ = [ - "row_key", - "cells", - "last_family", - "last_family_cells", - "last_qualifier", - "last_qualifier_cells", - "cell", - ] - - def __init__(self, row_key): - self.row_key = row_key - self.cells = OrderedDict() - - self.last_family = None - self.last_family_cells = OrderedDict() - self.last_qualifier = None - self.last_qualifier_cells = [] - - self.cell = None - - -class _PartialCell(object): - __slots__ = ["family", "qualifier", "timestamp", "labels", "value", "value_index"] - - def __init__(self): - self.family = None - self.qualifier = None - self.timestamp = None - self.labels = None - self.value = None - self.value_index = 0 - - -class _RowMerger(object): - """ - State machine to merge chunks from a response stream into logical rows. - - The implementation is a fairly linear state machine that is implemented as - a method for every state in the _State enum. In general the states flow - from top to bottom with some repetition. Each state handler will do some - sanity checks, update in progress data and set the next state. - - There can be multiple state transitions for each chunk, i.e. a single chunk - row will flow from ROW_START -> CELL_START -> CELL_COMPLETE -> ROW_COMPLETE - in a single iteration. - """ - - __slots__ = ["state", "last_seen_row_key", "row"] - - def __init__(self, last_seen_row=b""): - self.last_seen_row_key = last_seen_row - self.state = _State.ROW_START - self.row = None - - def process_chunks(self, response): - """ - Process the chunks in the given response and yield logical rows. - This class will maintain state across multiple response protos. - """ - if response.last_scanned_row_key: - if self.last_seen_row_key >= response.last_scanned_row_key: - raise InvalidChunk("Last scanned row key is out of order") - self.last_seen_row_key = response.last_scanned_row_key - - for chunk in response.chunks: - if chunk.reset_row: - self._handle_reset(chunk) - continue - - if self.state == _State.ROW_START: - self._handle_row_start(chunk) - - if self.state == _State.CELL_START: - self._handle_cell_start(chunk) - - if self.state == _State.CELL_IN_PROGRESS: - self._handle_cell_in_progress(chunk) - - if self.state == _State.CELL_COMPLETE: - self._handle_cell_complete(chunk) - - if self.state == _State.ROW_COMPLETE: - yield self._handle_row_complete(chunk) - elif chunk.commit_row: - raise InvalidChunk( - f"Chunk tried to commit row in wrong state (${self.state})" - ) - - def _handle_reset(self, chunk): - if self.state == _State.ROW_START: - raise InvalidChunk("Bare reset") - if chunk.row_key: - raise InvalidChunk("Reset chunk has a row key") - if chunk.HasField("family_name"): - raise InvalidChunk("Reset chunk has family_name") - if chunk.HasField("qualifier"): - raise InvalidChunk("Reset chunk has qualifier") - if chunk.timestamp_micros: - raise InvalidChunk("Reset chunk has a timestamp") - if chunk.labels: - raise InvalidChunk("Reset chunk has labels") - if chunk.value: - raise InvalidChunk("Reset chunk has a value") - - self.state = _State.ROW_START - self.row = None - - def _handle_row_start(self, chunk): - if not chunk.row_key: - raise InvalidChunk("New row is missing a row key") - if self.last_seen_row_key and self.last_seen_row_key >= chunk.row_key: - raise InvalidChunk("Out of order row keys") - - self.row = _PartialRow(chunk.row_key) - self.state = _State.CELL_START - - def _handle_cell_start(self, chunk): - # Ensure that all chunks after the first one either are missing a row - # key or the row is the same - if self.row.cells and chunk.row_key and chunk.row_key != self.row.row_key: - raise InvalidChunk("row key changed mid row") - - if not self.row.cell: - self.row.cell = _PartialCell() - - # Cells can inherit family/qualifier from previous cells - # However if the family changes, then qualifier must be specified as well - if chunk.HasField("family_name"): - self.row.cell.family = chunk.family_name.value - self.row.cell.qualifier = None - if not self.row.cell.family: - raise InvalidChunk("missing family for a new cell") - - if chunk.HasField("qualifier"): - self.row.cell.qualifier = chunk.qualifier.value - if self.row.cell.qualifier is None: - raise InvalidChunk("missing qualifier for a new cell") - - self.row.cell.timestamp = chunk.timestamp_micros - self.row.cell.labels = chunk.labels - - if chunk.value_size > 0: - # explicitly avoid pre-allocation as it seems that bytearray - # concatenation performs better than slice copies. - self.row.cell.value = bytearray() - self.state = _State.CELL_IN_PROGRESS - else: - self.row.cell.value = chunk.value - self.state = _State.CELL_COMPLETE - - def _handle_cell_in_progress(self, chunk): - # if this isn't the first cell chunk, make sure that everything except - # the value stayed constant. - if self.row.cell.value_index > 0: - if chunk.row_key: - raise InvalidChunk("found row key mid cell") - if chunk.HasField("family_name"): - raise InvalidChunk("In progress cell had a family name") - if chunk.HasField("qualifier"): - raise InvalidChunk("In progress cell had a qualifier") - if chunk.timestamp_micros: - raise InvalidChunk("In progress cell had a timestamp") - if chunk.labels: - raise InvalidChunk("In progress cell had labels") - - self.row.cell.value += chunk.value - self.row.cell.value_index += len(chunk.value) - - if chunk.value_size > 0: - self.state = _State.CELL_IN_PROGRESS - else: - self.row.cell.value = bytes(self.row.cell.value) - self.state = _State.CELL_COMPLETE - - def _handle_cell_complete(self, chunk): - # since we are guaranteed that all family & qualifier cells are - # contiguous, we can optimize away the dict lookup by caching the last - # family/qualifier and simply comparing and appending - family_changed = False - if self.row.last_family != self.row.cell.family: - family_changed = True - self.row.last_family = self.row.cell.family - self.row.cells[self.row.cell.family] = self.row.last_family_cells = ( - OrderedDict() - ) - - if family_changed or self.row.last_qualifier != self.row.cell.qualifier: - self.row.last_qualifier = self.row.cell.qualifier - self.row.last_family_cells[self.row.cell.qualifier] = ( - self.row.last_qualifier_cells - ) = [] - - self.row.last_qualifier_cells.append( - Cell( - self.row.cell.value, - self.row.cell.timestamp, - self.row.cell.labels, - ) - ) - - self.row.cell.timestamp = 0 - self.row.cell.value = None - self.row.cell.value_index = 0 - - if not chunk.commit_row: - self.state = _State.CELL_START - else: - self.state = _State.ROW_COMPLETE - - def _handle_row_complete(self, chunk): - new_row = PartialRowData(self.row.row_key) - new_row._cells = self.row.cells - - self.last_seen_row_key = new_row.row_key - self.row = None - self.state = _State.ROW_START - - return new_row - - def finalize(self): - """ - Must be called at the end of the stream to ensure there are no unmerged - rows. - """ - if self.row or self.state != _State.ROW_START: - raise ValueError("The row remains partial / is not committed.") diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py index 99dae972a6be..a7df6baf0d9a 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_set.py @@ -14,7 +14,13 @@ """User-friendly container for Google Cloud Bigtable RowSet""" -from google.cloud._helpers import _to_bytes # type: ignore +from google.cloud.bigtable.data.read_rows_query import ( + ReadRowsQuery, +) +from google.cloud.bigtable.data.read_rows_query import ( + RowRange as BaseRowRange, +) +from google.cloud.bigtable.helpers import _MappableAttributesMixin class RowSet(object): @@ -25,30 +31,25 @@ class RowSet(object): """ def __init__(self): - self.row_keys = [] - self.row_ranges = [] + self._read_rows_query = ReadRowsQuery() def __eq__(self, other): if not isinstance(other, self.__class__): return NotImplemented - if len(other.row_keys) != len(self.row_keys): - return False - - if len(other.row_ranges) != len(self.row_ranges): - return False - - if not set(other.row_keys) == set(self.row_keys): - return False - - if not set(other.row_ranges) == set(self.row_ranges): - return False - - return True + return self._read_rows_query == other._read_rows_query def __ne__(self, other): return not self == other + @property + def row_keys(self): + return self._read_rows_query.row_keys + + @property + def row_ranges(self): + return self._read_rows_query.row_ranges + def add_row_key(self, row_key): """Add row key to row_keys list. @@ -62,7 +63,7 @@ def add_row_key(self, row_key): :type row_key: bytes :param row_key: The key of a row to read """ - self.row_keys.append(row_key) + self._read_rows_query.add_key(row_key) def add_row_range(self, row_range): """Add row_range to row_ranges list. @@ -77,7 +78,7 @@ def add_row_range(self, row_range): :type row_range: class:`RowRange` :param row_range: The row range object having start and end key """ - self.row_ranges.append(row_range) + self._read_rows_query.add_range(row_range) def add_row_range_from_keys( self, start_key=None, end_key=None, start_inclusive=True, end_inclusive=False @@ -109,7 +110,7 @@ def add_row_range_from_keys( considered inclusive. The default is False (exclusive). """ row_range = RowRange(start_key, end_key, start_inclusive, end_inclusive) - self.row_ranges.append(row_range) + self._read_rows_query.add_range(row_range) def add_row_range_with_prefix(self, row_key_prefix): """Add row range to row_ranges list that start with the row_key_prefix from the row keys @@ -129,28 +130,15 @@ def add_row_range_with_prefix(self, row_key_prefix): row_key_prefix.encode("utf-8"), end_key.encode("utf-8") ) - def _update_message_request(self, message): - """Add row keys and row range to given request message - :type message: class:`data_messages_v2_pb2.ReadRowsRequest` - :param message: The ``ReadRowsRequest`` protobuf - """ - for each in self.row_keys: - message.rows.row_keys._pb.append(_to_bytes(each)) - - for each in self.row_ranges: - r_kwrags = each.get_range_kwargs() - message.rows.row_ranges.append(r_kwrags) - - -class RowRange(object): +class RowRange(_MappableAttributesMixin, BaseRowRange): """Convenience wrapper of google.bigtable.v2.RowRange - :type start_key: bytes + :type start_key: str | bytes :param start_key: (Optional) Start key of the row range. If left empty, will be interpreted as the empty string. - :type end_key: bytes + :type end_key: str | bytes :param end_key: (Optional) End key of the row range. If left empty, will be interpreted as the empty string and range will be unbounded on the high end. @@ -164,49 +152,26 @@ class RowRange(object): considered inclusive. The default is False (exclusive). """ - def __init__( - self, start_key=None, end_key=None, start_inclusive=True, end_inclusive=False - ): - self.start_key = start_key - self.start_inclusive = start_inclusive - self.end_key = end_key - self.end_inclusive = end_inclusive + _attribute_map = { + "start_inclusive": "start_is_inclusive", + "end_inclusive": "end_is_inclusive", + } def _key(self): - """A tuple key that uniquely describes this field. - - Used to compute this instance's hashcode and evaluate equality. - - Returns: - Tuple[str]: The contents of this :class:`.RowRange`. - """ - return (self.start_key, self.start_inclusive, self.end_key, self.end_inclusive) + return ( + self.start_key, + self.end_key, + self.start_is_inclusive, + self.end_is_inclusive, + ) def __hash__(self): return hash(self._key()) - def __eq__(self, other): - if not isinstance(other, self.__class__): - return NotImplemented - return self._key() == other._key() - - def __ne__(self, other): - return not self == other - def get_range_kwargs(self): """Convert row range object to dict which can be passed to google.bigtable.v2.RowRange add method. """ - range_kwargs = {} - if self.start_key is not None: - start_key_key = "start_key_open" - if self.start_inclusive: - start_key_key = "start_key_closed" - range_kwargs[start_key_key] = _to_bytes(self.start_key) - - if self.end_key is not None: - end_key_key = "end_key_open" - if self.end_inclusive: - end_key_key = "end_key_closed" - range_kwargs[end_key_key] = _to_bytes(self.end_key) - return range_kwargs + return { + descriptor.name: value for descriptor, value in self._pb._pb.ListFields() + } diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 4c2f2a703933..e0b2e27eb1f9 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -17,18 +17,17 @@ import warnings from typing import Set -from google.api_core import timeout from google.api_core.exceptions import ( Aborted, DeadlineExceeded, InternalServerError, NotFound, - RetryError, ServiceUnavailable, ) from google.api_core.gapic_v1.method import DEFAULT from google.api_core.retry import Retry, if_exception_type from google.cloud._helpers import _to_bytes # type: ignore +from google.rpc import code_pb2, status_pb2 from google.cloud.bigtable import enums from google.cloud.bigtable.backup import Backup @@ -38,21 +37,31 @@ MutationsBatcher, ) from google.cloud.bigtable.column_family import ColumnFamily, _gc_rule_from_pb +from google.cloud.bigtable.data._helpers import ( + TABLE_DEFAULT, + _get_statuses_from_mutations_exception_group, +) +from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup +from google.cloud.bigtable.data.mutations import RowMutationEntry +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.encryption_info import EncryptionInfo from google.cloud.bigtable.policy import Policy -from google.cloud.bigtable.row import AppendRow, ConditionalRow, DirectRow +from google.cloud.bigtable.row import ( + AppendRow, + ConditionalRow, + DirectRow, + PartialRowData, +) from google.cloud.bigtable.row_data import ( DEFAULT_RETRY_READ_ROWS, PartialRowsData, - _retriable_internal_server_error, ) -from google.cloud.bigtable.row_set import RowRange, RowSet +from google.cloud.bigtable.row_set import RowRange from google.cloud.bigtable_admin_v2 import BaseBigtableTableAdminClient from google.cloud.bigtable_admin_v2.types import ( bigtable_table_admin as table_admin_messages_v2_pb2, ) from google.cloud.bigtable_admin_v2.types import table as admin_messages_v2_pb2 -from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 # Maximum number of mutations in bulk (MutateRowsRequest message): # (https://cloud.google.com/bigtable/docs/reference/data/rpc/ @@ -568,18 +577,23 @@ def read_row(self, row_key, filter_=None, retry=DEFAULT_RETRY_READ_ROWS): :rtype: :class:`.PartialRowData`, :data:`NoneType ` :returns: The contents of the row if any chunks were returned in the response, otherwise :data:`None`. - :raises: :class:`ValueError ` if a commit row - chunk is never encountered. """ - row_set = RowSet() - row_set.add_row_key(row_key) - result_iter = iter( - self.read_rows(filter_=filter_, row_set=row_set, retry=retry) + attempt_timeout = ( + getattr(retry, "deadline", None) + or getattr(retry, "_deadline", None) + or TABLE_DEFAULT.READ_ROWS + ) + row = self._table_impl.read_row( + row_key, + row_filter=filter_, + operation_timeout=TABLE_DEFAULT.READ_ROWS, + attempt_timeout=attempt_timeout, + retryable_errors=TABLE_DEFAULT.READ_ROWS, ) - row = next(result_iter, None) - if next(result_iter, None) is not None: - raise ValueError("More than one row was returned.") - return row + if row is None: + return None + + return PartialRowData._from_data_client_row(row) def read_rows( self, @@ -640,18 +654,26 @@ def read_rows( :returns: A :class:`.PartialRowsData` a generator for consuming the streamed results. """ - request_pb = _create_row_request( - self.name, + attempt_timeout = ( + getattr(retry, "deadline", None) + or getattr(retry, "_deadline", None) + or TABLE_DEFAULT.READ_ROWS + ) + query = _create_row_request( start_key=start_key, end_key=end_key, filter_=filter_, limit=limit, end_inclusive=end_inclusive, - app_profile_id=self._app_profile_id, row_set=row_set, ) - data_client = self._instance._client.table_data_client - return PartialRowsData(data_client.read_rows, request_pb, retry) + generator = self._table_impl.read_rows_stream( + query, + operation_timeout=TABLE_DEFAULT.READ_ROWS, + attempt_timeout=attempt_timeout, + retryable_errors=TABLE_DEFAULT.READ_ROWS, + ) + return PartialRowsData._from_generator(generator) def yield_rows(self, **kwargs): """Read rows from this table. @@ -715,6 +737,9 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): specify a ``retry`` strategy of "do-nothing", a deadline of ``0.0`` can be specified. + If a deadline of ``None`` is specified, the deadline defaults to + a table-default of 600 seconds (10 minutes). + :type rows: list :param rows: List or other iterable of :class:`.DirectRow` instances. @@ -732,18 +757,54 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): :returns: A list of response statuses (`google.rpc.status_pb2.Status`) corresponding to success or failure of each row mutation sent. These will be in the same order as the `rows`. + + :raise: ValueError: If a row entry has no mutations, or too many mutations """ if timeout is DEFAULT: timeout = self.mutation_timeout - retryable_mutate_rows = _RetryableMutateRowsWorker( - self._instance._client, - self.name, - rows, - app_profile_id=self._app_profile_id, - timeout=timeout, - ) - return retryable_mutate_rows(retry=retry) + retryable_errors = RETRYABLE_MUTATION_ERRORS + + # The data client cannot take in zero or null values for deadline, so we set it to + # the default if that is the case. + if retry.deadline is None: + operation_timeout = TABLE_DEFAULT.MUTATE_ROWS + + # To adhere to the retry strategy of do-nothing being achievable with a deadline + # of 0.0, we modify the retryable errors to be empty if such a deadline is passed. + elif retry.deadline == 0: + operation_timeout = TABLE_DEFAULT.MUTATE_ROWS + retryable_errors = [] + else: + operation_timeout = retry.deadline + + attempt_timeout = timeout + mutation_entries = [ + RowMutationEntry(row.row_key, row._get_mutations()) for row in rows + ] + return_statuses = [ + status_pb2.Status(code=code_pb2.UNKNOWN) + for _ in range(len(mutation_entries)) + ] + + try: + self._table_impl.bulk_mutate_rows( + mutation_entries, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + retryable_errors=retryable_errors, + ) + except MutationsExceptionGroup as mut_exc_group: + return_statuses = _get_statuses_from_mutations_exception_group( + mut_exc_group, len(mutation_entries) + ) + else: + return_statuses = [ + status_pb2.Status(code=code_pb2.OK) + for _ in range(len(mutation_entries)) + ] + + return return_statuses def sample_row_keys(self): """Read a sample of row keys in the table. @@ -1078,133 +1139,6 @@ def restore(self, new_table_id, cluster_id=None, backup_id=None, backup_name=Non ) -class _RetryableMutateRowsWorker(object): - """A callable worker that can retry to mutate rows with transient errors. - - This class is a callable that can retry mutating rows that result in - transient errors. After all rows are successful or none of the rows - are retryable, any subsequent call on this callable will be a no-op. - """ - - def __init__(self, client, table_name, rows, app_profile_id=None, timeout=None): - self.client = client - self.table_name = table_name - self.rows = rows - self.app_profile_id = app_profile_id - self.responses_statuses = [None] * len(self.rows) - self.timeout = timeout - - def __call__(self, retry=DEFAULT_RETRY): - """Attempt to mutate all rows and retry rows with transient errors. - - Will retry the rows with transient errors until all rows succeed or - ``deadline`` specified in the `retry` is reached. - - :rtype: list - :returns: A list of response statuses (`google.rpc.status_pb2.Status`) - corresponding to success or failure of each row mutation - sent. These will be in the same order as the ``rows``. - """ - mutate_rows = self._do_mutate_retryable_rows - if retry: - mutate_rows = retry(self._do_mutate_retryable_rows) - - try: - mutate_rows() - except (_BigtableRetryableError, RetryError): - # - _BigtableRetryableError raised when no retry strategy is used - # and a retryable error on a mutation occurred. - # - RetryError raised when retry deadline is reached. - # In both cases, just return current `responses_statuses`. - pass - - return self.responses_statuses - - @staticmethod - def _is_retryable(status): - return status is None or status.code in RETRYABLE_CODES - - def _do_mutate_retryable_rows(self): - """Mutate all the rows that are eligible for retry. - - A row is eligible for retry if it has not been tried or if it resulted - in a transient error in a previous call. - - :rtype: list - :return: The responses statuses, which is a list of - :class:`~google.rpc.status_pb2.Status`. - :raises: One of the following: - - * :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 - """ - retryable_rows = [] - index_into_all_rows = [] - for index, status in enumerate(self.responses_statuses): - if self._is_retryable(status): - retryable_rows.append(self.rows[index]) - index_into_all_rows.append(index) - - if not retryable_rows: - # All mutations are either successful or non-retryable now. - return self.responses_statuses - - entries = _compile_mutation_entries(self.table_name, retryable_rows) - data_client = self.client.table_data_client - - kwargs = {} - if self.timeout is not None: - kwargs["timeout"] = timeout.ExponentialTimeout(deadline=self.timeout) - - try: - responses = data_client.mutate_rows( - table_name=self.table_name, - entries=entries, - app_profile_id=self.app_profile_id, - retry=None, - **kwargs, - ) - except RETRYABLE_MUTATION_ERRORS as exc: - # If an exception, considered retryable by `RETRYABLE_MUTATION_ERRORS`, is - # returned from the initial call, consider - # it to be retryable. Wrap as a Bigtable Retryable Error. - # For InternalServerError, it is only retriable if the message is related to RST Stream messages - if _retriable_internal_server_error(exc) or not isinstance( - exc, InternalServerError - ): - raise _BigtableRetryableError - else: - # re-raise the original exception - raise - - num_responses = 0 - num_retryable_responses = 0 - for response in responses: - for entry in response.entries: - num_responses += 1 - index = index_into_all_rows[entry.index] - self.responses_statuses[index] = entry.status - if self._is_retryable(entry.status): - num_retryable_responses += 1 - if entry.status.code == 0: - self.rows[index].clear() - - if len(retryable_rows) != num_responses: - raise RuntimeError( - "Unexpected number of responses", - num_responses, - "Expected", - len(retryable_rows), - ) - - if num_retryable_responses: - raise _BigtableRetryableError - - return self.responses_statuses - - class ClusterState(object): """Representation of a Cluster State. @@ -1281,13 +1215,11 @@ def __ne__(self, other): def _create_row_request( - table_name, start_key=None, end_key=None, filter_=None, limit=None, end_inclusive=False, - app_profile_id=None, row_set=None, ): """Creates a request to read rows in a table. @@ -1318,104 +1250,32 @@ def _create_row_request( :param end_inclusive: (Optional) Whether the ``end_key`` should be considered inclusive. The default is False (exclusive). - :type: app_profile_id: str - :param app_profile_id: (Optional) The unique name of the AppProfile. - :type row_set: :class:`.RowSet` :param row_set: (Optional) The row set containing multiple row keys and row_ranges. - :rtype: :class:`data_messages_v2_pb2.ReadRowsRequest` - :returns: The ``ReadRowsRequest`` protobuf corresponding to the inputs. + :rtype: :class:`ReadRowsQuery` + :returns: The `ReadRowsQuery` query object corresponding to the inputs. :raises: :class:`ValueError ` if both - ``row_set`` and one of ``start_key`` or ``end_key`` are set + ``row_set`` and one of ``start_key`` or ``end_key`` are set, or if + ``end_key`` is not greater than ``start_key`` """ - request_kwargs = {"table_name": table_name} if (start_key is not None or end_key is not None) and row_set is not None: raise ValueError("Row range and row set cannot be set simultaneously") - if filter_ is not None: - request_kwargs["filter"] = filter_._to_pb() - if limit is not None: - request_kwargs["rows_limit"] = limit - if app_profile_id is not None: - request_kwargs["app_profile_id"] = app_profile_id - - message = data_messages_v2_pb2.ReadRowsRequest(**request_kwargs) + query = ReadRowsQuery( + limit=limit, + row_filter=filter_, + ) if start_key is not None or end_key is not None: - row_set = RowSet() - row_set.add_row_range(RowRange(start_key, end_key, end_inclusive=end_inclusive)) + query.add_range(RowRange(start_key, end_key, end_inclusive=end_inclusive)) if row_set is not None: - row_set._update_message_request(message) + for row_key in row_set.row_keys: + query.add_key(row_key) - return message + for row_range in row_set.row_ranges: + query.add_range(row_range) - -def _compile_mutation_entries(table_name, rows): - """Create list of mutation entries - - :type table_name: str - :param table_name: The name of the table to write to. - - :type rows: list - :param rows: List or other iterable of :class:`.DirectRow` instances. - - :rtype: List[:class:`data_messages_v2_pb2.MutateRowsRequest.Entry`] - :returns: entries corresponding to the inputs. - :raises: :exc:`~.table.TooManyMutationsError` if the number of mutations is - greater than the max ({}) - """.format(_MAX_BULK_MUTATIONS) - entries = [] - mutations_count = 0 - entry_klass = data_messages_v2_pb2.MutateRowsRequest.Entry - - for row in rows: - _check_row_table_name(table_name, row) - _check_row_type(row) - mutations = row._get_mutations() - entries.append(entry_klass(row_key=row.row_key, mutations=mutations)) - mutations_count += len(mutations) - - if mutations_count > _MAX_BULK_MUTATIONS: - raise TooManyMutationsError( - "Maximum number of mutations is %s" % (_MAX_BULK_MUTATIONS,) - ) - return entries - - -def _check_row_table_name(table_name, row): - """Checks that a row belongs to a table. - - :type table_name: str - :param table_name: The name of the table. - - :type row: :class:`~google.cloud.bigtable.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.row.Row` - subclasses. - - :raises: :exc:`~.table.TableMismatchError` if the row does not belong to - the table. - """ - if row.table is not None and row.table.name != table_name: - raise TableMismatchError( - "Row %s is a part of %s table. Current table: %s" - % (row.row_key, row.table.name, table_name) - ) - - -def _check_row_type(row): - """Checks that a row is an instance of :class:`.DirectRow`. - - :type row: :class:`~google.cloud.bigtable.row.Row` - :param row: An instance of :class:`~google.cloud.bigtable.row.Row` - subclasses. - - :raises: :class:`TypeError ` if the row is not an - instance of DirectRow. - """ - if not isinstance(row, DirectRow): - raise TypeError( - "Bulk processing can not be applied for conditional or append mutations." - ) + return query diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py index 4ad84f01dda5..f88cf7719f90 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py @@ -512,6 +512,42 @@ async def test_mutations_batcher_timer_flush(self, client, target, temp_rows): # ensure cell is updated assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value + @pytest.mark.usefixtures("client") + @pytest.mark.usefixtures("target") + @CrossSync.Retry( + predicate=retry.if_exception_type(ClientError), initial=1, maximum=5 + ) + @CrossSync.pytest + async def test_mutations_batcher_completed_callback( + self, client, target, temp_rows + ): + """ + test batcher with batch completed callback. It should be called when the batcher flushes. + """ + import mock + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.mutations import RowMutationEntry + + callback = mock.Mock() + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await self._create_row_and_mutation( + target, temp_rows, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + flush_interval = 0.1 + async with target.mutations_batcher(flush_interval=flush_interval) as batcher: + batcher._user_batch_completed_callback = callback + await batcher.append(bulk_mutation) + await CrossSync.yield_to_event_loop() + assert len(batcher._staged_entries) == 1 + await CrossSync.sleep(flush_interval + 0.1) + assert len(batcher._staged_entries) == 0 + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + # ensure cell is updated + assert (await self._retrieve_cell_value(target, row_key)) == new_value + @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @CrossSync.Retry( diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py index 91844eb9bfce..23f069153eba 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py @@ -395,6 +395,35 @@ def test_mutations_batcher_timer_flush(self, client, target, temp_rows): assert len(batcher._staged_entries) == 0 assert temp_rows.retrieve_cell_value(target, row_key) == new_value + @pytest.mark.usefixtures("client") + @pytest.mark.usefixtures("target") + @CrossSync._Sync_Impl.Retry( + predicate=retry.if_exception_type(ClientError), initial=1, maximum=5 + ) + def test_mutations_batcher_completed_callback(self, client, target, temp_rows): + """test batcher with batch completed callback. It should be called when the batcher flushes.""" + import mock + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.mutations import RowMutationEntry + + callback = mock.Mock() + new_value = uuid.uuid4().hex.encode() + row_key, mutation = self._create_row_and_mutation( + target, temp_rows, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + flush_interval = 0.1 + with target.mutations_batcher(flush_interval=flush_interval) as batcher: + batcher._user_batch_completed_callback = callback + batcher.append(bulk_mutation) + CrossSync._Sync_Impl.yield_to_event_loop() + assert len(batcher._staged_entries) == 1 + CrossSync._Sync_Impl.sleep(flush_interval + 0.1) + assert len(batcher._staged_entries) == 0 + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + assert self._retrieve_cell_value(target, row_key) == new_value + @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @CrossSync._Sync_Impl.Retry( 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 cb90d62de934..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 @@ -13,7 +13,6 @@ # limitations under the License. import operator -import struct from datetime import datetime, timedelta, timezone import pytest @@ -46,6 +45,8 @@ INITIAL_ROW_SPLITS = [b"row_split_1", b"row_split_2", b"row_split_3"] JOY_EMOJI = "\N{FACE WITH TEARS OF JOY}" +GAP_MARGIN_OF_ERROR = 0.05 + PASS_ALL_FILTER = row_filters.PassAllFilter(True) BLOCK_ALL_FILTER = row_filters.BlockAllFilter(True) @@ -214,6 +215,37 @@ def test_table_read_rows_filter_millis(data_table): row_data.consume_all() +def test_table_direct_row_commit(data_table, rows_to_delete): + from google.rpc import code_pb2 + + row = data_table.direct_row(ROW_KEY) + + # Test set cell + row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) + row.set_cell(COLUMN_FAMILY_ID1, COL_NAME2, CELL_VAL1) + status = row.commit() + rows_to_delete.append(row) + assert status.code == code_pb2.Code.OK + row_data = data_table.read_row(ROW_KEY) + assert row_data.cells[COLUMN_FAMILY_ID1][COL_NAME1][0].value == CELL_VAL1 + assert row_data.cells[COLUMN_FAMILY_ID1][COL_NAME2][0].value == CELL_VAL1 + + # Test delete cell + row.delete_cell(COLUMN_FAMILY_ID1, COL_NAME1) + status = row.commit() + assert status.code == code_pb2.Code.OK + row_data = data_table.read_row(ROW_KEY) + assert COL_NAME1 not in row_data.cells[COLUMN_FAMILY_ID1] + assert row_data.cells[COLUMN_FAMILY_ID1][COL_NAME2][0].value == CELL_VAL1 + + # Test delete row + row.delete() + status = row.commit() + assert status.code == code_pb2.Code.OK + row_data = data_table.read_row(ROW_KEY) + assert row_data is None + + def test_table_mutate_rows(data_table, rows_to_delete): row1 = data_table.direct_row(ROW_KEY) row1.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) @@ -243,43 +275,12 @@ def test_table_mutate_rows(data_table, rows_to_delete): assert row2_data.cells[COLUMN_FAMILY_ID1][COL_NAME1][0].value == CELL_VAL4 -def _add_test_error_handler(retry): - """Overwrites the current on_error function to assert that backoff values are within expected bounds.""" - import time - - curr_time = time.monotonic() - times_triggered = 0 - - # Assert that the retry handler works properly. - def test_error_handler(exc): - nonlocal curr_time, times_triggered - next_time = time.monotonic() - if times_triggered >= 1: - gap = next_time - curr_time - - # Exponential backoff = uniform randomness from 0 to max_gap - max_gap = min( - retry._initial * retry._multiplier**times_triggered, - retry._maximum, - ) - # Allow a small tolerance margin (1.0s) for OS sleep scheduling latency - assert gap <= max_gap + 1.0 - times_triggered += 1 - curr_time = next_time - - retry._on_error = test_error_handler - - def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): - import copy - import mock - from google.api_core import retry as retries - from google.api_core.exceptions import InvalidArgument from google.rpc.code_pb2 import Code from google.rpc.status_pb2 import Status - from google.cloud.bigtable.table import DEFAULT_RETRY, _BigtableRetryableError + from google.cloud.bigtable.table import DEFAULT_RETRY from google.cloud.bigtable_v2 import MutateRowsResponse # Simulate a server error on row 2, and a normal response on row 1, followed by a bunch of error @@ -332,10 +333,7 @@ def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): rows_to_delete.append(row_2) row_2.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - # Testing the default retry - default_retry_copy = copy.copy(DEFAULT_RETRY) - _add_test_error_handler(default_retry_copy) - statuses = data_table.mutate_rows([row, row_2], retry=default_retry_copy) + statuses = data_table.mutate_rows([row, row_2]) assert statuses[0].code == Code.OK assert statuses[1].code == Code.OK @@ -355,28 +353,35 @@ def test_table_mutate_rows_retries_timeout(data_table, rows_to_delete): rows_to_delete.append(row_2) row_2.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - # Testing the default retry - default_retry_copy = copy.copy(DEFAULT_RETRY) - _add_test_error_handler(default_retry_copy) - statuses = data_table.mutate_rows([row, row_2], retry=default_retry_copy) + statuses = data_table.mutate_rows([row, row_2]) assert statuses[0].code == Code.OK - assert statuses[1].code == Code.INTERNAL + assert statuses[1].code == Code.DEADLINE_EXCEEDED - # Because of the way the retriable mutate worker class works, unusual things can happen - # when passing in custom retry predicates. - row = data_table.direct_row(ROW_KEY) - rows_to_delete.append(row) + # Retries with deadline 0 should do nothing. + with mock.patch.object( + data_table._instance._client.table_data_client, "mutate_rows" + ) as mutate_mock: + mutate_mock.side_effect = [ + initial_error_response, + followup_error_response, + followup_error_response, + final_success_response, + ] - row_2 = data_table.direct_row(ROW_KEY_ALT) - rows_to_delete.append(row_2) + row = data_table.direct_row(ROW_KEY) + rows_to_delete.append(row) + row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - retry = DEFAULT_RETRY.with_predicate( - retries.if_exception_type(_BigtableRetryableError, InvalidArgument) - ) - _add_test_error_handler(retry) - statuses = data_table.mutate_rows([row, row_2], retry=retry) - assert statuses[0] is None - assert statuses[1] is None + row_2 = data_table.direct_row(ROW_KEY_ALT) + rows_to_delete.append(row_2) + row_2.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) + + do_nothing_retry = DEFAULT_RETRY.with_deadline(0.0) + + statuses = data_table.mutate_rows([row, row_2], retry=do_nothing_retry) + assert statuses[0].code == Code.OK + assert statuses[1].code == Code.INTERNAL + mutate_mock.assert_called_once() def _populate_table( @@ -464,26 +469,23 @@ def test_table_mutate_rows_integers(data_table, rows_to_delete): def test_table_mutate_rows_input_errors(data_table, rows_to_delete): - from google.api_core.exceptions import InvalidArgument - - from google.cloud.bigtable.table import _MAX_BULK_MUTATIONS, TooManyMutationsError + from google.cloud.bigtable.table import _MAX_BULK_MUTATIONS row = data_table.direct_row(ROW_KEY) rows_to_delete.append(row) - # Mutate row with 0 mutations gives an API error from the service, not - # from the client library. - with pytest.raises(InvalidArgument): + # Mutate row with 0 mutations gives a ValueError from the client library. + with pytest.raises(ValueError): data_table.mutate_rows([row]) row.clear() - # Mutate row with >100k mutations gives a TooManyMutationsError from the + # Mutate row with >100k mutations gives a ValueError from the # client library. for _ in range(0, _MAX_BULK_MUTATIONS + 1): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - with pytest.raises(TooManyMutationsError): + with pytest.raises(ValueError): data_table.mutate_rows([row]) @@ -720,6 +722,27 @@ def _assert_data_table_read_rows_retry_correct(rows_data): ) +def test_table_read_rows_multiple_reads( + data_table_read_rows_retry_tests, +): + from types import SimpleNamespace + + rows_data = data_table_read_rows_retry_tests.read_rows() + first_iteration = SimpleNamespace() + first_iteration.rows = {} + + second_iteration = SimpleNamespace() + second_iteration.rows = {} + for item in rows_data: + first_iteration.rows[item.row_key] = item + + for item in rows_data: + second_iteration.rows[item.row_key] = item + + _assert_data_table_read_rows_retry_correct(first_iteration) + assert second_iteration.rows == {} + + def test_table_read_rows_retry_unretriable_error_establishing_stream( data_table_read_rows_retry_tests, ): @@ -727,11 +750,12 @@ def test_table_read_rows_retry_unretriable_error_establishing_stream( error_injector = data_table_read_rows_retry_tests.error_injector error_injector.errors_to_inject = [ - error_injector.make_exception(StatusCode.ABORTED, fail_mid_stream=False) + error_injector.make_exception(StatusCode.DATA_LOSS, fail_mid_stream=False) ] - with pytest.raises(exceptions.Aborted): - data_table_read_rows_retry_tests.read_rows() + rows_data = data_table_read_rows_retry_tests.read_rows() + with pytest.raises(exceptions.DataLoss): + rows_data.consume_all() def test_table_read_rows_retry_retriable_error_establishing_stream( @@ -792,25 +816,25 @@ def test_table_read_rows_retry_retriable_errors_mid_stream( def test_table_read_rows_retry_retriable_internal_errors_mid_stream( data_table_read_rows_retry_tests, ): - from google.cloud.bigtable.row_data import RETRYABLE_INTERNAL_ERROR_MESSAGES + from google.cloud.bigtable.data._helpers import _RETRYABLE_INTERNAL_ERROR_MESSAGES error_injector = data_table_read_rows_retry_tests.error_injector error_injector.errors_to_inject = [ error_injector.make_exception( StatusCode.INTERNAL, - message=RETRYABLE_INTERNAL_ERROR_MESSAGES[0], + message=_RETRYABLE_INTERNAL_ERROR_MESSAGES[0], fail_mid_stream=True, successes_before_fail=2, ), error_injector.make_exception( StatusCode.INTERNAL, - message=RETRYABLE_INTERNAL_ERROR_MESSAGES[1], + message=_RETRYABLE_INTERNAL_ERROR_MESSAGES[1], fail_mid_stream=True, successes_before_fail=1, ), error_injector.make_exception( StatusCode.INTERNAL, - message=RETRYABLE_INTERNAL_ERROR_MESSAGES[2], + message=_RETRYABLE_INTERNAL_ERROR_MESSAGES[2], fail_mid_stream=True, successes_before_fail=0, ), @@ -853,12 +877,12 @@ def test_table_read_rows_retry_retriable_error_mid_stream_unretriable_error_rees error_injector.make_exception( StatusCode.UNAVAILABLE, fail_mid_stream=True, successes_before_fail=5 ), - error_injector.make_exception(StatusCode.ABORTED, fail_mid_stream=False), + error_injector.make_exception(StatusCode.DATA_LOSS, fail_mid_stream=False), ] rows_data = data_table_read_rows_retry_tests.read_rows() - with pytest.raises(exceptions.Aborted): + with pytest.raises(exceptions.DataLoss): rows_data.consume_all() @@ -889,29 +913,32 @@ def test_table_read_rows_retry_timeout_mid_stream( from google.api_core import exceptions + from google.cloud.bigtable.data._helpers import _RETRYABLE_INTERNAL_ERROR_MESSAGES from google.cloud.bigtable.row_data import ( DEFAULT_RETRY_READ_ROWS, - RETRYABLE_INTERNAL_ERROR_MESSAGES, ) error_injector = data_table_read_rows_retry_tests.error_injector error_injector.errors_to_inject = [ error_injector.make_exception( StatusCode.INTERNAL, - message=RETRYABLE_INTERNAL_ERROR_MESSAGES[0], + message=_RETRYABLE_INTERNAL_ERROR_MESSAGES[0], fail_mid_stream=True, successes_before_fail=5, ), ] + [ error_injector.make_exception( StatusCode.INTERNAL, - message=RETRYABLE_INTERNAL_ERROR_MESSAGES[0], + message=_RETRYABLE_INTERNAL_ERROR_MESSAGES[0], fail_mid_stream=True, successes_before_fail=0, ), ] * 20 # Shorten the deadline so the timeout test is shorter. + data_table_read_rows_retry_tests._table_impl.default_read_rows_operation_timeout = ( + 10.0 + ) rows_data = data_table_read_rows_retry_tests.read_rows( retry=DEFAULT_RETRY_READ_ROWS.with_deadline(10.0) ) @@ -940,10 +967,14 @@ def test_table_read_rows_retry_timeout_establishing_stream( ] * 20 # Shorten the deadline so the timeout test is shorter. + data_table_read_rows_retry_tests._table_impl.default_read_rows_operation_timeout = ( + 10.0 + ) + rows_data = data_table_read_rows_retry_tests.read_rows( + retry=DEFAULT_RETRY_READ_ROWS.with_deadline(10.0) + ) with pytest.raises(exceptions.RetryError): - data_table_read_rows_retry_tests.read_rows( - retry=DEFAULT_RETRY_READ_ROWS.with_deadline(10.0) - ) + rows_data.consume_all() def test_table_check_and_mutate_rows(data_table, rows_to_delete): @@ -1032,8 +1063,6 @@ def test_table_sample_row_keys(data_table, skip_on_emulator): def test_table_direct_row_input_errors(data_table, rows_to_delete): - from google.api_core.exceptions import InvalidArgument - from google.cloud.bigtable.row import MAX_MUTATIONS row = data_table.direct_row(ROW_KEY) @@ -1047,22 +1076,11 @@ def test_table_direct_row_input_errors(data_table, rows_to_delete): with pytest.raises(TypeError): row.delete_cell(COLUMN_FAMILY_ID1, INT_COL_NAME) - # Unicode for column name and value does not get converted to bytes because - # internally we use to_bytes in ascii mode. - with pytest.raises(UnicodeEncodeError): - row.set_cell(COLUMN_FAMILY_ID1, JOY_EMOJI, CELL_VAL1) - - with pytest.raises(UnicodeEncodeError): - row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, JOY_EMOJI) - - with pytest.raises(UnicodeEncodeError): - row.delete_cell(COLUMN_FAMILY_ID1, JOY_EMOJI) - - # Various non int64s, we use struct to pack a Python int to bytes. - with pytest.raises(struct.error): + # Various non int64s + with pytest.raises(ValueError): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL) - with pytest.raises(struct.error): + with pytest.raises(ValueError): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL2) # Since floats aren't ints, they aren't converted to bytes via struct.pack, @@ -1070,20 +1088,18 @@ def test_table_direct_row_input_errors(data_table, rows_to_delete): with pytest.raises(TypeError): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, FLOAT_CELL_VAL) - # Can't have more than MAX_MUTATIONS mutations, but only enforced after - # a row.commit + # Can't have more than MAX_MUTATIONS mutations, enforced on server side now. row.clear() for _ in range(0, MAX_MUTATIONS + 1): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, CELL_VAL1) - with pytest.raises(ValueError): - row.commit() + resp = row.commit() + assert resp.code == StatusCode.INVALID_ARGUMENT.value[0] - # Not having any mutations gives a server error (InvalidArgument), not - # enforced on the client side. + # Not having any mutations raises a ValueError row.clear() - with pytest.raises(InvalidArgument): - row.commit() + with pytest.raises(ValueError): + resp = row.commit() def test_table_conditional_row_input_errors(data_table, rows_to_delete): @@ -1107,22 +1123,11 @@ def test_table_conditional_row_input_errors(data_table, rows_to_delete): with pytest.raises(TypeError): true_row.delete_cell(COLUMN_FAMILY_ID1, INT_COL_NAME) - # Unicode for column name and value does not get converted to bytes because - # internally we use to_bytes in ascii mode. - with pytest.raises(UnicodeEncodeError): - true_row.set_cell(COLUMN_FAMILY_ID1, JOY_EMOJI, CELL_VAL1) - - with pytest.raises(UnicodeEncodeError): - true_row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, JOY_EMOJI) - - with pytest.raises(UnicodeEncodeError): - true_row.delete_cell(COLUMN_FAMILY_ID1, JOY_EMOJI) - - # Various non int64s, we use struct to pack a Python int to bytes. - with pytest.raises(struct.error): + # Various non int64s + with pytest.raises(ValueError): true_row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL) - with pytest.raises(struct.error): + with pytest.raises(ValueError): true_row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL2) # Since floats aren't ints, they aren't converted to bytes via struct.pack, @@ -1171,39 +1176,17 @@ def test_table_append_row_input_errors(data_table, rows_to_delete): rows_to_delete.append(data_table.direct_row(ROW_KEY)) # Column names should be convertible to bytes (str or bytes) - with pytest.raises(TypeError): + with pytest.raises(AttributeError): row.append_cell_value(COLUMN_FAMILY_ID1, INT_COL_NAME, CELL_VAL1) - with pytest.raises(TypeError): + with pytest.raises(AttributeError): row.increment_cell_value(COLUMN_FAMILY_ID1, INT_COL_NAME, 1) - # Unicode for column name and value - with pytest.raises(UnicodeEncodeError): - row.append_cell_value(COLUMN_FAMILY_ID1, JOY_EMOJI, CELL_VAL1) - - with pytest.raises(UnicodeEncodeError): - row.append_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, JOY_EMOJI) - - with pytest.raises(UnicodeEncodeError): - row.increment_cell_value(COLUMN_FAMILY_ID1, JOY_EMOJI, 1) - - # Non-integer cell values for increment_cell_value with pytest.raises(ValueError): row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL) - # increment_cell_value does not do input validation on the int_value, instead using - # proto-plus to do validation. - row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, FLOAT_CELL_VAL) - row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME2, FLOAT_CELL_VAL2) - row.commit() - - row_data = data_table.read_row(ROW_KEY) - assert row_data.cells[COLUMN_FAMILY_ID1][COL_NAME1][0].value == int( - FLOAT_CELL_VAL - ).to_bytes(8, byteorder="big", signed=True) - assert row_data.cells[COLUMN_FAMILY_ID1][COL_NAME2][0].value == int( - FLOAT_CELL_VAL2 - ).to_bytes(8, byteorder="big", signed=True) + with pytest.raises(TypeError): + row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, FLOAT_CELL_VAL) # Can't have more than MAX_MUTATIONS mutations, but only enforced after # a row.commit @@ -1262,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_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index 815f26a367e3..7f7e45209313 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -1510,49 +1510,56 @@ def test_table_ctor_sync(self): @CrossSync.pytest # iterate over all retryable rpcs @pytest.mark.parametrize( - "fn_name,fn_args,is_stream,extra_retryables", + "fn_name,fn_args,is_read_rows_fn,is_stream,extra_retryables", [ ( "read_rows_stream", (ReadRowsQuery(),), True, + True, (), ), ( "read_rows", (ReadRowsQuery(),), True, + True, (), ), ( "read_row", (b"row_key",), True, + True, (), ), ( "read_rows_sharded", ([ReadRowsQuery()],), True, + True, (), ), ( "row_exists", (b"row_key",), True, + True, (), ), - ("sample_row_keys", (), False, ()), + ("sample_row_keys", (), False, False, ()), ( "mutate_row", (b"row_key", [DeleteAllFromRow()]), False, + False, (), ), ( "bulk_mutate_rows", ([mutations.RowMutationEntry(b"key", [DeleteAllFromRow()])],), False, + False, (_MutateRowsIncomplete,), ), ], @@ -1588,6 +1595,7 @@ async def test_customizable_retryable_errors( expected_retryables, fn_name, fn_args, + is_read_rows_fn, is_stream, extra_retryables, ): @@ -1600,8 +1608,18 @@ async def test_customizable_retryable_errors( retry_fn += "_stream" if CrossSync.is_async: retry_fn = f"CrossSync.{retry_fn}" + subpackage = "_async" else: retry_fn = f"CrossSync._Sync_Impl.{retry_fn}" + subpackage = "_sync_autogen" + + # Read Rows has its own custom predicate builder that also takes in + # a list of exceptions + if is_read_rows_fn: + predicate_builder = f"google.cloud.bigtable.data.{subpackage}._read_rows._rst_stream_aware_predicate" + else: + predicate_builder = "google.api_core.retry.if_exception_type" + with mock.patch( f"google.cloud.bigtable.data._cross_sync.{retry_fn}" ) as retry_fn_mock: @@ -1609,9 +1627,7 @@ async def test_customizable_retryable_errors( table = client.get_table("instance-id", "table-id") expected_predicate = expected_retryables.__contains__ retry_fn_mock.side_effect = RuntimeError("stop early") - with mock.patch( - "google.api_core.retry.if_exception_type" - ) as predicate_builder_mock: + with mock.patch(predicate_builder) as predicate_builder_mock: predicate_builder_mock.return_value = expected_predicate with pytest.raises(Exception): # we expect an exception from attempting to call the mock 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 2757e8a95cdb..aa10d8866f52 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 @@ -973,8 +979,85 @@ async def test__execute_mutate_rows_returns_errors(self): table.default_mutate_rows_attempt_timeout = 13 table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: + batch = [self._make_mutation(), self._make_mutation()] + result = await instance._execute_mutate_rows(batch, mock.Mock()) + assert len(result) == 2 + assert result[0] == err1 + assert result[1] == err2 + # indices should be set to None + assert result[0].index is None + assert result[1].index is None + + @CrossSync.pytest + async def test__execute_mutate_rows_batch_completed_callback(self): + from google.rpc import code_pb2, status_pb2 + + with mock.patch.object(CrossSync, "_MutateRowsOperation") as mutate_rows: + mutate_rows.return_value = CrossSync.Mock() + start_operation = mutate_rows().start + table = mock.Mock() + table.table_name = "test-table" + table.app_profile_id = "test-app-profile" + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + async with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback batch = [self._make_mutation()] result = await instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + assert start_operation.call_count == 1 + args, kwargs = mutate_rows.call_args + assert args[0] == table.client._gapic_client + assert args[1] == table + assert args[2] == batch + assert kwargs["operation_timeout"] == 17 + assert kwargs["attempt_timeout"] == 13 + assert result == [] + + @CrossSync.pytest + async def test__execute_mutate_rows_batch_completed_callback_errors(self): + from google.api_core import exceptions + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + with mock.patch.object(CrossSync._MutateRowsOperation, "start") as mutate_rows: + err1 = FailedMutationEntryError( + 1, mock.Mock(), exceptions.DataLoss("test error") + ) + err2 = FailedMutationEntryError( + 2, mock.Mock(), exceptions.DataLoss("test error") + ) + mutate_rows.side_effect = MutationsExceptionGroup([err1, err2], 10) + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + async with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [ + self._make_mutation(), + self._make_mutation(), + self._make_mutation(), + ] + result = await instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with( + [ + status_pb2.Status(code=code_pb2.OK), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + ] + ) assert len(result) == 2 assert result[0] == err1 assert result[1] == err2 @@ -982,6 +1065,22 @@ async def test__execute_mutate_rows_returns_errors(self): assert result[0].index is None assert result[1].index is None + @CrossSync.pytest + async def test__execute_mutate_rows_batch_completed_callback_exception(self): + with mock.patch.object(CrossSync, "_MutateRowsOperation") as mutate_rows: + mutate_rows.return_value = CrossSync.Mock() + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock(side_effect=RuntimeError("callback failed")) + async with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [self._make_mutation()] + result = await instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once() + assert result == [] + @CrossSync.pytest async def test__raise_exceptions(self): """Raise exceptions and reset error state""" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index 08d628f4aef7..899041e3055c 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -1263,19 +1263,20 @@ def test_ctor_invalid_timeout_values(self): client.close() @pytest.mark.parametrize( - "fn_name,fn_args,is_stream,extra_retryables", + "fn_name,fn_args,is_read_rows_fn,is_stream,extra_retryables", [ - ("read_rows_stream", (ReadRowsQuery(),), True, ()), - ("read_rows", (ReadRowsQuery(),), True, ()), - ("read_row", (b"row_key",), True, ()), - ("read_rows_sharded", ([ReadRowsQuery()],), True, ()), - ("row_exists", (b"row_key",), True, ()), - ("sample_row_keys", (), False, ()), - ("mutate_row", (b"row_key", [DeleteAllFromRow()]), False, ()), + ("read_rows_stream", (ReadRowsQuery(),), True, True, ()), + ("read_rows", (ReadRowsQuery(),), True, True, ()), + ("read_row", (b"row_key",), True, True, ()), + ("read_rows_sharded", ([ReadRowsQuery()],), True, True, ()), + ("row_exists", (b"row_key",), True, True, ()), + ("sample_row_keys", (), False, False, ()), + ("mutate_row", (b"row_key", [DeleteAllFromRow()]), False, False, ()), ( "bulk_mutate_rows", ([mutations.RowMutationEntry(b"key", [DeleteAllFromRow()])],), False, + False, (_MutateRowsIncomplete,), ), ], @@ -1310,6 +1311,7 @@ def test_customizable_retryable_errors( expected_retryables, fn_name, fn_args, + is_read_rows_fn, is_stream, extra_retryables, ): @@ -1319,6 +1321,11 @@ def test_customizable_retryable_errors( if is_stream: retry_fn += "_stream" retry_fn = f"CrossSync._Sync_Impl.{retry_fn}" + subpackage = "_sync_autogen" + if is_read_rows_fn: + predicate_builder = f"google.cloud.bigtable.data.{subpackage}._read_rows._rst_stream_aware_predicate" + else: + predicate_builder = "google.api_core.retry.if_exception_type" with mock.patch( f"google.cloud.bigtable.data._cross_sync.{retry_fn}" ) as retry_fn_mock: @@ -1326,9 +1333,7 @@ def test_customizable_retryable_errors( table = client.get_table("instance-id", "table-id") expected_predicate = expected_retryables.__contains__ retry_fn_mock.side_effect = RuntimeError("stop early") - with mock.patch( - "google.api_core.retry.if_exception_type" - ) as predicate_builder_mock: + with mock.patch(predicate_builder) as predicate_builder_mock: predicate_builder_mock.return_value = expected_predicate with pytest.raises(Exception): test_fn = table.__getattribute__(fn_name) 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 ad50533fe57c..7e6dbbdc767a 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 @@ -854,14 +860,109 @@ def test__execute_mutate_rows_returns_errors(self): table.default_mutate_rows_attempt_timeout = 13 table.default_mutate_rows_retryable_errors = () with self._make_one(table) as instance: + batch = [self._make_mutation(), self._make_mutation()] + result = instance._execute_mutate_rows(batch, mock.Mock()) + assert len(result) == 2 + assert result[0] == err1 + assert result[1] == err2 + assert result[0].index is None + assert result[1].index is None + + def test__execute_mutate_rows_batch_completed_callback(self): + from google.rpc import code_pb2, status_pb2 + + with mock.patch.object( + CrossSync._Sync_Impl, "_MutateRowsOperation" + ) as mutate_rows: + mutate_rows.return_value = CrossSync._Sync_Impl.Mock() + start_operation = mutate_rows().start + table = mock.Mock() + table.table_name = "test-table" + table.app_profile_id = "test-app-profile" + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback batch = [self._make_mutation()] result = instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with([status_pb2.Status(code=code_pb2.OK)]) + assert start_operation.call_count == 1 + args, kwargs = mutate_rows.call_args + assert args[0] == table.client._gapic_client + assert args[1] == table + assert args[2] == batch + assert kwargs["operation_timeout"] == 17 + assert kwargs["attempt_timeout"] == 13 + assert result == [] + + def test__execute_mutate_rows_batch_completed_callback_errors(self): + from google.api_core import exceptions + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + ) + + with mock.patch.object( + CrossSync._Sync_Impl._MutateRowsOperation, "start" + ) as mutate_rows: + err1 = FailedMutationEntryError( + 1, mock.Mock(), exceptions.DataLoss("test error") + ) + err2 = FailedMutationEntryError( + 2, mock.Mock(), exceptions.DataLoss("test error") + ) + mutate_rows.side_effect = MutationsExceptionGroup([err1, err2], 10) + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock() + with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [ + self._make_mutation(), + self._make_mutation(), + self._make_mutation(), + ] + result = instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once_with( + [ + status_pb2.Status(code=code_pb2.OK), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + status_pb2.Status( + code=code_pb2.DATA_LOSS, message="test error" + ), + ] + ) assert len(result) == 2 assert result[0] == err1 assert result[1] == err2 assert result[0].index is None assert result[1].index is None + def test__execute_mutate_rows_batch_completed_callback_exception(self): + with mock.patch.object( + CrossSync._Sync_Impl, "_MutateRowsOperation" + ) as mutate_rows: + mutate_rows.return_value = CrossSync._Sync_Impl.Mock() + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 17 + table.default_mutate_rows_attempt_timeout = 13 + table.default_mutate_rows_retryable_errors = () + callback = mock.Mock(side_effect=RuntimeError("callback failed")) + with self._make_one(table) as instance: + instance._user_batch_completed_callback = callback + batch = [self._make_mutation()] + result = instance._execute_mutate_rows(batch, mock.Mock()) + callback.assert_called_once() + assert result == [] + def test__raise_exceptions(self): """Raise exceptions and reset error state""" from google.cloud.bigtable.data import exceptions diff --git a/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py b/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py index 289b5ba7408b..3e44ec7c3d04 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py +++ b/packages/google-cloud-bigtable/tests/unit/data/test__helpers.py @@ -16,8 +16,10 @@ import mock import pytest from google.api_core import exceptions as core_exceptions +from google.rpc import code_pb2, status_pb2 import google.cloud.bigtable.data._helpers as _helpers +import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import TABLE_DEFAULT @@ -222,6 +224,134 @@ def test_get_timeouts_invalid(self, input_times): _helpers._align_timeouts(input_times[0], input_times[1]) +class TestRstStreamAwarePredicate: + @pytest.mark.parametrize( + "retryable_exceptions,exception,expected_is_retryable", + [ + ( + [core_exceptions.Aborted, core_exceptions.InternalServerError], + core_exceptions.InternalServerError("Sorry"), + True, + ), + ( + [core_exceptions.Aborted, core_exceptions.InternalServerError], + core_exceptions.DataLoss("Sorry"), + False, + ), + ( + [core_exceptions.ServiceUnavailable, core_exceptions.Aborted], + core_exceptions.InternalServerError("Sorry"), + False, + ), + ( + [core_exceptions.ServiceUnavailable, core_exceptions.Aborted], + core_exceptions.InternalServerError( + _helpers._RETRYABLE_INTERNAL_ERROR_MESSAGES[0] + ), + True, + ), + ( + [core_exceptions.InternalServerError, core_exceptions.Aborted], + core_exceptions.InternalServerError( + _helpers._RETRYABLE_INTERNAL_ERROR_MESSAGES[0] + ), + True, + ), + ], + ) + def test_rst_stream_aware_predicate( + self, retryable_exceptions, exception, expected_is_retryable + ): + predicate = _helpers._rst_stream_aware_predicate(*retryable_exceptions) + assert predicate(exception) is expected_is_retryable + + +class TestGetStatusesFromMutationsExceptionGroup: + @pytest.mark.parametrize( + "failed_idx,cause_exc,expected_status", + [ + ( + 0, + core_exceptions.DeadlineExceeded( + "Operation timed out after 40 seconds" + ), + status_pb2.Status( + code=code_pb2.DEADLINE_EXCEEDED, + message="Operation timed out after 40 seconds", + ), + ), + ( + 0, + RuntimeError("Something happened"), + status_pb2.Status(code=code_pb2.UNKNOWN, message="Something happened"), + ), + ( + 0, + bt_exceptions.RetryExceptionGroup( + excs=[ + core_exceptions.ServiceUnavailable("Service Unavailable"), + core_exceptions.ServiceUnavailable("Service Unavailable"), + core_exceptions.DeadlineExceeded( + "Operation timed out after 40 seconds" + ), + ] + ), + status_pb2.Status( + code=code_pb2.DEADLINE_EXCEEDED, + message="Operation timed out after 40 seconds", + ), + ), + ( + 0, + bt_exceptions.RetryExceptionGroup( + excs=[ + core_exceptions.ServiceUnavailable("Service Unavailable"), + core_exceptions.ServiceUnavailable("Service Unavailable"), + RuntimeError("Something happened"), + ] + ), + status_pb2.Status(code=code_pb2.UNKNOWN, message="Something happened"), + ), + ( + 0, + None, + status_pb2.Status( + code=code_pb2.UNKNOWN, message="An unknown error has occurred" + ), + ), + ( + 100, + RuntimeError("Something happened"), + status_pb2.Status(code=code_pb2.OK), + ), + ( + None, + RuntimeError("Something happened"), + status_pb2.Status(code=code_pb2.OK), + ), + ], + ) + def test_get_statuses_from_mutations_exception_group( + self, failed_idx, cause_exc, expected_status + ): + mutation_exception_group = bt_exceptions.MutationsExceptionGroup( + excs=[ + bt_exceptions.FailedMutationEntryError( + failed_idx=failed_idx, + failed_mutation_entry=mock.Mock(), + cause=cause_exc, + ) + ], + total_entries=1, + message="Mutations failed.", + ) + + statuses = _helpers._get_statuses_from_mutations_exception_group( + mutation_exception_group, 1 + ) + assert statuses[0] == expected_status + + class TestGetRetryableErrors: @pytest.mark.parametrize( "input_codes,input_table,expected", 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_row.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py index 5544327b3a97..7349f8be9b68 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py @@ -18,6 +18,8 @@ from ._testing import _make_credentials +_INSTANCE_ID = "test-instance" + def _make_client(*args, **kwargs): from google.cloud.bigtable.client import Client @@ -66,7 +68,7 @@ def test_direct_row_constructor(): row = _make_direct_row(row_key, table) assert row._row_key == row_key assert row._table is table - assert row._pb_mutations == [] + assert row._mutations == [] def test_direct_row_constructor_with_unicode(): @@ -89,10 +91,28 @@ def test_direct_row__get_mutations(): row_key = b"row_key" row = _make_direct_row(row_key, None) - row._pb_mutations = mutations = object() + row._mutations = mutations = object() assert mutations is row._get_mutations(None) +def test_direct_row__get_mutation_pbs(): + from google.cloud.bigtable.data.mutations import _SERVER_SIDE_TIMESTAMP, SetCell + + row_key = b"row_key" + row = _make_direct_row(row_key, None) + + mutation = SetCell( + family="column_family_id", + qualifier=b"column", + new_value=b"value", + timestamp_micros=_SERVER_SIDE_TIMESTAMP, + ) + + row._mutations = [mutation] + + assert row._get_mutation_pbs() == [mutation._to_pb()] + + def test_direct_row_get_mutations_size(): row_key = b"row_key" row = _make_direct_row(row_key, None) @@ -108,7 +128,7 @@ def test_direct_row_get_mutations_size(): row.set_cell(column_family_id2, column2, value) total_mutations_size = 0 - for mutation in row._get_mutations(): + for mutation in row._get_mutation_pbs(): total_mutations_size += mutation._pb.ByteSize() assert row.get_mutations_size() == total_mutations_size @@ -123,26 +143,27 @@ def _set_cell_helper( ): import struct + from google.cloud.bigtable.data.mutations import SetCell + row_key = b"row_key" column_family_id = "column_family_id" if column is None: column = b"column" table = object() row = _make_direct_row(row_key, table) - assert row._pb_mutations == [] + assert row._mutations == [] row.set_cell(column_family_id, column, value, timestamp=timestamp) if isinstance(value, int): value = struct.pack(">q", value) - expected_pb = _MutationPB( - set_cell=_MutationSetCellPB( - family_name=column_family_id, - column_qualifier=column_bytes or column, - timestamp_micros=timestamp_micros, - value=value, - ) + expected_mutation = SetCell( + family=column_family_id, + qualifier=column_bytes or column, + new_value=value, + timestamp_micros=timestamp_micros, ) - assert row._pb_mutations == [expected_pb] + + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_set_cell(): @@ -184,13 +205,15 @@ def test_direct_row_set_cell_with_non_null_timestamp(): def test_direct_row_delete(): + from google.cloud.bigtable.data.mutations import DeleteAllFromRow + row_key = b"row_key" row = _make_direct_row(row_key, object()) - assert row._pb_mutations == [] + assert row._mutations == [] row.delete() - expected_pb = _MutationPB(delete_from_row=_MutationDeleteFromRowPB()) - assert row._pb_mutations == [expected_pb] + expected_mutation = DeleteAllFromRow() + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_delete_cell(): @@ -214,14 +237,14 @@ def _delete_cells(self, *args, **kwargs): mock_row = MockRow(row_key, table) # Make sure no values are set before calling the method. - assert mock_row._pb_mutations == [] + assert mock_row._mutations == [] assert mock_row._args == [] assert mock_row._kwargs == [] # Actually make the request against the mock class. time_range = object() mock_row.delete_cell(column_family_id, column, time_range=time_range) - assert mock_row._pb_mutations == [] + assert mock_row._mutations == [] assert mock_row._args == [(column_family_id, [column])] assert mock_row._kwargs == [{"state": None, "time_range": time_range}] @@ -238,6 +261,7 @@ def test_direct_row_delete_cells_non_iterable(): def test_direct_row_delete_cells_all_columns(): + from google.cloud.bigtable.data.mutations import DeleteAllFromFamily from google.cloud.bigtable.row import DirectRow row_key = b"row_key" @@ -245,13 +269,11 @@ def test_direct_row_delete_cells_all_columns(): table = object() row = _make_direct_row(row_key, table) - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, DirectRow.ALL_COLUMNS) - expected_pb = _MutationPB( - delete_from_family=_MutationDeleteFromFamilyPB(family_name=column_family_id) - ) - assert row._pb_mutations == [expected_pb] + expected_mutation = DeleteAllFromFamily(family_to_delete=column_family_id) + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_delete_cells_no_columns(): @@ -261,12 +283,14 @@ def test_direct_row_delete_cells_no_columns(): row = _make_direct_row(row_key, table) columns = [] - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, columns) - assert row._pb_mutations == [] + assert row._mutations == [] def _delete_cells_helper(time_range=None): + from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn + row_key = b"row_key" column = b"column" column_family_id = "column_family_id" @@ -274,17 +298,17 @@ def _delete_cells_helper(time_range=None): row = _make_direct_row(row_key, table) columns = [column] - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, columns, time_range=time_range) - expected_pb = _MutationPB( - delete_from_column=_MutationDeleteFromColumnPB( - family_name=column_family_id, column_qualifier=column - ) - ) + expected_mutation = DeleteRangeFromColumn(family=column_family_id, qualifier=column) if time_range is not None: - expected_pb.delete_from_column.time_range._pb.CopyFrom(time_range._to_pb()._pb) - assert row._pb_mutations == [expected_pb] + timestamps = time_range._to_dict() + expected_mutation.start_timestamp_micros = timestamps.get( + "start_timestamp_micros" + ) + expected_mutation.end_timestamp_micros = timestamps.get("end_timestamp_micros") + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_delete_cells_no_time_range(): @@ -314,13 +338,15 @@ def test_direct_row_delete_cells_with_bad_column(): row = _make_direct_row(row_key, table) columns = [column, object()] - assert row._pb_mutations == [] + assert row._mutations == [] with pytest.raises(TypeError): row.delete_cells(column_family_id, columns) - assert row._pb_mutations == [] + assert row._mutations == [] def test_direct_row_delete_cells_with_string_columns(): + from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn + row_key = b"row_key" column_family_id = "column_family_id" column1 = "column1" @@ -331,61 +357,150 @@ def test_direct_row_delete_cells_with_string_columns(): row = _make_direct_row(row_key, table) columns = [column1, column2] - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, columns) - expected_pb1 = _MutationPB( - delete_from_column=_MutationDeleteFromColumnPB( - family_name=column_family_id, column_qualifier=column1_bytes - ) + expected_mutation1 = DeleteRangeFromColumn( + family=column_family_id, qualifier=column1_bytes ) - expected_pb2 = _MutationPB( - delete_from_column=_MutationDeleteFromColumnPB( - family_name=column_family_id, column_qualifier=column2_bytes - ) + expected_mutation2 = DeleteRangeFromColumn( + family=column_family_id, qualifier=column2_bytes ) - assert row._pb_mutations == [expected_pb1, expected_pb2] + _assert_mutations_equal(row._mutations, [expected_mutation1, expected_mutation2]) def test_direct_row_commit(): + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable_v2.services.bigtable import BigtableClient + project_id = "project-id" row_key = b"row_key" table_name = "projects/more-stuff" + app_profile_id = "app_profile_id" column_family_id = "column_family_id" column = b"column" credentials = _make_credentials() client = _make_client(project=project_id, credentials=credentials, admin=True) - table = _Table(table_name, client=client) + table = _Table(table_name, client=client, app_profile_id=app_profile_id) row = _make_direct_row(row_key, table) value = b"bytes-value" + # Set mock + api = mock.create_autospec(BigtableClient) + response_pb = _MutateRowResponsePB() + api.mutate_row.side_effect = [response_pb] + client.table_data_client + client._table_data_client._gapic_client = api + # Perform the method and check the result. row.set_cell(column_family_id, column, value) - row.commit() - assert table.mutated_rows == [row] + response = row.commit() + assert row._mutations == [] + assert response == status_pb2.Status(code=code_pb2.OK) + call_args = api.mutate_row.call_args + assert app_profile_id == call_args.app_profile_id[0] def test_direct_row_commit_with_exception(): - from google.rpc import status_pb2 + from google.api_core.exceptions import InternalServerError + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable_v2.services.bigtable import BigtableClient project_id = "project-id" row_key = b"row_key" table_name = "projects/more-stuff" + app_profile_id = "app_profile_id" column_family_id = "column_family_id" column = b"column" credentials = _make_credentials() client = _make_client(project=project_id, credentials=credentials, admin=True) - table = _Table(table_name, client=client) + table = _Table(table_name, client=client, app_profile_id=app_profile_id) row = _make_direct_row(row_key, table) value = b"bytes-value" + # Set mock + api = mock.create_autospec(BigtableClient) + exception_message = "Boom!" + exception = InternalServerError(exception_message) + api.mutate_row.side_effect = [exception] + client.table_data_client + client._table_data_client._gapic_client = api + # Perform the method and check the result. row.set_cell(column_family_id, column, value) result = row.commit() - expected = status_pb2.Status(code=0) - assert result == expected + assert row._mutations == [] + assert result == status_pb2.Status( + code=code_pb2.Code.INTERNAL, message=exception_message + ) + call_args = api.mutate_row.call_args + assert app_profile_id == call_args.app_profile_id[0] + + +def test_direct_row_commit_with_unknown_exception(): + from google.api_core.exceptions import GoogleAPICallError + from google.rpc import code_pb2, status_pb2 + + from google.cloud.bigtable_v2.services.bigtable import BigtableClient + + project_id = "project-id" + row_key = b"row_key" + table_name = "projects/more-stuff" + app_profile_id = "app_profile_id" + column_family_id = "column_family_id" + column = b"column" + + credentials = _make_credentials() + client = _make_client(project=project_id, credentials=credentials, admin=True) + table = _Table(table_name, client=client, app_profile_id=app_profile_id) + row = _make_direct_row(row_key, table) + value = b"bytes-value" + + # Set mock + api = mock.create_autospec(BigtableClient) + exception_message = "Boom!" + exception = GoogleAPICallError(message=exception_message) + api.mutate_row.side_effect = [exception] + client.table_data_client + client._table_data_client._gapic_client = api + + # Perform the method and check the result. + row.set_cell(column_family_id, column, value) + result = row.commit() + assert row._mutations == [] + assert result == status_pb2.Status( + code=code_pb2.Code.UNKNOWN, message=exception_message + ) + call_args = api.mutate_row.call_args + assert app_profile_id == call_args.app_profile_id[0] + + +def test_direct_row_commit_with_invalid_argument(): + from google.cloud.bigtable_v2.services.bigtable import BigtableClient + + project_id = "project-id" + row_key = b"row_key" + table_name = "projects/more-stuff" + app_profile_id = "app_profile_id" + + credentials = _make_credentials() + client = _make_client(project=project_id, credentials=credentials, admin=True) + table = _Table(table_name, client=client, app_profile_id=app_profile_id) + row = _make_direct_row(row_key, table) + + # Set mock + api = mock.create_autospec(BigtableClient) + client.table_data_client + client._table_data_client._gapic_client = api + + # Perform the method and check the result. + with pytest.raises(ValueError, match="No mutations provided"): + row.commit() + api.mutate_row.assert_not_called() def _make_conditional_row(*args, **kwargs): @@ -521,50 +636,54 @@ def test_append_row_constructor(): row = _make_append_row(row_key, table) assert row._row_key == row_key assert row._table is table - assert row._rule_pb_list == [] + assert row._rule_list == [] def test_append_row_clear(): row_key = b"row_key" table = object() row = _make_append_row(row_key, table) - row._rule_pb_list = [1, 2, 3] + row._rule_list = [1, 2, 3] row.clear() - assert row._rule_pb_list == [] + assert row._rule_list == [] def test_append_row_append_cell_value(): + from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule + table = object() row_key = b"row_key" row = _make_append_row(row_key, table) - assert row._rule_pb_list == [] + assert row._rule_list == [] column = b"column" column_family_id = "column_family_id" value = b"bytes-val" row.append_cell_value(column_family_id, column, value) - expected_pb = _ReadModifyWriteRulePB( - family_name=column_family_id, column_qualifier=column, append_value=value + expected_pb = AppendValueRule( + family=column_family_id, qualifier=column, append_value=value ) - assert row._rule_pb_list == [expected_pb] + _assert_mutations_equal(row._rule_list, [expected_pb]) def test_append_row_increment_cell_value(): + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + table = object() row_key = b"row_key" row = _make_append_row(row_key, table) - assert row._rule_pb_list == [] + assert row._rule_list == [] column = b"column" column_family_id = "column_family_id" int_value = 281330 row.increment_cell_value(column_family_id, column, int_value) - expected_pb = _ReadModifyWriteRulePB( - family_name=column_family_id, - column_qualifier=column, + expected_pb = IncrementRule( + family=column_family_id, + qualifier=column, increment_amount=int_value, ) - assert row._rule_pb_list == [expected_pb] + _assert_mutations_equal(row._rule_list, [expected_pb]) def test_append_row_commit(): @@ -610,7 +729,7 @@ def mock_parse_rmw_row_response(row_response): call_args = api.read_modify_write_row.call_args_list[0] assert app_profile_id == call_args.app_profile_id[0] assert result == expected_result - assert row._rule_pb_list == [] + assert row._rule_list == [] def test_append_row_commit_no_rules(): @@ -623,7 +742,7 @@ def test_append_row_commit_no_rules(): client = _make_client(project=project_id, credentials=credentials, admin=True) table = _Table(None, client=client) row = _make_append_row(row_key, table) - assert row._rule_pb_list == [] + assert row._rule_list == [] # Patch the stub used by the API method. stub = _FakeStub() @@ -643,8 +762,8 @@ def test_append_row_commit_too_many_mutations(): row_key = b"row_key" table = object() row = _make_append_row(row_key, table) - row._rule_pb_list = [1, 2, 3] - num_mutations = len(row._rule_pb_list) + row._rule_list = [1, 2, 3] + num_mutations = len(row._rule_list) with _Monkey(MUT, MAX_MUTATIONS=num_mutations - 1): with pytest.raises(ValueError): row.commit() @@ -653,6 +772,7 @@ def test_append_row_commit_too_many_mutations(): def test__parse_rmw_row_response(): from google.cloud._helpers import _datetime_from_microseconds + from google.cloud.bigtable.data.row import Row from google.cloud.bigtable.row import _parse_rmw_row_response col_fam1 = "col-fam-id" @@ -703,46 +823,14 @@ def test__parse_rmw_row_response(): ), ] ) - sample_input = _ReadModifyWriteRowResponsePB(row=response_row) + sample_input = Row._from_pb(response_row) assert expected_output == _parse_rmw_row_response(sample_input) -def test__parse_family_pb(): - from google.cloud._helpers import _datetime_from_microseconds - - from google.cloud.bigtable.row import _parse_family_pb - - col_fam1 = "col-fam-id" - col_name1 = b"col-name1" - col_name2 = b"col-name2" - cell_val1 = b"cell-val" - cell_val2 = b"cell-val-newer" - cell_val3 = b"altcol-cell-val" +def _MutateRowResponsePB(): + from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 - microseconds = 5554441037 - timestamp = _datetime_from_microseconds(microseconds) - expected_dict = { - col_name1: [(cell_val1, timestamp), (cell_val2, timestamp)], - col_name2: [(cell_val3, timestamp)], - } - expected_output = (col_fam1, expected_dict) - sample_input = _FamilyPB( - name=col_fam1, - columns=[ - _ColumnPB( - qualifier=col_name1, - cells=[ - _CellPB(value=cell_val1, timestamp_micros=microseconds), - _CellPB(value=cell_val2, timestamp_micros=microseconds), - ], - ), - _ColumnPB( - qualifier=col_name2, - cells=[_CellPB(value=cell_val3, timestamp_micros=microseconds)], - ), - ], - ) - assert expected_output == _parse_family_pb(sample_input) + return messages_v2_pb2.MutateRowResponse() def _CheckAndMutateRowResponsePB(*args, **kw): @@ -751,12 +839,6 @@ def _CheckAndMutateRowResponsePB(*args, **kw): return messages_v2_pb2.CheckAndMutateRowResponse(*args, **kw) -def _ReadModifyWriteRowResponsePB(*args, **kw): - from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 - - return messages_v2_pb2.ReadModifyWriteRowResponse(*args, **kw) - - def _CellPB(*args, **kw): from google.cloud.bigtable_v2.types import data as data_v2_pb2 @@ -775,46 +857,18 @@ def _FamilyPB(*args, **kw): return data_v2_pb2.Family(*args, **kw) -def _MutationPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation(*args, **kw) - - -def _MutationSetCellPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.SetCell(*args, **kw) - - -def _MutationDeleteFromColumnPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.DeleteFromColumn(*args, **kw) - - -def _MutationDeleteFromFamilyPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.DeleteFromFamily(*args, **kw) - - -def _MutationDeleteFromRowPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.DeleteFromRow(*args, **kw) - - def _RowPB(*args, **kw): from google.cloud.bigtable_v2.types import data as data_v2_pb2 return data_v2_pb2.Row(*args, **kw) -def _ReadModifyWriteRulePB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 +def _assert_mutations_equal(mutations_1, mutations_2): + assert len(mutations_1) == len(mutations_2) - return data_v2_pb2.ReadModifyWriteRule(*args, **kw) + for i in range(0, len(mutations_1)): + assert type(mutations_1[i]) is type(mutations_2[i]) + assert mutations_1[i]._to_pb() == mutations_2[i]._to_pb() class _Instance(object): @@ -828,10 +882,9 @@ def __init__(self, name, client=None, app_profile_id=None): self._instance = _Instance(client) self._app_profile_id = app_profile_id self.client = client - self.mutated_rows = [] - def mutate_rows(self, rows): - from google.rpc import status_pb2 - - self.mutated_rows.extend(rows) - return [status_pb2.Status(code=0)] + self._table_impl = self._instance._client._veneer_data_client.get_table( + _INSTANCE_ID, + self.name, + app_profile_id=self._app_profile_id, + ) diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py index 322135443ff6..bc7f8202d9b3 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_data.py @@ -13,10 +13,9 @@ # limitations under the License. -import mock import pytest -from ._testing import _make_credentials +from google.cloud.bigtable.data.row import Cell, Row TIMESTAMP_MICROS = 18738724000 # Make sure millis granularity ROW_KEY = b"row-key" @@ -24,6 +23,44 @@ QUALIFIER = b"qualifier" VALUE = b"value" TABLE_NAME = "table_name" +ROWS = [ + Row( + key=ROW_KEY, + cells=[ + Cell( + value=VALUE, + row_key=ROW_KEY, + family=FAMILY_NAME, + qualifier=QUALIFIER, + timestamp_micros=TIMESTAMP_MICROS, + ) + ], + ), + Row( + key=ROW_KEY + b"2", + cells=[ + Cell( + value=VALUE, + row_key=ROW_KEY + b"2", + family=FAMILY_NAME, + qualifier=QUALIFIER, + timestamp_micros=TIMESTAMP_MICROS, + ) + ], + ), + Row( + key=ROW_KEY + b"3", + cells=[ + Cell( + value=VALUE, + row_key=ROW_KEY + b"3", + family=FAMILY_NAME, + qualifier=QUALIFIER, + timestamp_micros=TIMESTAMP_MICROS, + ) + ], + ), +] def _make_cell(*args, **kwargs): @@ -67,6 +104,45 @@ def test_cell_from_pb_with_labels(): _cell_from_pb_test_helper(labels) +def test_cell__from_data_client_cell(): + from google.cloud.bigtable.data.row import Cell as DataCell + from google.cloud.bigtable.row_data import Cell + + data_cell = DataCell( + value=VALUE, + row_key=ROW_KEY, + family=FAMILY_NAME, + qualifier=QUALIFIER, + timestamp_micros=TIMESTAMP_MICROS, + ) + + cell = Cell._from_data_client_cell(data_cell) + + assert cell.value == VALUE + assert cell.timestamp_micros == TIMESTAMP_MICROS + assert cell.labels == [] + + +def test_cell__from_data_client_cell_with_labels(): + from google.cloud.bigtable.data.row import Cell as DataCell + from google.cloud.bigtable.row_data import Cell + + data_cell = DataCell( + value=VALUE, + row_key=ROW_KEY, + family=FAMILY_NAME, + qualifier=QUALIFIER, + timestamp_micros=TIMESTAMP_MICROS, + labels=["label1", "label2"], + ) + + cell = Cell._from_data_client_cell(data_cell) + + assert cell.value == VALUE + assert cell.timestamp_micros == TIMESTAMP_MICROS + assert cell.labels == ["label1", "label2"] + + def test_cell_constructor(): value = object() cell = _make_cell(value, TIMESTAMP_MICROS) @@ -114,6 +190,76 @@ def test_partial_row_data_constructor(): assert partial_row_data._cells == {} +def test_partial_row_data__from_data_client_row(): + from google.cloud.bigtable.data.row import Cell as DataCell + from google.cloud.bigtable.data.row import Row as DataRow + from google.cloud.bigtable.row_data import PartialRowData + + cells = [ + DataCell( + value=b"1", + row_key=ROW_KEY, + family="family1", + qualifier=b"qual1", + timestamp_micros=1, + ), + DataCell( + value=b"2", + row_key=ROW_KEY, + family="family1", + qualifier=b"qual1", + timestamp_micros=2, + ), + DataCell( + value=b"3", + row_key=ROW_KEY, + family="family1", + qualifier=b"qual2", + timestamp_micros=3, + ), + DataCell( + value=b"4", + row_key=ROW_KEY, + family="family2", + qualifier=b"qual1", + timestamp_micros=4, + ), + DataCell( + value=b"5", + row_key=ROW_KEY, + family="family2", + qualifier=b"qual2", + timestamp_micros=5, + ), + ] + + row = DataRow(ROW_KEY, cells) + partial_row_data = PartialRowData._from_data_client_row(row) + + expected_cells = { + "family1": { + b"qual1": [ + _make_cell(b"1", 1), + _make_cell(b"2", 2), + ], + b"qual2": [ + _make_cell(b"3", 3), + ], + }, + "family2": { + b"qual1": [ + _make_cell(b"4", 4), + ], + b"qual2": [ + _make_cell(b"5", 5), + ], + }, + } + + assert partial_row_data._row_key == ROW_KEY + assert partial_row_data._cells == expected_cells + + def test_partial_row_data___eq__(): row_key = object() partial_row_data1 = _make_partial_row_data(row_key) @@ -287,90 +433,6 @@ def trailing_metadata(self): return TestingException(exception) -def test__retry_read_rows_exception_miss(): - from google.api_core.exceptions import Conflict - - from google.cloud.bigtable.row_data import _retry_read_rows_exception - - exception = Conflict("testing") - assert not _retry_read_rows_exception(exception) - - -def test__retry_read_rows_exception_service_unavailable(): - from google.api_core.exceptions import ServiceUnavailable - - from google.cloud.bigtable.row_data import _retry_read_rows_exception - - exception = ServiceUnavailable("testing") - assert _retry_read_rows_exception(exception) - - -def test__retry_read_rows_exception_deadline_exceeded(): - from google.api_core.exceptions import DeadlineExceeded - - from google.cloud.bigtable.row_data import _retry_read_rows_exception - - exception = DeadlineExceeded("testing") - assert _retry_read_rows_exception(exception) - - -def test__retry_read_rows_exception_internal_server_not_retriable(): - from google.api_core.exceptions import InternalServerError - - from google.cloud.bigtable.row_data import ( - RETRYABLE_INTERNAL_ERROR_MESSAGES, - _retry_read_rows_exception, - ) - - err_message = "500 Error" - exception = InternalServerError(err_message) - assert err_message not in RETRYABLE_INTERNAL_ERROR_MESSAGES - assert not _retry_read_rows_exception(exception) - - -def test__retry_read_rows_exception_internal_server_retriable(): - from google.api_core.exceptions import InternalServerError - - from google.cloud.bigtable.row_data import ( - RETRYABLE_INTERNAL_ERROR_MESSAGES, - _retry_read_rows_exception, - ) - - for err_message in RETRYABLE_INTERNAL_ERROR_MESSAGES: - exception = InternalServerError(err_message) - assert _retry_read_rows_exception(exception) - - -def test__retry_read_rows_exception_miss_wrapped_in_grpc(): - from google.api_core.exceptions import Conflict - - from google.cloud.bigtable.row_data import _retry_read_rows_exception - - wrapped = Conflict("testing") - exception = _make_grpc_call_error(wrapped) - assert not _retry_read_rows_exception(exception) - - -def test__retry_read_rows_exception_service_unavailable_wrapped_in_grpc(): - from google.api_core.exceptions import ServiceUnavailable - - from google.cloud.bigtable.row_data import _retry_read_rows_exception - - wrapped = ServiceUnavailable("testing") - exception = _make_grpc_call_error(wrapped) - assert _retry_read_rows_exception(exception) - - -def test__retry_read_rows_exception_deadline_exceeded_wrapped_in_grpc(): - from google.api_core.exceptions import DeadlineExceeded - - from google.cloud.bigtable.row_data import _retry_read_rows_exception - - wrapped = DeadlineExceeded("testing") - exception = _make_grpc_call_error(wrapped) - assert _retry_read_rows_exception(exception) - - def test_partial_cell_data(): from google.cloud.bigtable.row_data import PartialCellData @@ -395,65 +457,32 @@ def test_partial_cell_data(): assert instance.value == added_value + added_value -def _make_partial_rows_data(*args, **kwargs): +def _make_partial_rows_data(generator): from google.cloud.bigtable.row_data import PartialRowsData - return PartialRowsData(*args, **kwargs) + return PartialRowsData._from_generator(generator) def _partial_rows_data_consume_all(yrd): return [row.row_key for row in yrd] -def _make_client(*args, **kwargs): - from google.cloud.bigtable.client import Client - - return Client(*args, **kwargs) +def _make_generator(rows, error=None): + for row in rows: + if error: + raise error + else: + yield row -def test_partial_rows_data_constructor(): - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS - - client = _Client() - client._data_stub = mock.MagicMock() - request = object() - partial_rows_data = _make_partial_rows_data(client._data_stub.ReadRows, request) - assert partial_rows_data.request is request - assert partial_rows_data.rows == {} - assert partial_rows_data.retry == DEFAULT_RETRY_READ_ROWS +def _assert_generator_closed(generator): + with pytest.raises(StopIteration): + next(generator) def test_partial_rows_data_consume_all(): - resp = _ReadRowsResponseV2( - [ - _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ), - _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY + b"2", - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ), - ] - ) - - call_count = 0 - iterator = _MockCancellableIterator(resp) - - def fake_read(*args, **kwargs): - nonlocal call_count - call_count += 1 - return iterator - - partial_rows_data = _make_partial_rows_data(fake_read, None) + generator = _make_generator(ROWS) + partial_rows_data = _make_partial_rows_data(generator) partial_rows_data.consume_all() row1 = _make_partial_row_data(ROW_KEY) @@ -464,772 +493,125 @@ def fake_read(*args, **kwargs): row2._cells[FAMILY_NAME] = { QUALIFIER: [_make_cell(value=VALUE, timestamp_micros=TIMESTAMP_MICROS)] } + row3 = _make_partial_row_data(ROW_KEY + b"3") + row3._cells[FAMILY_NAME] = { + QUALIFIER: [_make_cell(value=VALUE, timestamp_micros=TIMESTAMP_MICROS)] + } - assert partial_rows_data.rows == {row1.row_key: row1, row2.row_key: row2} - - -def test_partial_rows_data_constructor_with_retry(): - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS - - client = _Client() - client._data_stub = mock.MagicMock() - request = object() - retry = DEFAULT_RETRY_READ_ROWS - partial_rows_data = _make_partial_rows_data( - client._data_stub.ReadRows, request, retry - ) - partial_rows_data.read_method.assert_called_once_with( - request, - timeout=DEFAULT_RETRY_READ_ROWS.deadline + 1, - retry=DEFAULT_RETRY_READ_ROWS, - ) - assert partial_rows_data.request is request - assert partial_rows_data.rows == {} - assert partial_rows_data.retry == retry - - -def test_partial_rows_data___eq__(): - client = _Client() - client._data_stub = mock.MagicMock() - request = object() - partial_rows_data1 = _make_partial_rows_data(client._data_stub.ReadRows, request) - partial_rows_data2 = _make_partial_rows_data(client._data_stub.ReadRows, request) - assert partial_rows_data1.rows == partial_rows_data2.rows - - -def test_partial_rows_data___eq__type_differ(): - client = _Client() - client._data_stub = mock.MagicMock() - request = object() - partial_rows_data1 = _make_partial_rows_data(client._data_stub.ReadRows, request) - partial_rows_data2 = object() - assert not (partial_rows_data1 == partial_rows_data2) - - -def test_partial_rows_data___ne__same_value(): - client = _Client() - client._data_stub = mock.MagicMock() - request = object() - partial_rows_data1 = _make_partial_rows_data(client._data_stub.ReadRows, request) - partial_rows_data2 = _make_partial_rows_data(client._data_stub.ReadRows, request) - assert partial_rows_data1 != partial_rows_data2 - - -def test_partial_rows_data___ne__(): - client = _Client() - client._data_stub = mock.MagicMock() - request = object() - partial_rows_data1 = _make_partial_rows_data(client._data_stub.ReadRows, request) - partial_rows_data2 = _make_partial_rows_data(client._data_stub.ReadRows, request) - assert partial_rows_data1 != partial_rows_data2 - - -def test_partial_rows_data_rows_getter(): - client = _Client() - client._data_stub = mock.MagicMock() - request = object() - partial_rows_data = _make_partial_rows_data(client._data_stub.ReadRows, request) - partial_rows_data.rows = value = object() - assert partial_rows_data.rows is value - - -def test_partial_rows_data_state_start(): - client = _Client() - iterator = _MockCancellableIterator() - client._data_stub = mock.MagicMock() - client._data_stub.ReadRows.side_effect = [iterator] - request = object() - yrd = _make_partial_rows_data(client._data_stub.ReadRows, request) - assert yrd.state == yrd.NEW_ROW - - -def test_partial_rows_data_state_new_row_w_row(): - from google.cloud.bigtable_v2.services.bigtable import BigtableClient - - chunk = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - chunks = [chunk] - - response = _ReadRowsResponseV2(chunks) - iterator = _MockCancellableIterator(response) - - data_api = mock.create_autospec(BigtableClient) - - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - client._table_data_client = data_api - request = object() - - yrd = _make_partial_rows_data(client._table_data_client.read_rows, request) - assert yrd.retry._deadline == 60.0 - - yrd.response_iterator = iterator - rows = [row for row in yrd] - - result = rows[0] - assert result.row_key == ROW_KEY - assert yrd._counter == 1 - assert yrd.state == yrd.NEW_ROW - - -def test_partial_rows_data_multiple_chunks(): - from google.cloud.bigtable_v2.services.bigtable import BigtableClient - - chunk1 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=False, - ) - chunk2 = _ReadRowsResponseCellChunkPB( - qualifier=QUALIFIER + b"1", - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - chunks = [chunk1, chunk2] - - response = _ReadRowsResponseV2(chunks) - iterator = _MockCancellableIterator(response) - data_api = mock.create_autospec(BigtableClient) - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - client._table_data_client = data_api - request = object() - - yrd = _make_partial_rows_data(data_api.read_rows, request) + assert partial_rows_data.rows == { + row1.row_key: row1, + row2.row_key: row2, + row3.row_key: row3, + } - yrd.response_iterator = iterator - rows = [row for row in yrd] - result = rows[0] - assert result.row_key == ROW_KEY - assert yrd._counter == 1 - assert yrd.state == yrd.NEW_ROW + _assert_generator_closed(generator) def test_partial_rows_data_cancel(): - client = _Client() - response_iterator = _MockCancellableIterator() - client._data_stub = mock.MagicMock() - client._data_stub.ReadRows.side_effect = [response_iterator] - request = object() - yield_rows_data = _make_partial_rows_data(client._data_stub.ReadRows, request) - assert response_iterator.cancel_calls == 0 - yield_rows_data.cancel() - assert response_iterator.cancel_calls == 1 - assert list(yield_rows_data) == [] - - -def test_partial_rows_data_cancel_between_chunks(): - from google.cloud.bigtable_v2.services.bigtable import BigtableClient - - chunk1 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - chunk2 = _ReadRowsResponseCellChunkPB( - qualifier=QUALIFIER + b"1", - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - chunks = [chunk1, chunk2] - response = _ReadRowsResponseV2(chunks) - response_iterator = _MockCancellableIterator(response) - - client = _Client() - data_api = mock.create_autospec(BigtableClient) - client._table_data_client = data_api - request = object() - yrd = _make_partial_rows_data(data_api.read_rows, request) - yrd.response_iterator = response_iterator - - rows = [] - for row in yrd: - yrd.cancel() - rows.append(row) - - assert response_iterator.cancel_calls == 1 - assert list(yrd) == [] - - -def test_partial_rows_data_valid_last_scanned_row_key_on_start(): - client = _Client() - response = _ReadRowsResponseV2([], last_scanned_row_key=b"2.AFTER") - iterator = _MockCancellableIterator(response) - client._data_stub = mock.MagicMock() - client._data_stub.read_rows.side_effect = [iterator] - request = object() - yrd = _make_partial_rows_data(client._data_stub.read_rows, request) - yrd.last_scanned_row_key = b"1.BEFORE" - _partial_rows_data_consume_all(yrd) - assert yrd.last_scanned_row_key == b"2.AFTER" - - -def test_partial_rows_data_invalid_empty_chunk(): - from google.cloud.bigtable.row_data import InvalidChunk - from google.cloud.bigtable_v2.services.bigtable import BigtableClient - - client = _Client() - chunks = _generate_cell_chunks([""]) - response = _ReadRowsResponseV2(chunks) - iterator = _MockCancellableIterator(response) - client._data_stub = mock.create_autospec(BigtableClient) - client._data_stub.read_rows.side_effect = [iterator] - request = object() - yrd = _make_partial_rows_data(client._data_stub.read_rows, request) - with pytest.raises(InvalidChunk): - _partial_rows_data_consume_all(yrd) - - -def test_partial_rows_data_state_cell_in_progress(): - labels = ["L1", "L2"] - resp = _ReadRowsResponseV2( - [ - _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - value_size=(2 * len(VALUE)), - labels=labels, - ), - _ReadRowsResponseCellChunkPB(value=VALUE, commit_row=True), - ] - ) - - def fake_read(*args, **kwargs): - return iter([resp]) + generator = _make_generator(ROWS) + partial_rows_data = _make_partial_rows_data(generator) + row_data = [] - yrd = _make_partial_rows_data(fake_read, None) - yrd.consume_all() + count = 0 + for row in partial_rows_data: + row_data.append(row) + count += 1 + if count == 1: + partial_rows_data.cancel() - expected_row = _make_partial_row_data(ROW_KEY) - expected_row._cells = { - QUALIFIER: [ - _make_cell( - value=(VALUE + VALUE), timestamp_micros=TIMESTAMP_MICROS, labels=labels - ) - ] + row1 = _make_partial_row_data(ROW_KEY) + row1._cells[FAMILY_NAME] = { + QUALIFIER: [_make_cell(value=VALUE, timestamp_micros=TIMESTAMP_MICROS)] } + assert row_data == [row1] + # We should have closed the generator, so there should be no more + # elements left in there even though we haven't iterated through all the rows. + _assert_generator_closed(generator) -def test_partial_rows_data_yield_rows_data(): - from google.cloud.bigtable_v2.services.bigtable import BigtableClient - client = _Client() +def test_partial_rows_data_deadline_exceeded(): + from google.api_core import exceptions - chunk = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, + generator = _make_generator( + ROWS, error=exceptions.DeadlineExceeded("Operation timed out.") ) - chunks = [chunk] - - response = _ReadRowsResponseV2(chunks) - iterator = _MockCancellableIterator(response) - data_api = mock.create_autospec(BigtableClient) - client._data_stub = data_api - client._data_stub.read_rows.side_effect = [iterator] - - request = object() - yrd = _make_partial_rows_data(client._data_stub.read_rows, request) + partial_rows_data = _make_partial_rows_data(generator) + with pytest.raises(exceptions.RetryError): + list(partial_rows_data) - result = _partial_rows_data_consume_all(yrd)[0] + # An exception should close the generator. + _assert_generator_closed(generator) - assert result == ROW_KEY +def _make_cell_pb(value): + from google.cloud.bigtable import row_data -def test_partial_rows_data_yield_retry_rows_data(): - from google.api_core import retry + return row_data.Cell(value, TIMESTAMP_MICROS) - client = _Client() - retry_read_rows = retry.Retry(predicate=_read_rows_retry_exception) +def test_partial_rows_data_legacy_constructor_fallback(): + import warnings - chunk = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - chunks = [chunk] - - response = _ReadRowsResponseV2(chunks) - failure_iterator = _MockFailureIterator_1() - iterator = _MockCancellableIterator(response) - client._data_stub = mock.MagicMock() - client._data_stub.ReadRows.side_effect = [failure_iterator, iterator] + from google.cloud.bigtable.row_data import PartialRowsData + read_method = object() request = object() - yrd = _make_partial_rows_data(client._data_stub.ReadRows, request, retry_read_rows) - - result = _partial_rows_data_consume_all(yrd)[0] - - assert result == ROW_KEY - - -def _make_read_rows_request_manager(*args, **kwargs): - from google.cloud.bigtable.row_data import _ReadRowsRequestManager - - return _ReadRowsRequestManager(*args, **kwargs) - - -@pytest.fixture(scope="session") -def rrrm_data(): - from google.cloud.bigtable import row_set - - row_range1 = row_set.RowRange(b"row_key21", b"row_key29") - row_range2 = row_set.RowRange(b"row_key31", b"row_key39") - row_range3 = row_set.RowRange(b"row_key41", b"row_key49") - - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request.rows.row_ranges.append(row_range1.get_range_kwargs()) - request.rows.row_ranges.append(row_range2.get_range_kwargs()) - request.rows.row_ranges.append(row_range3.get_range_kwargs()) - - yield { - "row_range1": row_range1, - "row_range2": row_range2, - "row_range3": row_range3, - "request": request, - } - - -def test_RRRM_constructor(): - request = mock.Mock() - last_scanned_key = "last_key" - rows_read_so_far = 10 - - request_manager = _make_read_rows_request_manager( - request, last_scanned_key, rows_read_so_far - ) - assert request == request_manager.message - assert last_scanned_key == request_manager.last_scanned_key - assert rows_read_so_far == request_manager.rows_read_so_far - - -def test_RRRM__filter_row_key(): - table_name = "table_name" - request = _ReadRowsRequestPB(table_name=table_name) - request.rows.row_keys.extend([b"row_key1", b"row_key2", b"row_key3", b"row_key4"]) - - last_scanned_key = b"row_key2" - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - row_keys = request_manager._filter_rows_keys() - - expected_row_keys = [b"row_key3", b"row_key4"] - assert expected_row_keys == row_keys - - -def test_RRRM__filter_row_key_is_empty(): - table_name = "table_name" - request = _ReadRowsRequestPB(table_name=table_name) - request.rows.row_keys.extend([b"row_key1", b"row_key2", b"row_key3", b"row_key4"]) - - last_scanned_key = b"row_key4" - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 4) - row_keys = request_manager._filter_rows_keys() - - assert row_keys == [] - - -def test_RRRM__filter_row_ranges_all_ranges_added_back(rrrm_data): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - request = rrrm_data["request"] - last_scanned_key = b"row_key14" - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - row_ranges = request_manager._filter_row_ranges() - - exp_row_range1 = data_v2_pb2.RowRange( - start_key_closed=b"row_key21", end_key_open=b"row_key29" - ) - exp_row_range2 = data_v2_pb2.RowRange( - start_key_closed=b"row_key31", end_key_open=b"row_key39" - ) - exp_row_range3 = data_v2_pb2.RowRange( - start_key_closed=b"row_key41", end_key_open=b"row_key49" - ) - exp_row_ranges = [exp_row_range1, exp_row_range2, exp_row_range3] - - assert exp_row_ranges == row_ranges - - -def test_RRRM__filter_row_ranges_all_ranges_already_read(rrrm_data): - request = rrrm_data["request"] - last_scanned_key = b"row_key54" - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - row_ranges = request_manager._filter_row_ranges() - - assert row_ranges == [] - - -def test_RRRM__filter_row_ranges_all_ranges_already_read_open_closed(): - from google.cloud.bigtable import row_set - - last_scanned_key = b"row_key54" - - row_range1 = row_set.RowRange(b"row_key21", b"row_key29", False, True) - row_range2 = row_set.RowRange(b"row_key31", b"row_key39") - row_range3 = row_set.RowRange(b"row_key41", b"row_key49", False, True) - - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request.rows.row_ranges.append(row_range1.get_range_kwargs()) - request.rows.row_ranges.append(row_range2.get_range_kwargs()) - request.rows.row_ranges.append(row_range3.get_range_kwargs()) - - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - request_manager.new_message = _ReadRowsRequestPB(table_name=TABLE_NAME) - row_ranges = request_manager._filter_row_ranges() - - assert row_ranges == [] - - -def test_RRRM__filter_row_ranges_some_ranges_already_read(rrrm_data): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - request = rrrm_data["request"] - last_scanned_key = b"row_key22" - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - request_manager.new_message = _ReadRowsRequestPB(table_name=TABLE_NAME) - row_ranges = request_manager._filter_row_ranges() - - exp_row_range1 = data_v2_pb2.RowRange( - start_key_open=b"row_key22", end_key_open=b"row_key29" - ) - exp_row_range2 = data_v2_pb2.RowRange( - start_key_closed=b"row_key31", end_key_open=b"row_key39" - ) - exp_row_range3 = data_v2_pb2.RowRange( - start_key_closed=b"row_key41", end_key_open=b"row_key49" - ) - exp_row_ranges = [exp_row_range1, exp_row_range2, exp_row_range3] - - assert exp_row_ranges == row_ranges - - -def test_RRRM_build_updated_request(rrrm_data): - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable_v2 import types - - row_range1 = rrrm_data["row_range1"] - row_filter = RowSampleFilter(0.33) - last_scanned_key = b"row_key25" - request = _ReadRowsRequestPB( - filter=row_filter._to_pb(), - rows_limit=8, - table_name=TABLE_NAME, - app_profile_id="app-profile-id-1", - ) - request.rows.row_ranges.append(row_range1.get_range_kwargs()) - - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - - result = request_manager.build_updated_request() - - expected_result = _ReadRowsRequestPB( - table_name=TABLE_NAME, - filter=row_filter._to_pb(), - rows_limit=6, - app_profile_id="app-profile-id-1", - ) - - row_range1 = types.RowRange( - start_key_open=last_scanned_key, end_key_open=row_range1.end_key - ) - expected_result.rows.row_ranges.append(row_range1) - - assert expected_result == result - - -def test_RRRM_build_updated_request_full_table(): - from google.cloud.bigtable_v2 import types - - last_scanned_key = b"row_key14" - - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - - result = request_manager.build_updated_request() - expected_result = _ReadRowsRequestPB(table_name=TABLE_NAME) - row_range1 = types.RowRange(start_key_open=last_scanned_key) - expected_result.rows.row_ranges.append(row_range1) - assert expected_result == result - - -def test_RRRM_build_updated_request_no_start_key(): - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable_v2 import types - - row_filter = RowSampleFilter(0.33) - last_scanned_key = b"row_key25" - request = _ReadRowsRequestPB( - filter=row_filter._to_pb(), rows_limit=8, table_name=TABLE_NAME - ) - row_range1 = types.RowRange(end_key_open=b"row_key29") - request.rows.row_ranges.append(row_range1) - - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - - result = request_manager.build_updated_request() + with warnings.catch_warnings(record=True) as warned: + warnings.simplefilter("always") + prd = PartialRowsData(read_method, request) - expected_result = _ReadRowsRequestPB( - table_name=TABLE_NAME, filter=row_filter._to_pb(), rows_limit=6 + assert len(warned) == 1 + assert issubclass(warned[0].category, DeprecationWarning) + assert "Direct instantiation of `PartialRowsData` is deprecated" in str( + warned[0].message ) - - row_range2 = types.RowRange( - start_key_open=last_scanned_key, end_key_open=b"row_key29" - ) - expected_result.rows.row_ranges.append(row_range2) - - assert expected_result == result - - -def test_RRRM_build_updated_request_no_end_key(): - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable_v2 import types - - row_filter = RowSampleFilter(0.33) - last_scanned_key = b"row_key25" - request = _ReadRowsRequestPB( - filter=row_filter._to_pb(), rows_limit=8, table_name=TABLE_NAME - ) - - row_range1 = types.RowRange(start_key_closed=b"row_key20") - request.rows.row_ranges.append(row_range1) - - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - - result = request_manager.build_updated_request() - - expected_result = _ReadRowsRequestPB( - table_name=TABLE_NAME, filter=row_filter._to_pb(), rows_limit=6 - ) - row_range2 = types.RowRange(start_key_open=last_scanned_key) - expected_result.rows.row_ranges.append(row_range2) - - assert expected_result == result - - -def test_RRRM_build_updated_request_rows(): - from google.cloud.bigtable.row_filters import RowSampleFilter - - row_filter = RowSampleFilter(0.33) - last_scanned_key = b"row_key4" - request = _ReadRowsRequestPB( - filter=row_filter._to_pb(), rows_limit=5, table_name=TABLE_NAME - ) - request.rows.row_keys.extend( - [b"row_key1", b"row_key2", b"row_key4", b"row_key5", b"row_key7", b"row_key9"] - ) - - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 3) - - result = request_manager.build_updated_request() - - expected_result = _ReadRowsRequestPB( - table_name=TABLE_NAME, filter=row_filter._to_pb(), rows_limit=2 - ) - expected_result.rows.row_keys.extend([b"row_key5", b"row_key7", b"row_key9"]) - - assert expected_result == result + assert list(prd) == [] + assert prd.rows == {} + # cancel() should work cleanly on the generator fallback + prd.cancel() -def test_RRRM_build_updated_request_rows_limit(): - from google.cloud.bigtable_v2 import types +def test_partial_rows_data_keyword_generator(): + import warnings - last_scanned_key = b"row_key14" - - request = _ReadRowsRequestPB(table_name=TABLE_NAME, rows_limit=10) - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - - result = request_manager.build_updated_request() - expected_result = _ReadRowsRequestPB(table_name=TABLE_NAME, rows_limit=8) - row_range1 = types.RowRange(start_key_open=last_scanned_key) - expected_result.rows.row_ranges.append(row_range1) - assert expected_result == result - - -def test_RRRM__key_already_read(): - last_scanned_key = b"row_key14" - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request_manager = _make_read_rows_request_manager(request, last_scanned_key, 2) - - assert request_manager._key_already_read(b"row_key11") - assert not request_manager._key_already_read(b"row_key16") - - -def test_RRRM__rows_limit_reached(): - from google.cloud.bigtable.row_data import InvalidRetryRequest - - last_scanned_key = b"row_key14" - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request.rows_limit = 2 - request_manager = _make_read_rows_request_manager( - request, last_scanned_key=last_scanned_key, rows_read_so_far=2 - ) - with pytest.raises(InvalidRetryRequest): - request_manager.build_updated_request() - - -def test_RRRM_build_updated_request_last_row_read_raises_invalid_retry_request(): - from google.cloud.bigtable.row_data import InvalidRetryRequest - - last_scanned_key = b"row_key4" - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request.rows.row_keys.extend([b"row_key1", b"row_key2", b"row_key4"]) - - request_manager = _make_read_rows_request_manager( - request, last_scanned_key, rows_read_so_far=3 - ) - with pytest.raises(InvalidRetryRequest): - request_manager.build_updated_request() - - -def test_RRRM_build_updated_request_row_ranges_read_raises_invalid_retry_request(): - from google.cloud.bigtable import row_set - from google.cloud.bigtable.row_data import InvalidRetryRequest - - row_range1 = row_set.RowRange(b"row_key21", b"row_key29") - - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request.rows.row_ranges.append(row_range1.get_range_kwargs()) - - last_scanned_key = b"row_key4" - request = _ReadRowsRequestPB( - table_name=TABLE_NAME, - ) - request.rows.row_ranges.append(row_range1.get_range_kwargs()) - - request_manager = _make_read_rows_request_manager( - request, last_scanned_key, rows_read_so_far=2 - ) - with pytest.raises(InvalidRetryRequest): - request_manager.build_updated_request() - - -def test_RRRM_build_updated_request_row_ranges_valid(): - from google.cloud.bigtable import row_set - - row_range1 = row_set.RowRange(b"row_key21", b"row_key29") - - request = _ReadRowsRequestPB(table_name=TABLE_NAME) - request.rows.row_ranges.append(row_range1.get_range_kwargs()) - - last_scanned_key = b"row_key21" - request = _ReadRowsRequestPB( - table_name=TABLE_NAME, - ) - request.rows.row_ranges.append(row_range1.get_range_kwargs()) - - request_manager = _make_read_rows_request_manager( - request, last_scanned_key, rows_read_so_far=1 - ) - updated_request = request_manager.build_updated_request() - assert len(updated_request.rows.row_ranges) > 0 - - -class _MockCancellableIterator(object): - cancel_calls = 0 - - def __init__(self, *values): - self.iter_values = iter(values) - self.last_scanned_row_key = "" - - def cancel(self): - self.cancel_calls += 1 - - def next(self): - return next(self.iter_values) - - __next__ = next - - -class _MockFailureIterator_1(object): - def next(self): - from google.api_core.exceptions import DeadlineExceeded - - raise DeadlineExceeded("Failed to read from server") - - __next__ = next - - -def _ReadRowsResponseV2(chunks, last_scanned_row_key=b""): - from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 - - return messages_v2_pb2.ReadRowsResponse( - chunks=chunks, last_scanned_row_key=last_scanned_row_key - ) - - -def _generate_cell_chunks(chunk_text_pbs): - from google.protobuf.text_format import Merge - - from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse - - chunks = [] - - for chunk_text_pb in chunk_text_pbs: - chunk = ReadRowsResponse.CellChunk() - chunk._pb = Merge(chunk_text_pb, chunk._pb) - chunks.append(chunk) - - return chunks - - -def _ReadRowsResponseCellChunkPB(*args, **kw): - from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 - - family_name = kw.pop("family_name", None) - qualifier = kw.pop("qualifier", None) - message = messages_v2_pb2.ReadRowsResponse.CellChunk(*args, **kw) - - if family_name: - message.family_name = family_name - if qualifier: - message.qualifier = qualifier - - return message - - -def _make_cell_pb(value): - from google.cloud.bigtable import row_data - - return row_data.Cell(value, TIMESTAMP_MICROS) - - -def _ReadRowsRequestPB(*args, **kw): - from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 + from google.cloud.bigtable.row_data import PartialRowsData - return messages_v2_pb2.ReadRowsRequest(*args, **kw) + generator = _make_generator([]) + with warnings.catch_warnings(record=True) as warned: + warnings.simplefilter("always") + prd = PartialRowsData(generator=generator) + assert len(warned) == 0 + assert prd._generator is generator + assert prd.rows == {} -def _read_rows_retry_exception(exc): - from google.api_core.exceptions import DeadlineExceeded - return isinstance(exc, DeadlineExceeded) +def test_partial_rows_data_deprecated_properties(): + import warnings + from google.cloud.bigtable.row_data import PartialRowsData -class _Client(object): - data_stub = None + generator = _make_generator([]) + prd = PartialRowsData._from_generator(generator) + + for attr in [ + "state", + "last_scanned_row_key", + "read_method", + "request", + "retry", + "response_iterator", + ]: + with warnings.catch_warnings(record=True) as warned: + warnings.simplefilter("always") + val = getattr(prd, attr) + + assert val is None + assert len(warned) == 1 + assert issubclass(warned[0].category, DeprecationWarning) + assert f"The `{attr}` attribute on `PartialRowsData` is deprecated" in str( + warned[0].message + ) diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_merger.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_merger.py index cc9df9c70822..deac16a7f8fc 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_merger.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_merger.py @@ -1,12 +1,8 @@ import os -from itertools import zip_longest -from typing import List import proto -import pytest -from google.cloud.bigtable.row_data import InvalidChunk, PartialRowData, PartialRowsData -from google.cloud.bigtable.row_merger import _RowMerger +from google.cloud.bigtable.row_data import PartialRowData from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse @@ -59,172 +55,3 @@ def extract_results_from_row(row: PartialRowData): ) ) return results - - -@pytest.mark.parametrize( - "test_case", parse_readrows_acceptance_tests(), ids=lambda t: t.description -) -def test_scenario(test_case: ReadRowsTest): - def fake_read(*args, **kwargs): - return iter([ReadRowsResponse(chunks=test_case.chunks)]) - - actual_results: List[ReadRowsTest.Result] = [] - try: - for row in PartialRowsData(fake_read, request=None): - actual_results.extend(extract_results_from_row(row)) - except (InvalidChunk, ValueError): - actual_results.append(ReadRowsTest.Result(error=True)) - - for expected, actual in zip_longest(test_case.results, actual_results): - assert actual == expected - - -def test_out_of_order_rows(): - row_merger = _RowMerger(last_seen_row=b"z") - with pytest.raises(InvalidChunk): - list(row_merger.process_chunks(ReadRowsResponse(last_scanned_row_key=b"a"))) - - -def test_bare_reset(): - first_chunk = ReadRowsResponse.CellChunk( - ReadRowsResponse.CellChunk( - row_key=b"a", family_name="f", qualifier=b"q", value=b"v" - ) - ) - with pytest.raises(InvalidChunk): - _process_chunks( - first_chunk, - ReadRowsResponse.CellChunk( - ReadRowsResponse.CellChunk(reset_row=True, row_key=b"a") - ), - ) - with pytest.raises(InvalidChunk): - _process_chunks( - first_chunk, - ReadRowsResponse.CellChunk( - ReadRowsResponse.CellChunk(reset_row=True, family_name="f") - ), - ) - with pytest.raises(InvalidChunk): - _process_chunks( - first_chunk, - ReadRowsResponse.CellChunk( - ReadRowsResponse.CellChunk(reset_row=True, qualifier=b"q") - ), - ) - with pytest.raises(InvalidChunk): - _process_chunks( - first_chunk, - ReadRowsResponse.CellChunk( - ReadRowsResponse.CellChunk(reset_row=True, timestamp_micros=1000) - ), - ) - with pytest.raises(InvalidChunk): - _process_chunks( - first_chunk, - ReadRowsResponse.CellChunk( - ReadRowsResponse.CellChunk(reset_row=True, labels=["a"]) - ), - ) - with pytest.raises(InvalidChunk): - _process_chunks( - first_chunk, - ReadRowsResponse.CellChunk( - ReadRowsResponse.CellChunk(reset_row=True, value=b"v") - ), - ) - - -def test_missing_family(): - with pytest.raises(InvalidChunk): - _process_chunks( - ReadRowsResponse.CellChunk( - row_key=b"a", - qualifier=b"q", - timestamp_micros=1000, - value=b"v", - commit_row=True, - ) - ) - - -def test_mid_cell_row_key_change(): - with pytest.raises(InvalidChunk): - _process_chunks( - ReadRowsResponse.CellChunk( - row_key=b"a", - family_name="f", - qualifier=b"q", - timestamp_micros=1000, - value_size=2, - value=b"v", - ), - ReadRowsResponse.CellChunk(row_key=b"b", value=b"v", commit_row=True), - ) - - -def test_mid_cell_family_change(): - with pytest.raises(InvalidChunk): - _process_chunks( - ReadRowsResponse.CellChunk( - row_key=b"a", - family_name="f", - qualifier=b"q", - timestamp_micros=1000, - value_size=2, - value=b"v", - ), - ReadRowsResponse.CellChunk(family_name="f2", value=b"v", commit_row=True), - ) - - -def test_mid_cell_qualifier_change(): - with pytest.raises(InvalidChunk): - _process_chunks( - ReadRowsResponse.CellChunk( - row_key=b"a", - family_name="f", - qualifier=b"q", - timestamp_micros=1000, - value_size=2, - value=b"v", - ), - ReadRowsResponse.CellChunk(qualifier=b"q2", value=b"v", commit_row=True), - ) - - -def test_mid_cell_timestamp_change(): - with pytest.raises(InvalidChunk): - _process_chunks( - ReadRowsResponse.CellChunk( - row_key=b"a", - family_name="f", - qualifier=b"q", - timestamp_micros=1000, - value_size=2, - value=b"v", - ), - ReadRowsResponse.CellChunk( - timestamp_micros=2000, value=b"v", commit_row=True - ), - ) - - -def test_mid_cell_labels_change(): - with pytest.raises(InvalidChunk): - _process_chunks( - ReadRowsResponse.CellChunk( - row_key=b"a", - family_name="f", - qualifier=b"q", - timestamp_micros=1000, - value_size=2, - value=b"v", - ), - ReadRowsResponse.CellChunk(labels=["b"], value=b"v", commit_row=True), - ) - - -def _process_chunks(*chunks): - req = ReadRowsResponse.pb(ReadRowsResponse(chunks=chunks)) - return list(_RowMerger().process_chunks(req)) diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py index 4142348ee195..a79ceb717d0b 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_set.py @@ -48,38 +48,6 @@ def test_row_set__eq__type_differ(): assert not (row_set1 == row_set2) -def test_row_set__eq__len_row_keys_differ(): - from google.cloud.bigtable.row_set import RowSet - - row_key1 = b"row_key1" - row_key2 = b"row_key1" - - row_set1 = RowSet() - row_set2 = RowSet() - - row_set1.add_row_key(row_key1) - row_set1.add_row_key(row_key2) - row_set2.add_row_key(row_key2) - - assert not (row_set1 == row_set2) - - -def test_row_set__eq__len_row_ranges_differ(): - from google.cloud.bigtable.row_set import RowRange, RowSet - - row_range1 = RowRange(b"row_key4", b"row_key9") - row_range2 = RowRange(b"row_key4", b"row_key9") - - row_set1 = RowSet() - row_set2 = RowSet() - - row_set1.add_row_range(row_range1) - row_set1.add_row_range(row_range2) - row_set2.add_row_range(row_range2) - - assert not (row_set1 == row_set2) - - def test_row_set__eq__row_keys_differ(): from google.cloud.bigtable.row_set import RowSet @@ -195,36 +163,14 @@ def test_row_set_add_row_range_with_prefix(): assert row_set.row_ranges[0].end_key == b"rox" -def test_row_set__update_message_request(): - from google.cloud._helpers import _to_bytes - - from google.cloud.bigtable.row_set import RowRange, RowSet - - row_set = RowSet() - table_name = "table_name" - row_set.add_row_key("row_key1") - row_range1 = RowRange(b"row_key21", b"row_key29") - row_set.add_row_range(row_range1) - - request = _ReadRowsRequestPB(table_name=table_name) - row_set._update_message_request(request) - - expected_request = _ReadRowsRequestPB(table_name=table_name) - expected_request.rows.row_keys.append(_to_bytes("row_key1")) - - expected_request.rows.row_ranges.append(row_range1.get_range_kwargs()) - - assert request == expected_request - - def test_row_range_constructor(): from google.cloud.bigtable.row_set import RowRange start_key = "row_key1" end_key = "row_key9" row_range = RowRange(start_key, end_key) - assert start_key == row_range.start_key - assert end_key == row_range.end_key + assert start_key == row_range.start_key.decode() + assert end_key == row_range.end_key.decode() assert row_range.start_inclusive assert not row_range.end_inclusive 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 55cd5eb9a134..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 @@ -20,6 +20,8 @@ from google.api_core.exceptions import DeadlineExceeded from grpc import StatusCode +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from ._testing import _make_credentials PROJECT_ID = "project-id" @@ -48,105 +50,7 @@ RETRYABLES = (RETRYABLE_1, RETRYABLE_2, RETRYABLE_3) NON_RETRYABLE = StatusCode.CANCELLED.value[0] STATUS_INTERNAL = StatusCode.INTERNAL.value[0] - - -@mock.patch("google.cloud.bigtable.table._MAX_BULK_MUTATIONS", new=3) -def test__compile_mutation_entries_w_too_many_mutations(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import ( - TooManyMutationsError, - _compile_mutation_entries, - ) - - table = mock.Mock(name="table", spec=["name"]) - table.name = "table" - rows = [ - DirectRow(row_key=b"row_key", table=table), - DirectRow(row_key=b"row_key_2", table=table), - ] - rows[0].set_cell("cf1", b"c1", 1) - rows[0].set_cell("cf1", b"c1", 2) - rows[1].set_cell("cf1", b"c1", 3) - rows[1].set_cell("cf1", b"c1", 4) - - with pytest.raises(TooManyMutationsError): - _compile_mutation_entries("table", rows) - - -def test__compile_mutation_entries_normal(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _compile_mutation_entries - from google.cloud.bigtable_v2.types import MutateRowsRequest, data - - table = mock.Mock(spec=["name"]) - table.name = "table" - rows = [ - DirectRow(row_key=b"row_key", table=table), - DirectRow(row_key=b"row_key_2"), - ] - rows[0].set_cell("cf1", b"c1", b"1") - rows[1].set_cell("cf1", b"c1", b"2") - - result = _compile_mutation_entries("table", rows) - - entry_1 = MutateRowsRequest.Entry() - entry_1.row_key = b"row_key" - mutations_1 = data.Mutation() - mutations_1.set_cell.family_name = "cf1" - mutations_1.set_cell.column_qualifier = b"c1" - mutations_1.set_cell.timestamp_micros = -1 - mutations_1.set_cell.value = b"1" - entry_1.mutations.append(mutations_1) - - entry_2 = MutateRowsRequest.Entry() - entry_2.row_key = b"row_key_2" - mutations_2 = data.Mutation() - mutations_2.set_cell.family_name = "cf1" - mutations_2.set_cell.column_qualifier = b"c1" - mutations_2.set_cell.timestamp_micros = -1 - mutations_2.set_cell.value = b"2" - entry_2.mutations.append(mutations_2) - assert result == [entry_1, entry_2] - - -def test__check_row_table_name_w_wrong_table_name(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import TableMismatchError, _check_row_table_name - - table = mock.Mock(name="table", spec=["name"]) - table.name = "table" - row = DirectRow(row_key=b"row_key", table=table) - - with pytest.raises(TableMismatchError): - _check_row_table_name("other_table", row) - - -def test__check_row_table_name_w_right_table_name(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _check_row_table_name - - table = mock.Mock(name="table", spec=["name"]) - table.name = "table" - row = DirectRow(row_key=b"row_key", table=table) - - assert not _check_row_table_name("table", row) - - -def test__check_row_type_w_wrong_row_type(): - from google.cloud.bigtable.row import ConditionalRow - from google.cloud.bigtable.table import _check_row_type - - row = ConditionalRow(row_key=b"row_key", table="table", filter_=None) - with pytest.raises(TypeError): - _check_row_type(row) - - -def test__check_row_type_w_right_row_type(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _check_row_type - - row = DirectRow(row_key=b"row_key", table="table") - assert not _check_row_type(row) +STATUS_UNKNOWN = StatusCode.UNKNOWN.value[0] def _make_client(*args, **kwargs): @@ -683,141 +587,25 @@ def _make_gapic_api(client): return gapic_client_mock -def _table_read_row_helper(chunks, expected_result, app_profile_id=None): - from google.cloud._testing import _Monkey - - from google.cloud.bigtable import table as MUT - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS - from google.cloud.bigtable.row_filters import RowSampleFilter - from google.cloud.bigtable.row_set import RowSet - - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance, app_profile_id=app_profile_id) - - # Create request_pb - request_pb = object() # Returned by our mock. - mock_created = [] - - def mock_create_row_request(table_name, **kwargs): - mock_created.append((table_name, kwargs)) - return request_pb - - # Create response_iterator - if chunks is None: - response_iterator = iter(()) # no responses at all - else: - response_pb = _ReadRowsResponsePB(chunks=chunks) - response_iterator = iter([response_pb]) - - gapic_api = _make_gapic_api(client) - gapic_api.read_rows.return_value = response_iterator - - filter_obj = RowSampleFilter(0.33) - - with _Monkey(MUT, _create_row_request=mock_create_row_request): - result = table.read_row(ROW_KEY, filter_=filter_obj) - - row_set = RowSet() - row_set.add_row_key(ROW_KEY) - expected_request = [ - ( - table.name, - { - "end_inclusive": False, - "row_set": row_set, - "app_profile_id": app_profile_id, - "end_key": None, - "limit": None, - "start_key": None, - "filter_": filter_obj, - }, - ) - ] - assert result == expected_result - assert mock_created == expected_request - - gapic_api.read_rows.assert_called_once_with( - request_pb, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS - ) - - -def test_table_read_row_miss_no__responses(): - _table_read_row_helper(None, None) - - -def test_table_read_row_miss_no_chunks_in_response(): - chunks = [] - _table_read_row_helper(chunks, None) - - -def test_table_read_row_complete(): - from google.cloud.bigtable.row_data import Cell, PartialRowData - - app_profile_id = "app-profile-id" - chunk = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - chunks = [chunk] - expected_result = PartialRowData(row_key=ROW_KEY) - family = expected_result._cells.setdefault(FAMILY_NAME, {}) - column = family.setdefault(QUALIFIER, []) - column.append(Cell.from_pb(chunk)) - - _table_read_row_helper(chunks, expected_result, app_profile_id) - - -def test_table_read_row_more_than_one_row_returned(): - app_profile_id = "app-profile-id" - chunk_1 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - )._pb - chunk_2 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_2, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - )._pb - - chunks = [chunk_1, chunk_2] - - with pytest.raises(ValueError): - _table_read_row_helper(chunks, None, app_profile_id) - - -def test_table_read_row_still_partial(): - chunk = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - ) - chunks = [chunk] # No "commit row". - - with pytest.raises(ValueError): - _table_read_row_helper(chunks, None) - - def _table_mutate_rows_helper( - mutation_timeout=None, app_profile_id=None, retry=None, timeout=None + mutation_timeout=None, + app_profile_id=None, + retry=None, + timeout=None, + expected_operation_timeout=None, + expected_attempt_timeout=None, + expected_retryable_errors=None, ): - from google.rpc.status_pb2 import Status + from google.api_core import exceptions as api_exceptions + from google.rpc import status_pb2 - from google.cloud.bigtable.table import DEFAULT_RETRY + from google.cloud.bigtable.data.exceptions import ( + FailedMutationEntryError, + MutationsExceptionGroup, + RetryExceptionGroup, + ) + from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.table import DEFAULT_RETRY, RETRYABLE_MUTATION_ERRORS credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) @@ -830,15 +618,20 @@ def _table_mutate_rows_helper( if app_profile_id is not None: ctor_kwargs["app_profile_id"] = app_profile_id - table = _make_table(TABLE_ID, instance, **ctor_kwargs) + if expected_operation_timeout is None: + expected_operation_timeout = DEFAULT_RETRY.deadline - rows = [mock.MagicMock(), mock.MagicMock()] - response = [Status(code=0), Status(code=1)] - instance_mock = mock.Mock(return_value=response) - klass_mock = mock.patch( - "google.cloud.bigtable.table._RetryableMutateRowsWorker", - new=mock.MagicMock(return_value=instance_mock), - ) + if expected_retryable_errors is None: + expected_retryable_errors = RETRYABLE_MUTATION_ERRORS + + rows = [ + _MockRow(ROW_KEY), + _MockRow(ROW_KEY_1), + _MockRow(ROW_KEY_2), + _MockRow(ROW_KEY_3), + ] + + table = _make_table(TABLE_ID, instance, **ctor_kwargs) call_kwargs = {} @@ -846,29 +639,84 @@ def _table_mutate_rows_helper( call_kwargs["retry"] = retry if timeout is not None: - expected_timeout = call_kwargs["timeout"] = timeout - else: - expected_timeout = mutation_timeout + call_kwargs["timeout"] = timeout + + with mock.patch.object(table._table_impl, "bulk_mutate_rows") as mutate_rows_mock: + # First entry = success + # Second entry = api error + # Third entry = non-api error + # Fourth entry = retryexceptiongroup + mutate_rows_mock.side_effect = MutationsExceptionGroup( + excs=[ + FailedMutationEntryError( + failed_idx=1, + failed_mutation_entry=RowMutationEntry( + ROW_KEY_1, [mock.MagicMock()] + ), + cause=api_exceptions.InternalServerError("Failure"), + ), + FailedMutationEntryError( + failed_idx=2, + failed_mutation_entry=RowMutationEntry( + ROW_KEY_2, [mock.MagicMock()] + ), + cause=ValueError("Invalid argument"), + ), + FailedMutationEntryError( + failed_idx=3, + failed_mutation_entry=RowMutationEntry( + ROW_KEY_3, [mock.MagicMock()] + ), + cause=RetryExceptionGroup( + [ + api_exceptions.InternalServerError("First failure"), + OSError("Out of memory"), + api_exceptions.InternalServerError("Final failure"), + ] + ), + ), + ], + total_entries=4, + ) - with klass_mock: statuses = table.mutate_rows(rows, **call_kwargs) - result = [status.code for status in statuses] - expected_result = [0, 1] - assert result == expected_result - - klass_mock.new.assert_called_once_with( - client, - TABLE_NAME, - rows, - app_profile_id=app_profile_id, - timeout=expected_timeout, - ) + assert statuses == [ + status_pb2.Status( + code=SUCCESS, + message="", + ), + status_pb2.Status( + code=STATUS_INTERNAL, + message="Failure", + ), + status_pb2.Status( + code=STATUS_UNKNOWN, + message="Invalid argument", + ), + status_pb2.Status( + code=STATUS_INTERNAL, + message="Final failure", + ), + ] - if retry is not None: - instance_mock.assert_called_once_with(retry=retry) - else: - instance_mock.assert_called_once_with(retry=DEFAULT_RETRY) + # Check all call args other than mutation_entries + mutate_rows_mock.assert_called_once_with( + mock.ANY, + operation_timeout=expected_operation_timeout, + attempt_timeout=expected_attempt_timeout, + retryable_errors=expected_retryable_errors, + ) + + # Check that mutation entries are in order + mutation_entries = mutate_rows_mock.call_args.args[0] + mutation_entry_keys = [row.row_key for row in mutation_entries] + assert mutation_entry_keys == [ + ROW_KEY, + ROW_KEY_1, + ROW_KEY_2, + ROW_KEY_3, + ] def test_table_mutate_rows_w_default_mutation_timeout_app_profile_id(): @@ -876,8 +724,10 @@ def test_table_mutate_rows_w_default_mutation_timeout_app_profile_id(): def test_table_mutate_rows_w_mutation_timeout(): - mutation_timeout = 123 - _table_mutate_rows_helper(mutation_timeout=mutation_timeout) + mutation_timeout = 50 + _table_mutate_rows_helper( + mutation_timeout=mutation_timeout, expected_attempt_timeout=mutation_timeout + ) def test_table_mutate_rows_w_app_profile_id(): @@ -886,54 +736,106 @@ def test_table_mutate_rows_w_app_profile_id(): def test_table_mutate_rows_w_retry(): + deadline = 456.0 + retry = mock.Mock() + retry.deadline = deadline + _table_mutate_rows_helper(retry=retry, expected_operation_timeout=deadline) + + +def test_table_mutate_rows_w_zero_deadline_retry(): + from google.cloud.bigtable.data._helpers import TABLE_DEFAULT + + deadline = 0.0 retry = mock.Mock() - _table_mutate_rows_helper(retry=retry) + retry.deadline = deadline + _table_mutate_rows_helper( + retry=retry, + expected_operation_timeout=TABLE_DEFAULT.MUTATE_ROWS, + expected_retryable_errors=[], + ) + + +def test_table_mutate_rows_w_none_deadline_retry(): + from google.cloud.bigtable.data._helpers import TABLE_DEFAULT + + deadline = None + retry = mock.Mock() + retry.deadline = deadline + _table_mutate_rows_helper( + retry=retry, expected_operation_timeout=TABLE_DEFAULT.MUTATE_ROWS + ) def test_table_mutate_rows_w_timeout_arg(): - timeout = 123 - _table_mutate_rows_helper(timeout=timeout) + timeout = 40 + _table_mutate_rows_helper(timeout=timeout, expected_attempt_timeout=timeout) def test_table_mutate_rows_w_mutation_timeout_and_timeout_arg(): - mutation_timeout = 123 - timeout = 456 - _table_mutate_rows_helper(mutation_timeout=mutation_timeout, timeout=timeout) + mutation_timeout = 50 + timeout = 100 + _table_mutate_rows_helper( + mutation_timeout=mutation_timeout, + timeout=timeout, + expected_attempt_timeout=timeout, + ) def test_table_read_rows(): - from google.cloud._testing import _Monkey - - from google.cloud.bigtable import table as MUT + from google.cloud.bigtable.data._helpers import TABLE_DEFAULT + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row import Cell, Row + from google.cloud.bigtable.row import Cell as PartialRowDataCell + from google.cloud.bigtable.row import PartialRowData from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS, PartialRowsData + from google.cloud.bigtable.row_set import RowRange credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) - gapic_api = _make_gapic_api(client) instance = client.instance(instance_id=INSTANCE_ID) app_profile_id = "app-profile-id" table = _make_table(TABLE_ID, instance, app_profile_id=app_profile_id) - # Create request_pb - request_pb = object() # Returned by our mock. - retry = DEFAULT_RETRY_READ_ROWS - mock_created = [] - - def mock_create_row_request(table_name, **kwargs): - mock_created.append((table_name, kwargs)) - return request_pb + # Create read_rows return value + rows = [ + Row( + key=ROW_KEY, + cells=[ + Cell( + value=VALUE, + row_key=ROW_KEY, + family=FAMILY_NAME, + qualifier=QUALIFIER, + timestamp_micros=TIMESTAMP_MICROS, + ) + ], + ), + Row( + key=ROW_KEY_1, + cells=[ + Cell( + value=VALUE, + row_key=ROW_KEY_1, + family=FAMILY_NAME, + qualifier=QUALIFIER, + timestamp_micros=TIMESTAMP_MICROS, + ) + ], + ), + ] + generator = (r for r in rows) - # Create expected_result. - expected_result = PartialRowsData( - client._table_data_client._gapic_client.transport.read_rows, request_pb, retry - ) + # Create expected result. + expected_result = PartialRowsData._from_generator(generator) # Perform the method and check the result. - start_key = b"start-key" + start_key = b"begin-key" end_key = b"end-key" filter_obj = object() limit = 22 - with _Monkey(MUT, _create_row_request=mock_create_row_request): + retry = DEFAULT_RETRY_READ_ROWS + with mock.patch.object(table._table_impl, "read_rows_stream") as read_rows_mock: + read_rows_mock.return_value = generator result = table.read_rows( start_key=start_key, end_key=end_key, @@ -943,264 +845,42 @@ def mock_create_row_request(table_name, **kwargs): ) assert result.rows == expected_result.rows - assert result.retry == expected_result.retry - created_kwargs = { - "start_key": start_key, - "end_key": end_key, - "filter_": filter_obj, - "limit": limit, - "end_inclusive": False, - "app_profile_id": app_profile_id, - "row_set": None, - } - assert mock_created == [(table.name, created_kwargs)] - - gapic_api.read_rows.assert_called_once_with(request_pb, timeout=61.0, retry=retry) - - -def test_table_read_retry_rows(): - from google.api_core import retry + assert result._generator == expected_result._generator - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - gapic_api = _make_gapic_api(client) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - - retry_read_rows = retry.Retry(predicate=_read_rows_retry_exception) - - # Create response_iterator - chunk_1 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_1, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - - chunk_2 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_2, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, + expected_read_rows_query = ReadRowsQuery( + row_ranges=RowRange(start_key=start_key, end_key=end_key), + row_filter=filter_obj, + limit=limit, ) - response_1 = _ReadRowsResponseV2([chunk_1]) - response_2 = _ReadRowsResponseV2([chunk_2]) - response_failure_iterator_1 = _MockFailureIterator_1() - response_failure_iterator_2 = _MockFailureIterator_2([response_1]) - response_iterator = _MockReadRowsIterator(response_2) - - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" + read_rows_mock.assert_called_once_with( + expected_read_rows_query, + operation_timeout=TABLE_DEFAULT.READ_ROWS, + attempt_timeout=retry.deadline, + retryable_errors=TABLE_DEFAULT.READ_ROWS, ) - gapic_api.read_rows.side_effect = [ - response_failure_iterator_1, - response_failure_iterator_2, - response_iterator, - ] - - rows = [ - row - for row in table.read_rows( - start_key=ROW_KEY_1, end_key=ROW_KEY_2, retry=retry_read_rows - ) - ] - - result = rows[1] - assert result.row_key == ROW_KEY_2 - - assert len(gapic_api.read_rows.mock_calls) == 3 - - -def test_table_read_retry_rows_no_full_table_scan(): - from google.api_core import retry - - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - gapic_api = _make_gapic_api(client) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - - retry_read_rows = retry.Retry(predicate=_read_rows_retry_exception) - - # Create response_iterator - chunk_1 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_2, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - - response_1 = _ReadRowsResponseV2([chunk_1]) - response_failure_iterator_2 = _MockFailureIterator_2([response_1]) - - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - - gapic_api.read_rows.side_effect = [ - response_failure_iterator_2, - ] - - rows = [ - row - for row in table.read_rows( - start_key="doesn't matter", end_key=ROW_KEY_2, retry=retry_read_rows - ) - ] - assert len(rows) == 1 - result = rows[0] - assert result.row_key == ROW_KEY_2 - - assert len(gapic_api.read_rows.mock_calls) == 1 - assert ( - len(gapic_api.read_rows.mock_calls[0].args[0].rows.row_ranges) > 0 - ) # not empty row_ranges - - -def test_table_yield_retry_rows(): - from google.cloud.bigtable.table import _create_row_request - - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - - # Create response_iterator - chunk_1 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_1, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - - chunk_2 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_2, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - - response_1 = _ReadRowsResponseV2([chunk_1]) - response_2 = _ReadRowsResponseV2([chunk_2]) - response_failure_iterator_1 = _MockFailureIterator_1() - response_failure_iterator_2 = _MockFailureIterator_2([response_1]) - response_iterator = _MockReadRowsIterator(response_2) - - gapic_api = _make_gapic_api(client) - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - gapic_api.read_rows.side_effect = [ - response_failure_iterator_1, - response_failure_iterator_2, - response_iterator, - ] - - rows = [] - with warnings.catch_warnings(record=True) as warned: - for row in table.yield_rows(start_key=ROW_KEY_1, end_key=ROW_KEY_2): - rows.append(row) - - assert len(warned) >= 1 - assert DeprecationWarning in [w.category for w in warned] - - result = rows[1] - assert result.row_key == ROW_KEY_2 - - expected_request = _create_row_request( - table.name, - start_key=ROW_KEY_1, - end_key=ROW_KEY_2, - ) - gapic_api.read_rows.mock_calls = [expected_request] * 3 - - -def test_table_yield_rows_with_row_set(): - from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS - from google.cloud.bigtable.row_set import RowRange, RowSet - from google.cloud.bigtable.table import _create_row_request - - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - - # Create response_iterator - chunk_1 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_1, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - - chunk_2 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_2, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - - chunk_3 = _ReadRowsResponseCellChunkPB( - row_key=ROW_KEY_3, - family_name=FAMILY_NAME, - qualifier=QUALIFIER, - timestamp_micros=TIMESTAMP_MICROS, - value=VALUE, - commit_row=True, - ) - - response_1 = _ReadRowsResponseV2([chunk_1]) - response_2 = _ReadRowsResponseV2([chunk_2]) - response_3 = _ReadRowsResponseV2([chunk_3]) - response_iterator = _MockReadRowsIterator(response_1, response_2, response_3) - - gapic_api = _make_gapic_api(client) - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - gapic_api.read_rows.side_effect = [response_iterator] - - rows = [] - row_set = RowSet() - row_set.add_row_range(RowRange(start_key=ROW_KEY_1, end_key=ROW_KEY_2)) - row_set.add_row_key(ROW_KEY_3) - - with warnings.catch_warnings(record=True) as warned: - for row in table.yield_rows(row_set=row_set): - rows.append(row) - - assert len(warned) >= 1 - assert DeprecationWarning in [w.category for w in warned] - - assert rows[0].row_key == ROW_KEY_1 - assert rows[1].row_key == ROW_KEY_2 - assert rows[2].row_key == ROW_KEY_3 + # Test that the correct rows get returned. + partial_row_data = PartialRowData(ROW_KEY) + partial_row_data._cells = { + FAMILY_NAME: { + QUALIFIER: [ + PartialRowDataCell(value=VALUE, timestamp_micros=TIMESTAMP_MICROS), + ], + }, + } + partial_row_data_1 = PartialRowData(ROW_KEY_1) + partial_row_data_1._cells = { + FAMILY_NAME: { + QUALIFIER: [ + PartialRowDataCell(value=VALUE, timestamp_micros=TIMESTAMP_MICROS), + ], + }, + } + expected_row_data = [partial_row_data, partial_row_data_1] - expected_request = _create_row_request( - table.name, - start_key=ROW_KEY_1, - end_key=ROW_KEY_2, - ) - expected_request.rows.row_keys.append(ROW_KEY_3) - gapic_api.read_rows.assert_called_once_with( - expected_request, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS - ) + row_data = list(result) + assert row_data == expected_row_data def test_table_sample_row_keys(): @@ -1300,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(): @@ -1558,518 +1238,6 @@ def test_table_restore_table_w_backup_name(): _table_restore_helper(backup_name=BACKUP_NAME) -def _make_worker(*args, **kwargs): - from google.cloud.bigtable.table import _RetryableMutateRowsWorker - - return _RetryableMutateRowsWorker(*args, **kwargs) - - -def _make_responses_statuses(codes): - from google.rpc.status_pb2 import Status - - response = [Status(code=code) for code in codes] - return response - - -def _make_responses(codes): - from google.rpc.status_pb2 import Status - - from google.cloud.bigtable_v2.types.bigtable import MutateRowsResponse - - entries = [ - MutateRowsResponse.Entry(index=i, status=Status(code=codes[i])) - for i in range(len(codes)) - ] - return MutateRowsResponse(entries=entries) - - -def test_rmrw_callable_empty_rows(): - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - gapic_api = _make_gapic_api(client) - gapic_api.mutate_rows.return_value = [] - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - - worker = _make_worker(client, table.name, []) - statuses = worker() - - assert len(statuses) == 0 - - -def test_rmrw_callable_no_retry_strategy(): - from google.cloud.bigtable.row import DirectRow - - # Setup: - # - Mutate 3 rows. - # Action: - # - Attempt to mutate the rows w/o any retry strategy. - # Expectation: - # - Since no retry, should return statuses as they come back. - # - Even if there are retryable errors, no retry attempt is made. - # - State of responses_statuses should be - # [success, retryable, non-retryable] - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - - row_1 = DirectRow(row_key=b"row_key", table=table) - row_1.set_cell("cf", b"col", b"value1") - row_2 = DirectRow(row_key=b"row_key_2", table=table) - row_2.set_cell("cf", b"col", b"value2") - row_3 = DirectRow(row_key=b"row_key_3", table=table) - row_3.set_cell("cf", b"col", b"value3") - - response_codes = [SUCCESS, RETRYABLE_1, NON_RETRYABLE] - response = _make_responses(response_codes) - - gapic_api = _make_gapic_api(client) - gapic_api.mutate_rows.return_value = [response] - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - worker = _make_worker(client, table.name, [row_1, row_2, row_3]) - - statuses = worker(retry=None) - - result = [status.code for status in statuses] - assert result == response_codes - - gapic_api.mutate_rows.assert_called_once() - - -def test_rmrw_callable_retry(): - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import DEFAULT_RETRY - - # Setup: - # - Mutate 3 rows. - # Action: - # - Initial attempt will mutate all 3 rows. - # Expectation: - # - First attempt will result in one retryable error. - # - Second attempt will result in success for the retry-ed row. - # - Check MutateRows is called twice. - # - State of responses_statuses should be - # [success, success, non-retryable] - - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - row_1 = DirectRow(row_key=b"row_key", table=table) - row_1.set_cell("cf", b"col", b"value1") - row_2 = DirectRow(row_key=b"row_key_2", table=table) - row_2.set_cell("cf", b"col", b"value2") - row_3 = DirectRow(row_key=b"row_key_3", table=table) - row_3.set_cell("cf", b"col", b"value3") - - response_1 = _make_responses([SUCCESS, RETRYABLE_1, NON_RETRYABLE]) - response_2 = _make_responses([SUCCESS]) - gapic_api = _make_gapic_api(client) - gapic_api.mutate_rows.side_effect = [[response_1], [response_2]] - gapic_api.table_path.return_value = ( - f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" - ) - worker = _make_worker(client, table.name, [row_1, row_2, row_3]) - retry = DEFAULT_RETRY.with_delay(initial=0.1) - - statuses = worker(retry=retry) - - result = [status.code for status in statuses] - - assert result == [SUCCESS, SUCCESS, NON_RETRYABLE] - - assert client._table_data_client._gapic_client.mutate_rows.call_count == 2 - - -def _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=None, - expected_result=None, - raising_retry=False, - retryable_error=False, - timeout=None, - mutate_rows_side_effect=None, -): - from google.api_core.exceptions import ServiceUnavailable - - from google.cloud.bigtable.row import DirectRow - from google.cloud.bigtable.table import _BigtableRetryableError - from google.cloud.bigtable_v2.types import bigtable as data_messages_v2_pb2 - - # Setup: - # - Mutate 2 rows. - # Action: - # - Initial attempt will mutate all 2 rows. - # Expectation: - # - Expect [success, non-retryable] - - credentials = _make_credentials() - client = _make_client(project="project-id", credentials=credentials, admin=True) - instance = client.instance(instance_id=INSTANCE_ID) - table = _make_table(TABLE_ID, instance) - - rows = [] - for row_key, cell_data in row_cells: - row = DirectRow(row_key=row_key, table=table) - row.set_cell(*cell_data) - rows.append(row) - - response = _make_responses(responses) - - gapic_api = _make_gapic_api(client) - if retryable_error: - if mutate_rows_side_effect is not None: - gapic_api.mutate_rows.side_effect = mutate_rows_side_effect - else: - gapic_api.mutate_rows.side_effect = ServiceUnavailable("testing") - else: - if mutate_rows_side_effect is not None: - gapic_api.mutate_rows.side_effect = mutate_rows_side_effect - gapic_api.mutate_rows.return_value = [response] - - worker = _make_worker(client, table.name, rows=rows) - - if prior_statuses is not None: - assert len(prior_statuses) == len(rows) - worker.responses_statuses = _make_responses_statuses(prior_statuses) - - expected_entries = [] - for row, prior_status in zip(rows, worker.responses_statuses): - if prior_status is None or prior_status.code in RETRYABLES: - mutations = row._get_mutations().copy() # row clears on success - entry = data_messages_v2_pb2.MutateRowsRequest.Entry( - row_key=row.row_key, - mutations=mutations, - ) - expected_entries.append(entry) - - expected_kwargs = {} - if timeout is not None: - worker.timeout = timeout - expected_kwargs["timeout"] = mock.ANY - - if retryable_error or raising_retry: - with pytest.raises(_BigtableRetryableError): - worker._do_mutate_retryable_rows() - statuses = worker.responses_statuses - else: - statuses = worker._do_mutate_retryable_rows() - - if not retryable_error: - result = [status.code for status in statuses] - - if expected_result is None: - expected_result = responses - - assert result == expected_result - - if len(responses) == 0 and not retryable_error: - gapic_api.mutate_rows.assert_not_called() - else: - gapic_api.mutate_rows.assert_called_once_with( - table_name=table.name, - entries=expected_entries, - app_profile_id=None, - retry=None, - **expected_kwargs, - ) - if timeout is not None: - called = gapic_api.mutate_rows.mock_calls[0] - assert called.kwargs["timeout"]._deadline == timeout - - -def test_rmrw_do_mutate_retryable_rows_empty_rows(): - # - # Setup: - # - No mutated rows. - # Action: - # - No API call made. - # Expectation: - # - No change. - # - row_cells = [] - responses = [] - - _do_mutate_retryable_rows_helper(row_cells, responses) - - -def test_rmrw_do_mutate_retryable_rows_w_timeout(): - # - # Setup: - # - Mutate 2 rows. - # Action: - # - Initial attempt will mutate all 2 rows. - # Expectation: - # - No retryable error codes, so don't expect a raise. - # - State of responses_statuses should be [success, non-retryable]. - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = [SUCCESS, NON_RETRYABLE] - - timeout = 5 # seconds - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - timeout=timeout, - ) - - -def test_rmrw_do_mutate_retryable_rows_w_retryable_error(): - # - # Setup: - # - Mutate 2 rows. - # Action: - # - Initial attempt will mutate all 2 rows. - # Expectation: - # - No retryable error codes, so don't expect a raise. - # - State of responses_statuses should be [success, non-retryable]. - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = () - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - retryable_error=True, - ) - - -def test_rmrw_do_mutate_retryable_rows_w_retryable_error_internal_rst_stream_error(): - # Mutate two rows - # Raise internal server error with RST STREAM error messages - # There should be no error raised and that the request is retried - from google.api_core.exceptions import InternalServerError - - from google.cloud.bigtable.row_data import RETRYABLE_INTERNAL_ERROR_MESSAGES - - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - responses = () - - for retryable_internal_error_message in RETRYABLE_INTERNAL_ERROR_MESSAGES: - for message in [ - retryable_internal_error_message, - retryable_internal_error_message.upper(), - ]: - _do_mutate_retryable_rows_helper( - row_cells, - responses, - retryable_error=True, - mutate_rows_side_effect=InternalServerError(message), - ) - - -def test_rmrw_do_mutate_rows_w_retryable_error_internal_not_retryable(): - # Mutate two rows - # Raise internal server error but not RST STREAM error messages - # mutate_rows should raise Internal Server Error - from google.api_core.exceptions import InternalServerError - - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - responses = () - - with pytest.raises(InternalServerError): - _do_mutate_retryable_rows_helper( - row_cells, - responses, - mutate_rows_side_effect=InternalServerError("Error not retryable."), - ) - - -def test_rmrw_do_mutate_retryable_rows_retry(): - # - # Setup: - # - Mutate 3 rows. - # Action: - # - Initial attempt will mutate all 3 rows. - # Expectation: - # - Second row returns retryable error code, so expect a raise. - # - State of responses_statuses should be - # [success, retryable, non-retryable] - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - (b"row_key_3", ("cf", b"col", b"value3")), - ] - - responses = [SUCCESS, RETRYABLE_1, NON_RETRYABLE] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - raising_retry=True, - ) - - -def test_rmrw_do_mutate_retryable_rows_second_retry(): - # - # Setup: - # - Mutate 4 rows. - # - First try results: - # [success, retryable, non-retryable, retryable] - # Action: - # - Second try should re-attempt the 'retryable' rows. - # Expectation: - # - After second try: - # [success, success, non-retryable, retryable] - # - One of the rows tried second time returns retryable error code, - # so expect a raise. - # - Exception contains response whose index should be '3' even though - # only two rows were retried. - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - (b"row_key_3", ("cf", b"col", b"value3")), - (b"row_key_4", ("cf", b"col", b"value4")), - ] - - responses = [SUCCESS, RETRYABLE_1] - - prior_statuses = [ - SUCCESS, - RETRYABLE_1, - NON_RETRYABLE, - RETRYABLE_2, - ] - - expected_result = [ - SUCCESS, - SUCCESS, - NON_RETRYABLE, - RETRYABLE_1, - ] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=prior_statuses, - expected_result=expected_result, - raising_retry=True, - ) - - -def test_rmrw_do_mutate_retryable_rows_second_try(): - # - # Setup: - # - Mutate 4 rows. - # - First try results: - # [success, retryable, non-retryable, retryable] - # Action: - # - Second try should re-attempt the 'retryable' rows. - # Expectation: - # - After second try: - # [success, non-retryable, non-retryable, success] - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - (b"row_key_3", ("cf", b"col", b"value3")), - (b"row_key_4", ("cf", b"col", b"value4")), - ] - - responses = [NON_RETRYABLE, SUCCESS] - - prior_statuses = [ - SUCCESS, - RETRYABLE_1, - NON_RETRYABLE, - RETRYABLE_2, - ] - - expected_result = [ - SUCCESS, - NON_RETRYABLE, - NON_RETRYABLE, - SUCCESS, - ] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=prior_statuses, - expected_result=expected_result, - ) - - -def test_rmrw_do_mutate_retryable_rows_second_try_no_retryable(): - # - # Setup: - # - Mutate 2 rows. - # - First try results: [success, non-retryable] - # Action: - # - Second try has no row to retry. - # Expectation: - # - After second try: [success, non-retryable] - # - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = [] # no calls will be made - - prior_statuses = [ - SUCCESS, - NON_RETRYABLE, - ] - - expected_result = [ - SUCCESS, - NON_RETRYABLE, - ] - - _do_mutate_retryable_rows_helper( - row_cells, - responses, - prior_statuses=prior_statuses, - expected_result=expected_result, - ) - - -def test_rmrw_do_mutate_retryable_rows_mismatch_num_responses(): - row_cells = [ - (b"row_key_1", ("cf", b"col", b"value1")), - (b"row_key_2", ("cf", b"col", b"value2")), - ] - - responses = [SUCCESS] - - with pytest.raises(RuntimeError): - _do_mutate_retryable_rows_helper(row_cells, responses) - - -def test__create_row_request_table_name_only(): - from google.cloud.bigtable.table import _create_row_request - - table_name = "table_name" - result = _create_row_request(table_name) - expected_result = _ReadRowsRequestPB(table_name=table_name) - assert result == expected_result - - def test__create_row_request_row_range_row_set_conflict(): from google.cloud.bigtable.table import _create_row_request @@ -2081,12 +1249,9 @@ def test__create_row_request_row_range_start_key(): from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange - table_name = "table_name" - start_key = b"start_key" - result = _create_row_request(table_name, start_key=start_key) - expected_result = _ReadRowsRequestPB(table_name=table_name) - row_range = RowRange(start_key_closed=start_key) - expected_result.rows.row_ranges.append(row_range) + start_key = b"begin_key" + result = _create_row_request(start_key=start_key) + expected_result = ReadRowsQuery(row_ranges=RowRange(start_key_closed=start_key)) assert result == expected_result @@ -2094,12 +1259,9 @@ def test__create_row_request_row_range_end_key(): from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange - table_name = "table_name" - end_key = b"end_key" - result = _create_row_request(table_name, end_key=end_key) - expected_result = _ReadRowsRequestPB(table_name=table_name) - row_range = RowRange(end_key_open=end_key) - expected_result.rows.row_ranges.append(row_range) + end_key = b"begin_key" + result = _create_row_request(end_key=end_key) + expected_result = ReadRowsQuery(row_ranges=RowRange(end_key_open=end_key)) assert result == expected_result @@ -2107,13 +1269,12 @@ def test__create_row_request_row_range_both_keys(): from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange - table_name = "table_name" - start_key = b"start_key" + start_key = b"begin_key" end_key = b"end_key" - result = _create_row_request(table_name, start_key=start_key, end_key=end_key) - row_range = RowRange(start_key_closed=start_key, end_key_open=end_key) - expected_result = _ReadRowsRequestPB(table_name=table_name) - expected_result.rows.row_ranges.append(row_range) + result = _create_row_request(start_key=start_key, end_key=end_key) + expected_result = ReadRowsQuery( + row_ranges=RowRange(start_key_closed=start_key, end_key_open=end_key) + ) assert result == expected_result @@ -2121,15 +1282,14 @@ def test__create_row_request_row_range_both_keys_inclusive(): from google.cloud.bigtable.table import _create_row_request from google.cloud.bigtable_v2.types import RowRange - table_name = "table_name" - start_key = b"start_key" + start_key = b"begin_key" end_key = b"end_key" result = _create_row_request( - table_name, start_key=start_key, end_key=end_key, end_inclusive=True + start_key=start_key, end_key=end_key, end_inclusive=True + ) + expected_result = ReadRowsQuery( + row_ranges=RowRange(start_key_closed=start_key, end_key_closed=end_key) ) - expected_result = _ReadRowsRequestPB(table_name=table_name) - row_range = RowRange(start_key_closed=start_key, end_key_closed=end_key) - expected_result.rows.row_ranges.append(row_range) assert result == expected_result @@ -2137,22 +1297,18 @@ def test__create_row_request_with_filter(): from google.cloud.bigtable.row_filters import RowSampleFilter from google.cloud.bigtable.table import _create_row_request - table_name = "table_name" row_filter = RowSampleFilter(0.33) - result = _create_row_request(table_name, filter_=row_filter) - expected_result = _ReadRowsRequestPB( - table_name=table_name, filter=row_filter._to_pb() - ) + result = _create_row_request(filter_=row_filter) + expected_result = ReadRowsQuery(row_filter=row_filter) assert result == expected_result def test__create_row_request_with_limit(): from google.cloud.bigtable.table import _create_row_request - table_name = "table_name" limit = 1337 - result = _create_row_request(table_name, limit=limit) - expected_result = _ReadRowsRequestPB(table_name=table_name, rows_limit=limit) + result = _create_row_request(limit=limit) + expected_result = ReadRowsQuery(limit=limit) assert result == expected_result @@ -2160,32 +1316,12 @@ def test__create_row_request_with_row_set(): from google.cloud.bigtable.row_set import RowSet from google.cloud.bigtable.table import _create_row_request - table_name = "table_name" row_set = RowSet() - result = _create_row_request(table_name, row_set=row_set) - expected_result = _ReadRowsRequestPB(table_name=table_name) - assert result == expected_result - - -def test__create_row_request_with_app_profile_id(): - from google.cloud.bigtable.table import _create_row_request - - table_name = "table_name" - limit = 1337 - app_profile_id = "app-profile-id" - result = _create_row_request(table_name, limit=limit, app_profile_id=app_profile_id) - expected_result = _ReadRowsRequestPB( - table_name=table_name, rows_limit=limit, app_profile_id=app_profile_id - ) + result = _create_row_request(row_set=row_set) + expected_result = ReadRowsQuery() assert result == expected_result -def _ReadRowsRequestPB(*args, **kw): - from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 - - return messages_v2_pb2.ReadRowsRequest(*args, **kw) - - def test_cluster_state___eq__(): from google.cloud.bigtable.enums import Table as enum_table from google.cloud.bigtable.table import ClusterState @@ -2282,6 +1418,14 @@ def _ReadRowsResponsePB(*args, **kw): return messages_v2_pb2.ReadRowsResponse(*args, **kw) +class _MockRow(object): + def __init__(self, row_key): + self.row_key = row_key + + def _get_mutations(self): + return [mock.MagicMock()] + + class _MockReadRowsIterator(object): def __init__(self, *values): self.iter_values = iter(values)