diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 8e6833bca..a4eb93bfd 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -31,6 +31,7 @@ if TYPE_CHECKING: from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.data._metrics import ActiveOperationMetric if CrossSync.is_async: from google.cloud.bigtable_v2.services.bigtable.async_client import ( @@ -68,6 +69,8 @@ class _MutateRowsOperationAsync: operation_timeout: the timeout to use for the entire operation, in seconds. attempt_timeout: the timeout to use for each mutate_rows attempt, in seconds. If not specified, the request will run until operation_timeout is reached. + metric: the metric object representing the active operation + retryable_exceptions: a list of exceptions that should be retried """ @CrossSync.convert @@ -78,6 +81,7 @@ def __init__( mutation_entries: list["RowMutationEntry"], operation_timeout: float, attempt_timeout: float | None, + metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]] = (), ): # check that mutations are within limits @@ -97,13 +101,13 @@ def __init__( # Entry level errors bt_exceptions._MutateRowsIncomplete, ) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) self._operation = lambda: CrossSync.retry_target( self._run_attempt, self.is_retryable, - sleep_generator, + metric.backoff_generator, operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=metric.track_terminal_error(_retry_exception_factory), + on_error=metric.track_retryable_error, ) # initialize state self.timeout_generator = _attempt_timeout_generator( @@ -112,6 +116,8 @@ def __init__( self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries] self.remaining_indices = list(range(len(self.mutations))) self.errors: dict[int, list[Exception]] = {} + # set up metrics + self._operation_metric = metric @CrossSync.convert async def start(self): @@ -121,34 +127,35 @@ async def start(self): Raises: MutationsExceptionGroup: if any mutations failed """ - try: - # trigger mutate_rows - await self._operation() - except Exception as exc: - # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations - incomplete_indices = self.remaining_indices.copy() - for idx in incomplete_indices: - self._handle_entry_error(idx, exc) - finally: - # raise exception detailing incomplete mutations - all_errors: list[Exception] = [] - for idx, exc_list in self.errors.items(): - if len(exc_list) == 0: - raise core_exceptions.ClientError( - f"Mutation {idx} failed with no associated errors" + with self._operation_metric: + try: + # trigger mutate_rows + await self._operation() + except Exception as exc: + # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations + incomplete_indices = self.remaining_indices.copy() + for idx in incomplete_indices: + self._handle_entry_error(idx, exc) + finally: + # raise exception detailing incomplete mutations + all_errors: list[Exception] = [] + for idx, exc_list in self.errors.items(): + if len(exc_list) == 0: + raise core_exceptions.ClientError( + f"Mutation {idx} failed with no associated errors" + ) + elif len(exc_list) == 1: + cause_exc = exc_list[0] + else: + cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) + entry = self.mutations[idx].entry + all_errors.append( + bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) + ) + if all_errors: + raise bt_exceptions.MutationsExceptionGroup( + all_errors, len(self.mutations) ) - elif len(exc_list) == 1: - cause_exc = exc_list[0] - else: - cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) - entry = self.mutations[idx].entry - all_errors.append( - bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) - ) - if all_errors: - raise bt_exceptions.MutationsExceptionGroup( - all_errors, len(self.mutations) - ) @CrossSync.convert async def _run_attempt(self): @@ -160,6 +167,8 @@ async def _run_attempt(self): retry after the attempt is complete GoogleAPICallError: if the gapic rpc fails """ + # register attempt start + self._operation_metric.start_attempt() request_entries = [self.mutations[idx].proto for idx in self.remaining_indices] # track mutations in this request that have not been finalized yet active_request_indices = { diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 8787bfa71..35b2e44e9 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -17,6 +17,10 @@ from typing import Sequence, TYPE_CHECKING +import time + +from grpc import StatusCode + from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB from google.cloud.bigtable_v2.types import RowSet as RowSetPB @@ -31,11 +35,12 @@ from google.cloud.bigtable.data._helpers import _retry_exception_factory from google.api_core import retry as retries -from google.api_core.retry import exponential_sleep_generator from google.cloud.bigtable.data._cross_sync import CrossSync if TYPE_CHECKING: + from google.cloud.bigtable.data._metrics import ActiveOperationMetric + if CrossSync.is_async: from google.cloud.bigtable.data._async.client import ( _DataApiTargetAsync as TargetType, @@ -64,6 +69,7 @@ class _ReadRowsOperationAsync: target: The table or view to send the request to operation_timeout: The total time to allow for the operation, in seconds attempt_timeout: The time to allow for each individual attempt, in seconds + metric: the metric object representing the active operation retryable_exceptions: A list of exceptions that should trigger a retry """ @@ -75,6 +81,7 @@ class _ReadRowsOperationAsync: "_predicate", "_last_yielded_row_key", "_remaining_count", + "_operation_metric", ) def __init__( @@ -83,6 +90,7 @@ def __init__( target: TargetType, operation_timeout: float, attempt_timeout: float, + metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]] = (), ): self.attempt_timeout_gen = _attempt_timeout_generator( @@ -101,6 +109,7 @@ def __init__( self._predicate = retries.if_exception_type(*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 def start_operation(self) -> CrossSync.Iterable[Row]: """ @@ -112,9 +121,12 @@ def start_operation(self) -> CrossSync.Iterable[Row]: return CrossSync.retry_target_stream( self._read_rows_attempt, self._predicate, - exponential_sleep_generator(0.01, 60, multiplier=2), + self._operation_metric.backoff_generator, self.operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=self._operation_metric.track_terminal_error( + _retry_exception_factory + ), + on_error=self._operation_metric.track_retryable_error, ) def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: @@ -127,6 +139,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ + self._operation_metric.start_attempt() # revise request keys and ranges between attempts if self._last_yielded_row_key is not None: # if this is a retry, try to trim down the request to avoid ones we've already processed @@ -204,12 +217,11 @@ async def chunk_stream( raise InvalidChunk("emit count exceeds row limit") current_key = None - @staticmethod @CrossSync.convert( replace_symbols={"__aiter__": "__iter__", "__anext__": "__next__"}, ) async def merge_rows( - chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None, + self, chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None ) -> CrossSync.Iterable[Row]: """ Merge chunks into rows @@ -219,108 +231,125 @@ async def merge_rows( Yields: Row: the next row in the stream """ - if chunks is None: - return - it = chunks.__aiter__() - # For each row - while True: - try: - c = await it.__anext__() - except CrossSync.StopIteration: - # stream complete + try: + if chunks is None: + self._operation_metric.end_with_success() return - row_key = c.row_key - - if not row_key: - raise InvalidChunk("first row chunk is missing key") - - cells = [] - - # shared per cell storage - family: str | None = None - qualifier: bytes | None = None - - try: - # for each cell - while True: - if c.reset_row: - raise _ResetRow(c) - k = c.row_key - f = c.family_name.value - q = c.qualifier.value if c.HasField("qualifier") else None - if k and k != row_key: - raise InvalidChunk("unexpected new row key") - if f: - family = f - if q is not None: - qualifier = q - else: - raise InvalidChunk("new family without qualifier") - elif family is None: - raise InvalidChunk("missing family") - elif q is not None: - if family is None: - raise InvalidChunk("new qualifier without family") - qualifier = q - elif qualifier is None: - raise InvalidChunk("missing qualifier") - - ts = c.timestamp_micros - labels = c.labels if c.labels else [] - value = c.value - - # merge split cells - if c.value_size > 0: - buffer = [value] - while c.value_size > 0: - # throws when premature end - c = await it.__anext__() - - t = c.timestamp_micros - cl = c.labels - k = c.row_key - if ( - c.HasField("family_name") - and c.family_name.value != family - ): - raise InvalidChunk("family changed mid cell") - if ( - c.HasField("qualifier") - and c.qualifier.value != qualifier - ): - raise InvalidChunk("qualifier changed mid cell") - if t and t != ts: - raise InvalidChunk("timestamp changed mid cell") - if cl and cl != labels: - raise InvalidChunk("labels changed mid cell") - if k and k != row_key: - raise InvalidChunk("row key changed mid cell") - - if c.reset_row: - raise _ResetRow(c) - buffer.append(c.value) - value = b"".join(buffer) - cells.append( - Cell(value, row_key, family, qualifier, ts, list(labels)) - ) - if c.commit_row: - yield Row(row_key, cells) - break + it = chunks.__aiter__() + # For each row + while True: + try: c = await it.__anext__() - except _ResetRow as e: - c = e.chunk - if ( - c.row_key - or c.HasField("family_name") - or c.HasField("qualifier") - or c.timestamp_micros - or c.labels - or c.value - ): - raise InvalidChunk("reset row with data") - continue - except CrossSync.StopIteration: - raise InvalidChunk("premature end of stream") + except CrossSync.StopIteration: + # stream complete + self._operation_metric.end_with_success() + return + row_key = c.row_key + + if not row_key: + raise InvalidChunk("first row chunk is missing key") + + cells = [] + + # shared per cell storage + family: str | None = None + qualifier: bytes | None = None + + try: + # for each cell + while True: + if c.reset_row: + raise _ResetRow(c) + k = c.row_key + f = c.family_name.value + q = c.qualifier.value if c.HasField("qualifier") else None + if k and k != row_key: + raise InvalidChunk("unexpected new row key") + if f: + family = f + if q is not None: + qualifier = q + else: + raise InvalidChunk("new family without qualifier") + elif family is None: + raise InvalidChunk("missing family") + elif q is not None: + if family is None: + raise InvalidChunk("new qualifier without family") + qualifier = q + elif qualifier is None: + raise InvalidChunk("missing qualifier") + + ts = c.timestamp_micros + labels = c.labels if c.labels else [] + value = c.value + + # merge split cells + if c.value_size > 0: + buffer = [value] + while c.value_size > 0: + # throws when premature end + c = await it.__anext__() + + t = c.timestamp_micros + cl = c.labels + k = c.row_key + if ( + c.HasField("family_name") + and c.family_name.value != family + ): + raise InvalidChunk("family changed mid cell") + if ( + c.HasField("qualifier") + and c.qualifier.value != qualifier + ): + raise InvalidChunk("qualifier changed mid cell") + if t and t != ts: + raise InvalidChunk("timestamp changed mid cell") + if cl and cl != labels: + raise InvalidChunk("labels changed mid cell") + if k and k != row_key: + raise InvalidChunk("row key changed mid cell") + + if c.reset_row: + raise _ResetRow(c) + buffer.append(c.value) + value = b"".join(buffer) + cells.append( + Cell(value, row_key, family, qualifier, ts, list(labels)) + ) + if c.commit_row: + block_time = time.monotonic_ns() + yield Row(row_key, cells) + # most metric operations use setters, but this one updates + # the value directly to avoid extra overhead + if self._operation_metric.active_attempt is not None: + self._operation_metric.active_attempt.application_blocking_time_ns += ( # type: ignore + time.monotonic_ns() - block_time + ) + break + c = await it.__anext__() + except _ResetRow as e: + c = e.chunk + if ( + c.row_key + or c.HasField("family_name") + or c.HasField("qualifier") + or c.timestamp_micros + or c.labels + or c.value + ): + raise InvalidChunk("reset row with data") + continue + except CrossSync.StopIteration: + raise InvalidChunk("premature end of stream") + except GeneratorExit as close_exception: + # handle aclose() + self._operation_metric.end_with_status(StatusCode.CANCELLED) + raise close_exception + except Exception as generic_exception: + # handle exceptions in retry wrapper + raise generic_exception @staticmethod def _revise_request_rowset( diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index f8c7b287d..a40190898 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -86,6 +86,7 @@ from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter from google.cloud.bigtable.data.row_filters import RowFilterChain from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController +from google.cloud.bigtable.data._metrics import OperationType from google.cloud.bigtable.data._cross_sync import CrossSync @@ -1045,6 +1046,9 @@ async def read_rows_stream( self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=True + ), retryable_exceptions=retryable_excs, ) return row_merger.start_operation() @@ -1137,15 +1141,28 @@ async def read_row( if row_key is None: raise ValueError("row_key must be string or bytes") query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1) - results = await self.read_rows( + + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self + ) + retryable_excs = _get_retryable_errors(retryable_errors, self) + + row_merger = CrossSync._ReadRowsOperation( query, + self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, - retryable_errors=retryable_errors, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=False + ), + retryable_exceptions=retryable_excs, ) - if len(results) == 0: + results_generator = row_merger.start_operation() + try: + results = [a async for a in results_generator] + return results[0] + except IndexError: return None - return results[0] @CrossSync.convert async def read_rows_sharded( @@ -1284,20 +1301,17 @@ async def row_exists( from any retries that failed google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error """ - if row_key is None: - raise ValueError("row_key must be string or bytes") - strip_filter = StripValueTransformerFilter(flag=True) limit_filter = CellsRowLimitFilter(1) chain_filter = RowFilterChain(filters=[limit_filter, strip_filter]) - query = ReadRowsQuery(row_keys=row_key, limit=1, row_filter=chain_filter) - results = await self.read_rows( - query, + result = await self.read_row( + row_key=row_key, + row_filter=chain_filter, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, retryable_errors=retryable_errors, ) - return len(results) > 0 + return result is not None @CrossSync.convert async def sample_row_keys( @@ -1349,26 +1363,31 @@ async def sample_row_keys( retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - - @CrossSync.convert - async def execute_rpc(): - results = await self.client._gapic_client.sample_row_keys( - request=SampleRowKeysRequest( - app_profile_id=self.app_profile_id, **self._request_path + with self._metrics.create_operation( + OperationType.SAMPLE_ROW_KEYS + ) as operation_metric: + + @CrossSync.convert + async def execute_rpc(): + results = await self.client._gapic_client.sample_row_keys( + request=SampleRowKeysRequest( + app_profile_id=self.app_profile_id, **self._request_path + ), + timeout=next(attempt_timeout_gen), + retry=None, + ) + return [(s.row_key, s.offset_bytes) async for s in results] + + return await CrossSync.retry_target( + execute_rpc, + predicate, + operation_metric.backoff_generator, + operation_timeout, + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory ), - timeout=next(attempt_timeout_gen), - retry=None, + on_error=operation_metric.track_retryable_error, ) - return [(s.row_key, s.offset_bytes) async for s in results] - - return await CrossSync.retry_target( - execute_rpc, - predicate, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, - ) @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"}) def mutations_batcher( @@ -1479,28 +1498,32 @@ async def mutate_row( # mutations should not be retried predicate = retries.if_exception_type() - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - - target = partial( - self.client._gapic_client.mutate_row, - request=MutateRowRequest( - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - mutations=[mutation._to_pb() for mutation in mutations_list], - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=attempt_timeout, - retry=None, - ) - return await CrossSync.retry_target( - target, - predicate, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, - ) + with self._metrics.create_operation( + OperationType.MUTATE_ROW + ) as operation_metric: + target = partial( + self.client._gapic_client.mutate_row, + request=MutateRowRequest( + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + mutations=[mutation._to_pb() for mutation in mutations_list], + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=attempt_timeout, + retry=None, + ) + return await CrossSync.retry_target( + target, + predicate, + operation_metric.backoff_generator, + operation_timeout, + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory + ), + on_error=operation_metric.track_retryable_error, + ) @CrossSync.convert async def bulk_mutate_rows( @@ -1554,6 +1577,7 @@ async def bulk_mutate_rows( mutation_entries, operation_timeout, attempt_timeout, + metric=self._metrics.create_operation(OperationType.BULK_MUTATE_ROWS), retryable_exceptions=retryable_excs, ) await operation.start() @@ -1611,21 +1635,25 @@ async def check_and_mutate_row( ): false_case_mutations = [false_case_mutations] false_case_list = [m._to_pb() for m in false_case_mutations or []] - result = await self.client._gapic_client.check_and_mutate_row( - request=CheckAndMutateRowRequest( - true_mutations=true_case_list, - false_mutations=false_case_list, - predicate_filter=predicate._to_pb() if predicate is not None else None, - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=operation_timeout, - retry=None, - ) - return result.predicate_matched + + with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE): + result = await self.client._gapic_client.check_and_mutate_row( + request=CheckAndMutateRowRequest( + true_mutations=true_case_list, + false_mutations=false_case_list, + predicate_filter=predicate._to_pb() + if predicate is not None + else None, + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=operation_timeout, + retry=None, + ) + return result.predicate_matched @CrossSync.convert async def read_modify_write_row( @@ -1665,20 +1693,22 @@ async def read_modify_write_row( rules = [rules] if not rules: raise ValueError("rules must contain at least one item") - result = await self.client._gapic_client.read_modify_write_row( - request=ReadModifyWriteRowRequest( - rules=[rule._to_pb() for rule in rules], - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=operation_timeout, - retry=None, - ) - # construct Row from result - return Row._from_pb(result.row) + + with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE): + result = await self.client._gapic_client.read_modify_write_row( + request=ReadModifyWriteRowRequest( + rules=[rule._to_pb() for rule in rules], + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=operation_timeout, + retry=None, + ) + # construct Row from result + return Row._from_pb(result.row) @CrossSync.convert async def close(self): diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 0bd401a78..dad9ee602 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -92,10 +92,6 @@ async def _get_metadata(source) -> dict[str, str | bytes] | None: class AsyncBigtableMetricsInterceptor( UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor, MetricsHandler ): - """ - An async gRPC interceptor to add client metadata and print server metadata. - """ - @CrossSync.convert @_with_operation_from_metadata async def intercept_unary_unary( diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index a8e99ea9e..8b8571b27 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -16,6 +16,7 @@ from typing import Sequence, TYPE_CHECKING, cast import atexit +import time import warnings from collections import deque import concurrent.futures @@ -25,6 +26,8 @@ from google.cloud.bigtable.data._helpers import _get_retryable_errors from google.cloud.bigtable.data._helpers import _get_timeouts from google.cloud.bigtable.data._helpers import TABLE_DEFAULT +from google.cloud.bigtable.data._metrics import OperationType +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data.mutations import ( _MUTATE_ROWS_REQUEST_MUTATION_LIMIT, @@ -35,6 +38,7 @@ if TYPE_CHECKING: from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController if CrossSync.is_async: from google.cloud.bigtable.data._async.client import ( @@ -177,6 +181,24 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry] ) yield mutations[start_idx:end_idx] + @CrossSync.convert(replace_symbols={"__anext__": "__next__"}) + async def add_to_flow_with_metrics( + self, + mutations: RowMutationEntry | list[RowMutationEntry], + metrics_controller: BigtableClientSideMetricsController, + ): + inner_generator = self.add_to_flow(mutations) + while True: + # start a new metric + metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS) + flow_start_time = time.monotonic_ns() + try: + value = await inner_generator.__anext__() + except CrossSync.StopIteration: + return + metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time + yield value, metric + @CrossSync.convert_class(sync_name="MutationsBatcher") class MutationsBatcherAsync: @@ -353,9 +375,14 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]): """ # flush new entries in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = [] - async for batch in self._flow_control.add_to_flow(new_entries): + async for batch, metric in self._flow_control.add_to_flow_with_metrics( + new_entries, self._target._metrics + ): batch_task = CrossSync.create_task( - self._execute_mutate_rows, batch, sync_executor=self._sync_rpc_executor + self._execute_mutate_rows, + batch, + metric, + sync_executor=self._sync_rpc_executor, ) in_process_requests.append(batch_task) # wait for all inflight requests to complete @@ -366,7 +393,7 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]): @CrossSync.convert async def _execute_mutate_rows( - self, batch: list[RowMutationEntry] + self, batch: list[RowMutationEntry], metric: ActiveOperationMetric ) -> list[FailedMutationEntryError]: """ Helper to execute mutation operation on a batch @@ -387,6 +414,7 @@ async def _execute_mutate_rows( batch, operation_timeout=self._operation_timeout, attempt_timeout=self._attempt_timeout, + metric=metric, retryable_exceptions=self._retryable_errors, ) await operation.start() diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index e848ebc6f..13bcfcc29 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -103,6 +103,7 @@ def _retry_exception_factory( tuple[Exception, Exception|None]: tuple of the exception to raise, and a cause exception if applicable """ + exc_list = exc_list.copy() if reason == RetryFailureReason.TIMEOUT: timeout_val_str = f"of {timeout_val:0.1f}s " if timeout_val is not None else "" # if failed due to timeout, raise deadline exceeded as primary exception diff --git a/google/cloud/bigtable/data/_metrics/handlers/_stdout.py b/google/cloud/bigtable/data/_metrics/handlers/_stdout.py new file mode 100644 index 000000000..e0a4bc8da --- /dev/null +++ b/google/cloud/bigtable/data/_metrics/handlers/_stdout.py @@ -0,0 +1,48 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric + + +class _StdoutMetricsHandler(MetricsHandler): + """ + Prints a table of metric data after each operation, for debugging purposes. + """ + + def __init__(self, **kwargs): + self._completed_ops = {} + + def on_operation_complete(self, op: CompletedOperationMetric) -> None: + """ + After each operation, update the state and print the metrics table. + """ + current_list = self._completed_ops.setdefault(op.op_type, []) + current_list.append(op) + self.print() + + def print(self): + """ + Print the current state of the metrics table. + """ + print("Bigtable Metrics:") + for ops_type, ops_list in self._completed_ops.items(): + count = len(ops_list) + total_latency = sum([op.duration_ns // 1_000_000 for op in ops_list]) + total_attempts = sum([len(op.completed_attempts) for op in ops_list]) + avg_latency = total_latency / count + avg_attempts = total_attempts / count + print( + f"{ops_type}: count: {count}, avg latency: {avg_latency:.2f} ms, avg attempts: {avg_attempts:.1f}" + ) + print() diff --git a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index 3bf7b562f..81a5be4cf 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -29,6 +29,7 @@ if TYPE_CHECKING: from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable_v2.services.bigtable.client import ( BigtableClient as GapicClientType, ) @@ -54,6 +55,8 @@ class _MutateRowsOperation: operation_timeout: the timeout to use for the entire operation, in seconds. attempt_timeout: the timeout to use for each mutate_rows attempt, in seconds. If not specified, the request will run until operation_timeout is reached. + metric: the metric object representing the active operation + retryable_exceptions: a list of exceptions that should be retried """ def __init__( @@ -63,6 +66,7 @@ def __init__( mutation_entries: list["RowMutationEntry"], operation_timeout: float, attempt_timeout: float | None, + metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]] = (), ): total_mutations = sum((len(entry.mutations) for entry in mutation_entries)) @@ -75,13 +79,13 @@ def __init__( self.is_retryable = retries.if_exception_type( *retryable_exceptions, bt_exceptions._MutateRowsIncomplete ) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) self._operation = lambda: CrossSync._Sync_Impl.retry_target( self._run_attempt, self.is_retryable, - sleep_generator, + metric.backoff_generator, operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=metric.track_terminal_error(_retry_exception_factory), + on_error=metric.track_retryable_error, ) self.timeout_generator = _attempt_timeout_generator( attempt_timeout, operation_timeout @@ -89,37 +93,39 @@ def __init__( self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries] self.remaining_indices = list(range(len(self.mutations))) self.errors: dict[int, list[Exception]] = {} + self._operation_metric = metric def start(self): """Start the operation, and run until completion Raises: MutationsExceptionGroup: if any mutations failed""" - try: - self._operation() - except Exception as exc: - incomplete_indices = self.remaining_indices.copy() - for idx in incomplete_indices: - self._handle_entry_error(idx, exc) - finally: - all_errors: list[Exception] = [] - for idx, exc_list in self.errors.items(): - if len(exc_list) == 0: - raise core_exceptions.ClientError( - f"Mutation {idx} failed with no associated errors" + with self._operation_metric: + try: + self._operation() + except Exception as exc: + incomplete_indices = self.remaining_indices.copy() + for idx in incomplete_indices: + self._handle_entry_error(idx, exc) + finally: + all_errors: list[Exception] = [] + for idx, exc_list in self.errors.items(): + if len(exc_list) == 0: + raise core_exceptions.ClientError( + f"Mutation {idx} failed with no associated errors" + ) + elif len(exc_list) == 1: + cause_exc = exc_list[0] + else: + cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) + entry = self.mutations[idx].entry + all_errors.append( + bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) + ) + if all_errors: + raise bt_exceptions.MutationsExceptionGroup( + all_errors, len(self.mutations) ) - elif len(exc_list) == 1: - cause_exc = exc_list[0] - else: - cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) - entry = self.mutations[idx].entry - all_errors.append( - bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) - ) - if all_errors: - raise bt_exceptions.MutationsExceptionGroup( - all_errors, len(self.mutations) - ) def _run_attempt(self): """Run a single attempt of the mutate_rows rpc. @@ -128,6 +134,7 @@ def _run_attempt(self): _MutateRowsIncomplete: if there are failed mutations eligible for retry after the attempt is complete GoogleAPICallError: if the gapic rpc fails""" + self._operation_metric.start_attempt() request_entries = [self.mutations[idx].proto for idx in self.remaining_indices] active_request_indices = { req_idx: orig_idx diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index 3593475a9..adbe819eb 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -18,6 +18,8 @@ from __future__ import annotations from typing import Sequence, TYPE_CHECKING +import time +from grpc import StatusCode from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB from google.cloud.bigtable_v2.types import RowSet as RowSetPB @@ -30,10 +32,10 @@ from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory from google.api_core import retry as retries -from google.api_core.retry import exponential_sleep_generator from google.cloud.bigtable.data._cross_sync import CrossSync if TYPE_CHECKING: + from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data._sync_autogen.client import ( _DataApiTarget as TargetType, ) @@ -56,6 +58,7 @@ class _ReadRowsOperation: target: The table or view to send the request to operation_timeout: The total time to allow for the operation, in seconds attempt_timeout: The time to allow for each individual attempt, in seconds + metric: the metric object representing the active operation retryable_exceptions: A list of exceptions that should trigger a retry """ @@ -67,6 +70,7 @@ class _ReadRowsOperation: "_predicate", "_last_yielded_row_key", "_remaining_count", + "_operation_metric", ) def __init__( @@ -75,6 +79,7 @@ def __init__( target: TargetType, operation_timeout: float, attempt_timeout: float, + metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]] = (), ): self.attempt_timeout_gen = _attempt_timeout_generator( @@ -91,6 +96,7 @@ def __init__( self._predicate = retries.if_exception_type(*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 def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]: """Start the read_rows operation, retrying on retryable errors. @@ -100,9 +106,12 @@ def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]: return CrossSync._Sync_Impl.retry_target_stream( self._read_rows_attempt, self._predicate, - exponential_sleep_generator(0.01, 60, multiplier=2), + self._operation_metric.backoff_generator, self.operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=self._operation_metric.track_terminal_error( + _retry_exception_factory + ), + on_error=self._operation_metric.track_retryable_error, ) def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: @@ -113,6 +122,7 @@ def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: Yields: Row: The next row in the stream""" + self._operation_metric.start_attempt() if self._last_yielded_row_key is not None: try: self.request.rows = self._revise_request_rowset( @@ -174,9 +184,8 @@ def chunk_stream( raise InvalidChunk("emit count exceeds row limit") current_key = None - @staticmethod def merge_rows( - chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None, + self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None ) -> CrossSync._Sync_Impl.Iterable[Row]: """Merge chunks into rows @@ -184,94 +193,107 @@ def merge_rows( chunks: the chunk stream to merge Yields: Row: the next row in the stream""" - if chunks is None: - return - it = chunks.__iter__() - while True: - try: - c = it.__next__() - except CrossSync._Sync_Impl.StopIteration: + try: + if chunks is None: + self._operation_metric.end_with_success() return - row_key = c.row_key - if not row_key: - raise InvalidChunk("first row chunk is missing key") - cells = [] - family: str | None = None - qualifier: bytes | None = None - try: - while True: - if c.reset_row: - raise _ResetRow(c) - k = c.row_key - f = c.family_name.value - q = c.qualifier.value if c.HasField("qualifier") else None - if k and k != row_key: - raise InvalidChunk("unexpected new row key") - if f: - family = f - if q is not None: - qualifier = q - else: - raise InvalidChunk("new family without qualifier") - elif family is None: - raise InvalidChunk("missing family") - elif q is not None: - if family is None: - raise InvalidChunk("new qualifier without family") - qualifier = q - elif qualifier is None: - raise InvalidChunk("missing qualifier") - ts = c.timestamp_micros - labels = c.labels if c.labels else [] - value = c.value - if c.value_size > 0: - buffer = [value] - while c.value_size > 0: - c = it.__next__() - t = c.timestamp_micros - cl = c.labels - k = c.row_key - if ( - c.HasField("family_name") - and c.family_name.value != family - ): - raise InvalidChunk("family changed mid cell") - if ( - c.HasField("qualifier") - and c.qualifier.value != qualifier - ): - raise InvalidChunk("qualifier changed mid cell") - if t and t != ts: - raise InvalidChunk("timestamp changed mid cell") - if cl and cl != labels: - raise InvalidChunk("labels changed mid cell") - if k and k != row_key: - raise InvalidChunk("row key changed mid cell") - if c.reset_row: - raise _ResetRow(c) - buffer.append(c.value) - value = b"".join(buffer) - cells.append( - Cell(value, row_key, family, qualifier, ts, list(labels)) - ) - if c.commit_row: - yield Row(row_key, cells) - break + it = chunks.__iter__() + while True: + try: c = it.__next__() - except _ResetRow as e: - c = e.chunk - if ( - c.row_key - or c.HasField("family_name") - or c.HasField("qualifier") - or c.timestamp_micros - or c.labels - or c.value - ): - raise InvalidChunk("reset row with data") - continue - except CrossSync._Sync_Impl.StopIteration: - raise InvalidChunk("premature end of stream") + except CrossSync._Sync_Impl.StopIteration: + self._operation_metric.end_with_success() + return + row_key = c.row_key + if not row_key: + raise InvalidChunk("first row chunk is missing key") + cells = [] + family: str | None = None + qualifier: bytes | None = None + try: + while True: + if c.reset_row: + raise _ResetRow(c) + k = c.row_key + f = c.family_name.value + q = c.qualifier.value if c.HasField("qualifier") else None + if k and k != row_key: + raise InvalidChunk("unexpected new row key") + if f: + family = f + if q is not None: + qualifier = q + else: + raise InvalidChunk("new family without qualifier") + elif family is None: + raise InvalidChunk("missing family") + elif q is not None: + if family is None: + raise InvalidChunk("new qualifier without family") + qualifier = q + elif qualifier is None: + raise InvalidChunk("missing qualifier") + ts = c.timestamp_micros + labels = c.labels if c.labels else [] + value = c.value + if c.value_size > 0: + buffer = [value] + while c.value_size > 0: + c = it.__next__() + t = c.timestamp_micros + cl = c.labels + k = c.row_key + if ( + c.HasField("family_name") + and c.family_name.value != family + ): + raise InvalidChunk("family changed mid cell") + if ( + c.HasField("qualifier") + and c.qualifier.value != qualifier + ): + raise InvalidChunk("qualifier changed mid cell") + if t and t != ts: + raise InvalidChunk("timestamp changed mid cell") + if cl and cl != labels: + raise InvalidChunk("labels changed mid cell") + if k and k != row_key: + raise InvalidChunk("row key changed mid cell") + if c.reset_row: + raise _ResetRow(c) + buffer.append(c.value) + value = b"".join(buffer) + cells.append( + Cell(value, row_key, family, qualifier, ts, list(labels)) + ) + if c.commit_row: + block_time = time.monotonic_ns() + yield Row(row_key, cells) + if self._operation_metric.active_attempt is not None: + self._operation_metric.active_attempt.application_blocking_time_ns += ( + time.monotonic_ns() - block_time + ) + break + c = it.__next__() + except _ResetRow as e: + c = e.chunk + if ( + c.row_key + or c.HasField("family_name") + or c.HasField("qualifier") + or c.timestamp_micros + or c.labels + or c.value + ): + raise InvalidChunk("reset row with data") + continue + except CrossSync._Sync_Impl.StopIteration: + raise InvalidChunk("premature end of stream") + except GeneratorExit as close_exception: + self._operation_metric.end_with_status(StatusCode.CANCELLED) + raise close_exception + except Exception as generic_exception: + raise generic_exception @staticmethod def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSetPB: diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index 86993cf56..b1c3cec47 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -73,6 +73,7 @@ from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter from google.cloud.bigtable.data.row_filters import RowFilterChain from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController +from google.cloud.bigtable.data._metrics import OperationType from google.cloud.bigtable.data._cross_sync import CrossSync from typing import Iterable from grpc import insecure_channel @@ -826,6 +827,9 @@ def read_rows_stream( self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=True + ), retryable_exceptions=retryable_excs, ) return row_merger.start_operation() @@ -914,15 +918,26 @@ def read_row( if row_key is None: raise ValueError("row_key must be string or bytes") query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1) - results = self.read_rows( + (operation_timeout, attempt_timeout) = _get_timeouts( + operation_timeout, attempt_timeout, self + ) + retryable_excs = _get_retryable_errors(retryable_errors, self) + row_merger = CrossSync._Sync_Impl._ReadRowsOperation( query, + self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, - retryable_errors=retryable_errors, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=False + ), + retryable_exceptions=retryable_excs, ) - if len(results) == 0: + results_generator = row_merger.start_operation() + try: + results = [a for a in results_generator] + return results[0] + except IndexError: return None - return results[0] def read_rows_sharded( self, @@ -1045,19 +1060,17 @@ def row_exists( from any retries that failed google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error """ - if row_key is None: - raise ValueError("row_key must be string or bytes") strip_filter = StripValueTransformerFilter(flag=True) limit_filter = CellsRowLimitFilter(1) chain_filter = RowFilterChain(filters=[limit_filter, strip_filter]) - query = ReadRowsQuery(row_keys=row_key, limit=1, row_filter=chain_filter) - results = self.read_rows( - query, + result = self.read_row( + row_key=row_key, + row_filter=chain_filter, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, retryable_errors=retryable_errors, ) - return len(results) > 0 + return result is not None def sample_row_keys( self, @@ -1104,25 +1117,30 @@ def sample_row_keys( ) retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - - def execute_rpc(): - results = self.client._gapic_client.sample_row_keys( - request=SampleRowKeysRequest( - app_profile_id=self.app_profile_id, **self._request_path + with self._metrics.create_operation( + OperationType.SAMPLE_ROW_KEYS + ) as operation_metric: + + def execute_rpc(): + results = self.client._gapic_client.sample_row_keys( + request=SampleRowKeysRequest( + app_profile_id=self.app_profile_id, **self._request_path + ), + timeout=next(attempt_timeout_gen), + retry=None, + ) + return [(s.row_key, s.offset_bytes) for s in results] + + return CrossSync._Sync_Impl.retry_target( + execute_rpc, + predicate, + operation_metric.backoff_generator, + operation_timeout, + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory ), - timeout=next(attempt_timeout_gen), - retry=None, + on_error=operation_metric.track_retryable_error, ) - return [(s.row_key, s.offset_bytes) for s in results] - - return CrossSync._Sync_Impl.retry_target( - execute_rpc, - predicate, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, - ) def mutations_batcher( self, @@ -1223,27 +1241,32 @@ def mutate_row( ) else: predicate = retries.if_exception_type() - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - target = partial( - self.client._gapic_client.mutate_row, - request=MutateRowRequest( - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - mutations=[mutation._to_pb() for mutation in mutations_list], - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=attempt_timeout, - retry=None, - ) - return CrossSync._Sync_Impl.retry_target( - target, - predicate, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, - ) + with self._metrics.create_operation( + OperationType.MUTATE_ROW + ) as operation_metric: + target = partial( + self.client._gapic_client.mutate_row, + request=MutateRowRequest( + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + mutations=[mutation._to_pb() for mutation in mutations_list], + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=attempt_timeout, + retry=None, + ) + return CrossSync._Sync_Impl.retry_target( + target, + predicate, + operation_metric.backoff_generator, + operation_timeout, + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory + ), + on_error=operation_metric.track_retryable_error, + ) def bulk_mutate_rows( self, @@ -1293,6 +1316,7 @@ def bulk_mutate_rows( mutation_entries, operation_timeout, attempt_timeout, + metric=self._metrics.create_operation(OperationType.BULK_MUTATE_ROWS), retryable_exceptions=retryable_excs, ) operation.start() @@ -1347,21 +1371,24 @@ def check_and_mutate_row( ): false_case_mutations = [false_case_mutations] false_case_list = [m._to_pb() for m in false_case_mutations or []] - result = self.client._gapic_client.check_and_mutate_row( - request=CheckAndMutateRowRequest( - true_mutations=true_case_list, - false_mutations=false_case_list, - predicate_filter=predicate._to_pb() if predicate is not None else None, - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=operation_timeout, - retry=None, - ) - return result.predicate_matched + with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE): + result = self.client._gapic_client.check_and_mutate_row( + request=CheckAndMutateRowRequest( + true_mutations=true_case_list, + false_mutations=false_case_list, + predicate_filter=predicate._to_pb() + if predicate is not None + else None, + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=operation_timeout, + retry=None, + ) + return result.predicate_matched def read_modify_write_row( self, @@ -1398,19 +1425,20 @@ def read_modify_write_row( rules = [rules] if not rules: raise ValueError("rules must contain at least one item") - result = self.client._gapic_client.read_modify_write_row( - request=ReadModifyWriteRowRequest( - rules=[rule._to_pb() for rule in rules], - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=operation_timeout, - retry=None, - ) - return Row._from_pb(result.row) + with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE): + result = self.client._gapic_client.read_modify_write_row( + request=ReadModifyWriteRowRequest( + rules=[rule._to_pb() for rule in rules], + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=operation_timeout, + retry=None, + ) + return Row._from_pb(result.row) def close(self): """Called to close the Table instance and release any resources held by it.""" diff --git a/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py b/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py index 84f0ba8c0..9f47eb39a 100644 --- a/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py +++ b/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py @@ -18,6 +18,7 @@ from __future__ import annotations from typing import Sequence, TYPE_CHECKING, cast import atexit +import time import warnings from collections import deque import concurrent.futures @@ -26,12 +27,15 @@ from google.cloud.bigtable.data._helpers import _get_retryable_errors from google.cloud.bigtable.data._helpers import _get_timeouts from google.cloud.bigtable.data._helpers import TABLE_DEFAULT +from google.cloud.bigtable.data._metrics import OperationType +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT from google.cloud.bigtable.data.mutations import Mutation from google.cloud.bigtable.data._cross_sync import CrossSync if TYPE_CHECKING: from google.cloud.bigtable.data.mutations import RowMutationEntry + from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController from google.cloud.bigtable.data._sync_autogen.client import ( _DataApiTarget as TargetType, ) @@ -147,6 +151,22 @@ def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]): ) yield mutations[start_idx:end_idx] + def add_to_flow_with_metrics( + self, + mutations: RowMutationEntry | list[RowMutationEntry], + metrics_controller: BigtableClientSideMetricsController, + ): + inner_generator = self.add_to_flow(mutations) + while True: + metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS) + flow_start_time = time.monotonic_ns() + try: + value = inner_generator.__next__() + except CrossSync._Sync_Impl.StopIteration: + return + metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time + yield (value, metric) + class MutationsBatcher: """ @@ -302,9 +322,14 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]): in_process_requests: list[ CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]] ] = [] - for batch in self._flow_control.add_to_flow(new_entries): + for batch, metric in self._flow_control.add_to_flow_with_metrics( + new_entries, self._target._metrics + ): batch_task = CrossSync._Sync_Impl.create_task( - self._execute_mutate_rows, batch, sync_executor=self._sync_rpc_executor + self._execute_mutate_rows, + batch, + metric, + sync_executor=self._sync_rpc_executor, ) in_process_requests.append(batch_task) found_exceptions = self._wait_for_batch_results(*in_process_requests) @@ -312,7 +337,7 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]): self._add_exceptions(found_exceptions) def _execute_mutate_rows( - self, batch: list[RowMutationEntry] + self, batch: list[RowMutationEntry], metric: ActiveOperationMetric ) -> list[FailedMutationEntryError]: """Helper to execute mutation operation on a batch @@ -331,6 +356,7 @@ def _execute_mutate_rows( batch, operation_timeout=self._operation_timeout, attempt_timeout=self._attempt_timeout, + metric=metric, retryable_exceptions=self._retryable_errors, ) operation.start() diff --git a/tests/system/conftest.py b/tests/system/conftest.py index 39480942d..f7eccbe00 100644 --- a/tests/system/conftest.py +++ b/tests/system/conftest.py @@ -23,10 +23,6 @@ script_path = os.path.dirname(os.path.realpath(__file__)) sys.path.append(script_path) -pytest_plugins = [ - "data.setup_fixtures", -] - @pytest.fixture(scope="session") def event_loop(): diff --git a/tests/system/data/__init__.py b/tests/system/data/__init__.py index 2b35cea8f..6f836fb96 100644 --- a/tests/system/data/__init__.py +++ b/tests/system/data/__init__.py @@ -13,7 +13,242 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import pytest +import os +import uuid TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" TEST_AGGREGATE_FAMILY = "test-aggregate-family" + +# authorized view subset to allow all qualifiers +ALLOW_ALL = "" +ALL_QUALIFIERS = {"qualifier_prefixes": [ALLOW_ALL]} + + +class SystemTestRunner: + """ + configures a system test class with configuration for clusters/tables/etc + + used by standard system tests, and metrics tests + """ + + @pytest.fixture(scope="session") + def init_table_id(self): + """ + The table_id to use when creating a new test table + """ + return f"test-table-{uuid.uuid4().hex}" + + @pytest.fixture(scope="session") + def cluster_config(self, project_id): + """ + Configuration for the clusters to use when creating a new instance + """ + from google.cloud.bigtable_admin_v2 import types + + cluster = { + "test-cluster": types.Cluster( + location=f"projects/{project_id}/locations/us-central1-b", + serve_nodes=1, + ) + } + return cluster + + @pytest.fixture(scope="session") + def column_family_config(self): + """ + specify column families to create when creating a new test table + """ + from google.cloud.bigtable_admin_v2 import types + + int_aggregate_type = types.Type.Aggregate( + input_type=types.Type(int64_type={"encoding": {"big_endian_bytes": {}}}), + sum={}, + ) + return { + TEST_FAMILY: types.ColumnFamily(), + TEST_FAMILY_2: types.ColumnFamily(), + TEST_AGGREGATE_FAMILY: types.ColumnFamily( + value_type=types.Type(aggregate_type=int_aggregate_type) + ), + } + + @pytest.fixture(scope="session") + def admin_client(self): + """ + Client for interacting with Table and Instance admin APIs + """ + from google.cloud.bigtable.client import Client + + client = Client(admin=True) + yield client + + @pytest.fixture(scope="session") + def instance_id(self, admin_client, project_id, cluster_config): + """ + Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session + """ + from google.cloud.bigtable_admin_v2 import types + from google.api_core import exceptions + from google.cloud.environment_vars import BIGTABLE_EMULATOR + + # use user-specified instance if available + user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") + if user_specified_instance: + print("Using user-specified instance: {}".format(user_specified_instance)) + yield user_specified_instance + return + + # create a new temporary test instance + instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" + if os.getenv(BIGTABLE_EMULATOR): + # don't create instance if in emulator mode + yield instance_id + else: + try: + operation = admin_client.instance_admin_client.create_instance( + parent=f"projects/{project_id}", + instance_id=instance_id, + instance=types.Instance( + display_name="Test Instance", + # labels={"python-system-test": "true"}, + ), + clusters=cluster_config, + ) + operation.result(timeout=240) + except exceptions.AlreadyExists: + pass + yield instance_id + admin_client.instance_admin_client.delete_instance( + name=f"projects/{project_id}/instances/{instance_id}" + ) + + @pytest.fixture(scope="session") + def column_split_config(self): + """ + specify initial splits to create when creating a new test table + """ + return [(num * 1000).to_bytes(8, "big") for num in range(1, 10)] + + @pytest.fixture(scope="session") + def table_id( + self, + admin_client, + project_id, + instance_id, + column_family_config, + init_table_id, + column_split_config, + ): + """ + Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session + + Args: + - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture. + - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. + - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. + - init_column_families: A list of column families to initialize the table with, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. + Supplied by the init_column_families fixture. + - init_table_id: The table ID to give to the test table, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. + Supplied by the init_table_id fixture. + - column_split_config: A list of row keys to use as initial splits when creating the test table. + """ + from google.api_core import exceptions + from google.api_core import retry + + # use user-specified instance if available + user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") + if user_specified_table: + print("Using user-specified table: {}".format(user_specified_table)) + yield user_specified_table + return + + retry = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + try: + parent_path = f"projects/{project_id}/instances/{instance_id}" + print(f"Creating table: {parent_path}/tables/{init_table_id}") + admin_client.table_admin_client.create_table( + request={ + "parent": parent_path, + "table_id": init_table_id, + "table": {"column_families": column_family_config}, + "initial_splits": [{"key": key} for key in column_split_config], + }, + retry=retry, + ) + except exceptions.AlreadyExists: + pass + yield init_table_id + print(f"Deleting table: {parent_path}/tables/{init_table_id}") + try: + admin_client.table_admin_client.delete_table( + name=f"{parent_path}/tables/{init_table_id}" + ) + except exceptions.NotFound: + print(f"Table {init_table_id} not found, skipping deletion") + + @pytest.fixture(scope="session") + def authorized_view_id( + self, + admin_client, + project_id, + instance_id, + table_id, + ): + """ + Creates and returns a new temporary authorized view for the test session + + Args: + - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture. + - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. + - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. + - table_id: The ID of the table to create the authorized view for. Supplied by the table_id fixture. + """ + from google.api_core import exceptions + from google.api_core import retry + + retry = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + new_view_id = uuid.uuid4().hex[:8] + parent_path = f"projects/{project_id}/instances/{instance_id}/tables/{table_id}" + new_path = f"{parent_path}/authorizedViews/{new_view_id}" + try: + print(f"Creating view: {new_path}") + admin_client.table_admin_client.create_authorized_view( + request={ + "parent": parent_path, + "authorized_view_id": new_view_id, + "authorized_view": { + "subset_view": { + "row_prefixes": [ALLOW_ALL], + "family_subsets": { + TEST_FAMILY: ALL_QUALIFIERS, + TEST_FAMILY_2: ALL_QUALIFIERS, + TEST_AGGREGATE_FAMILY: ALL_QUALIFIERS, + }, + }, + }, + }, + retry=retry, + ) + except exceptions.AlreadyExists: + pass + except exceptions.MethodNotImplemented: + # will occur when run in emulator. Pass empty id + new_view_id = None + yield new_view_id + if new_view_id: + print(f"Deleting view: {new_path}") + try: + admin_client.table_admin_client.delete_authorized_view(name=new_path) + except exceptions.NotFound: + print(f"View {new_view_id} not found, skipping deletion") + + @pytest.fixture(scope="session") + def project_id(self, client): + """Returns the project ID from the client.""" + yield client.project diff --git a/tests/system/data/setup_fixtures.py b/tests/system/data/setup_fixtures.py deleted file mode 100644 index 169e2396b..000000000 --- a/tests/system/data/setup_fixtures.py +++ /dev/null @@ -1,210 +0,0 @@ -# Copyright 2023 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Contains a set of pytest fixtures for setting up and populating a -Bigtable database for testing purposes. -""" - -import pytest -import os -import uuid - -from . import TEST_FAMILY, TEST_FAMILY_2, TEST_AGGREGATE_FAMILY - -# authorized view subset to allow all qualifiers -ALLOW_ALL = "" -ALL_QUALIFIERS = {"qualifier_prefixes": [ALLOW_ALL]} - - -@pytest.fixture(scope="session") -def admin_client(): - """ - Client for interacting with Table and Instance admin APIs - """ - from google.cloud.bigtable.client import Client - - client = Client(admin=True) - yield client - - -@pytest.fixture(scope="session") -def instance_id(admin_client, project_id, cluster_config): - """ - Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session - """ - from google.cloud.bigtable_admin_v2 import types - from google.api_core import exceptions - from google.cloud.environment_vars import BIGTABLE_EMULATOR - - # use user-specified instance if available - user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") - if user_specified_instance: - print("Using user-specified instance: {}".format(user_specified_instance)) - yield user_specified_instance - return - - # create a new temporary test instance - instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" - if os.getenv(BIGTABLE_EMULATOR): - # don't create instance if in emulator mode - yield instance_id - else: - try: - operation = admin_client.instance_admin_client.create_instance( - parent=f"projects/{project_id}", - instance_id=instance_id, - instance=types.Instance( - display_name="Test Instance", - # labels={"python-system-test": "true"}, - ), - clusters=cluster_config, - ) - operation.result(timeout=240) - except exceptions.AlreadyExists: - pass - yield instance_id - admin_client.instance_admin_client.delete_instance( - name=f"projects/{project_id}/instances/{instance_id}" - ) - - -@pytest.fixture(scope="session") -def column_split_config(): - """ - specify initial splits to create when creating a new test table - """ - return [(num * 1000).to_bytes(8, "big") for num in range(1, 10)] - - -@pytest.fixture(scope="session") -def table_id( - admin_client, - project_id, - instance_id, - column_family_config, - init_table_id, - column_split_config, -): - """ - Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session - - Args: - - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture. - - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. - - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. - - init_column_families: A list of column families to initialize the table with, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. - Supplied by the init_column_families fixture. - - init_table_id: The table ID to give to the test table, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. - Supplied by the init_table_id fixture. - - column_split_config: A list of row keys to use as initial splits when creating the test table. - """ - from google.api_core import exceptions - from google.api_core import retry - - # use user-specified instance if available - user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") - if user_specified_table: - print("Using user-specified table: {}".format(user_specified_table)) - yield user_specified_table - return - - retry = retry.Retry( - predicate=retry.if_exception_type(exceptions.FailedPrecondition) - ) - try: - parent_path = f"projects/{project_id}/instances/{instance_id}" - print(f"Creating table: {parent_path}/tables/{init_table_id}") - admin_client.table_admin_client.create_table( - request={ - "parent": parent_path, - "table_id": init_table_id, - "table": {"column_families": column_family_config}, - "initial_splits": [{"key": key} for key in column_split_config], - }, - retry=retry, - ) - except exceptions.AlreadyExists: - pass - yield init_table_id - print(f"Deleting table: {parent_path}/tables/{init_table_id}") - try: - admin_client.table_admin_client.delete_table( - name=f"{parent_path}/tables/{init_table_id}" - ) - except exceptions.NotFound: - print(f"Table {init_table_id} not found, skipping deletion") - - -@pytest.fixture(scope="session") -def authorized_view_id( - admin_client, - project_id, - instance_id, - table_id, -): - """ - Creates and returns a new temporary authorized view for the test session - - Args: - - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture. - - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. - - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. - - table_id: The ID of the table to create the authorized view for. Supplied by the table_id fixture. - """ - from google.api_core import exceptions - from google.api_core import retry - - retry = retry.Retry( - predicate=retry.if_exception_type(exceptions.FailedPrecondition) - ) - new_view_id = uuid.uuid4().hex[:8] - parent_path = f"projects/{project_id}/instances/{instance_id}/tables/{table_id}" - new_path = f"{parent_path}/authorizedViews/{new_view_id}" - try: - print(f"Creating view: {new_path}") - admin_client.table_admin_client.create_authorized_view( - request={ - "parent": parent_path, - "authorized_view_id": new_view_id, - "authorized_view": { - "subset_view": { - "row_prefixes": [ALLOW_ALL], - "family_subsets": { - TEST_FAMILY: ALL_QUALIFIERS, - TEST_FAMILY_2: ALL_QUALIFIERS, - TEST_AGGREGATE_FAMILY: ALL_QUALIFIERS, - }, - }, - }, - }, - retry=retry, - ) - except exceptions.AlreadyExists: - pass - except exceptions.MethodNotImplemented: - # will occur when run in emulator. Pass empty id - new_view_id = None - yield new_view_id - if new_view_id: - print(f"Deleting view: {new_path}") - try: - admin_client.table_admin_client.delete_authorized_view(name=new_path) - except exceptions.NotFound: - print(f"View {new_view_id} not found, skipping deletion") - - -@pytest.fixture(scope="session") -def project_id(client): - """Returns the project ID from the client.""" - yield client.project diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py new file mode 100644 index 000000000..af4735d97 --- /dev/null +++ b/tests/system/data/test_metrics_async.py @@ -0,0 +1,2169 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import os +import pytest +import uuid + +from grpc import StatusCode + +from google.api_core.exceptions import Aborted +from google.api_core.exceptions import GoogleAPICallError +from google.api_core.exceptions import PermissionDenied +from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data._metrics.data_model import ( + CompletedOperationMetric, + CompletedAttemptMetric, +) +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery +from google.cloud.bigtable_v2.types import ResponseParams + +from google.cloud.bigtable.data._cross_sync import CrossSync + +from . import TEST_FAMILY, SystemTestRunner + +if CrossSync.is_async: + from grpc.aio import UnaryUnaryClientInterceptor + from grpc.aio import UnaryStreamClientInterceptor + from grpc.aio import AioRpcError + from grpc.aio import Metadata +else: + from grpc import UnaryUnaryClientInterceptor + from grpc import UnaryStreamClientInterceptor + from grpc import RpcError + from grpc import intercept_channel + +__CROSS_SYNC_OUTPUT__ = "tests.system.data.test_metrics_autogen" + + +class _MetricsTestHandler(MetricsHandler): + """ + Store completed metrics events in internal lists for testing + """ + + def __init__(self, **kwargs): + self.completed_operations = [] + self.completed_attempts = [] + + def on_operation_complete(self, op): + self.completed_operations.append(op) + + def on_attempt_complete(self, attempt, _): + self.completed_attempts.append(attempt) + + def clear(self): + self.completed_operations.clear() + self.completed_attempts.clear() + + def __repr__(self): + return f"{self.__class__}(completed_operations={len(self.completed_operations)}, completed_attempts={len(self.completed_attempts)}" + + +@CrossSync.convert_class +class _ErrorInjectorInterceptor( + UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor +): + """ + Gprc interceptor used to inject errors into rpc calls, to test failures + """ + + def __init__(self): + self._exc_list = [] + self.fail_mid_stream = False + + def push(self, exc: Exception): + self._exc_list.append(exc) + + def clear(self): + self._exc_list.clear() + self.fail_mid_stream = False + + @CrossSync.convert + async def intercept_unary_unary(self, continuation, client_call_details, request): + if self._exc_list: + raise self._exc_list.pop(0) + return await continuation(client_call_details, request) + + @CrossSync.convert + async def intercept_unary_stream(self, continuation, client_call_details, request): + if not self.fail_mid_stream and self._exc_list: + raise self._exc_list.pop(0) + + response = await continuation(client_call_details, request) + + if self.fail_mid_stream and self._exc_list: + exc = self._exc_list.pop(0) + + class CallWrapper: + def __init__(self, call, exc_to_raise): + self._call = call + self._exc = exc_to_raise + self._raised = False + + @CrossSync.convert(sync_name="__iter__") + def __aiter__(self): + return self + + @CrossSync.convert( + sync_name="__next__", replace_symbols={"__anext__": "__next__"} + ) + async def __anext__(self): + if not self._raised: + self._raised = True + if self._exc: + raise self._exc + return await self._call.__anext__() + + def __getattr__(self, name): + return getattr(self._call, name) + + return CallWrapper(response, exc) + + return response + + +@CrossSync.convert_class(sync_name="TestMetrics") +class TestMetricsAsync(SystemTestRunner): + def _make_client(self): + project = os.getenv("GOOGLE_CLOUD_PROJECT") or None + return CrossSync.DataClient(project=project) + + def _make_exception(self, status, cluster_id=None, zone_id=None): + if cluster_id or zone_id: + metadata = ( + "x-goog-ext-425905942-bin", + ResponseParams.serialize( + ResponseParams(cluster_id=cluster_id, zone_id=zone_id) + ), + ) + else: + metadata = None + if CrossSync.is_async: + metadata = Metadata(metadata) if metadata else Metadata() + return AioRpcError(status, Metadata(), metadata) + else: + exc = RpcError(status) + exc.trailing_metadata = lambda: [metadata] if metadata else [] + exc.initial_metadata = lambda: [] + exc.code = lambda: status + exc.details = lambda: None + + def _result(): + raise exc + + exc.result = _result + return exc + + @pytest.fixture(scope="session") + def handler(self): + return _MetricsTestHandler() + + @pytest.fixture(scope="session") + def error_injector(self): + return _ErrorInjectorInterceptor() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function", autouse=True) + async def _clear_state(self, handler, error_injector): + """Clear handler and interceptor between each test""" + handler.clear() + error_injector.clear() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def client(self, error_injector): + async with self._make_client() as client: + if CrossSync.is_async: + client.transport.grpc_channel._unary_unary_interceptors.append( + error_injector + ) + client.transport.grpc_channel._unary_stream_interceptors.append( + error_injector + ) + else: + # inject interceptor after bigtable metrics interceptors + metrics_channel = client.transport._grpc_channel._channel._channel + client.transport._grpc_channel._channel._channel = intercept_channel( + metrics_channel, error_injector + ) + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function") + async def temp_rows(self, table): + builder = CrossSync.TempRowBuilder(table) + yield builder + await builder.delete_rows() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def table(self, client, table_id, instance_id, handler): + async with client.get_table(instance_id, table_id) as table: + table._metrics.add_handler(handler) + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def authorized_view( + self, client, table_id, instance_id, authorized_view_id, handler + ): + async with client.get_authorized_view( + instance_id, table_id, authorized_view_id + ) as table: + table._metrics.add_handler(handler) + yield table + + @CrossSync.pytest + async def test_read_rows(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + row_list = await table.read_rows(ReadRowsQuery()) + assert len(row_list) == 2 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_read_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + await table.read_rows(ReadRowsQuery(), operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.read_rows( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + # full table scan + generator = await table.read_rows_stream(ReadRowsQuery()) + row_list = [r async for r in generator] + assert len(row_list) == 2 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + @CrossSync.convert(replace_symbols={"__anext__": "__next__", "aclose": "close"}) + async def test_read_rows_stream_failure_closed( + self, table, temp_rows, handler, error_injector + ): + """ + Test how metrics collection handles closed generator + """ + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + generator = await table.read_rows_stream(ReadRowsQuery()) + await generator.__anext__() + await generator.aclose() + with pytest.raises(CrossSync.StopIteration): + await generator.__anext__() + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "CANCELLED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "CANCELLED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) + with pytest.raises(PermissionDenied): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + generator = await table.read_rows_stream( + ReadRowsQuery(), operation_timeout=0.001 + ) + with pytest.raises(GoogleAPICallError): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = await authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + [_ async for _ in generator] + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_stream_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = await authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + [_ async for _ in generator] + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + @CrossSync.pytest + async def test_read_rows_stream_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) + with pytest.raises(PermissionDenied): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 2 + # validate retried attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + # validate final attempt + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + @CrossSync.pytest + async def test_read_row(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + handler.clear() + await table.read_row(b"row_key_1") + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns > 0 + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_read_row_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.read_row(b"row_key_1", retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_row_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + await table.read_row(b"row_key_1", operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.read_row( + b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) + handler.clear() + row_list = await table.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + # validate counts + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # validate operations + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + attempt = operation.completed_attempts[0] + assert attempt in handler.completed_attempts + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + assert isinstance(attempt, CompletedAttemptMetric) + assert ( + attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + ) + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 + and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_read_rows_sharded_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + + error_injector.push(self._make_exception(StatusCode.ABORTED)) + await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + # validate operations + for op in handler.completed_operations: + assert op.final_status.name == "OK" + assert op.op_type.value == "ReadRows" + assert op.is_streaming is True + # validate attempts + assert ( + len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) + == 2 + ) + assert ( + len( + [ + a + for a in handler.completed_attempts + if a.end_status.name == "ABORTED" + ] + ) + == 1 + ) + + @CrossSync.pytest + async def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.api_core.exceptions import DeadlineExceeded + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2], operation_timeout=0.005) + assert len(e.value.exceptions) == 2 + for sub_exc in e.value.exceptions: + assert isinstance(sub_exc.__cause__, DeadlineExceeded) + # both shards should fail + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # validate operations + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = operation.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_sharded_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await authorized_view.read_rows_sharded([query1, query2]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + # one shard will fail, the other will succeed + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # sort operations by status + failed_op = next( + op for op in handler.completed_operations if op.final_status.name != "OK" + ) + success_op = next( + op for op in handler.completed_operations if op.final_status.name == "OK" + ) + # validate failed operation + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert failed_op.cluster_id == next(iter(cluster_config.keys())) + assert ( + failed_op.zone + == cluster_config[failed_op.cluster_id].location.split("/")[-1] + ) + # validate failed attempt + failed_attempt = failed_op.completed_attempts[0] + assert failed_attempt.end_status.name == "PERMISSION_DENIED" + assert ( + failed_attempt.gfe_latency_ns >= 0 + and failed_attempt.gfe_latency_ns < failed_op.duration_ns + ) + # validate successful operation + assert success_op.final_status.name == "OK" + assert success_op.op_type.value == "ReadRows" + assert success_op.is_streaming is True + assert len(success_op.completed_attempts) == 1 + # validate successful attempt + success_attempt = success_op.completed_attempts[0] + assert success_attempt.end_status.name == "OK" + + @CrossSync.pytest + async def test_read_rows_sharded_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + # one shard will fail, the other will succeed + # the failing shard will have one retry + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + # sort operations by status + failed_op = next( + op for op in handler.completed_operations if op.final_status.name != "OK" + ) + success_op = next( + op for op in handler.completed_operations if op.final_status.name == "OK" + ) + # validate failed operation + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + # validate successful operation + assert success_op.final_status.name == "OK" + assert len(success_op.completed_attempts) == 2 + # validate failed attempt + attempt = failed_op.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + # validate retried attempt + retried_attempt = success_op.completed_attempts[0] + assert retried_attempt.end_status.name == "ABORTED" + # validate successful attempt + success_attempt = success_op.completed_attempts[-1] + assert success_attempt.end_status.name == "OK" + + @CrossSync.pytest + async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + + handler.clear() + await table.bulk_mutate_rows([bulk_mutation]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + await table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup): + await table.bulk_mutate_rows([entry], operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup): + await authorized_view.bulk_mutate_rows([entry]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + + For bulk_mutate, the rpc returns success, with failures returned in the response. + For this reason, We expect the attempts to be marked as successful, even though + the underlying mutation is retried + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup) as e: + await authorized_view.bulk_mutate_rows( + [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 + ) + assert len(e.value.exceptions) == 1 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] + + @CrossSync.pytest + async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + row_key2, mutation2 = await temp_rows.create_row_and_mutation( + table, new_value=new_value2 + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + + handler.clear() + async with table.mutations_batcher() as batcher: + await batcher.append(bulk_mutation) + await batcher.append(bulk_mutation2) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # bacher expects to cancel staged operation on close + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert ( + operation.flow_throttling_time_ns > 0 + and operation.flow_throttling_time_ns < operation.duration_ns + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_with_retries( + self, table, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + async with table.mutations_batcher( + batch_retryable_errors=[Aborted] + ) as batcher: + await batcher.append(entry) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup): + async with table.mutations_batcher( + batch_operation_timeout=0.001 + ) as batcher: + await batcher.append(entry) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup) as e: + async with authorized_view.mutations_batcher() as batcher: + await batcher.append(entry) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_mutate_row(self, table, temp_rows, handler, cluster_config): + row_key = b"mutate" + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + handler.clear() + await table.mutate_row(row_key, mutation) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_mutate_row_failure_with_retries( + self, table, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.mutate_row(row_key, [mutation], retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_row_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + + with pytest.raises(GoogleAPICallError): + await table.mutate_row(row_key, [mutation], operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.mutate_row(row_key, [mutation]) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_mutate_row_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.mutate_row( + row_key, + [mutation], + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + @CrossSync.pytest + async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): + await table.sample_row_keys() + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "SampleRowKeys" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.drop + @CrossSync.pytest + async def test_sample_row_keys_failure_cancelled( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + test with retryable errors, then a terminal one + + No headers expected + """ + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(asyncio.CancelledError) + with pytest.raises(asyncio.CancelledError): + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "UNKNOWN" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "UNKNOWN" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_sample_row_keys_failure_with_retries( + self, table, temp_rows, handler, error_injector, cluster_config + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a success + """ + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(self._make_exception(StatusCode.ABORTED)) + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "OK" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "OK" + assert ( + final_attempt.gfe_latency_ns > 0 + and final_attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_sample_row_keys_failure_timeout(self, table, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + with pytest.raises(GoogleAPICallError): + await table.sample_row_keys(operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_sample_row_keys_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(PermissionDenied): + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 2 + # validate retried attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + # validate final attempt + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + @CrossSync.pytest + async def test_read_modify_write(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + await table.read_modify_write_row(row_key, rule) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadModifyWriteRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.drop + @CrossSync.pytest + async def test_read_modify_write_failure_cancelled( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + + # trigger an exception + exc = asyncio.CancelledError("injected") + error_injector.push(exc) + with pytest.raises(asyncio.CancelledError): + await table.read_modify_write_row(row_key, rule) + + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "UNKNOWN" + assert operation.is_streaming is False + assert operation.op_type.value == "ReadModifyWriteRow" + assert len(operation.completed_attempts) == len(handler.completed_attempts) + assert operation.completed_attempts == handler.completed_attempts + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 + assert attempt.end_status.name == "UNKNOWN" + assert attempt.backoff_before_attempt_ns == 0 + assert attempt.gfe_latency_ns is None + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_read_modify_write_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + with pytest.raises(GoogleAPICallError): + await table.read_modify_write_row(row_key, rule, operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_modify_write_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + rule = IncrementRule("unauthorized", qualifier, 1) + with pytest.raises(GoogleAPICallError): + await authorized_view.read_modify_write_row(row_key, rule) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_check_and_mutate_row( + self, table, temp_rows, handler, cluster_config + ): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + predicate = ValueRangeFilter(0, 2) + await table.check_and_mutate_row( + row_key, + predicate, + true_case_mutations=true_mutation, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "CheckAndMutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.drop + @CrossSync.pytest + async def test_check_and_mutate_row_failure_cancelled( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + + # trigger an exception + exc = asyncio.CancelledError("injected") + error_injector.push(exc) + with pytest.raises(asyncio.CancelledError): + await table.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "UNKNOWN" + assert operation.is_streaming is False + assert operation.op_type.value == "CheckAndMutateRow" + assert len(operation.completed_attempts) == len(handler.completed_attempts) + assert operation.completed_attempts == handler.completed_attempts + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 + assert attempt.end_status.name == "UNKNOWN" + assert attempt.backoff_before_attempt_ns == 0 + assert attempt.gfe_latency_ns is None + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + + @CrossSync.pytest + async def test_check_and_mutate_row_failure_timeout( + self, table, temp_rows, handler + ): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + with pytest.raises(GoogleAPICallError): + await table.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=true_mutation, + operation_timeout=0.001, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_check_and_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + mutation_value = b"true-mutation-value" + mutation = SetCell( + family="unauthorized", qualifier=qualifier, new_value=mutation_value + ) + with pytest.raises(GoogleAPICallError): + await authorized_view.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=mutation, + false_case_mutations=mutation, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) diff --git a/tests/system/data/test_metrics_autogen.py b/tests/system/data/test_metrics_autogen.py new file mode 100644 index 000000000..54a9f2256 --- /dev/null +++ b/tests/system/data/test_metrics_autogen.py @@ -0,0 +1,1680 @@ +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This file is automatically generated by CrossSync. Do not edit manually. + +import os +import pytest +import uuid +from grpc import StatusCode +from google.api_core.exceptions import Aborted +from google.api_core.exceptions import GoogleAPICallError +from google.api_core.exceptions import PermissionDenied +from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data._metrics.data_model import ( + CompletedOperationMetric, + CompletedAttemptMetric, +) +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery +from google.cloud.bigtable_v2.types import ResponseParams +from google.cloud.bigtable.data._cross_sync import CrossSync +from . import TEST_FAMILY, SystemTestRunner +from grpc import UnaryUnaryClientInterceptor +from grpc import UnaryStreamClientInterceptor +from grpc import RpcError +from grpc import intercept_channel + + +class _MetricsTestHandler(MetricsHandler): + """ + Store completed metrics events in internal lists for testing + """ + + def __init__(self, **kwargs): + self.completed_operations = [] + self.completed_attempts = [] + + def on_operation_complete(self, op): + self.completed_operations.append(op) + + def on_attempt_complete(self, attempt, _): + self.completed_attempts.append(attempt) + + def clear(self): + self.completed_operations.clear() + self.completed_attempts.clear() + + def __repr__(self): + return f"{self.__class__}(completed_operations={len(self.completed_operations)}, completed_attempts={len(self.completed_attempts)}" + + +class _ErrorInjectorInterceptor( + UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor +): + """ + Gprc interceptor used to inject errors into rpc calls, to test failures + """ + + def __init__(self): + self._exc_list = [] + self.fail_mid_stream = False + + def push(self, exc: Exception): + self._exc_list.append(exc) + + def clear(self): + self._exc_list.clear() + self.fail_mid_stream = False + + def intercept_unary_unary(self, continuation, client_call_details, request): + if self._exc_list: + raise self._exc_list.pop(0) + return continuation(client_call_details, request) + + def intercept_unary_stream(self, continuation, client_call_details, request): + if not self.fail_mid_stream and self._exc_list: + raise self._exc_list.pop(0) + response = continuation(client_call_details, request) + if self.fail_mid_stream and self._exc_list: + exc = self._exc_list.pop(0) + + class CallWrapper: + def __init__(self, call, exc_to_raise): + self._call = call + self._exc = exc_to_raise + self._raised = False + + def __iter__(self): + return self + + def __next__(self): + if not self._raised: + self._raised = True + if self._exc: + raise self._exc + return self._call.__next__() + + def __getattr__(self, name): + return getattr(self._call, name) + + return CallWrapper(response, exc) + return response + + +class TestMetrics(SystemTestRunner): + def _make_client(self): + project = os.getenv("GOOGLE_CLOUD_PROJECT") or None + return CrossSync._Sync_Impl.DataClient(project=project) + + def _make_exception(self, status, cluster_id=None, zone_id=None): + if cluster_id or zone_id: + metadata = ( + "x-goog-ext-425905942-bin", + ResponseParams.serialize( + ResponseParams(cluster_id=cluster_id, zone_id=zone_id) + ), + ) + else: + metadata = None + exc = RpcError(status) + exc.trailing_metadata = lambda: [metadata] if metadata else [] + exc.initial_metadata = lambda: [] + exc.code = lambda: status + exc.details = lambda: None + + def _result(): + raise exc + + exc.result = _result + return exc + + @pytest.fixture(scope="session") + def handler(self): + return _MetricsTestHandler() + + @pytest.fixture(scope="session") + def error_injector(self): + return _ErrorInjectorInterceptor() + + @pytest.fixture(scope="function", autouse=True) + def _clear_state(self, handler, error_injector): + """Clear handler and interceptor between each test""" + handler.clear() + error_injector.clear() + + @pytest.fixture(scope="session") + def client(self, error_injector): + with self._make_client() as client: + metrics_channel = client.transport._grpc_channel._channel._channel + client.transport._grpc_channel._channel._channel = intercept_channel( + metrics_channel, error_injector + ) + yield client + + @pytest.fixture(scope="function") + def temp_rows(self, table): + builder = CrossSync._Sync_Impl.TempRowBuilder(table) + yield builder + builder.delete_rows() + + @pytest.fixture(scope="session") + def table(self, client, table_id, instance_id, handler): + with client.get_table(instance_id, table_id) as table: + table._metrics.add_handler(handler) + yield table + + @pytest.fixture(scope="session") + def authorized_view( + self, client, table_id, instance_id, authorized_view_id, handler + ): + with client.get_authorized_view( + instance_id, table_id, authorized_view_id + ) as table: + table._metrics.add_handler(handler) + yield table + + def test_read_rows(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + row_list = table.read_rows(ReadRowsQuery()) + assert len(row_list) == 2 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 + + def test_read_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_rows_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + table.read_rows(ReadRowsQuery(), operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + authorized_view.read_rows( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery()) + row_list = [r for r in generator] + assert len(row_list) == 2 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 + + def test_read_rows_stream_failure_closed( + self, table, temp_rows, handler, error_injector + ): + """Test how metrics collection handles closed generator""" + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery()) + generator.__next__() + generator.close() + with pytest.raises(CrossSync._Sync_Impl.StopIteration): + generator.__next__() + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "CANCELLED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "CANCELLED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + with pytest.raises(PermissionDenied): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery(), operation_timeout=0.001) + with pytest.raises(GoogleAPICallError): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + [_ for _ in generator] + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_stream_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + [_ for _ in generator] + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + def test_read_rows_stream_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + temp_rows.add_row(b"row_key_1") + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + with pytest.raises(PermissionDenied): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 2 + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + def test_read_row(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + handler.clear() + table.read_row(b"row_key_1") + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns > 0 + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 + + def test_read_row_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.read_row(b"row_key_1", retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_row_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + table.read_row(b"row_key_1", operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + authorized_view.read_row( + b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + temp_rows.add_row(b"c") + temp_rows.add_row(b"d") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) + handler.clear() + row_list = table.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + attempt = operation.completed_attempts[0] + assert attempt in handler.completed_attempts + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + assert isinstance(attempt, CompletedAttemptMetric) + assert ( + attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + ) + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 + and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + assert attempt.grpc_throttling_time_ns == 0 + + def test_read_rows_sharded_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors""" + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.push(self._make_exception(StatusCode.ABORTED)) + table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + for op in handler.completed_operations: + assert op.final_status.name == "OK" + assert op.op_type.value == "ReadRows" + assert op.is_streaming is True + assert ( + len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) + == 2 + ) + assert ( + len( + [ + a + for a in handler.completed_attempts + if a.end_status.name == "ABORTED" + ] + ) + == 1 + ) + + def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.api_core.exceptions import DeadlineExceeded + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + table.read_rows_sharded([query1, query2], operation_timeout=0.005) + assert len(e.value.exceptions) == 2 + for sub_exc in e.value.exceptions: + assert isinstance(sub_exc.__cause__, DeadlineExceeded) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = operation.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_sharded_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + authorized_view.read_rows_sharded([query1, query2]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + failed_op = next( + (op for op in handler.completed_operations if op.final_status.name != "OK") + ) + success_op = next( + (op for op in handler.completed_operations if op.final_status.name == "OK") + ) + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert failed_op.cluster_id == next(iter(cluster_config.keys())) + assert ( + failed_op.zone + == cluster_config[failed_op.cluster_id].location.split("/")[-1] + ) + failed_attempt = failed_op.completed_attempts[0] + assert failed_attempt.end_status.name == "PERMISSION_DENIED" + assert ( + failed_attempt.gfe_latency_ns >= 0 + and failed_attempt.gfe_latency_ns < failed_op.duration_ns + ) + assert success_op.final_status.name == "OK" + assert success_op.op_type.value == "ReadRows" + assert success_op.is_streaming is True + assert len(success_op.completed_attempts) == 1 + success_attempt = success_op.completed_attempts[0] + assert success_attempt.end_status.name == "OK" + + def test_read_rows_sharded_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + failed_op = next( + (op for op in handler.completed_operations if op.final_status.name != "OK") + ) + success_op = next( + (op for op in handler.completed_operations if op.final_status.name == "OK") + ) + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert success_op.final_status.name == "OK" + assert len(success_op.completed_attempts) == 2 + attempt = failed_op.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + retried_attempt = success_op.completed_attempts[0] + assert retried_attempt.end_status.name == "ABORTED" + success_attempt = success_op.completed_attempts[-1] + assert success_attempt.end_status.name == "OK" + + def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + (row_key, mutation) = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + handler.clear() + table.bulk_mutate_rows([bulk_mutation]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 + + def test_bulk_mutate_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup): + table.bulk_mutate_rows([entry], operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_bulk_mutate_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup): + authorized_view.bulk_mutate_rows([entry]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_bulk_mutate_rows_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out + + For bulk_mutate, the rpc returns success, with failures returned in the response. + For this reason, We expect the attempts to be marked as successful, even though + the underlying mutation is retried""" + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup) as e: + authorized_view.bulk_mutate_rows( + [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 + ) + assert len(e.value.exceptions) == 1 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] + + def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + (new_value, new_value2) = [uuid.uuid4().hex.encode() for _ in range(2)] + (row_key, mutation) = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + (row_key2, mutation2) = temp_rows.create_row_and_mutation( + table, new_value=new_value2 + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + handler.clear() + with table.mutations_batcher() as batcher: + batcher.append(bulk_mutation) + batcher.append(bulk_mutation2) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert ( + operation.flow_throttling_time_ns > 0 + and operation.flow_throttling_time_ns < operation.duration_ns + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 + + def test_mutate_rows_batcher_failure_with_retries( + self, table, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + with table.mutations_batcher(batch_retryable_errors=[Aborted]) as batcher: + batcher.append(entry) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + with pytest.raises(MutationsExceptionGroup): + with table.mutations_batcher(batch_operation_timeout=0.001) as batcher: + batcher.append(entry) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_mutate_rows_batcher_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + with pytest.raises(MutationsExceptionGroup) as e: + with authorized_view.mutations_batcher() as batcher: + batcher.append(entry) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_mutate_row(self, table, temp_rows, handler, cluster_config): + row_key = b"mutate" + new_value = uuid.uuid4().hex.encode() + (row_key, mutation) = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + handler.clear() + table.mutate_row(row_key, mutation) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 + + def test_mutate_row_failure_with_retries(self, table, handler, error_injector): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.mutate_row(row_key, [mutation], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_mutate_row_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + with pytest.raises(GoogleAPICallError): + table.mutate_row(row_key, [mutation], operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + with pytest.raises(GoogleAPICallError) as e: + authorized_view.mutate_row(row_key, [mutation]) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_mutate_row_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + with pytest.raises(GoogleAPICallError) as e: + authorized_view.mutate_row( + row_key, + [mutation], + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): + table.sample_row_keys() + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "SampleRowKeys" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 + + def test_sample_row_keys_failure_with_retries( + self, table, temp_rows, handler, error_injector, cluster_config + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a success""" + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(self._make_exception(StatusCode.ABORTED)) + table.sample_row_keys(retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "OK" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "OK" + assert ( + final_attempt.gfe_latency_ns > 0 + and final_attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_sample_row_keys_failure_timeout(self, table, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + with pytest.raises(GoogleAPICallError): + table.sample_row_keys(operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_sample_row_keys_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(PermissionDenied): + table.sample_row_keys(retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 2 + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + def test_read_modify_write(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + table.read_modify_write_row(row_key, rule) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadModifyWriteRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 + + def test_read_modify_write_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + with pytest.raises(GoogleAPICallError): + table.read_modify_write_row(row_key, rule, operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + def test_read_modify_write_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + rule = IncrementRule("unauthorized", qualifier, 1) + with pytest.raises(GoogleAPICallError): + authorized_view.read_modify_write_row(row_key, rule) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_check_and_mutate_row(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + predicate = ValueRangeFilter(0, 2) + table.check_and_mutate_row( + row_key, predicate, true_case_mutations=true_mutation + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "CheckAndMutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + assert attempt.grpc_throttling_time_ns == 0 + + def test_check_and_mutate_row_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + with pytest.raises(GoogleAPICallError): + table.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=true_mutation, + operation_timeout=0.001, + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + def test_check_and_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + mutation_value = b"true-mutation-value" + mutation = SetCell( + family="unauthorized", qualifier=qualifier, new_value=mutation_value + ) + with pytest.raises(GoogleAPICallError): + authorized_view.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=mutation, + false_case_mutations=mutation, + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) diff --git a/tests/system/data/test_system_async.py b/tests/system/data/test_system_async.py index 39c454996..4f9c1ae4e 100644 --- a/tests/system/data/test_system_async.py +++ b/tests/system/data/test_system_async.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import asyncio import pytest import datetime import uuid @@ -26,7 +27,7 @@ from google.cloud.bigtable.data._cross_sync import CrossSync -from . import TEST_FAMILY, TEST_FAMILY_2, TEST_AGGREGATE_FAMILY +from . import TEST_FAMILY, TEST_FAMILY_2, TEST_AGGREGATE_FAMILY, SystemTestRunner if CrossSync.is_async: from google.cloud.bigtable_v2.services.bigtable.transports.grpc_asyncio import ( @@ -116,9 +117,43 @@ async def delete_rows(self): } await self.target.client._gapic_client.mutate_rows(request) + @CrossSync.convert + async def retrieve_cell_value(self, target, row_key): + """ + Helper to read an individual row + """ + from google.cloud.bigtable.data import ReadRowsQuery + + row_list = await target.read_rows(ReadRowsQuery(row_keys=row_key)) + assert len(row_list) == 1 + row = row_list[0] + cell = row.cells[0] + return cell.value + + @CrossSync.convert + async def create_row_and_mutation( + self, table, *, start_value=b"start", new_value=b"new_value" + ): + """ + Helper to create a new row, and a sample set_cell mutation to change its value + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = uuid.uuid4().hex.encode() + family = TEST_FAMILY + qualifier = b"test-qualifier" + await self.add_row( + row_key, family=family, qualifier=qualifier, value=start_value + ) + # ensure cell is initialized + assert await self.retrieve_cell_value(table, row_key) == start_value + + mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) + return row_key, mutation + @CrossSync.convert_class(sync_name="TestSystem") -class TestSystemAsync: +class TestSystemAsync(SystemTestRunner): def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None return CrossSync.DataClient(project=project) @@ -148,82 +183,6 @@ async def target(self, client, table_id, authorized_view_id, instance_id, reques else: raise ValueError(f"unknown target type: {request.param}") - @pytest.fixture(scope="session") - def column_family_config(self): - """ - specify column families to create when creating a new test table - """ - from google.cloud.bigtable_admin_v2 import types - - int_aggregate_type = types.Type.Aggregate( - input_type=types.Type(int64_type={"encoding": {"big_endian_bytes": {}}}), - sum={}, - ) - return { - TEST_FAMILY: types.ColumnFamily(), - TEST_FAMILY_2: types.ColumnFamily(), - TEST_AGGREGATE_FAMILY: types.ColumnFamily( - value_type=types.Type(aggregate_type=int_aggregate_type) - ), - } - - @pytest.fixture(scope="session") - def init_table_id(self): - """ - The table_id to use when creating a new test table - """ - return f"test-table-{uuid.uuid4().hex}" - - @pytest.fixture(scope="session") - def cluster_config(self, project_id): - """ - Configuration for the clusters to use when creating a new instance - """ - from google.cloud.bigtable_admin_v2 import types - - cluster = { - "test-cluster": types.Cluster( - location=f"projects/{project_id}/locations/us-central1-b", - serve_nodes=1, - ) - } - return cluster - - @CrossSync.convert - @pytest.mark.usefixtures("target") - async def _retrieve_cell_value(self, target, row_key): - """ - Helper to read an individual row - """ - from google.cloud.bigtable.data import ReadRowsQuery - - row_list = await target.read_rows(ReadRowsQuery(row_keys=row_key)) - assert len(row_list) == 1 - row = row_list[0] - cell = row.cells[0] - return cell.value - - @CrossSync.convert - async def _create_row_and_mutation( - self, table, temp_rows, *, start_value=b"start", new_value=b"new_value" - ): - """ - Helper to create a new row, and a sample set_cell mutation to change its value - """ - from google.cloud.bigtable.data.mutations import SetCell - - row_key = uuid.uuid4().hex.encode() - family = TEST_FAMILY - qualifier = b"test-qualifier" - await temp_rows.add_row( - row_key, family=family, qualifier=qualifier, value=start_value - ) - # ensure cell is initialized - assert await self._retrieve_cell_value(table, row_key) == start_value - - mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) - return row_key, mutation - @CrossSync.convert @CrossSync.pytest_fixture(scope="function") async def temp_rows(self, target): @@ -321,13 +280,13 @@ async def test_mutation_set_cell(self, target, temp_rows): """ row_key = b"bulk_mutate" new_value = uuid.uuid4().hex.encode() - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value ) await target.mutate_row(row_key, mutation) # ensure cell is updated - assert (await self._retrieve_cell_value(target, row_key)) == new_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value @CrossSync.pytest @pytest.mark.usefixtures("target") @@ -349,14 +308,14 @@ async def test_mutation_add_to_cell(self, target, temp_rows): await target.mutate_row( row_key, AddToCell(family, qualifier, 1, timestamp_micros=0) ) - encoded_result = await self._retrieve_cell_value(target, row_key) + encoded_result = await temp_rows.retrieve_cell_value(target, row_key) int_result = int.from_bytes(encoded_result, byteorder="big") assert int_result == 1 # update again await target.mutate_row( row_key, AddToCell(family, qualifier, 9, timestamp_micros=0) ) - encoded_result = await self._retrieve_cell_value(target, row_key) + encoded_result = await temp_rows.retrieve_cell_value(target, row_key) int_result = int.from_bytes(encoded_result, byteorder="big") assert int_result == 10 @@ -398,15 +357,15 @@ async def test_bulk_mutations_set_cell(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry new_value = uuid.uuid4().hex.encode() - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) await target.bulk_mutate_rows([bulk_mutation]) # ensure cell is updated - assert (await self._retrieve_cell_value(target, row_key)) == new_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value @CrossSync.pytest async def test_bulk_mutations_raise_exception(self, client, target): @@ -444,11 +403,11 @@ async def test_mutations_batcher_context_manager(self, client, target, temp_rows from google.cloud.bigtable.data.mutations import RowMutationEntry new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value ) - row_key2, mutation2 = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value2 + row_key2, mutation2 = await temp_rows.create_row_and_mutation( + target, new_value=new_value2 ) bulk_mutation = RowMutationEntry(row_key, [mutation]) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) @@ -457,7 +416,7 @@ async def test_mutations_batcher_context_manager(self, client, target, temp_rows await batcher.append(bulk_mutation) await batcher.append(bulk_mutation2) # ensure cell is updated - assert (await self._retrieve_cell_value(target, row_key)) == new_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value assert len(batcher._staged_entries) == 0 @pytest.mark.usefixtures("client") @@ -473,8 +432,8 @@ async def test_mutations_batcher_timer_flush(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry new_value = uuid.uuid4().hex.encode() - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) flush_interval = 0.1 @@ -485,7 +444,7 @@ async def test_mutations_batcher_timer_flush(self, client, target, temp_rows): await CrossSync.sleep(flush_interval + 0.1) assert len(batcher._staged_entries) == 0 # ensure cell is updated - assert (await self._retrieve_cell_value(target, row_key)) == new_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -500,12 +459,12 @@ async def test_mutations_batcher_count_flush(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) - row_key2, mutation2 = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value2 + row_key2, mutation2 = await temp_rows.create_row_and_mutation( + target, new_value=new_value2 ) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) @@ -525,8 +484,8 @@ async def test_mutations_batcher_count_flush(self, client, target, temp_rows): assert len(batcher._staged_entries) == 0 assert len(batcher._flush_jobs) == 0 # ensure cells were updated - assert (await self._retrieve_cell_value(target, row_key)) == new_value - assert (await self._retrieve_cell_value(target, row_key2)) == new_value2 + assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value + assert (await temp_rows.retrieve_cell_value(target, row_key2)) == new_value2 @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -541,12 +500,12 @@ async def test_mutations_batcher_bytes_flush(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) - row_key2, mutation2 = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value2 + row_key2, mutation2 = await temp_rows.create_row_and_mutation( + target, new_value=new_value2 ) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) @@ -566,8 +525,8 @@ async def test_mutations_batcher_bytes_flush(self, client, target, temp_rows): # for sync version: grab result future.result() # ensure cells were updated - assert (await self._retrieve_cell_value(target, row_key)) == new_value - assert (await self._retrieve_cell_value(target, row_key2)) == new_value2 + assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value + assert (await temp_rows.retrieve_cell_value(target, row_key2)) == new_value2 @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -580,12 +539,12 @@ async def test_mutations_batcher_no_flush(self, client, target, temp_rows): new_value = uuid.uuid4().hex.encode() start_value = b"unchanged" - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, start_value=start_value, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, start_value=start_value, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) - row_key2, mutation2 = await self._create_row_and_mutation( - target, temp_rows, start_value=start_value, new_value=new_value + row_key2, mutation2 = await temp_rows.create_row_and_mutation( + target, start_value=start_value, new_value=new_value ) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) @@ -602,8 +561,10 @@ async def test_mutations_batcher_no_flush(self, client, target, temp_rows): assert len(batcher._staged_entries) == 2 assert len(batcher._flush_jobs) == 0 # ensure cells were not updated - assert (await self._retrieve_cell_value(target, row_key)) == start_value - assert (await self._retrieve_cell_value(target, row_key2)) == start_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == start_value + assert ( + await temp_rows.retrieve_cell_value(target, row_key2) + ) == start_value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -674,7 +635,7 @@ async def test_read_modify_write_row_increment( assert result[0].qualifier == qualifier assert int(result[0]) == expected # ensure that reading from server gives same value - assert (await self._retrieve_cell_value(target, row_key)) == result[0].value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == result[0].value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -714,7 +675,7 @@ async def test_read_modify_write_row_append( assert result[0].qualifier == qualifier assert result[0].value == expected # ensure that reading from server gives same value - assert (await self._retrieve_cell_value(target, row_key)) == result[0].value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == result[0].value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -751,7 +712,7 @@ async def test_read_modify_write_row_chained(self, client, target, temp_rows): + b"helloworld!" ) # ensure that reading from server gives same value - assert (await self._retrieve_cell_value(target, row_key)) == result[0].value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == result[0].value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -800,7 +761,7 @@ async def test_check_and_mutate( expected_value = ( true_mutation_value if expected_result else false_mutation_value ) - assert (await self._retrieve_cell_value(target, row_key)) == expected_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == expected_value @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), diff --git a/tests/system/data/test_system_autogen.py b/tests/system/data/test_system_autogen.py index 37c00f2ae..66ca27a66 100644 --- a/tests/system/data/test_system_autogen.py +++ b/tests/system/data/test_system_autogen.py @@ -26,7 +26,7 @@ from google.cloud.environment_vars import BIGTABLE_EMULATOR from google.type import date_pb2 from google.cloud.bigtable.data._cross_sync import CrossSync -from . import TEST_FAMILY, TEST_FAMILY_2, TEST_AGGREGATE_FAMILY +from . import TEST_FAMILY, TEST_FAMILY_2, TEST_AGGREGATE_FAMILY, SystemTestRunner from google.cloud.bigtable_v2.services.bigtable.transports.grpc import ( _LoggingClientInterceptor as GapicInterceptor, ) @@ -100,8 +100,32 @@ def delete_rows(self): } self.target.client._gapic_client.mutate_rows(request) + def retrieve_cell_value(self, target, row_key): + """Helper to read an individual row""" + from google.cloud.bigtable.data import ReadRowsQuery + + row_list = target.read_rows(ReadRowsQuery(row_keys=row_key)) + assert len(row_list) == 1 + row = row_list[0] + cell = row.cells[0] + return cell.value -class TestSystem: + def create_row_and_mutation( + self, table, *, start_value=b"start", new_value=b"new_value" + ): + """Helper to create a new row, and a sample set_cell mutation to change its value""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = uuid.uuid4().hex.encode() + family = TEST_FAMILY + qualifier = b"test-qualifier" + self.add_row(row_key, family=family, qualifier=qualifier, value=start_value) + assert self.retrieve_cell_value(table, row_key) == start_value + mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) + return (row_key, mutation) + + +class TestSystem(SystemTestRunner): def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None return CrossSync._Sync_Impl.DataClient(project=project) @@ -127,67 +151,6 @@ def target(self, client, table_id, authorized_view_id, instance_id, request): else: raise ValueError(f"unknown target type: {request.param}") - @pytest.fixture(scope="session") - def column_family_config(self): - """specify column families to create when creating a new test table""" - from google.cloud.bigtable_admin_v2 import types - - int_aggregate_type = types.Type.Aggregate( - input_type=types.Type(int64_type={"encoding": {"big_endian_bytes": {}}}), - sum={}, - ) - return { - TEST_FAMILY: types.ColumnFamily(), - TEST_FAMILY_2: types.ColumnFamily(), - TEST_AGGREGATE_FAMILY: types.ColumnFamily( - value_type=types.Type(aggregate_type=int_aggregate_type) - ), - } - - @pytest.fixture(scope="session") - def init_table_id(self): - """The table_id to use when creating a new test table""" - return f"test-table-{uuid.uuid4().hex}" - - @pytest.fixture(scope="session") - def cluster_config(self, project_id): - """Configuration for the clusters to use when creating a new instance""" - from google.cloud.bigtable_admin_v2 import types - - cluster = { - "test-cluster": types.Cluster( - location=f"projects/{project_id}/locations/us-central1-b", serve_nodes=1 - ) - } - return cluster - - @pytest.mark.usefixtures("target") - def _retrieve_cell_value(self, target, row_key): - """Helper to read an individual row""" - from google.cloud.bigtable.data import ReadRowsQuery - - row_list = target.read_rows(ReadRowsQuery(row_keys=row_key)) - assert len(row_list) == 1 - row = row_list[0] - cell = row.cells[0] - return cell.value - - def _create_row_and_mutation( - self, table, temp_rows, *, start_value=b"start", new_value=b"new_value" - ): - """Helper to create a new row, and a sample set_cell mutation to change its value""" - from google.cloud.bigtable.data.mutations import SetCell - - row_key = uuid.uuid4().hex.encode() - family = TEST_FAMILY - qualifier = b"test-qualifier" - temp_rows.add_row( - row_key, family=family, qualifier=qualifier, value=start_value - ) - assert self._retrieve_cell_value(table, row_key) == start_value - mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) - return (row_key, mutation) - @pytest.fixture(scope="function") def temp_rows(self, target): builder = CrossSync._Sync_Impl.TempRowBuilder(target) @@ -260,11 +223,11 @@ def test_mutation_set_cell(self, target, temp_rows): """Ensure cells can be set properly""" row_key = b"bulk_mutate" new_value = uuid.uuid4().hex.encode() - (row_key, mutation) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value + (row_key, mutation) = temp_rows.create_row_and_mutation( + target, new_value=new_value ) target.mutate_row(row_key, mutation) - assert self._retrieve_cell_value(target, row_key) == new_value + assert temp_rows.retrieve_cell_value(target, row_key) == new_value @pytest.mark.usefixtures("target") @CrossSync._Sync_Impl.Retry( @@ -279,11 +242,11 @@ def test_mutation_add_to_cell(self, target, temp_rows): qualifier = b"test-qualifier" temp_rows.add_aggregate_row(row_key, family=family, qualifier=qualifier) target.mutate_row(row_key, AddToCell(family, qualifier, 1, timestamp_micros=0)) - encoded_result = self._retrieve_cell_value(target, row_key) + encoded_result = temp_rows.retrieve_cell_value(target, row_key) int_result = int.from_bytes(encoded_result, byteorder="big") assert int_result == 1 target.mutate_row(row_key, AddToCell(family, qualifier, 9, timestamp_micros=0)) - encoded_result = self._retrieve_cell_value(target, row_key) + encoded_result = temp_rows.retrieve_cell_value(target, row_key) int_result = int.from_bytes(encoded_result, byteorder="big") assert int_result == 10 @@ -314,12 +277,12 @@ def test_bulk_mutations_set_cell(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry new_value = uuid.uuid4().hex.encode() - (row_key, mutation) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value + (row_key, mutation) = temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) target.bulk_mutate_rows([bulk_mutation]) - assert self._retrieve_cell_value(target, row_key) == new_value + assert temp_rows.retrieve_cell_value(target, row_key) == new_value def test_bulk_mutations_raise_exception(self, client, target): """If an invalid mutation is passed, an exception should be raised""" @@ -350,18 +313,18 @@ def test_mutations_batcher_context_manager(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry (new_value, new_value2) = [uuid.uuid4().hex.encode() for _ in range(2)] - (row_key, mutation) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value + (row_key, mutation) = temp_rows.create_row_and_mutation( + target, new_value=new_value ) - (row_key2, mutation2) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value2 + (row_key2, mutation2) = temp_rows.create_row_and_mutation( + target, new_value=new_value2 ) bulk_mutation = RowMutationEntry(row_key, [mutation]) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) with target.mutations_batcher() as batcher: batcher.append(bulk_mutation) batcher.append(bulk_mutation2) - assert self._retrieve_cell_value(target, row_key) == new_value + assert temp_rows.retrieve_cell_value(target, row_key) == new_value assert len(batcher._staged_entries) == 0 @pytest.mark.usefixtures("client") @@ -374,8 +337,8 @@ def test_mutations_batcher_timer_flush(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry new_value = uuid.uuid4().hex.encode() - (row_key, mutation) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value + (row_key, mutation) = temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) flush_interval = 0.1 @@ -385,7 +348,7 @@ def test_mutations_batcher_timer_flush(self, client, target, temp_rows): assert len(batcher._staged_entries) == 1 CrossSync._Sync_Impl.sleep(flush_interval + 0.1) assert len(batcher._staged_entries) == 0 - assert self._retrieve_cell_value(target, row_key) == new_value + assert temp_rows.retrieve_cell_value(target, row_key) == new_value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -397,12 +360,12 @@ def test_mutations_batcher_count_flush(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry (new_value, new_value2) = [uuid.uuid4().hex.encode() for _ in range(2)] - (row_key, mutation) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value + (row_key, mutation) = temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) - (row_key2, mutation2) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value2 + (row_key2, mutation2) = temp_rows.create_row_and_mutation( + target, new_value=new_value2 ) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) with target.mutations_batcher(flush_limit_mutation_count=2) as batcher: @@ -416,8 +379,8 @@ def test_mutations_batcher_count_flush(self, client, target, temp_rows): future.result() assert len(batcher._staged_entries) == 0 assert len(batcher._flush_jobs) == 0 - assert self._retrieve_cell_value(target, row_key) == new_value - assert self._retrieve_cell_value(target, row_key2) == new_value2 + assert temp_rows.retrieve_cell_value(target, row_key) == new_value + assert temp_rows.retrieve_cell_value(target, row_key2) == new_value2 @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -429,12 +392,12 @@ def test_mutations_batcher_bytes_flush(self, client, target, temp_rows): from google.cloud.bigtable.data.mutations import RowMutationEntry (new_value, new_value2) = [uuid.uuid4().hex.encode() for _ in range(2)] - (row_key, mutation) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value + (row_key, mutation) = temp_rows.create_row_and_mutation( + target, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) - (row_key2, mutation2) = self._create_row_and_mutation( - target, temp_rows, new_value=new_value2 + (row_key2, mutation2) = temp_rows.create_row_and_mutation( + target, new_value=new_value2 ) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) flush_limit = bulk_mutation.size() + bulk_mutation2.size() - 1 @@ -448,8 +411,8 @@ def test_mutations_batcher_bytes_flush(self, client, target, temp_rows): for future in list(batcher._flush_jobs): future future.result() - assert self._retrieve_cell_value(target, row_key) == new_value - assert self._retrieve_cell_value(target, row_key2) == new_value2 + assert temp_rows.retrieve_cell_value(target, row_key) == new_value + assert temp_rows.retrieve_cell_value(target, row_key2) == new_value2 @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -459,12 +422,12 @@ def test_mutations_batcher_no_flush(self, client, target, temp_rows): new_value = uuid.uuid4().hex.encode() start_value = b"unchanged" - (row_key, mutation) = self._create_row_and_mutation( - target, temp_rows, start_value=start_value, new_value=new_value + (row_key, mutation) = temp_rows.create_row_and_mutation( + target, start_value=start_value, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) - (row_key2, mutation2) = self._create_row_and_mutation( - target, temp_rows, start_value=start_value, new_value=new_value + (row_key2, mutation2) = temp_rows.create_row_and_mutation( + target, start_value=start_value, new_value=new_value ) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) size_limit = bulk_mutation.size() + bulk_mutation2.size() + 1 @@ -478,8 +441,8 @@ def test_mutations_batcher_no_flush(self, client, target, temp_rows): CrossSync._Sync_Impl.yield_to_event_loop() assert len(batcher._staged_entries) == 2 assert len(batcher._flush_jobs) == 0 - assert self._retrieve_cell_value(target, row_key) == start_value - assert self._retrieve_cell_value(target, row_key2) == start_value + assert temp_rows.retrieve_cell_value(target, row_key) == start_value + assert temp_rows.retrieve_cell_value(target, row_key2) == start_value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -537,7 +500,7 @@ def test_read_modify_write_row_increment( assert result[0].family == family assert result[0].qualifier == qualifier assert int(result[0]) == expected - assert self._retrieve_cell_value(target, row_key) == result[0].value + assert temp_rows.retrieve_cell_value(target, row_key) == result[0].value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -570,7 +533,7 @@ def test_read_modify_write_row_append( assert result[0].family == family assert result[0].qualifier == qualifier assert result[0].value == expected - assert self._retrieve_cell_value(target, row_key) == result[0].value + assert temp_rows.retrieve_cell_value(target, row_key) == result[0].value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -602,7 +565,7 @@ def test_read_modify_write_row_chained(self, client, target, temp_rows): == (start_amount + increment_amount).to_bytes(8, "big", signed=True) + b"helloworld!" ) - assert self._retrieve_cell_value(target, row_key) == result[0].value + assert temp_rows.retrieve_cell_value(target, row_key) == result[0].value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") @@ -640,7 +603,7 @@ def test_check_and_mutate( expected_value = ( true_mutation_value if expected_result else false_mutation_value ) - assert self._retrieve_cell_value(target, row_key) == expected_value + assert temp_rows.retrieve_cell_value(target, row_key) == expected_value @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index f14fa6dee..a43b0a35f 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -17,6 +17,7 @@ from google.cloud.bigtable_v2.types import MutateRowsResponse from google.cloud.bigtable.data.mutations import RowMutationEntry from google.cloud.bigtable.data.mutations import DeleteAllFromRow +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.rpc import status_pb2 from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import Forbidden @@ -48,6 +49,9 @@ def _make_one(self, *args, **kwargs): kwargs["attempt_timeout"] = kwargs.pop("attempt_timeout", 0.1) kwargs["retryable_exceptions"] = kwargs.pop("retryable_exceptions", ()) kwargs["mutation_entries"] = kwargs.pop("mutation_entries", []) + kwargs["metric"] = kwargs.pop( + "metric", ActiveOperationMetric("MUTATE_ROWS") + ) return self._target_class()(*args, **kwargs) def _make_mutation(self, count=1, size=1): @@ -90,6 +94,7 @@ def test_ctor(self): entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 attempt_timeout = 0.01 + metric = mock.Mock() retryable_exceptions = () instance = self._make_one( client, @@ -97,6 +102,7 @@ def test_ctor(self): entries, operation_timeout, attempt_timeout, + metric, retryable_exceptions, ) # running gapic_fn should trigger a client call with baked-in args @@ -116,6 +122,7 @@ def test_ctor(self): assert instance.is_retryable(RuntimeError("")) is False assert instance.remaining_indices == list(range(len(entries))) assert instance.errors == {} + assert instance._operation_metric == metric def test_ctor_too_many_entries(self): """ @@ -139,6 +146,7 @@ def test_ctor_too_many_entries(self): entries, operation_timeout, attempt_timeout, + mock.Mock(), ) assert "mutate_rows requests can contain at most 100000 mutations" in str( e.value @@ -152,6 +160,7 @@ async def test_mutate_rows_operation(self): """ client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 cls = self._target_class() @@ -159,7 +168,7 @@ async def test_mutate_rows_operation(self): f"{cls.__module__}.{cls.__name__}._run_attempt", CrossSync.Mock() ) as attempt_mock: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) await instance.start() assert attempt_mock.call_count == 1 @@ -173,6 +182,7 @@ async def test_mutate_rows_attempt_exception(self, exc_type): client = CrossSync.Mock() table = mock.Mock() table._request_path = {"table_name": "table"} + metric = ActiveOperationMetric("MUTATE_ROWS") table.app_profile_id = None entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 @@ -181,7 +191,7 @@ async def test_mutate_rows_attempt_exception(self, exc_type): found_exc = None try: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) await instance._run_attempt() except Exception as e: @@ -203,6 +213,7 @@ async def test_mutate_rows_exception(self, exc_type): client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 expected_cause = exc_type("abort") @@ -215,7 +226,7 @@ async def test_mutate_rows_exception(self, exc_type): found_exc = None try: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) await instance.start() except MutationsExceptionGroup as e: @@ -239,6 +250,7 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation()] operation_timeout = 1 expected_cause = exc_type("retry") @@ -255,6 +267,7 @@ async def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): entries, operation_timeout, operation_timeout, + metric, retryable_exceptions=(exc_type,), ) await instance.start() @@ -271,6 +284,7 @@ async def test_mutate_rows_incomplete_ignored(self): client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation()] operation_timeout = 0.05 with mock.patch.object( @@ -282,7 +296,7 @@ async def test_mutate_rows_incomplete_ignored(self): found_exc = None try: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) await instance.start() except MutationsExceptionGroup as e: diff --git a/tests/unit/data/_async/test__read_rows.py b/tests/unit/data/_async/test__read_rows.py index c43f46d5a..4bf0f933e 100644 --- a/tests/unit/data/_async/test__read_rows.py +++ b/tests/unit/data/_async/test__read_rows.py @@ -15,6 +15,7 @@ import pytest from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data._metrics import ActiveOperationMetric # try/except added for compatibility with python < 3.8 try: @@ -59,6 +60,7 @@ def test_ctor(self): expected_operation_timeout = 42 expected_request_timeout = 44 time_gen_mock = mock.Mock() + expected_metric = mock.Mock() subpath = "_async" if CrossSync.is_async else "_sync_autogen" with mock.patch( f"google.cloud.bigtable.data.{subpath}._read_rows._attempt_timeout_generator", @@ -69,6 +71,7 @@ def test_ctor(self): table, operation_timeout=expected_operation_timeout, attempt_timeout=expected_request_timeout, + metric=expected_metric, ) assert time_gen_mock.call_count == 1 time_gen_mock.assert_called_once_with( @@ -81,6 +84,7 @@ def test_ctor(self): assert instance.request.table_name == "test_table" assert instance.request.app_profile_id == table.app_profile_id assert instance.request.rows_limit == row_limit + assert instance._operation_metric == expected_metric @pytest.mark.parametrize( "in_keys,last_key,expected", @@ -269,7 +273,9 @@ async def mock_stream(): table = mock.Mock() table._request_path = {"table_name": "table_name"} table.app_profile_id = "app_profile_id" - instance = self._make_one(query, table, 10, 10) + instance = self._make_one( + query, table, 10, 10, ActiveOperationMetric("READ_ROWS") + ) assert instance._remaining_count == start_limit # read emit_num rows async for val in instance.chunk_stream(awaitable_stream()): @@ -308,7 +314,9 @@ async def mock_stream(): table = mock.Mock() table._request_path = {"table_name": "table_name"} table.app_profile_id = "app_profile_id" - instance = self._make_one(query, table, 10, 10) + instance = self._make_one( + query, table, 10, 10, ActiveOperationMetric("READ_ROWS") + ) assert instance._remaining_count == start_limit with pytest.raises(InvalidChunk) as e: # read emit_num rows @@ -334,7 +342,9 @@ async def mock_stream(): with mock.patch.object( self._get_target_class(), "_read_rows_attempt" ) as mock_attempt: - instance = self._make_one(mock.Mock(), mock.Mock(), 1, 1) + instance = self._make_one( + mock.Mock(), mock.Mock(), 1, 1, ActiveOperationMetric("READ_ROWS") + ) wrapped_gen = mock_stream() mock_attempt.return_value = wrapped_gen gen = instance.start_operation() diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 2cae7a08c..b24f1f47e 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1971,9 +1971,21 @@ async def test_read_row(self): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() expected_result = object() - read_rows.side_effect = lambda *args, **kwargs: [expected_result] + + if CrossSync.is_async: + + async def mock_generator(): + yield expected_result + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = [expected_result] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 row = await table.read_row( @@ -1982,16 +1994,17 @@ async def test_read_row(self): attempt_timeout=expected_req_timeout, ) assert row == expected_result - assert read_rows.call_count == 1 - args, kwargs = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + args, kwargs = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout - assert len(args) == 1 + assert len(args) == 2 assert isinstance(args[0], ReadRowsQuery) query = args[0] assert query.row_keys == [row_key] assert query.row_ranges == [] assert query.limit == 1 + assert args[1] is table @CrossSync.pytest async def test_read_row_w_filter(self): @@ -1999,14 +2012,24 @@ async def test_read_row_w_filter(self): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() expected_result = object() - read_rows.side_effect = lambda *args, **kwargs: [expected_result] + + if CrossSync.is_async: + + async def mock_generator(): + yield expected_result + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = [expected_result] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 - mock_filter = mock.Mock() - expected_filter = {"filter": "mock filter"} - mock_filter._to_dict.return_value = expected_filter + expected_filter = mock.Mock() row = await table.read_row( row_key, operation_timeout=expected_op_timeout, @@ -2014,11 +2037,11 @@ async def test_read_row_w_filter(self): row_filter=expected_filter, ) assert row == expected_result - assert read_rows.call_count == 1 - args, kwargs = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + args, kwargs = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout - assert len(args) == 1 + assert len(args) == 2 assert isinstance(args[0], ReadRowsQuery) query = args[0] assert query.row_keys == [row_key] @@ -2032,9 +2055,21 @@ async def test_read_row_no_response(self): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: - # return no rows - read_rows.side_effect = lambda *args, **kwargs: [] + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() + + if CrossSync.is_async: + + async def mock_generator(): + if False: + yield + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = [] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 result = await table.read_row( @@ -2043,8 +2078,8 @@ async def test_read_row_no_response(self): attempt_timeout=expected_req_timeout, ) assert result is None - assert read_rows.call_count == 1 - args, kwargs = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + args, kwargs = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout assert isinstance(args[0], ReadRowsQuery) @@ -2067,22 +2102,36 @@ async def test_row_exists(self, return_value, expected_result): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: - # return no rows - read_rows.side_effect = lambda *args, **kwargs: return_value - expected_op_timeout = 1 - expected_req_timeout = 2 + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() + if CrossSync.is_async: + + async def mock_generator(): + for item in return_value: + yield item + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = return_value + mock_op_constructor.return_value = mock_op + expected_op_timeout = 2 + expected_req_timeout = 1 result = await table.row_exists( row_key, operation_timeout=expected_op_timeout, attempt_timeout=expected_req_timeout, ) assert expected_result == result - assert read_rows.call_count == 1 - args, kwargs = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + args, kwargs = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout - assert isinstance(args[0], ReadRowsQuery) + query = args[0] + assert isinstance(query, ReadRowsQuery) + assert query.row_keys == [row_key] + assert query.limit == 1 expected_filter = { "chain": { "filters": [ @@ -2091,10 +2140,6 @@ async def test_row_exists(self, return_value, expected_result): ] } } - query = args[0] - assert query.row_keys == [row_key] - assert query.row_ranges == [] - assert query.limit == 1 assert query.filter._to_dict() == expected_filter diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index 29f2f1026..d94b1e98c 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -307,6 +307,9 @@ def _get_target_class(self): def _make_one(self, table=None, **kwargs): from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import ServiceUnavailable + from google.cloud.bigtable.data._metrics import ( + BigtableClientSideMetricsController, + ) if table is None: table = mock.Mock() @@ -318,6 +321,7 @@ def _make_one(self, table=None, **kwargs): DeadlineExceeded, ServiceUnavailable, ) + table._metrics = BigtableClientSideMetricsController([]) return self._get_target_class()(table, **kwargs) @@ -935,14 +939,16 @@ async def test__execute_mutate_rows(self): table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: batch = [self._make_mutation()] - result = await instance._execute_mutate_rows(batch) + expected_metric = mock.Mock() + result = await instance._execute_mutate_rows(batch, expected_metric) 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 - kwargs["operation_timeout"] == 17 - kwargs["attempt_timeout"] == 13 + assert kwargs["operation_timeout"] == 17 + assert kwargs["attempt_timeout"] == 13 + assert kwargs["metric"] == expected_metric assert result == [] @CrossSync.pytest @@ -963,7 +969,7 @@ async def test__execute_mutate_rows_returns_errors(self): table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: batch = [self._make_mutation()] - result = await instance._execute_mutate_rows(batch) + result = await instance._execute_mutate_rows(batch, mock.Mock()) assert len(result) == 2 assert result[0] == err1 assert result[1] == err2 @@ -1093,7 +1099,9 @@ async def test_timeout_args_passed(self): assert instance._operation_timeout == expected_operation_timeout assert instance._attempt_timeout == expected_attempt_timeout # make simulated gapic call - await instance._execute_mutate_rows([self._make_mutation()]) + await instance._execute_mutate_rows( + [self._make_mutation()], mock.Mock() + ) assert mutate_rows.call_count == 1 kwargs = mutate_rows.call_args[1] assert kwargs["operation_timeout"] == expected_operation_timeout @@ -1191,6 +1199,8 @@ async def test_customizable_retryable_errors( Test that retryable functions support user-configurable arguments, and that the configured retryables are passed down to the gapic layer. """ + from google.cloud.bigtable.data._metrics import ActiveOperationMetric + with mock.patch.object( google.api_core.retry, "if_exception_type" ) as predicate_builder_mock: @@ -1206,7 +1216,9 @@ async def test_customizable_retryable_errors( predicate_builder_mock.return_value = expected_predicate retry_fn_mock.side_effect = RuntimeError("stop early") mutation = self._make_mutation(count=1, size=1) - await instance._execute_mutate_rows([mutation]) + await instance._execute_mutate_rows( + [mutation], ActiveOperationMetric("MUTATE_ROWS") + ) # passed in errors should be used to build the predicate predicate_builder_mock.assert_called_once_with( *expected_retryables, _MutateRowsIncomplete diff --git a/tests/unit/data/_async/test_read_rows_acceptance.py b/tests/unit/data/_async/test_read_rows_acceptance.py index ab9502223..1fbbd7c82 100644 --- a/tests/unit/data/_async/test_read_rows_acceptance.py +++ b/tests/unit/data/_async/test_read_rows_acceptance.py @@ -24,6 +24,7 @@ from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable.data.row import Row +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from ...v2_client.test_row_merger import ReadRowsTest, TestFile @@ -39,8 +40,11 @@ class TestReadRowsAcceptanceAsync: @staticmethod @CrossSync.convert - def _get_operation_class(): - return CrossSync._ReadRowsOperation + def _make_operation(): + metric = ActiveOperationMetric("READ_ROWS") + op = CrossSync._ReadRowsOperation(mock.Mock(), mock.Mock(), 5, 5, metric) + op._remaining_count = None + return op @staticmethod @CrossSync.convert @@ -83,13 +87,8 @@ async def _process_chunks(self, *chunks): async def _row_stream(): yield ReadRowsResponse(chunks=chunks) - instance = mock.Mock() - instance._remaining_count = None - instance._last_yielded_row_key = None - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_row_stream()) - ) - merger = self._get_operation_class().merge_rows(chunker) + chunker = self._make_operation().chunk_stream(self._coro_wrapper(_row_stream())) + merger = self._make_operation().merge_rows(chunker) results = [] async for row in merger: results.append(row) @@ -106,13 +105,10 @@ async def _scenerio_stream(): try: results = [] - instance = mock.Mock() - instance._last_yielded_row_key = None - instance._remaining_count = None - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_scenerio_stream()) + chunker = self._make_operation().chunk_stream( + self._coro_wrapper(_scenerio_stream()) ) - merger = self._get_operation_class().merge_rows(chunker) + merger = self._make_operation().merge_rows(chunker) async for row in merger: for cell in row: cell_result = ReadRowsTest.Result( @@ -199,13 +195,10 @@ async def test_out_of_order_rows(self): async def _row_stream(): yield ReadRowsResponse(last_scanned_row_key=b"a") - instance = mock.Mock() - instance._remaining_count = None - instance._last_yielded_row_key = b"b" - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_row_stream()) - ) - merger = self._get_operation_class().merge_rows(chunker) + op = self._make_operation() + op._last_yielded_row_key = b"b" + chunker = op.chunk_stream(self._coro_wrapper(_row_stream())) + merger = self._make_operation().merge_rows(chunker) with pytest.raises(InvalidChunk): async for _ in merger: pass diff --git a/tests/unit/data/_sync_autogen/test__mutate_rows.py b/tests/unit/data/_sync_autogen/test__mutate_rows.py index b198df01b..248c01a74 100644 --- a/tests/unit/data/_sync_autogen/test__mutate_rows.py +++ b/tests/unit/data/_sync_autogen/test__mutate_rows.py @@ -19,6 +19,7 @@ from google.cloud.bigtable_v2.types import MutateRowsResponse from google.cloud.bigtable.data.mutations import RowMutationEntry from google.cloud.bigtable.data.mutations import DeleteAllFromRow +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.rpc import status_pb2 from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import Forbidden @@ -45,6 +46,9 @@ def _make_one(self, *args, **kwargs): kwargs["attempt_timeout"] = kwargs.pop("attempt_timeout", 0.1) kwargs["retryable_exceptions"] = kwargs.pop("retryable_exceptions", ()) kwargs["mutation_entries"] = kwargs.pop("mutation_entries", []) + kwargs["metric"] = kwargs.pop( + "metric", ActiveOperationMetric("MUTATE_ROWS") + ) return self._target_class()(*args, **kwargs) def _make_mutation(self, count=1, size=1): @@ -84,6 +88,7 @@ def test_ctor(self): entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 attempt_timeout = 0.01 + metric = mock.Mock() retryable_exceptions = () instance = self._make_one( client, @@ -91,6 +96,7 @@ def test_ctor(self): entries, operation_timeout, attempt_timeout, + metric, retryable_exceptions, ) assert client.mutate_rows.call_count == 0 @@ -106,6 +112,7 @@ def test_ctor(self): assert instance.is_retryable(RuntimeError("")) is False assert instance.remaining_indices == list(range(len(entries))) assert instance.errors == {} + assert instance._operation_metric == metric def test_ctor_too_many_entries(self): """should raise an error if an operation is created with more than 100,000 entries""" @@ -120,7 +127,9 @@ def test_ctor_too_many_entries(self): operation_timeout = 0.05 attempt_timeout = 0.01 with pytest.raises(ValueError) as e: - self._make_one(client, table, entries, operation_timeout, attempt_timeout) + self._make_one( + client, table, entries, operation_timeout, attempt_timeout, mock.Mock() + ) assert "mutate_rows requests can contain at most 100000 mutations" in str( e.value ) @@ -130,6 +139,7 @@ def test_mutate_rows_operation(self): """Test successful case of mutate_rows_operation""" client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 cls = self._target_class() @@ -137,7 +147,7 @@ def test_mutate_rows_operation(self): f"{cls.__module__}.{cls.__name__}._run_attempt", CrossSync._Sync_Impl.Mock() ) as attempt_mock: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) instance.start() assert attempt_mock.call_count == 1 @@ -148,6 +158,7 @@ def test_mutate_rows_attempt_exception(self, exc_type): client = CrossSync._Sync_Impl.Mock() table = mock.Mock() table._request_path = {"table_name": "table"} + metric = ActiveOperationMetric("MUTATE_ROWS") table.app_profile_id = None entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 @@ -156,7 +167,7 @@ def test_mutate_rows_attempt_exception(self, exc_type): found_exc = None try: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) instance._run_attempt() except Exception as e: @@ -175,6 +186,7 @@ def test_mutate_rows_exception(self, exc_type): client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation(), self._make_mutation()] operation_timeout = 0.05 expected_cause = exc_type("abort") @@ -185,7 +197,7 @@ def test_mutate_rows_exception(self, exc_type): found_exc = None try: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) instance.start() except MutationsExceptionGroup as e: @@ -202,6 +214,7 @@ def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): """If an exception fails but eventually passes, it should not raise an exception""" client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation()] operation_timeout = 1 expected_cause = exc_type("retry") @@ -216,6 +229,7 @@ def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): entries, operation_timeout, operation_timeout, + metric, retryable_exceptions=(exc_type,), ) instance.start() @@ -229,6 +243,7 @@ def test_mutate_rows_incomplete_ignored(self): client = mock.Mock() table = mock.Mock() + metric = ActiveOperationMetric("MUTATE_ROWS") entries = [self._make_mutation()] operation_timeout = 0.05 with mock.patch.object( @@ -238,7 +253,7 @@ def test_mutate_rows_incomplete_ignored(self): found_exc = None try: instance = self._make_one( - client, table, entries, operation_timeout, operation_timeout + client, table, entries, operation_timeout, operation_timeout, metric ) instance.start() except MutationsExceptionGroup as e: diff --git a/tests/unit/data/_sync_autogen/test__read_rows.py b/tests/unit/data/_sync_autogen/test__read_rows.py index a545142d3..e6c00f848 100644 --- a/tests/unit/data/_sync_autogen/test__read_rows.py +++ b/tests/unit/data/_sync_autogen/test__read_rows.py @@ -17,6 +17,7 @@ import pytest from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data._metrics import ActiveOperationMetric try: from unittest import mock @@ -53,6 +54,7 @@ def test_ctor(self): expected_operation_timeout = 42 expected_request_timeout = 44 time_gen_mock = mock.Mock() + expected_metric = mock.Mock() subpath = "_async" if CrossSync._Sync_Impl.is_async else "_sync_autogen" with mock.patch( f"google.cloud.bigtable.data.{subpath}._read_rows._attempt_timeout_generator", @@ -63,6 +65,7 @@ def test_ctor(self): table, operation_timeout=expected_operation_timeout, attempt_timeout=expected_request_timeout, + metric=expected_metric, ) assert time_gen_mock.call_count == 1 time_gen_mock.assert_called_once_with( @@ -75,6 +78,7 @@ def test_ctor(self): assert instance.request.table_name == "test_table" assert instance.request.app_profile_id == table.app_profile_id assert instance.request.rows_limit == row_limit + assert instance._operation_metric == expected_metric @pytest.mark.parametrize( "in_keys,last_key,expected", @@ -254,7 +258,9 @@ def mock_stream(): table = mock.Mock() table._request_path = {"table_name": "table_name"} table.app_profile_id = "app_profile_id" - instance = self._make_one(query, table, 10, 10) + instance = self._make_one( + query, table, 10, 10, ActiveOperationMetric("READ_ROWS") + ) assert instance._remaining_count == start_limit for val in instance.chunk_stream(awaitable_stream()): pass @@ -289,7 +295,9 @@ def mock_stream(): table = mock.Mock() table._request_path = {"table_name": "table_name"} table.app_profile_id = "app_profile_id" - instance = self._make_one(query, table, 10, 10) + instance = self._make_one( + query, table, 10, 10, ActiveOperationMetric("READ_ROWS") + ) assert instance._remaining_count == start_limit with pytest.raises(InvalidChunk) as e: for val in instance.chunk_stream(awaitable_stream()): @@ -307,7 +315,9 @@ def mock_stream(): with mock.patch.object( self._get_target_class(), "_read_rows_attempt" ) as mock_attempt: - instance = self._make_one(mock.Mock(), mock.Mock(), 1, 1) + instance = self._make_one( + mock.Mock(), mock.Mock(), 1, 1, ActiveOperationMetric("READ_ROWS") + ) wrapped_gen = mock_stream() mock_attempt.return_value = wrapped_gen gen = instance.start_operation() diff --git a/tests/unit/data/_sync_autogen/test_client.py b/tests/unit/data/_sync_autogen/test_client.py index 42f5388ee..ab624130c 100644 --- a/tests/unit/data/_sync_autogen/test_client.py +++ b/tests/unit/data/_sync_autogen/test_client.py @@ -1632,9 +1632,13 @@ def test_read_row(self): with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: + with mock.patch.object( + CrossSync._Sync_Impl, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() expected_result = object() - read_rows.side_effect = lambda *args, **kwargs: [expected_result] + mock_op.start_operation.return_value = [expected_result] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 row = table.read_row( @@ -1643,30 +1647,33 @@ def test_read_row(self): attempt_timeout=expected_req_timeout, ) assert row == expected_result - assert read_rows.call_count == 1 - (args, kwargs) = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + (args, kwargs) = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout - assert len(args) == 1 + assert len(args) == 2 assert isinstance(args[0], ReadRowsQuery) query = args[0] assert query.row_keys == [row_key] assert query.row_ranges == [] assert query.limit == 1 + assert args[1] is table def test_read_row_w_filter(self): """Test reading a single row with an added filter""" with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: + with mock.patch.object( + CrossSync._Sync_Impl, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() expected_result = object() - read_rows.side_effect = lambda *args, **kwargs: [expected_result] + mock_op.start_operation.return_value = [expected_result] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 - mock_filter = mock.Mock() - expected_filter = {"filter": "mock filter"} - mock_filter._to_dict.return_value = expected_filter + expected_filter = mock.Mock() row = table.read_row( row_key, operation_timeout=expected_op_timeout, @@ -1674,11 +1681,11 @@ def test_read_row_w_filter(self): row_filter=expected_filter, ) assert row == expected_result - assert read_rows.call_count == 1 - (args, kwargs) = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + (args, kwargs) = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout - assert len(args) == 1 + assert len(args) == 2 assert isinstance(args[0], ReadRowsQuery) query = args[0] assert query.row_keys == [row_key] @@ -1691,8 +1698,12 @@ def test_read_row_no_response(self): with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: - read_rows.side_effect = lambda *args, **kwargs: [] + with mock.patch.object( + CrossSync._Sync_Impl, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() + mock_op.start_operation.return_value = [] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 result = table.read_row( @@ -1701,8 +1712,8 @@ def test_read_row_no_response(self): attempt_timeout=expected_req_timeout, ) assert result is None - assert read_rows.call_count == 1 - (args, kwargs) = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + (args, kwargs) = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout assert isinstance(args[0], ReadRowsQuery) @@ -1720,21 +1731,28 @@ def test_row_exists(self, return_value, expected_result): with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: - read_rows.side_effect = lambda *args, **kwargs: return_value - expected_op_timeout = 1 - expected_req_timeout = 2 + with mock.patch.object( + CrossSync._Sync_Impl, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() + mock_op.start_operation.return_value = return_value + mock_op_constructor.return_value = mock_op + expected_op_timeout = 2 + expected_req_timeout = 1 result = table.row_exists( row_key, operation_timeout=expected_op_timeout, attempt_timeout=expected_req_timeout, ) assert expected_result == result - assert read_rows.call_count == 1 - (args, kwargs) = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + (args, kwargs) = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout - assert isinstance(args[0], ReadRowsQuery) + query = args[0] + assert isinstance(query, ReadRowsQuery) + assert query.row_keys == [row_key] + assert query.limit == 1 expected_filter = { "chain": { "filters": [ @@ -1743,10 +1761,6 @@ def test_row_exists(self, return_value, expected_result): ] } } - query = args[0] - assert query.row_keys == [row_key] - assert query.row_ranges == [] - assert query.limit == 1 assert query.filter._to_dict() == expected_filter diff --git a/tests/unit/data/_sync_autogen/test_mutations_batcher.py b/tests/unit/data/_sync_autogen/test_mutations_batcher.py index 72db64146..cb16ecfff 100644 --- a/tests/unit/data/_sync_autogen/test_mutations_batcher.py +++ b/tests/unit/data/_sync_autogen/test_mutations_batcher.py @@ -257,6 +257,9 @@ def _get_target_class(self): def _make_one(self, table=None, **kwargs): from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import ServiceUnavailable + from google.cloud.bigtable.data._metrics import ( + BigtableClientSideMetricsController, + ) if table is None: table = mock.Mock() @@ -268,6 +271,7 @@ def _make_one(self, table=None, **kwargs): DeadlineExceeded, ServiceUnavailable, ) + table._metrics = BigtableClientSideMetricsController([]) return self._get_target_class()(table, **kwargs) @staticmethod @@ -815,14 +819,16 @@ def test__execute_mutate_rows(self): table.default_mutate_rows_retryable_errors = () with self._make_one(table) as instance: batch = [self._make_mutation()] - result = instance._execute_mutate_rows(batch) + expected_metric = mock.Mock() + result = instance._execute_mutate_rows(batch, expected_metric) 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 - kwargs["operation_timeout"] == 17 - kwargs["attempt_timeout"] == 13 + assert kwargs["operation_timeout"] == 17 + assert kwargs["attempt_timeout"] == 13 + assert kwargs["metric"] == expected_metric assert result == [] def test__execute_mutate_rows_returns_errors(self): @@ -844,7 +850,7 @@ def test__execute_mutate_rows_returns_errors(self): table.default_mutate_rows_retryable_errors = () with self._make_one(table) as instance: batch = [self._make_mutation()] - result = instance._execute_mutate_rows(batch) + result = instance._execute_mutate_rows(batch, mock.Mock()) assert len(result) == 2 assert result[0] == err1 assert result[1] == err2 @@ -952,7 +958,7 @@ def test_timeout_args_passed(self): ) as instance: assert instance._operation_timeout == expected_operation_timeout assert instance._attempt_timeout == expected_attempt_timeout - instance._execute_mutate_rows([self._make_mutation()]) + instance._execute_mutate_rows([self._make_mutation()], mock.Mock()) assert mutate_rows.call_count == 1 kwargs = mutate_rows.call_args[1] assert kwargs["operation_timeout"] == expected_operation_timeout @@ -1038,6 +1044,8 @@ def test__add_exceptions(self, limit, in_e, start_e, end_e): def test_customizable_retryable_errors(self, input_retryables, expected_retryables): """Test that retryable functions support user-configurable arguments, and that the configured retryables are passed down to the gapic layer.""" + from google.cloud.bigtable.data._metrics import ActiveOperationMetric + with mock.patch.object( google.api_core.retry, "if_exception_type" ) as predicate_builder_mock: @@ -1055,7 +1063,9 @@ def test_customizable_retryable_errors(self, input_retryables, expected_retryabl predicate_builder_mock.return_value = expected_predicate retry_fn_mock.side_effect = RuntimeError("stop early") mutation = self._make_mutation(count=1, size=1) - instance._execute_mutate_rows([mutation]) + instance._execute_mutate_rows( + [mutation], ActiveOperationMetric("MUTATE_ROWS") + ) predicate_builder_mock.assert_called_once_with( *expected_retryables, _MutateRowsIncomplete ) diff --git a/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py b/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py index 8ceb0daf7..40515a9b1 100644 --- a/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py +++ b/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py @@ -23,14 +23,20 @@ from google.cloud.bigtable_v2 import ReadRowsResponse from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable.data.row import Row +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from ...v2_client.test_row_merger import ReadRowsTest, TestFile from google.cloud.bigtable.data._cross_sync import CrossSync class TestReadRowsAcceptance: @staticmethod - def _get_operation_class(): - return CrossSync._Sync_Impl._ReadRowsOperation + def _make_operation(): + metric = ActiveOperationMetric("READ_ROWS") + op = CrossSync._Sync_Impl._ReadRowsOperation( + mock.Mock(), mock.Mock(), 5, 5, metric + ) + op._remaining_count = None + return op @staticmethod def _get_client_class(): @@ -68,13 +74,8 @@ def _process_chunks(self, *chunks): def _row_stream(): yield ReadRowsResponse(chunks=chunks) - instance = mock.Mock() - instance._remaining_count = None - instance._last_yielded_row_key = None - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_row_stream()) - ) - merger = self._get_operation_class().merge_rows(chunker) + chunker = self._make_operation().chunk_stream(self._coro_wrapper(_row_stream())) + merger = self._make_operation().merge_rows(chunker) results = [] for row in merger: results.append(row) @@ -90,13 +91,10 @@ def _scenerio_stream(): try: results = [] - instance = mock.Mock() - instance._last_yielded_row_key = None - instance._remaining_count = None - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_scenerio_stream()) + chunker = self._make_operation().chunk_stream( + self._coro_wrapper(_scenerio_stream()) ) - merger = self._get_operation_class().merge_rows(chunker) + merger = self._make_operation().merge_rows(chunker) for row in merger: for cell in row: cell_result = ReadRowsTest.Result( @@ -179,13 +177,10 @@ def test_out_of_order_rows(self): def _row_stream(): yield ReadRowsResponse(last_scanned_row_key=b"a") - instance = mock.Mock() - instance._remaining_count = None - instance._last_yielded_row_key = b"b" - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_row_stream()) - ) - merger = self._get_operation_class().merge_rows(chunker) + op = self._make_operation() + op._last_yielded_row_key = b"b" + chunker = op.chunk_stream(self._coro_wrapper(_row_stream())) + merger = self._make_operation().merge_rows(chunker) with pytest.raises(InvalidChunk): for _ in merger: pass