From 0d938893246e4f4cd7b976278e377ff5b1ba7fce Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 4 Aug 2025 16:16:45 -0700 Subject: [PATCH 01/78] added stdout handler to test --- google/cloud/bigtable/data/_async/client.py | 2 + .../data/_metrics/handlers/_stdout.py | 48 +++++++++++++++++++ 2 files changed, 50 insertions(+) create mode 100644 google/cloud/bigtable/data/_metrics/handlers/_stdout.py diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index d63909282..62ec54b76 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.handlers._stdout import _StdoutMetricsHandler from google.cloud.bigtable.data._cross_sync import CrossSync @@ -942,6 +943,7 @@ def __init__( self._metrics = BigtableClientSideMetricsController( client._metrics_interceptor, + handlers=[_StdoutMetricsHandler()], project_id=self.client.project, instance_id=instance_id, table_id=table_id, 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..a2a3624ed --- /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 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}, avg attempts: {avg_attempts:.1f}" + ) + print() \ No newline at end of file From a580fa28614d2494bc92bd95b15da2818256acef Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 4 Aug 2025 16:17:02 -0700 Subject: [PATCH 02/78] instrumented check_and_mutate --- google/cloud/bigtable/data/_async/client.py | 35 ++++++++++++--------- 1 file changed, 20 insertions(+), 15 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 62ec54b76..1de67f786 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._metrics.handlers._stdout import _StdoutMetricsHandler from google.cloud.bigtable.data._cross_sync import CrossSync @@ -1589,21 +1590,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 + + async 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( From bc13b468629508ebbe29a07bc7ef664331c52826 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 4 Aug 2025 16:21:52 -0700 Subject: [PATCH 03/78] added instrumentation to read_modify_write --- google/cloud/bigtable/data/_async/client.py | 32 +++++++++++-------- .../data/_metrics/handlers/_stdout.py | 4 +-- 2 files changed, 20 insertions(+), 16 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 1de67f786..cc4fd9217 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1648,20 +1648,24 @@ 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) + + async 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/_metrics/handlers/_stdout.py b/google/cloud/bigtable/data/_metrics/handlers/_stdout.py index a2a3624ed..fcec94938 100644 --- a/google/cloud/bigtable/data/_metrics/handlers/_stdout.py +++ b/google/cloud/bigtable/data/_metrics/handlers/_stdout.py @@ -38,11 +38,11 @@ def print(self): print("Bigtable Metrics:") for ops_type, ops_list in self._completed_ops.items(): count = len(ops_list) - total_latency = sum([op.duration_ns for op in 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}, avg attempts: {avg_attempts:.1f}" + f"{ops_type}: count: {count}, avg latency: {avg_latency:.2f} ms, avg attempts: {avg_attempts:.1f}" ) print() \ No newline at end of file From 6c3be4641d51510a61707782653514459c591274 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 4 Aug 2025 16:27:19 -0700 Subject: [PATCH 04/78] added instrumentation to sample_row_keys --- google/cloud/bigtable/data/_async/client.py | 40 +++++++++++---------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index cc4fd9217..609436cab 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -78,6 +78,7 @@ 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 _attempt_timeout_generator +from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule @@ -1328,26 +1329,29 @@ 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) + sleep_generator = TrackedBackoffGenerator(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 - ), - timeout=next(attempt_timeout_gen), - retry=None, + async with self._metrics.create_operation( + OperationType.SAMPLE_ROW_KEYS, backoff_generator=sleep_generator + ): + @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, + sleep_generator, + operation_timeout, + exception_factory=_retry_exception_factory, ) - 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( From ca386151d1a25daad0f63e2b2718d4af58cf4c54 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 4 Aug 2025 16:34:30 -0700 Subject: [PATCH 05/78] instrumented mutate_row --- google/cloud/bigtable/data/_async/client.py | 46 +++++++++++---------- 1 file changed, 24 insertions(+), 22 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 609436cab..361769e9a 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1462,28 +1462,30 @@ 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, - ) + sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) + async with self._metrics.create_operation( + OperationType.MUTATE_ROW, backoff_generator=sleep_generator + ): + 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, + ) @CrossSync.convert async def bulk_mutate_rows( From eb82ae9d9dcf0e8a14f8351a3c70f39e8934f188 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 6 Aug 2025 15:02:07 -0700 Subject: [PATCH 06/78] added instrumentation to mutate_rows --- .../bigtable/data/_async/_mutate_rows.py | 67 +++++++++++-------- google/cloud/bigtable/data/_async/client.py | 1 + .../bigtable/data/_async/mutations_batcher.py | 27 +++++++- 3 files changed, 64 insertions(+), 31 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 8e6833bca..518249022 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -22,6 +22,7 @@ import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator # mutate_rows requests are limited to this number of mutations from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT @@ -31,6 +32,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 +70,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 +82,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,7 +102,7 @@ def __init__( # Entry level errors bt_exceptions._MutateRowsIncomplete, ) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) + sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) self._operation = lambda: CrossSync.retry_target( self._run_attempt, self.is_retryable, @@ -112,6 +117,9 @@ 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 + metric.backoff_generator = sleep_generator + self._operation_metric = metric @CrossSync.convert async def start(self): @@ -121,34 +129,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" + async 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 +169,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/client.py b/google/cloud/bigtable/data/_async/client.py index 361769e9a..f7ca76b67 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1539,6 +1539,7 @@ async def bulk_mutate_rows( mutation_entries, operation_timeout, attempt_timeout, + operation=self._metrics.create_operation(OperationType.BULK_MUTATE_ROWS), retryable_exceptions=retryable_excs, ) await operation.start() diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index a8e99ea9e..e679e25f7 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,22 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry] ) yield mutations[start_idx:end_idx] + 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() + try: + value = await inner_generator.__anext__() + except StopAsyncIteration: + metric.cancel() + return + metric.flow_throttling_time = time.monotonic() - flow_start_time + yield value, metric + @CrossSync.convert_class(sync_name="MutationsBatcher") class MutationsBatcherAsync: @@ -353,9 +373,9 @@ 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 +386,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 +407,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() From 996acf29c98d5fd0aec60cd83cc61c385543dc95 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 7 Aug 2025 14:44:49 -0700 Subject: [PATCH 07/78] instrumented read_rows --- .../cloud/bigtable/data/_async/_read_rows.py | 33 +++++++- google/cloud/bigtable/data/_async/client.py | 78 +++++++++++-------- 2 files changed, 76 insertions(+), 35 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 8787bfa71..e085fa833 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -17,6 +17,8 @@ from typing import Sequence, TYPE_CHECKING +import time + 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 @@ -29,6 +31,7 @@ from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.api_core import retry as retries from google.api_core.retry import exponential_sleep_generator @@ -36,6 +39,8 @@ 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]: """ @@ -109,10 +118,13 @@ def start_operation(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ + self._operation_metric.backoff_generator = TrackedBackoffGenerator( + 0.01, 60, multiplier=2 + ) 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, ) @@ -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 @@ -137,12 +150,12 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: ) except _RowSetComplete: # if we've already seen all the rows, we're done - return self.merge_rows(None) + return self.merge_rows(None, self._operation_metric) # revise the limit based on number of rows already yielded if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: - return self.merge_rows(None) + return self.merge_rows(None, self._operation_metric) # create and return a new row merger gapic_stream = self.target.client._gapic_client.read_rows( self.request, @@ -150,7 +163,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: retry=None, ) chunked_stream = self.chunk_stream(gapic_stream) - return self.merge_rows(chunked_stream) + return self.merge_rows(chunked_stream, self._operation_metric) @CrossSync.convert() async def chunk_stream( @@ -210,6 +223,7 @@ async def chunk_stream( ) async def merge_rows( chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None, + operation_metric: ActiveOperationMetric, ) -> CrossSync.Iterable[Row]: """ Merge chunks into rows @@ -222,6 +236,7 @@ async def merge_rows( if chunks is None: return it = chunks.__aiter__() + is_first_row = True # For each row while True: try: @@ -304,7 +319,17 @@ async def merge_rows( Cell(value, row_key, family, qualifier, ts, list(labels)) ) if c.commit_row: + if is_first_row: + # record first row latency in metrics + is_first_row = False + operation_metric.attempt_first_response() + block_time = time.monotonic() yield Row(row_key, cells) + # most metric operations use setters, but this one updates + # the value directly to avoid extra overhead + operation_metric.active_attempt.application_blocking_time_ms += ( # type: ignore + time.monotonic() - block_time + ) * 1000 break c = await it.__anext__() except _ResetRow as e: diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 8c33c360e..05c26e92f 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -125,6 +125,7 @@ if TYPE_CHECKING: from google.cloud.bigtable.data._helpers import RowKeySamples from google.cloud.bigtable.data._helpers import ShardedQuery + from google.cloud.bigtable.data._metrics import ActiveOperationMetric if CrossSync.is_async: from google.cloud.bigtable.data._async.mutations_batcher import ( @@ -1020,14 +1021,18 @@ async def read_rows_stream( ) retryable_excs = _get_retryable_errors(retryable_errors, self) - row_merger = CrossSync._ReadRowsOperation( - query, - self, - operation_timeout=operation_timeout, - attempt_timeout=attempt_timeout, - retryable_exceptions=retryable_excs, - ) - return row_merger.start_operation() + with self._metrics.create_operation( + OperationType.READ_ROWS, streaming=True + ) as operation_metric: + row_merger = CrossSync._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=operation_metric, + retryable_exceptions=retryable_excs, + ) + return row_merger.start_operation() @CrossSync.convert async def read_rows( @@ -1077,7 +1082,9 @@ async def read_rows( ) return [row async for row in row_generator] - @CrossSync.convert + @CrossSync.convert( + replace_symbols={"__anext__": "__next__", "StopAsyncIteration": "StopIteration"} + ) async def read_row( self, row_key: str | bytes, @@ -1117,15 +1124,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( - query, - operation_timeout=operation_timeout, - attempt_timeout=attempt_timeout, - retryable_errors=retryable_errors, + + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self ) - if len(results) == 0: - return None - return results[0] + retryable_excs = _get_retryable_errors(retryable_errors, self) + + with self._metrics.create_operation( + OperationType.READ_ROWS, streaming=False + ) as operation_metric: + row_merger = CrossSync._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=operation_metric, + retryable_exceptions=retryable_excs, + ) + results_generator = row_merger.start_operation() + try: + return results_generator.__anext__() + except StopAsyncIteration: + return None @CrossSync.convert async def read_rows_sharded( @@ -1264,20 +1284,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( @@ -1334,6 +1351,7 @@ async def sample_row_keys( async with self._metrics.create_operation( OperationType.SAMPLE_ROW_KEYS, backoff_generator=sleep_generator ): + @CrossSync.convert async def execute_rpc(): results = await self.client._gapic_client.sample_row_keys( @@ -1598,14 +1616,14 @@ 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 []] - async with self._metrics.create_operation( - OperationType.CHECK_AND_MUTATE - ): + async 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, + 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, @@ -1656,9 +1674,7 @@ async def read_modify_write_row( if not rules: raise ValueError("rules must contain at least one item") - async with self._metrics.create_operation( - OperationType.READ_MODIFY_WRITE - ): + async 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], From 3eae4aa0c5e07d5a9cfbf87c48d8c64222e392bd Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 7 Aug 2025 14:59:02 -0700 Subject: [PATCH 08/78] fixed lint and mypy --- .../bigtable/data/_async/_mutate_rows.py | 2 +- .../cloud/bigtable/data/_async/_read_rows.py | 20 +- google/cloud/bigtable/data/_async/client.py | 5 +- .../bigtable/data/_async/mutations_batcher.py | 17 +- .../bigtable/data/_metrics/data_model.py | 7 +- .../bigtable/data/_metrics/handlers/_base.py | 2 +- .../data/_metrics/handlers/_stdout.py | 2 +- .../data/_sync_autogen/_mutate_rows.py | 59 ++--- .../bigtable/data/_sync_autogen/_read_rows.py | 30 ++- .../bigtable/data/_sync_autogen/client.py | 204 ++++++++++-------- .../data/_sync_autogen/mutations_batcher.py | 33 ++- 11 files changed, 236 insertions(+), 145 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 518249022..c170ff041 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -169,7 +169,7 @@ async def _run_attempt(self): retry after the attempt is complete GoogleAPICallError: if the gapic rpc fails """ - # register attempt start + # 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 diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index e085fa833..7e5ea4136 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -34,11 +34,11 @@ from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator 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 ActiveAttemptMetric from google.cloud.bigtable.data._metrics import ActiveOperationMetric if CrossSync.is_async: @@ -139,7 +139,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ - self._operation_metric.start_attempt() + attempt_metric = 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 @@ -150,12 +150,12 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: ) except _RowSetComplete: # if we've already seen all the rows, we're done - return self.merge_rows(None, self._operation_metric) + return self.merge_rows(None, attempt_metric) # revise the limit based on number of rows already yielded if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: - return self.merge_rows(None, self._operation_metric) + return self.merge_rows(None, attempt_metric) # create and return a new row merger gapic_stream = self.target.client._gapic_client.read_rows( self.request, @@ -163,7 +163,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: retry=None, ) chunked_stream = self.chunk_stream(gapic_stream) - return self.merge_rows(chunked_stream, self._operation_metric) + return self.merge_rows(chunked_stream, attempt_metric) @CrossSync.convert() async def chunk_stream( @@ -223,7 +223,7 @@ async def chunk_stream( ) async def merge_rows( chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None, - operation_metric: ActiveOperationMetric, + attempt_metric: ActiveAttemptMetric, ) -> CrossSync.Iterable[Row]: """ Merge chunks into rows @@ -322,13 +322,13 @@ async def merge_rows( if is_first_row: # record first row latency in metrics is_first_row = False - operation_metric.attempt_first_response() - block_time = time.monotonic() + attempt_metric.attempt_first_response() + 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 - operation_metric.active_attempt.application_blocking_time_ms += ( # type: ignore - time.monotonic() - block_time + attempt_metric.active_attempt.application_blocking_time_ns += ( # type: ignore + time.monotonic_ns() - block_time ) * 1000 break c = await it.__anext__() diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 05c26e92f..969957b62 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -125,7 +125,6 @@ if TYPE_CHECKING: from google.cloud.bigtable.data._helpers import RowKeySamples from google.cloud.bigtable.data._helpers import ShardedQuery - from google.cloud.bigtable.data._metrics import ActiveOperationMetric if CrossSync.is_async: from google.cloud.bigtable.data._async.mutations_batcher import ( @@ -1021,7 +1020,7 @@ async def read_rows_stream( ) retryable_excs = _get_retryable_errors(retryable_errors, self) - with self._metrics.create_operation( + async with self._metrics.create_operation( OperationType.READ_ROWS, streaming=True ) as operation_metric: row_merger = CrossSync._ReadRowsOperation( @@ -1130,7 +1129,7 @@ async def read_row( ) retryable_excs = _get_retryable_errors(retryable_errors, self) - with self._metrics.create_operation( + async with self._metrics.create_operation( OperationType.READ_ROWS, streaming=False ) as operation_metric: row_merger = CrossSync._ReadRowsOperation( diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index e679e25f7..42fbc4d3b 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -182,19 +182,21 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry] yield mutations[start_idx:end_idx] async def add_to_flow_with_metrics( - self, mutations: RowMutationEntry | list[RowMutationEntry], metrics_controller: BigtableClientSideMetricsController + 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() + flow_start_time = time.monotonic_ns() try: value = await inner_generator.__anext__() except StopAsyncIteration: metric.cancel() return - metric.flow_throttling_time = time.monotonic() - flow_start_time + metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time yield value, metric @@ -373,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, metric in self._flow_control.add_to_flow_with_metrics(new_entries, self._target._metrics): + 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, metric, 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 diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 9ff8b11ab..ba6731570 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -13,7 +13,7 @@ # limitations under the License. from __future__ import annotations -from typing import Callable, Tuple, cast, TYPE_CHECKING +from typing import Tuple, cast, TYPE_CHECKING import time import re @@ -184,7 +184,7 @@ def start(self) -> None: return self._handle_error(INVALID_STATE_ERROR.format("start", self.state)) self.start_time_ns = time.monotonic_ns() - def start_attempt(self) -> None: + def start_attempt(self) -> ActiveAttemptMetric: """ Called to initiate a new attempt for the operation. @@ -210,6 +210,7 @@ def start_attempt(self) -> None: backoff_ns = 0 self.active_attempt = ActiveAttemptMetric(backoff_before_attempt_ns=backoff_ns) + return self.active_attempt def add_response_metadata(self, metadata: dict[str, bytes | str]) -> None: """ @@ -428,4 +429,4 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): if exc_val is None: self.end_with_success() else: - self.end_with_status(exc_val) + self.end_with_status(exc_val) \ No newline at end of file diff --git a/google/cloud/bigtable/data/_metrics/handlers/_base.py b/google/cloud/bigtable/data/_metrics/handlers/_base.py index 05132d618..64cc89b05 100644 --- a/google/cloud/bigtable/data/_metrics/handlers/_base.py +++ b/google/cloud/bigtable/data/_metrics/handlers/_base.py @@ -35,4 +35,4 @@ def on_operation_cancelled(self, op: ActiveOperationMetric) -> None: def on_attempt_complete( self, attempt: CompletedAttemptMetric, op: ActiveOperationMetric ) -> None: - pass \ No newline at end of file + pass diff --git a/google/cloud/bigtable/data/_metrics/handlers/_stdout.py b/google/cloud/bigtable/data/_metrics/handlers/_stdout.py index fcec94938..e0a4bc8da 100644 --- a/google/cloud/bigtable/data/_metrics/handlers/_stdout.py +++ b/google/cloud/bigtable/data/_metrics/handlers/_stdout.py @@ -45,4 +45,4 @@ def print(self): print( f"{ops_type}: count: {count}, avg latency: {avg_latency:.2f} ms, avg attempts: {avg_attempts:.1f}" ) - print() \ No newline at end of file + 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..55c2d8668 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -23,12 +23,14 @@ import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT from google.cloud.bigtable.data.mutations import _EntryWithProto 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 ActiveOperationMetric from google.cloud.bigtable_v2.services.bigtable.client import ( BigtableClient as GapicClientType, ) @@ -54,6 +56,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 +67,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,7 +80,7 @@ def __init__( self.is_retryable = retries.if_exception_type( *retryable_exceptions, bt_exceptions._MutateRowsIncomplete ) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) + sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) self._operation = lambda: CrossSync._Sync_Impl.retry_target( self._run_attempt, self.is_retryable, @@ -89,37 +94,40 @@ 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]] = {} + metric.backoff_generator = sleep_generator + 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 +136,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..a704912e9 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -18,6 +18,7 @@ from __future__ import annotations from typing import Sequence, TYPE_CHECKING +import time 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 @@ -29,11 +30,13 @@ from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator 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 ActiveAttemptMetric + from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data._sync_autogen.client import ( _DataApiTarget as TargetType, ) @@ -56,6 +59,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 +71,7 @@ class _ReadRowsOperation: "_predicate", "_last_yielded_row_key", "_remaining_count", + "_operation_metric", ) def __init__( @@ -75,6 +80,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,16 +97,20 @@ 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. Yields: Row: The next row in the stream""" + self._operation_metric.backoff_generator = TrackedBackoffGenerator( + 0.01, 60, multiplier=2 + ) 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, ) @@ -113,6 +123,7 @@ def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: Yields: Row: The next row in the stream""" + attempt_metric = self._operation_metric.start_attempt() if self._last_yielded_row_key is not None: try: self.request.rows = self._revise_request_rowset( @@ -120,16 +131,16 @@ def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: last_seen_row_key=self._last_yielded_row_key, ) except _RowSetComplete: - return self.merge_rows(None) + return self.merge_rows(None, attempt_metric) if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: - return self.merge_rows(None) + return self.merge_rows(None, attempt_metric) gapic_stream = self.target.client._gapic_client.read_rows( self.request, timeout=next(self.attempt_timeout_gen), retry=None ) chunked_stream = self.chunk_stream(gapic_stream) - return self.merge_rows(chunked_stream) + return self.merge_rows(chunked_stream, attempt_metric) def chunk_stream( self, @@ -177,6 +188,7 @@ def chunk_stream( @staticmethod def merge_rows( chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None, + attempt_metric: ActiveAttemptMetric, ) -> CrossSync._Sync_Impl.Iterable[Row]: """Merge chunks into rows @@ -187,6 +199,7 @@ def merge_rows( if chunks is None: return it = chunks.__iter__() + is_first_row = True while True: try: c = it.__next__() @@ -255,7 +268,14 @@ def merge_rows( Cell(value, row_key, family, qualifier, ts, list(labels)) ) if c.commit_row: + if is_first_row: + is_first_row = False + attempt_metric.attempt_first_response() + block_time = time.monotonic_ns() yield Row(row_key, cells) + attempt_metric.active_attempt.application_blocking_time_ns += ( + time.monotonic_ns() - block_time + ) * 1000 break c = it.__next__() except _ResetRow as e: diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index a7a9b16b4..e28d6462a 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -66,6 +66,7 @@ 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 _attempt_timeout_generator +from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule from google.cloud.bigtable.data.row_filters import RowFilter @@ -73,6 +74,8 @@ 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._metrics.handlers._stdout import _StdoutMetricsHandler from google.cloud.bigtable.data._cross_sync import CrossSync from typing import Iterable from grpc import insecure_channel @@ -733,6 +736,7 @@ def __init__( ) self._metrics = BigtableClientSideMetricsController( client._metrics_interceptor, + handlers=[_StdoutMetricsHandler()], project_id=self.client.project, instance_id=instance_id, table_id=table_id, @@ -801,14 +805,18 @@ def read_rows_stream( 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_exceptions=retryable_excs, - ) - return row_merger.start_operation() + with self._metrics.create_operation( + OperationType.READ_ROWS, streaming=True + ) as operation_metric: + row_merger = CrossSync._Sync_Impl._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=operation_metric, + retryable_exceptions=retryable_excs, + ) + return row_merger.start_operation() def read_rows( self, @@ -894,15 +902,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( - query, - operation_timeout=operation_timeout, - attempt_timeout=attempt_timeout, - retryable_errors=retryable_errors, + (operation_timeout, attempt_timeout) = _get_timeouts( + operation_timeout, attempt_timeout, self ) - if len(results) == 0: - return None - return results[0] + retryable_excs = _get_retryable_errors(retryable_errors, self) + with self._metrics.create_operation( + OperationType.READ_ROWS, streaming=False + ) as operation_metric: + row_merger = CrossSync._Sync_Impl._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=operation_metric, + retryable_exceptions=retryable_excs, + ) + results_generator = row_merger.start_operation() + try: + return results_generator.__next__() + except StopIteration: + return None def read_rows_sharded( self, @@ -1025,19 +1044,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, @@ -1084,25 +1101,28 @@ 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) + sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) + with self._metrics.create_operation( + OperationType.SAMPLE_ROW_KEYS, backoff_generator=sleep_generator + ): - 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, + 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, + sleep_generator, + operation_timeout, + exception_factory=_retry_exception_factory, ) - 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, @@ -1203,27 +1223,30 @@ 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, - ) + sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) + with self._metrics.create_operation( + OperationType.MUTATE_ROW, backoff_generator=sleep_generator + ): + 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, + ) def bulk_mutate_rows( self, @@ -1273,6 +1296,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() @@ -1327,21 +1351,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, @@ -1378,19 +1405,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..3f4c61bc3 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,23 @@ def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]): ) yield mutations[start_idx:end_idx] + async 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 = await inner_generator.__anext__() + except StopAsyncIteration: + metric.cancel() + return + metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time + yield (value, metric) + class MutationsBatcher: """ @@ -302,9 +323,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 +338,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 +357,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() From 4684a7290bd70c6efa8eba399062766b67426a68 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 7 Aug 2025 15:11:52 -0700 Subject: [PATCH 09/78] use default backoff instance --- google/cloud/bigtable/data/_async/_mutate_rows.py | 4 +--- google/cloud/bigtable/data/_async/_read_rows.py | 3 --- google/cloud/bigtable/data/_async/client.py | 15 ++++++--------- google/cloud/bigtable/data/_metrics/data_model.py | 2 +- 4 files changed, 8 insertions(+), 16 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index c170ff041..3386a68a7 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -102,11 +102,10 @@ def __init__( # Entry level errors bt_exceptions._MutateRowsIncomplete, ) - sleep_generator = TrackedBackoffGenerator(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, ) @@ -118,7 +117,6 @@ def __init__( self.remaining_indices = list(range(len(self.mutations))) self.errors: dict[int, list[Exception]] = {} # set up metrics - metric.backoff_generator = sleep_generator self._operation_metric = metric @CrossSync.convert diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 7e5ea4136..871d39271 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -118,9 +118,6 @@ def start_operation(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ - self._operation_metric.backoff_generator = TrackedBackoffGenerator( - 0.01, 60, multiplier=2 - ) return CrossSync.retry_target_stream( self._read_rows_attempt, self._predicate, diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 969957b62..eed5f50dc 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1345,11 +1345,9 @@ async def sample_row_keys( retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) - async with self._metrics.create_operation( - OperationType.SAMPLE_ROW_KEYS, backoff_generator=sleep_generator - ): + OperationType.SAMPLE_ROW_KEYS + ) as operation_metric: @CrossSync.convert async def execute_rpc(): @@ -1365,7 +1363,7 @@ async def execute_rpc(): return await CrossSync.retry_target( execute_rpc, predicate, - sleep_generator, + operation_metric.backoff_generator, operation_timeout, exception_factory=_retry_exception_factory, ) @@ -1479,10 +1477,9 @@ async def mutate_row( # mutations should not be retried predicate = retries.if_exception_type() - sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) async with self._metrics.create_operation( - OperationType.MUTATE_ROW, backoff_generator=sleep_generator - ): + OperationType.MUTATE_ROW + ) as operation_metric: target = partial( self.client._gapic_client.mutate_row, request=MutateRowRequest( @@ -1499,7 +1496,7 @@ async def mutate_row( return await CrossSync.retry_target( target, predicate, - sleep_generator, + operation_metric.backoff_generator, operation_timeout, exception_factory=_retry_exception_factory, ) diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index ba6731570..07cd52429 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -429,4 +429,4 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): if exc_val is None: self.end_with_success() else: - self.end_with_status(exc_val) \ No newline at end of file + self.end_with_status(exc_val) From 98e400699afabbff2f2e4748d5be617360352b3b Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 7 Aug 2025 15:14:21 -0700 Subject: [PATCH 10/78] remove async with for metrics --- google/cloud/bigtable/data/_async/_mutate_rows.py | 2 +- google/cloud/bigtable/data/_async/client.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 3386a68a7..290bba23d 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -127,7 +127,7 @@ async def start(self): Raises: MutationsExceptionGroup: if any mutations failed """ - async with self._operation_metric: + with self._operation_metric: try: # trigger mutate_rows await self._operation() diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index eed5f50dc..a13fc53d5 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1020,7 +1020,7 @@ async def read_rows_stream( ) retryable_excs = _get_retryable_errors(retryable_errors, self) - async with self._metrics.create_operation( + with self._metrics.create_operation( OperationType.READ_ROWS, streaming=True ) as operation_metric: row_merger = CrossSync._ReadRowsOperation( @@ -1129,7 +1129,7 @@ async def read_row( ) retryable_excs = _get_retryable_errors(retryable_errors, self) - async with self._metrics.create_operation( + with self._metrics.create_operation( OperationType.READ_ROWS, streaming=False ) as operation_metric: row_merger = CrossSync._ReadRowsOperation( @@ -1345,7 +1345,7 @@ async def sample_row_keys( retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - async with self._metrics.create_operation( + with self._metrics.create_operation( OperationType.SAMPLE_ROW_KEYS ) as operation_metric: @@ -1477,7 +1477,7 @@ async def mutate_row( # mutations should not be retried predicate = retries.if_exception_type() - async with self._metrics.create_operation( + with self._metrics.create_operation( OperationType.MUTATE_ROW ) as operation_metric: target = partial( @@ -1612,7 +1612,7 @@ 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 []] - async with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE): + 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, @@ -1670,7 +1670,7 @@ async def read_modify_write_row( if not rules: raise ValueError("rules must contain at least one item") - async with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE): + 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], From fa865cb438c586e2dc55696ca03b02e13a7c21a1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 7 Aug 2025 15:17:38 -0700 Subject: [PATCH 11/78] followed same operation management pattern for read_rows as mutate_rows --- .../cloud/bigtable/data/_async/_read_rows.py | 15 +++--- google/cloud/bigtable/data/_async/client.py | 54 +++++++++---------- 2 files changed, 34 insertions(+), 35 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 871d39271..92fc9c2cf 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -118,13 +118,14 @@ def start_operation(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ - return CrossSync.retry_target_stream( - self._read_rows_attempt, - self._predicate, - self._operation_metric.backoff_generator, - self.operation_timeout, - exception_factory=_retry_exception_factory, - ) + with self._operation_metric: + return CrossSync.retry_target_stream( + self._read_rows_attempt, + self._predicate, + self._operation_metric.backoff_generator, + self.operation_timeout, + exception_factory=_retry_exception_factory, + ) def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: """ diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index a13fc53d5..42692033d 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1020,18 +1020,17 @@ async def read_rows_stream( ) retryable_excs = _get_retryable_errors(retryable_errors, self) - with self._metrics.create_operation( - OperationType.READ_ROWS, streaming=True - ) as operation_metric: - row_merger = CrossSync._ReadRowsOperation( - query, - self, - operation_timeout=operation_timeout, - attempt_timeout=attempt_timeout, - metric=operation_metric, - retryable_exceptions=retryable_excs, - ) - return row_merger.start_operation() + row_merger = CrossSync._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, streaming=True + ), + retryable_exceptions=retryable_excs, + ) + return row_merger.start_operation() @CrossSync.convert async def read_rows( @@ -1129,22 +1128,21 @@ async def read_row( ) retryable_excs = _get_retryable_errors(retryable_errors, self) - with self._metrics.create_operation( - OperationType.READ_ROWS, streaming=False - ) as operation_metric: - row_merger = CrossSync._ReadRowsOperation( - query, - self, - operation_timeout=operation_timeout, - attempt_timeout=attempt_timeout, - metric=operation_metric, - retryable_exceptions=retryable_excs, - ) - results_generator = row_merger.start_operation() - try: - return results_generator.__anext__() - except StopAsyncIteration: - return None + row_merger = CrossSync._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, streaming=False + ), + retryable_exceptions=retryable_excs, + ) + results_generator = row_merger.start_operation() + try: + return results_generator.__anext__() + except StopAsyncIteration: + return None @CrossSync.convert async def read_rows_sharded( From ff7e68150abb89f247612a300444598df7fa7c5f Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 7 Aug 2025 15:30:36 -0700 Subject: [PATCH 12/78] fixed lint --- google/cloud/bigtable/data/_async/_mutate_rows.py | 1 - google/cloud/bigtable/data/_async/_read_rows.py | 1 - google/cloud/bigtable/data/_async/client.py | 1 - google/cloud/bigtable/data/_metrics/data_model.py | 8 ++++---- 4 files changed, 4 insertions(+), 7 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 290bba23d..f6ed2ace3 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -22,7 +22,6 @@ import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory -from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator # mutate_rows requests are limited to this number of mutations from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 92fc9c2cf..9367a67cd 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -31,7 +31,6 @@ from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory -from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.api_core import retry as retries diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 42692033d..987d4b8ee 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -78,7 +78,6 @@ 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 _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 380018fd3..b4f848344 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -146,7 +146,9 @@ class ActiveOperationMetric: uuid: str = field(default_factory=lambda: str(uuid.uuid4())) # create a default backoff generator, initialized with standard default backoff values backoff_generator: TrackedBackoffGenerator = field( - default_factory=lambda: TrackedBackoffGenerator(initial=0.01, maximum=60, multiplier=2) + default_factory=lambda: TrackedBackoffGenerator( + initial=0.01, maximum=60, multiplier=2 + ) ) # keep monotonic timestamps for active operations start_time_ns: int = field(default_factory=time.monotonic_ns) @@ -204,9 +206,7 @@ def start_attempt(self) -> ActiveAttemptMetric: try: # find backoff value before this attempt prev_attempt_idx = len(self.completed_attempts) - 1 - backoff = self.backoff_generator.get_attempt_backoff( - prev_attempt_idx - ) + backoff = self.backoff_generator.get_attempt_backoff(prev_attempt_idx) # generator will return the backoff time in seconds, so convert to nanoseconds backoff_ns = int(backoff * 1e9) except IndexError: From 410cfb8b8954d4b07ab2c8c7819f139d96feb7ac Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 8 Aug 2025 10:59:02 -0700 Subject: [PATCH 13/78] fixed mypy --- .../cloud/bigtable/data/_async/_read_rows.py | 14 ++-- .../data/_sync_autogen/_mutate_rows.py | 5 +- .../bigtable/data/_sync_autogen/_read_rows.py | 32 ++++----- .../bigtable/data/_sync_autogen/client.py | 69 +++++++++---------- .../data/_sync_autogen/metrics_interceptor.py | 10 ++- 5 files changed, 60 insertions(+), 70 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 9367a67cd..7f310fe50 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -220,7 +220,7 @@ async def chunk_stream( ) async def merge_rows( chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None, - attempt_metric: ActiveAttemptMetric, + attempt_metric: ActiveAttemptMetric | None, ) -> CrossSync.Iterable[Row]: """ Merge chunks into rows @@ -233,7 +233,6 @@ async def merge_rows( if chunks is None: return it = chunks.__aiter__() - is_first_row = True # For each row while True: try: @@ -316,17 +315,14 @@ async def merge_rows( Cell(value, row_key, family, qualifier, ts, list(labels)) ) if c.commit_row: - if is_first_row: - # record first row latency in metrics - is_first_row = False - attempt_metric.attempt_first_response() 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 - attempt_metric.active_attempt.application_blocking_time_ns += ( # type: ignore - time.monotonic_ns() - block_time - ) * 1000 + if attempt_metric is not None: + attempt_metric.application_blocking_time_ns += ( # type: ignore + time.monotonic_ns() - block_time + ) * 1000 break c = await it.__anext__() except _ResetRow as e: diff --git a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index 55c2d8668..e7df85431 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -23,7 +23,6 @@ import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory -from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT from google.cloud.bigtable.data.mutations import _EntryWithProto from google.cloud.bigtable.data._cross_sync import CrossSync @@ -80,11 +79,10 @@ def __init__( self.is_retryable = retries.if_exception_type( *retryable_exceptions, bt_exceptions._MutateRowsIncomplete ) - sleep_generator = TrackedBackoffGenerator(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, ) @@ -94,7 +92,6 @@ 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]] = {} - metric.backoff_generator = sleep_generator self._operation_metric = metric def start(self): diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index a704912e9..eeb20e437 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -30,7 +30,6 @@ from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator from google.cloud.bigtable.data._helpers import _retry_exception_factory -from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.api_core import retry as retries from google.cloud.bigtable.data._cross_sync import CrossSync @@ -104,16 +103,14 @@ def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]: Yields: Row: The next row in the stream""" - self._operation_metric.backoff_generator = TrackedBackoffGenerator( - 0.01, 60, multiplier=2 - ) - return CrossSync._Sync_Impl.retry_target_stream( - self._read_rows_attempt, - self._predicate, - self._operation_metric.backoff_generator, - self.operation_timeout, - exception_factory=_retry_exception_factory, - ) + with self._operation_metric: + return CrossSync._Sync_Impl.retry_target_stream( + self._read_rows_attempt, + self._predicate, + self._operation_metric.backoff_generator, + self.operation_timeout, + exception_factory=_retry_exception_factory, + ) def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: """Attempt a single read_rows rpc call. @@ -188,7 +185,7 @@ def chunk_stream( @staticmethod def merge_rows( chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None, - attempt_metric: ActiveAttemptMetric, + attempt_metric: ActiveAttemptMetric | None, ) -> CrossSync._Sync_Impl.Iterable[Row]: """Merge chunks into rows @@ -199,7 +196,6 @@ def merge_rows( if chunks is None: return it = chunks.__iter__() - is_first_row = True while True: try: c = it.__next__() @@ -268,14 +264,12 @@ def merge_rows( Cell(value, row_key, family, qualifier, ts, list(labels)) ) if c.commit_row: - if is_first_row: - is_first_row = False - attempt_metric.attempt_first_response() block_time = time.monotonic_ns() yield Row(row_key, cells) - attempt_metric.active_attempt.application_blocking_time_ns += ( - time.monotonic_ns() - block_time - ) * 1000 + if attempt_metric is not None: + attempt_metric.application_blocking_time_ns += ( + time.monotonic_ns() - block_time + ) * 1000 break c = it.__next__() except _ResetRow as e: diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index e28d6462a..48ac397ee 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -66,7 +66,6 @@ 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 _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule from google.cloud.bigtable.data.row_filters import RowFilter @@ -805,18 +804,17 @@ def read_rows_stream( operation_timeout, attempt_timeout, self ) retryable_excs = _get_retryable_errors(retryable_errors, self) - with self._metrics.create_operation( - OperationType.READ_ROWS, streaming=True - ) as operation_metric: - row_merger = CrossSync._Sync_Impl._ReadRowsOperation( - query, - self, - operation_timeout=operation_timeout, - attempt_timeout=attempt_timeout, - metric=operation_metric, - retryable_exceptions=retryable_excs, - ) - return row_merger.start_operation() + row_merger = CrossSync._Sync_Impl._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, streaming=True + ), + retryable_exceptions=retryable_excs, + ) + return row_merger.start_operation() def read_rows( self, @@ -906,22 +904,21 @@ def read_row( operation_timeout, attempt_timeout, self ) retryable_excs = _get_retryable_errors(retryable_errors, self) - with self._metrics.create_operation( - OperationType.READ_ROWS, streaming=False - ) as operation_metric: - row_merger = CrossSync._Sync_Impl._ReadRowsOperation( - query, - self, - operation_timeout=operation_timeout, - attempt_timeout=attempt_timeout, - metric=operation_metric, - retryable_exceptions=retryable_excs, - ) - results_generator = row_merger.start_operation() - try: - return results_generator.__next__() - except StopIteration: - return None + row_merger = CrossSync._Sync_Impl._ReadRowsOperation( + query, + self, + operation_timeout=operation_timeout, + attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, streaming=False + ), + retryable_exceptions=retryable_excs, + ) + results_generator = row_merger.start_operation() + try: + return results_generator.__next__() + except StopIteration: + return None def read_rows_sharded( self, @@ -1101,10 +1098,9 @@ def sample_row_keys( ) retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) with self._metrics.create_operation( - OperationType.SAMPLE_ROW_KEYS, backoff_generator=sleep_generator - ): + OperationType.SAMPLE_ROW_KEYS + ) as operation_metric: def execute_rpc(): results = self.client._gapic_client.sample_row_keys( @@ -1119,7 +1115,7 @@ def execute_rpc(): return CrossSync._Sync_Impl.retry_target( execute_rpc, predicate, - sleep_generator, + operation_metric.backoff_generator, operation_timeout, exception_factory=_retry_exception_factory, ) @@ -1223,10 +1219,9 @@ def mutate_row( ) else: predicate = retries.if_exception_type() - sleep_generator = TrackedBackoffGenerator(0.01, 2, 60) with self._metrics.create_operation( - OperationType.MUTATE_ROW, backoff_generator=sleep_generator - ): + OperationType.MUTATE_ROW + ) as operation_metric: target = partial( self.client._gapic_client.mutate_row, request=MutateRowRequest( @@ -1243,7 +1238,7 @@ def mutate_row( return CrossSync._Sync_Impl.retry_target( target, predicate, - sleep_generator, + operation_metric.backoff_generator, operation_timeout, exception_factory=_retry_exception_factory, ) diff --git a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py index 21dd6752a..b1e4a6b8f 100644 --- a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py @@ -15,6 +15,7 @@ # This file is automatically generated by CrossSync. Do not edit manually. from __future__ import annotations +import time from functools import wraps from google.cloud.bigtable.data._metrics.data_model import ( OPERATION_INTERCEPTOR_METADATA_KEY, @@ -76,7 +77,8 @@ def register_operation(self, operation): operation.handlers.append(self) def on_operation_complete(self, op): - del self.operation_map[op.uuid] + if op.uuid in self.operation_map: + del self.operation_map[op.uuid] def on_operation_cancelled(self, op): self.on_operation_complete(op) @@ -105,9 +107,15 @@ def intercept_unary_stream( self, operation, continuation, client_call_details, request ): def response_wrapper(call): + has_first_response = operation.first_response_latency is not None encountered_exc = None try: for response in call: + if not has_first_response: + operation.first_response_latency_ns = ( + time.monotonic_ns() - operation.start_time_ns + ) + has_first_response = True yield response except Exception as e: encountered_exc = e From 2adec5ac9f0550826f57353f9c509b7feb508e43 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 8 Aug 2025 11:12:28 -0700 Subject: [PATCH 14/78] fixed errors --- google/cloud/bigtable/data/_async/client.py | 6 +++--- google/cloud/bigtable/data/_metrics/data_model.py | 2 +- google/cloud/bigtable/data/_sync_autogen/client.py | 4 ++-- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 987d4b8ee..e981270b7 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1025,7 +1025,7 @@ async def read_rows_stream( operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, metric=self._metrics.create_operation( - OperationType.READ_ROWS, streaming=True + OperationType.READ_ROWS, is_streaming=True ), retryable_exceptions=retryable_excs, ) @@ -1133,13 +1133,13 @@ async def read_row( operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, metric=self._metrics.create_operation( - OperationType.READ_ROWS, streaming=False + OperationType.READ_ROWS, is_streaming=False ), retryable_exceptions=retryable_excs, ) results_generator = row_merger.start_operation() try: - return results_generator.__anext__() + return await results_generator.__anext__() except StopAsyncIteration: return None diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 3a7754be3..2bcd40021 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -28,11 +28,11 @@ import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable_v2.types.response_params import ResponseParams +from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.protobuf.message import DecodeError if TYPE_CHECKING: from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler - from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator LOGGER = logging.getLogger(__name__) diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index 48ac397ee..6ed90e41f 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -810,7 +810,7 @@ def read_rows_stream( operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, metric=self._metrics.create_operation( - OperationType.READ_ROWS, streaming=True + OperationType.READ_ROWS, is_streaming=True ), retryable_exceptions=retryable_excs, ) @@ -910,7 +910,7 @@ def read_row( operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, metric=self._metrics.create_operation( - OperationType.READ_ROWS, streaming=False + OperationType.READ_ROWS, is_streaming=False ), retryable_exceptions=retryable_excs, ) From 14d252bffae494adf85f1f525f0ccce3c6fa1f64 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 8 Aug 2025 12:06:08 -0700 Subject: [PATCH 15/78] fixed operation end for read_rows --- .../cloud/bigtable/data/_async/_read_rows.py | 252 +++++++++--------- .../bigtable/data/_sync_autogen/_read_rows.py | 220 +++++++-------- 2 files changed, 245 insertions(+), 227 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 7f310fe50..86f95e278 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -15,7 +15,7 @@ from __future__ import annotations -from typing import Sequence, TYPE_CHECKING +from typing import Callable, Sequence, TYPE_CHECKING import time @@ -117,14 +117,13 @@ def start_operation(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ - with self._operation_metric: - return CrossSync.retry_target_stream( - self._read_rows_attempt, - self._predicate, - self._operation_metric.backoff_generator, - self.operation_timeout, - exception_factory=_retry_exception_factory, - ) + return CrossSync.retry_target_stream( + self._read_rows_attempt, + self._predicate, + self._operation_metric.backoff_generator, + self.operation_timeout, + exception_factory=_retry_exception_factory, + ) def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: """ @@ -136,7 +135,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ - attempt_metric = self._operation_metric.start_attempt() + 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 @@ -147,12 +146,12 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: ) except _RowSetComplete: # if we've already seen all the rows, we're done - return self.merge_rows(None, attempt_metric) + return self.merge_rows(None, self._operation_metric, self._predicate) # revise the limit based on number of rows already yielded if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: - return self.merge_rows(None, attempt_metric) + return self.merge_rows(None, self._operation_metric, self._predicate) # create and return a new row merger gapic_stream = self.target.client._gapic_client.read_rows( self.request, @@ -160,7 +159,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: retry=None, ) chunked_stream = self.chunk_stream(gapic_stream) - return self.merge_rows(chunked_stream, attempt_metric) + return self.merge_rows(chunked_stream, self._operation_metric, self._predicate) @CrossSync.convert() async def chunk_stream( @@ -220,7 +219,8 @@ async def chunk_stream( ) async def merge_rows( chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None, - attempt_metric: ActiveAttemptMetric | None, + operation_metric: ActiveOperationMetric, + retryable_predicate: Callable[[Exception], bool], ) -> CrossSync.Iterable[Row]: """ Merge chunks into rows @@ -230,115 +230,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: + 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 attempt_metric is not None: - attempt_metric.application_blocking_time_ns += ( # type: ignore - time.monotonic_ns() - block_time - ) * 1000 - 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 + 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 operation_metric.active_attempt is not None: + operation_metric.active_attempt.application_blocking_time_ns += ( # type: ignore + time.monotonic_ns() - block_time + ) * 1000 + 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 Exception as generic_exception: + if not retryable_predicate(generic_exception): + operation_metric.end_attempt_with_status(generic_exception) + raise generic_exception + else: + operation_metric.end_with_success() + @staticmethod def _revise_request_rowset( diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index eeb20e437..82deb4322 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -17,7 +17,7 @@ # This file is automatically generated by CrossSync. Do not edit manually. from __future__ import annotations -from typing import Sequence, TYPE_CHECKING +from typing import Callable, Sequence, TYPE_CHECKING import time from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB @@ -34,7 +34,6 @@ from google.cloud.bigtable.data._cross_sync import CrossSync if TYPE_CHECKING: - from google.cloud.bigtable.data._metrics import ActiveAttemptMetric from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data._sync_autogen.client import ( _DataApiTarget as TargetType, @@ -103,14 +102,13 @@ def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]: Yields: Row: The next row in the stream""" - with self._operation_metric: - return CrossSync._Sync_Impl.retry_target_stream( - self._read_rows_attempt, - self._predicate, - self._operation_metric.backoff_generator, - self.operation_timeout, - exception_factory=_retry_exception_factory, - ) + return CrossSync._Sync_Impl.retry_target_stream( + self._read_rows_attempt, + self._predicate, + self._operation_metric.backoff_generator, + self.operation_timeout, + exception_factory=_retry_exception_factory, + ) def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: """Attempt a single read_rows rpc call. @@ -120,7 +118,7 @@ def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: Yields: Row: The next row in the stream""" - attempt_metric = self._operation_metric.start_attempt() + self._operation_metric.start_attempt() if self._last_yielded_row_key is not None: try: self.request.rows = self._revise_request_rowset( @@ -128,16 +126,16 @@ def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: last_seen_row_key=self._last_yielded_row_key, ) except _RowSetComplete: - return self.merge_rows(None, attempt_metric) + return self.merge_rows(None, self._operation_metric, self._predicate) if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: - return self.merge_rows(None, attempt_metric) + return self.merge_rows(None, self._operation_metric, self._predicate) gapic_stream = self.target.client._gapic_client.read_rows( self.request, timeout=next(self.attempt_timeout_gen), retry=None ) chunked_stream = self.chunk_stream(gapic_stream) - return self.merge_rows(chunked_stream, attempt_metric) + return self.merge_rows(chunked_stream, self._operation_metric, self._predicate) def chunk_stream( self, @@ -185,7 +183,8 @@ def chunk_stream( @staticmethod def merge_rows( chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None, - attempt_metric: ActiveAttemptMetric | None, + operation_metric: ActiveOperationMetric, + retryable_predicate: Callable[[Exception], bool], ) -> CrossSync._Sync_Impl.Iterable[Row]: """Merge chunks into rows @@ -193,99 +192,108 @@ 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: + 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 attempt_metric is not None: - attempt_metric.application_blocking_time_ns += ( - time.monotonic_ns() - block_time - ) * 1000 - 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: + 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 operation_metric.active_attempt is not None: + operation_metric.active_attempt.application_blocking_time_ns += ( + time.monotonic_ns() - block_time + ) * 1000 + 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 Exception as generic_exception: + if not retryable_predicate(generic_exception): + operation_metric.end_attempt_with_status(generic_exception) + raise generic_exception + else: + operation_metric.end_with_success() @staticmethod def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSetPB: From ab30b028b22c298fa64c230a1136d6cb9264c8a0 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 8 Aug 2025 15:16:14 -0700 Subject: [PATCH 16/78] fixed lint --- google/cloud/bigtable/data/_async/_read_rows.py | 2 -- google/cloud/bigtable/data/_metrics/data_model.py | 1 + 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 86f95e278..19848c070 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -37,7 +37,6 @@ from google.cloud.bigtable.data._cross_sync import CrossSync if TYPE_CHECKING: - from google.cloud.bigtable.data._metrics import ActiveAttemptMetric from google.cloud.bigtable.data._metrics import ActiveOperationMetric if CrossSync.is_async: @@ -349,7 +348,6 @@ async def merge_rows( else: operation_metric.end_with_success() - @staticmethod def _revise_request_rowset( row_set: RowSetPB, diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 2bcd40021..981f4455c 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -396,6 +396,7 @@ def _handle_error(message: str) -> None: """ full_message = f"Error in Bigtable Metrics: {message}" LOGGER.warning(full_message) + raise RuntimeError(full_message) def __enter__(self): """ From bb55c4604fe40cbff6867214cb81783360360798 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 8 Aug 2025 15:31:54 -0700 Subject: [PATCH 17/78] made merge_rows into an instance method --- .../cloud/bigtable/data/_async/_read_rows.py | 26 +++++++++---------- google/cloud/bigtable/data/_async/client.py | 3 +-- .../bigtable/data/_metrics/data_model.py | 1 - .../bigtable/data/_sync_autogen/_read_rows.py | 26 +++++++++---------- .../bigtable/data/_sync_autogen/client.py | 3 +-- 5 files changed, 26 insertions(+), 33 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 19848c070..4c28dac30 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -15,7 +15,7 @@ from __future__ import annotations -from typing import Callable, Sequence, TYPE_CHECKING +from typing import Sequence, TYPE_CHECKING import time @@ -145,12 +145,12 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: ) except _RowSetComplete: # if we've already seen all the rows, we're done - return self.merge_rows(None, self._operation_metric, self._predicate) + return self.merge_rows(None) # revise the limit based on number of rows already yielded if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: - return self.merge_rows(None, self._operation_metric, self._predicate) + return self.merge_rows(None) # create and return a new row merger gapic_stream = self.target.client._gapic_client.read_rows( self.request, @@ -158,7 +158,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: retry=None, ) chunked_stream = self.chunk_stream(gapic_stream) - return self.merge_rows(chunked_stream, self._operation_metric, self._predicate) + return self.merge_rows(chunked_stream) @CrossSync.convert() async def chunk_stream( @@ -217,9 +217,7 @@ async def chunk_stream( replace_symbols={"__aiter__": "__iter__", "__anext__": "__next__"}, ) async def merge_rows( - chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None, - operation_metric: ActiveOperationMetric, - retryable_predicate: Callable[[Exception], bool], + self, chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None ) -> CrossSync.Iterable[Row]: """ Merge chunks into rows @@ -231,7 +229,7 @@ async def merge_rows( """ try: if chunks is None: - operation_metric.end_with_success() + self._operation_metric.end_with_success() return it = chunks.__aiter__() # For each row @@ -240,7 +238,7 @@ async def merge_rows( c = await it.__anext__() except CrossSync.StopIteration: # stream complete - operation_metric.end_with_success() + self._operation_metric.end_with_success() return row_key = c.row_key @@ -321,8 +319,8 @@ async def merge_rows( yield Row(row_key, cells) # most metric operations use setters, but this one updates # the value directly to avoid extra overhead - if operation_metric.active_attempt is not None: - operation_metric.active_attempt.application_blocking_time_ns += ( # type: ignore + 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 ) * 1000 break @@ -342,11 +340,11 @@ async def merge_rows( except CrossSync.StopIteration: raise InvalidChunk("premature end of stream") except Exception as generic_exception: - if not retryable_predicate(generic_exception): - operation_metric.end_attempt_with_status(generic_exception) + if not self._predicate(generic_exception): + self._operation_metric.end_attempt_with_status(generic_exception) raise generic_exception else: - operation_metric.end_with_success() + self._operation_metric.end_with_success() @staticmethod def _revise_request_rowset( diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index e981270b7..762abd206 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -87,7 +87,6 @@ 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._metrics.handlers._stdout import _StdoutMetricsHandler from google.cloud.bigtable.data._cross_sync import CrossSync @@ -944,7 +943,7 @@ def __init__( self._metrics = BigtableClientSideMetricsController( client._metrics_interceptor, - handlers=[_StdoutMetricsHandler()], + handlers=[], project_id=self.client.project, instance_id=instance_id, table_id=table_id, diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 981f4455c..2bcd40021 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -396,7 +396,6 @@ def _handle_error(message: str) -> None: """ full_message = f"Error in Bigtable Metrics: {message}" LOGGER.warning(full_message) - raise RuntimeError(full_message) def __enter__(self): """ diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index 82deb4322..939b222ab 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -17,7 +17,7 @@ # This file is automatically generated by CrossSync. Do not edit manually. from __future__ import annotations -from typing import Callable, Sequence, TYPE_CHECKING +from typing import Sequence, TYPE_CHECKING import time from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB @@ -126,16 +126,16 @@ def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: last_seen_row_key=self._last_yielded_row_key, ) except _RowSetComplete: - return self.merge_rows(None, self._operation_metric, self._predicate) + return self.merge_rows(None) if self._remaining_count is not None: self.request.rows_limit = self._remaining_count if self._remaining_count == 0: - return self.merge_rows(None, self._operation_metric, self._predicate) + return self.merge_rows(None) gapic_stream = self.target.client._gapic_client.read_rows( self.request, timeout=next(self.attempt_timeout_gen), retry=None ) chunked_stream = self.chunk_stream(gapic_stream) - return self.merge_rows(chunked_stream, self._operation_metric, self._predicate) + return self.merge_rows(chunked_stream) def chunk_stream( self, @@ -182,9 +182,7 @@ def chunk_stream( @staticmethod def merge_rows( - chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None, - operation_metric: ActiveOperationMetric, - retryable_predicate: Callable[[Exception], bool], + self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None ) -> CrossSync._Sync_Impl.Iterable[Row]: """Merge chunks into rows @@ -194,14 +192,14 @@ def merge_rows( Row: the next row in the stream""" try: if chunks is None: - operation_metric.end_with_success() + self._operation_metric.end_with_success() return it = chunks.__iter__() while True: try: c = it.__next__() except CrossSync._Sync_Impl.StopIteration: - operation_metric.end_with_success() + self._operation_metric.end_with_success() return row_key = c.row_key if not row_key: @@ -268,8 +266,8 @@ def merge_rows( if c.commit_row: block_time = time.monotonic_ns() yield Row(row_key, cells) - if operation_metric.active_attempt is not None: - operation_metric.active_attempt.application_blocking_time_ns += ( + if self._operation_metric.active_attempt is not None: + self._operation_metric.active_attempt.application_blocking_time_ns += ( time.monotonic_ns() - block_time ) * 1000 break @@ -289,11 +287,11 @@ def merge_rows( except CrossSync._Sync_Impl.StopIteration: raise InvalidChunk("premature end of stream") except Exception as generic_exception: - if not retryable_predicate(generic_exception): - operation_metric.end_attempt_with_status(generic_exception) + if not self._predicate(generic_exception): + self._operation_metric.end_attempt_with_status(generic_exception) raise generic_exception else: - operation_metric.end_with_success() + self._operation_metric.end_with_success() @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 6ed90e41f..039048744 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -74,7 +74,6 @@ 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._metrics.handlers._stdout import _StdoutMetricsHandler from google.cloud.bigtable.data._cross_sync import CrossSync from typing import Iterable from grpc import insecure_channel @@ -735,7 +734,7 @@ def __init__( ) self._metrics = BigtableClientSideMetricsController( client._metrics_interceptor, - handlers=[_StdoutMetricsHandler()], + handlers=[], project_id=self.client.project, instance_id=instance_id, table_id=table_id, From a4746f099cfa85bd04e5776ca73645e11ad115ac Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 12:22:24 -0700 Subject: [PATCH 18/78] added new file for system tests --- tests/system/data/test_metrics_async.py | 151 ++++++++++++++++++++++++ 1 file changed, 151 insertions(+) create mode 100644 tests/system/data/test_metrics_async.py diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py new file mode 100644 index 000000000..1e838394d --- /dev/null +++ b/tests/system/data/test_metrics_async.py @@ -0,0 +1,151 @@ +# 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 google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric + +from google.cloud.bigtable.data._cross_sync import CrossSync + + +from . import TEST_FAMILY + +__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 = [] + self.cancelled_operations = [] + + def on_operation_complete(self, op): + self.completed_operations.append(op) + + def on_operation_cancelled(self, op): + self.cancelled_operations.append(op) + + def on_attempt_complete(self, attempt, op): + self.completed_attempts.append((attempt, op)) + + def total(self): + return len(self.completed_operations) + len(self._ancelled_operations) + len(self.completed_attempts) + + def clear(self): + self.cancelled_operations.clear() + self.completed_operations.clear() + self.completed_attempts.clear() + + +@CrossSync.convert_class(sync_name="TestSystem") +class TestMetricsAsync: + + @CrossSync.drop + @pytest.fixture(scope="session") + def event_loop(self): + loop = asyncio.get_event_loop() + yield loop + loop.stop() + loop.close() + + @pytest.fixture(scope="session") + def init_table_id(self): + """ + The table_id to use when creating a new test table + """ + return f"test-metrics-{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 + + return {TEST_FAMILY: types.ColumnFamily()} + + def _make_client(self): + project = os.getenv("GOOGLE_CLOUD_PROJECT") or None + return CrossSync.DataClient(project=project) + + @pytest.fixture(scope="session") + def handler(self): + return _MetricsTestHandler() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function", autouse=True) + async def _clear_handler(self, handler): + """Clear handler between each test""" + handler.clear() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def client(self): + async with self._make_client() as client: + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function") + async def temp_rows(self, target): + builder = CrossSync.TempRowBuilder(target) + yield builder + await builder.delete_rows() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def target(self, client, table_id, instance_id, handler): + """ + This fixture runs twice: once for a standard table, and once with an authorized view + + Note: emulator doesn't support authorized views. Only use target + """ + async with client.get_table(instance_id, table_id) as table: + table._metrics.add_handler(handler) + yield table + + async def test_read_modify_write(self, target, temp_rows): + 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) + result = await target.read_modify_write_row(row_key, rule) + breakpoint() + print(result) From 417faf0505f0372a94ed606d436aef24d7a49184 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 13:07:22 -0700 Subject: [PATCH 19/78] got tests to run --- tests/system/data/__init__.py | 60 +++++++++++++++++++++++++ tests/system/data/test_metrics_async.py | 59 +++++++----------------- tests/system/data/test_system_async.py | 53 +--------------------- 3 files changed, 77 insertions(+), 95 deletions(-) diff --git a/tests/system/data/__init__.py b/tests/system/data/__init__.py index 2b35cea8f..a30174cd6 100644 --- a/tests/system/data/__init__.py +++ b/tests/system/data/__init__.py @@ -13,7 +13,67 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import asyncio +import pytest +import uuid +from google.cloud.bigtable.data._cross_sync import CrossSync TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" TEST_AGGREGATE_FAMILY = "test-aggregate-family" + +class SystemTestRunner: + """ + configures a system test class with configuration for clusters/tables/etc + + used by standard system tests, and metrics tests + """ + + @CrossSync.drop + @pytest.fixture(scope="session") + def event_loop(self): + loop = asyncio.get_event_loop() + yield loop + loop.stop() + loop.close() + + @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) + ), + } \ No newline at end of file diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 1e838394d..f3edb8afd 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -22,7 +22,7 @@ from google.cloud.bigtable.data._cross_sync import CrossSync -from . import TEST_FAMILY +from . import TEST_FAMILY, SystemTestRunner __CROSS_SYNC_OUTPUT__ = "tests.system.data.test_metrics_autogen" @@ -54,48 +54,12 @@ def clear(self): self.completed_operations.clear() self.completed_attempts.clear() + def __repr__(self): + return f"{self.__class__}(completed_operations={len(self.completed_operations)}, cancelled_operations={len(self.cancelled_operations)}, completed_attempts={len(self.completed_attempts)}" -@CrossSync.convert_class(sync_name="TestSystem") -class TestMetricsAsync: - @CrossSync.drop - @pytest.fixture(scope="session") - def event_loop(self): - loop = asyncio.get_event_loop() - yield loop - loop.stop() - loop.close() - - @pytest.fixture(scope="session") - def init_table_id(self): - """ - The table_id to use when creating a new test table - """ - return f"test-metrics-{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 - - return {TEST_FAMILY: types.ColumnFamily()} +@CrossSync.convert_class(sync_name="TestMetrics") +class TestMetricsAsync(SystemTestRunner): def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None @@ -136,7 +100,8 @@ async def target(self, client, table_id, instance_id, handler): table._metrics.add_handler(handler) yield table - async def test_read_modify_write(self, target, temp_rows): + @CrossSync.pytest + async def test_read_modify_write(self, target, temp_rows, handler): from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule row_key = b"test-row-key" @@ -146,6 +111,12 @@ async def test_read_modify_write(self, target, temp_rows): row_key, value=0, family=family, qualifier=qualifier ) rule = IncrementRule(family, qualifier, 1) - result = await target.read_modify_write_row(row_key, rule) + await target.read_modify_write_row(row_key, rule) breakpoint() - print(result) + assert handler.total() == 1 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + operation = handler.completed_operations[0] + assert operation.final_status.value[0] == 0 + assert operation.final_status.value[0] == 0 diff --git a/tests/system/data/test_system_async.py b/tests/system/data/test_system_async.py index b4e661e6c..6d1938cc3 100644 --- a/tests/system/data/test_system_async.py +++ b/tests/system/data/test_system_async.py @@ -27,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 ( @@ -119,7 +119,7 @@ async def delete_rows(self): @CrossSync.convert_class(sync_name="TestSystem") -class TestSystemAsync: +class TestSystemAsync(SystemTestRunner): @CrossSync.convert @CrossSync.pytest_fixture(scope="session") async def client(self): @@ -146,55 +146,6 @@ async def target(self, client, table_id, authorized_view_id, instance_id, reques else: raise ValueError(f"unknown target type: {request.param}") - @CrossSync.drop - @pytest.fixture(scope="session") - def event_loop(self): - loop = asyncio.get_event_loop() - yield loop - loop.stop() - loop.close() - - @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): From 4c8ffe5e99e89fec572141beb63b96a874196405 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 13:43:59 -0700 Subject: [PATCH 20/78] added test for read_modify_write --- google/cloud/bigtable/data/_async/client.py | 3 +- tests/system/data/test_metrics_async.py | 34 +++++++++++++++------ 2 files changed, 26 insertions(+), 11 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 762abd206..5b4f9017c 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1666,7 +1666,7 @@ async def read_modify_write_row( if not rules: raise ValueError("rules must contain at least one item") - with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE): + with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op: result = await self.client._gapic_client.read_modify_write_row( request=ReadModifyWriteRowRequest( rules=[rule._to_pb() for rule in rules], @@ -1678,6 +1678,7 @@ async def read_modify_write_row( ), timeout=operation_timeout, retry=None, + metadata=[op.interceptor_metadata], ) # construct Row from result return Row._from_pb(result.row) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index f3edb8afd..123339fb4 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -17,7 +17,7 @@ import uuid from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler -from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric +from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric from google.cloud.bigtable.data._cross_sync import CrossSync @@ -43,11 +43,8 @@ def on_operation_complete(self, op): def on_operation_cancelled(self, op): self.cancelled_operations.append(op) - def on_attempt_complete(self, attempt, op): - self.completed_attempts.append((attempt, op)) - - def total(self): - return len(self.completed_operations) + len(self._ancelled_operations) + len(self.completed_attempts) + def on_attempt_complete(self, attempt, _): + self.completed_attempts.append(attempt) def clear(self): self.cancelled_operations.clear() @@ -101,7 +98,7 @@ async def target(self, client, table_id, instance_id, handler): yield table @CrossSync.pytest - async def test_read_modify_write(self, target, temp_rows, handler): + async def test_read_modify_write(self, target, temp_rows, handler, cluster_config): from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule row_key = b"test-row-key" @@ -112,11 +109,28 @@ async def test_read_modify_write(self, target, temp_rows, handler): ) rule = IncrementRule(family, qualifier, 1) await target.read_modify_write_row(row_key, rule) - breakpoint() - assert handler.total() == 1 + # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 assert len(handler.cancelled_operations) == 0 + # validate operation operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 - 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 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 From 542dfb9ce87c2f1cfc044c186cc4c376c5e6fa8a Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 13:49:13 -0700 Subject: [PATCH 21/78] added test for check_and_mutate --- google/cloud/bigtable/data/_async/client.py | 3 +- tests/system/data/test_metrics_async.py | 54 +++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 5b4f9017c..9b8ffcc1d 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1608,7 +1608,7 @@ 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 []] - with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE): + with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op: result = await self.client._gapic_client.check_and_mutate_row( request=CheckAndMutateRowRequest( true_mutations=true_case_list, @@ -1624,6 +1624,7 @@ async def check_and_mutate_row( ), timeout=operation_timeout, retry=None, + metadata=[op.interceptor_metadata], ) return result.predicate_matched diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 123339fb4..5a5084ab6 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -134,3 +134,57 @@ async def test_read_modify_write(self, target, temp_rows, handler, cluster_confi 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_check_and_mutate_row(self, target, 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 + ) + + + false_mutation_value = b"false-mutation-value" + false_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=false_mutation_value + ) + 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 target.check_and_mutate_row( + row_key, + predicate, + true_case_mutations=true_mutation, + false_case_mutations=false_mutation, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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 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 From b815674c5dd1518ab9ed3b8dc01b9906db2ec4ce Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 14:12:30 -0700 Subject: [PATCH 22/78] added test for sample_row_keys --- google/cloud/bigtable/data/_async/client.py | 1 + .../data/_async/metrics_interceptor.py | 4 ++- tests/system/data/test_metrics_async.py | 29 +++++++++++++++++++ 3 files changed, 33 insertions(+), 1 deletion(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 9b8ffcc1d..4cd531054 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1353,6 +1353,7 @@ async def execute_rpc(): ), timeout=next(attempt_timeout_gen), retry=None, + metadata=[operation_metric.interceptor_metadata], ) return [(s.row_key, s.offset_bytes) async for s in results] diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index e5777ea0f..5e206f3b7 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -20,6 +20,7 @@ ) from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric from google.cloud.bigtable.data._metrics.data_model import OperationState +from google.cloud.bigtable.data._metrics.data_model import OperationType from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from google.cloud.bigtable.data._cross_sync import CrossSync @@ -129,7 +130,8 @@ async def intercept_unary_stream( ): # TODO: benchmark async def response_wrapper(call): - has_first_response = operation.first_response_latency is not None + # only track has_first response for READ_ROWS + has_first_response = operation.first_response_latency_ns is not None or operation.op_type != OperationType.READ_ROWS encountered_exc = None try: async for response in call: diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 5a5084ab6..ac3ea2925 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -97,6 +97,35 @@ async def target(self, client, table_id, instance_id, handler): table._metrics.add_handler(handler) yield table + @CrossSync.pytest + async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config): + await target.sample_row_keys() + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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 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_read_modify_write(self, target, temp_rows, handler, cluster_config): from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule From f84151f17fe1b3cf37ecc24ea4d69ff65d3cf79d Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 14:14:44 -0700 Subject: [PATCH 23/78] pass down metadata in other rpcs --- google/cloud/bigtable/data/_async/_mutate_rows.py | 1 + google/cloud/bigtable/data/_async/_read_rows.py | 1 + google/cloud/bigtable/data/_async/client.py | 1 + 3 files changed, 3 insertions(+) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index f6ed2ace3..9c1fc30b3 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -187,6 +187,7 @@ async def _run_attempt(self): ), timeout=next(self.timeout_generator), retry=None, + metadata=[self._operation_metric.interceptor_metadata], ) async for result_list in result_generator: for result in result_list.entries: diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 4c28dac30..899e05d7a 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -156,6 +156,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: self.request, timeout=next(self.attempt_timeout_gen), retry=None, + metadata=[self._operation_metric.interceptor_metadata], ) chunked_stream = self.chunk_stream(gapic_stream) return self.merge_rows(chunked_stream) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 4cd531054..142c3c392 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1489,6 +1489,7 @@ async def mutate_row( ), timeout=attempt_timeout, retry=None, + metadata=[operation_metric.interceptor_metadata], ) return await CrossSync.retry_target( target, From b57ab242213b4c9971c366bedd02d6e5f8906136 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 14:14:58 -0700 Subject: [PATCH 24/78] added stubs for other rpcs --- tests/system/data/test_metrics_async.py | 28 +++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index ac3ea2925..3da4f18e7 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -97,6 +97,34 @@ async def target(self, client, table_id, instance_id, handler): table._metrics.add_handler(handler) yield table + @CrossSync.pytest + async def test_read_rows(self, target, temp_rows, handler, cluster_config): + pass + + @CrossSync.pytest + async def test_read_rows_stream(self, target, temp_rows, handler, cluster_config): + pass + + @CrossSync.pytest + async def test_read_row(self, target, temp_rows, handler, cluster_config): + pass + + @CrossSync.pytest + async def test_read_rows_sharded(self, target, temp_rows, handler, cluster_config): + pass + + @CrossSync.pytest + async def test_bulk_mutate_rows(self, target, temp_rows, handler, cluster_config): + pass + + @CrossSync.pytest + async def test_row_batcher(self, target, temp_rows, handler, cluster_config): + pass + + @CrossSync.pytest + async def test_mutate_row(self, target, temp_rows, handler, cluster_config): + pass + @CrossSync.pytest async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config): await target.sample_row_keys() From 55ff4d1c148cdedc04bdd23fcdf339754a5db8b6 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 14:49:32 -0700 Subject: [PATCH 25/78] fixed read_rows --- google/cloud/bigtable/data/_async/_read_rows.py | 1 - google/cloud/bigtable/data/_sync_autogen/_read_rows.py | 1 - 2 files changed, 2 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 899e05d7a..b935798ec 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -213,7 +213,6 @@ async def chunk_stream( raise InvalidChunk("emit count exceeds row limit") current_key = None - @staticmethod @CrossSync.convert( replace_symbols={"__aiter__": "__iter__", "__anext__": "__next__"}, ) diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index 939b222ab..1a40b2fe5 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -180,7 +180,6 @@ def chunk_stream( raise InvalidChunk("emit count exceeds row limit") current_key = None - @staticmethod def merge_rows( self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None ) -> CrossSync._Sync_Impl.Iterable[Row]: From 614c32a17bb07263617dce208e9c99189a5acb59 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 14:54:31 -0700 Subject: [PATCH 26/78] refactored system tests --- tests/system/data/test_system_async.py | 144 ++++++++++++------------- 1 file changed, 71 insertions(+), 73 deletions(-) diff --git a/tests/system/data/test_system_async.py b/tests/system/data/test_system_async.py index 93d64fe53..82939ea00 100644 --- a/tests/system/data/test_system_async.py +++ b/tests/system/data/test_system_async.py @@ -116,6 +116,39 @@ 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(SystemTestRunner): @@ -148,41 +181,6 @@ async def target(self, client, table_id, authorized_view_id, instance_id, reques else: raise ValueError(f"unknown target type: {request.param}") - @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): @@ -280,13 +278,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") @@ -308,14 +306,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 @@ -357,15 +355,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): @@ -403,11 +401,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]) @@ -416,7 +414,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") @@ -432,8 +430,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 @@ -444,7 +442,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") @@ -459,12 +457,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]) @@ -484,8 +482,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") @@ -500,12 +498,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]) @@ -525,8 +523,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") @@ -539,12 +537,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]) @@ -561,8 +559,8 @@ 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") @@ -633,7 +631,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") @@ -673,7 +671,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") @@ -710,7 +708,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") @@ -759,7 +757,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)), From d3f8dbf8383da968da89009566d8c52031a51e4a Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 14:55:51 -0700 Subject: [PATCH 27/78] moved event loop back into test files --- tests/system/data/__init__.py | 9 --------- tests/system/data/test_metrics_async.py | 8 ++++++++ tests/system/data/test_system_async.py | 10 ++++++++++ 3 files changed, 18 insertions(+), 9 deletions(-) diff --git a/tests/system/data/__init__.py b/tests/system/data/__init__.py index a30174cd6..d123f43fe 100644 --- a/tests/system/data/__init__.py +++ b/tests/system/data/__init__.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import asyncio import pytest import uuid from google.cloud.bigtable.data._cross_sync import CrossSync @@ -29,14 +28,6 @@ class SystemTestRunner: used by standard system tests, and metrics tests """ - @CrossSync.drop - @pytest.fixture(scope="session") - def event_loop(self): - loop = asyncio.get_event_loop() - yield loop - loop.stop() - loop.close() - @pytest.fixture(scope="session") def init_table_id(self): """ diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 3da4f18e7..3f3a32666 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -58,6 +58,14 @@ def __repr__(self): @CrossSync.convert_class(sync_name="TestMetrics") class TestMetricsAsync(SystemTestRunner): + @CrossSync.drop + @pytest.fixture(scope="session") + def event_loop(self): + loop = asyncio.get_event_loop() + yield loop + loop.stop() + loop.close() + def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None return CrossSync.DataClient(project=project) diff --git a/tests/system/data/test_system_async.py b/tests/system/data/test_system_async.py index 82939ea00..44eed7a42 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 @@ -152,6 +153,15 @@ async def create_row_and_mutation( @CrossSync.convert_class(sync_name="TestSystem") class TestSystemAsync(SystemTestRunner): + + @CrossSync.drop + @pytest.fixture(scope="session") + def event_loop(self): + loop = asyncio.get_event_loop() + yield loop + loop.stop() + loop.close() + def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None return CrossSync.DataClient(project=project) From 2a93a28113fe0b196f804e6f945b2a1ce29013fe Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 15 Sep 2025 15:46:32 -0700 Subject: [PATCH 28/78] implemented full success tests for rpcs --- google/cloud/bigtable/data/_async/client.py | 9 +- tests/system/data/test_metrics_async.py | 264 ++++++++++++++++++-- 2 files changed, 254 insertions(+), 19 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 46acdc139..9cbf99a4f 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1108,9 +1108,7 @@ async def read_rows( ) return [row async for row in row_generator] - @CrossSync.convert( - replace_symbols={"__anext__": "__next__", "StopAsyncIteration": "StopIteration"} - ) + @CrossSync.convert async def read_row( self, row_key: str | bytes, @@ -1168,8 +1166,9 @@ async def read_row( ) results_generator = row_merger.start_operation() try: - return await results_generator.__anext__() - except StopAsyncIteration: + results = [a async for a in results_generator] + return results[0] + except IndexError: return None @CrossSync.convert diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 3f3a32666..108193aba 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -17,7 +17,7 @@ import uuid 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._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric, ActiveOperationMetric, OperationState from google.cloud.bigtable.data._cross_sync import CrossSync @@ -96,42 +96,275 @@ async def temp_rows(self, target): @CrossSync.convert @CrossSync.pytest_fixture(scope="session") async def target(self, client, table_id, instance_id, handler): - """ - This fixture runs twice: once for a standard table, and once with an authorized view - - Note: emulator doesn't support authorized views. Only use target - """ async with client.get_table(instance_id, table_id) as table: table._metrics.add_handler(handler) yield table @CrossSync.pytest async def test_read_rows(self, target, temp_rows, handler, cluster_config): - pass + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + row_list = await target.read_rows({}) + assert len(row_list) == 2 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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_stream(self, target, temp_rows, handler, cluster_config): - pass + 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 target.read_rows_stream({}) + 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 + assert len(handler.cancelled_operations) == 0 + # 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_row(self, target, temp_rows, handler, cluster_config): - pass + await temp_rows.add_row(b"row_key_1") + handler.clear() + await target.read_row(b"row_key_1") + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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_rows_sharded(self, target, temp_rows, handler, cluster_config): - pass + 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 target.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + # validate counts + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + assert len(handler.cancelled_operations) == 0 + # 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_bulk_mutate_rows(self, target, temp_rows, handler, cluster_config): - pass + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + + handler.clear() + await target.bulk_mutate_rows([bulk_mutation]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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_row_batcher(self, target, temp_rows, handler, cluster_config): - pass + async def test_mutate_rows_batcher(self, target, 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( + target, new_value=new_value + ) + 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]) + + handler.clear() + async with target.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 + assert len(handler.cancelled_operations) == 1 + cancelled = handler.cancelled_operations[0] + assert isinstance(cancelled, ActiveOperationMetric) + assert cancelled.state == OperationState.CREATED + assert not cancelled.completed_attempts + # 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_row(self, target, temp_rows, handler, cluster_config): - pass + row_key = b"bulk_mutate" + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value + ) + handler.clear() + await target.mutate_row(row_key, mutation) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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_sample_row_keys(self, target, temp_rows, handler, cluster_config): @@ -155,6 +388,7 @@ async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config) 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 @@ -193,6 +427,7 @@ async def test_read_modify_write(self, target, temp_rows, handler, cluster_confi 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 @@ -247,6 +482,7 @@ async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_co 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 From 3fb4f135daed8b2735eb8b201e89ec45263af4d1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 11:09:21 -0700 Subject: [PATCH 29/78] added failure tests for checK_and_mutate --- tests/system/data/test_metrics_async.py | 188 ++++++++++++++++++++++-- 1 file changed, 178 insertions(+), 10 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 108193aba..d4b16b84c 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -15,15 +15,26 @@ import os import pytest import uuid +from grpc import RpcError +from grpc.aio import AioRpcError +from grpc.aio import Metadata +from google.api_core.exceptions import GoogleAPICallError +from google.cloud.bigtable_v2.types import ResponseParams from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric, ActiveOperationMetric, OperationState 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 +else: + from grpc import UnaryUnaryClientInterceptor + from grpc import UnaryStreamClientInterceptor + __CROSS_SYNC_OUTPUT__ = "tests.system.data.test_metrics_autogen" @@ -55,6 +66,27 @@ def __repr__(self): return f"{self.__class__}(completed_operations={len(self.completed_operations)}, cancelled_operations={len(self.cancelled_operations)}, completed_attempts={len(self.completed_attempts)}" +class _ErrorInjectorInterceptor(UnaryUnaryClientInterceptor): + """ + Gprc interceptor used to inject errors into rpc calls, to test failures + """ + + def __init__(self): + self._exc_list = [] + + def push(self, exc: Exception): + self._exc_list.append(exc) + + def clear(self): + self._exc_list.clear() + + 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_class(sync_name="TestMetrics") class TestMetricsAsync(SystemTestRunner): @@ -74,16 +106,23 @@ def _make_client(self): 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_handler(self, handler): - """Clear handler between each test""" + 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): + 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) yield client @CrossSync.convert @@ -447,11 +486,6 @@ async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_co row_key, value=1, family=family, qualifier=qualifier ) - - false_mutation_value = b"false-mutation-value" - false_mutation = SetCell( - family=TEST_FAMILY, qualifier=qualifier, new_value=false_mutation_value - ) true_mutation_value = b"true-mutation-value" true_mutation = SetCell( family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value @@ -461,7 +495,6 @@ async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_co row_key, predicate, true_case_mutations=true_mutation, - false_case_mutations=false_mutation, ) # validate counts assert len(handler.completed_operations) == 1 @@ -489,3 +522,138 @@ async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_co 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_check_and_mutate_row_failure_grpc( + self, target, 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.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) + + # trigger an exception + exc = RuntimeError("injected") + error_injector.push(exc) + with pytest.raises(RuntimeError): + await target.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0,2), + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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_invalid_argument( + self, target, temp_rows, handler + ): + """ + Test failure on backend by passing invalid argument + + We expect a server-timing header, but no cluster/zone info + """ + 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) + + predicate = ValueRangeFilter(-1, -1) + with pytest.raises(GoogleAPICallError): + await target.check_and_mutate_row( + row_key, + predicate, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "INVALID_ARGUMENT" + assert operation.cluster_id == "unspecified" + assert operation.zone == "global" + # 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_failure_timeout( + self, target, temp_rows, handler, error_injector + ): + """ + 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 target.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 + assert len(handler.cancelled_operations) == 0 + # 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 \ No newline at end of file From 66d3254eb134deee9537d2a7d891e8b01ebc4ee1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 11:15:47 -0700 Subject: [PATCH 30/78] added unauthenticated error test --- .../data/_async/metrics_interceptor.py | 4 +- tests/system/data/test_metrics_async.py | 120 +++++++++--------- 2 files changed, 65 insertions(+), 59 deletions(-) diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 9404aa563..b4b6da642 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -89,8 +89,8 @@ def _end_attempt(operation, exc, metadata): async def _get_metadata(source): """Helper to extract metadata from a call or RpcError""" try: - return (await source.trailing_metadata() or []) + ( - await source.initial_metadata() or [] + return (await source.trailing_metadata() or {}) + ( + await source.initial_metadata() or {} ) except Exception: # ignore errors while fetching metadata diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index d4b16b84c..49fe44895 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -127,24 +127,31 @@ async def client(self, error_injector): @CrossSync.convert @CrossSync.pytest_fixture(scope="function") - async def temp_rows(self, target): - builder = CrossSync.TempRowBuilder(target) + 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 target(self, client, table_id, instance_id, handler): + 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, target, temp_rows, handler, cluster_config): + 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 target.read_rows({}) + row_list = await table.read_rows({}) assert len(row_list) == 2 # validate counts assert len(handler.completed_operations) == 1 @@ -174,12 +181,12 @@ async def test_read_rows(self, target, temp_rows, handler, cluster_config): assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_rows_stream(self, target, temp_rows, handler, cluster_config): + 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 target.read_rows_stream({}) + generator = await table.read_rows_stream({}) row_list = [r async for r in generator] assert len(row_list) == 2 # validate counts @@ -210,10 +217,10 @@ async def test_read_rows_stream(self, target, temp_rows, handler, cluster_config assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_row(self, target, temp_rows, handler, cluster_config): + async def test_read_row(self, table, temp_rows, handler, cluster_config): await temp_rows.add_row(b"row_key_1") handler.clear() - await target.read_row(b"row_key_1") + await table.read_row(b"row_key_1") # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 @@ -242,7 +249,7 @@ async def test_read_row(self, target, temp_rows, handler, cluster_config): assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_rows_sharded(self, target, temp_rows, handler, cluster_config): + 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") @@ -251,7 +258,7 @@ async def test_read_rows_sharded(self, target, temp_rows, handler, cluster_confi query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) handler.clear() - row_list = await target.read_rows_sharded([query1, query2]) + row_list = await table.read_rows_sharded([query1, query2]) assert len(row_list) == 4 # validate counts assert len(handler.completed_operations) == 2 @@ -281,17 +288,17 @@ async def test_read_rows_sharded(self, target, temp_rows, handler, cluster_confi assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_bulk_mutate_rows(self, target, temp_rows, handler, cluster_config): + 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( - target, new_value=new_value + table, new_value=new_value ) bulk_mutation = RowMutationEntry(row_key, [mutation]) handler.clear() - await target.bulk_mutate_rows([bulk_mutation]) + await table.bulk_mutate_rows([bulk_mutation]) # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 @@ -320,21 +327,21 @@ async def test_bulk_mutate_rows(self, target, temp_rows, handler, cluster_config assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_mutate_rows_batcher(self, target, temp_rows, handler, cluster_config): + 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( - target, new_value=new_value + table, new_value=new_value ) row_key2, mutation2 = await temp_rows.create_row_and_mutation( - target, new_value=new_value2 + table, new_value=new_value2 ) bulk_mutation = RowMutationEntry(row_key, [mutation]) bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) handler.clear() - async with target.mutations_batcher() as batcher: + async with table.mutations_batcher() as batcher: await batcher.append(bulk_mutation) await batcher.append(bulk_mutation2) # validate counts @@ -370,14 +377,14 @@ async def test_mutate_rows_batcher(self, target, temp_rows, handler, cluster_con assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_mutate_row(self, target, temp_rows, handler, cluster_config): + async def test_mutate_row(self, table, temp_rows, handler, cluster_config): row_key = b"bulk_mutate" new_value = uuid.uuid4().hex.encode() row_key, mutation = await temp_rows.create_row_and_mutation( - target, new_value=new_value + table, new_value=new_value ) handler.clear() - await target.mutate_row(row_key, mutation) + await table.mutate_row(row_key, mutation) # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 @@ -406,8 +413,8 @@ async def test_mutate_row(self, target, temp_rows, handler, cluster_config): assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config): - await target.sample_row_keys() + 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 @@ -436,7 +443,7 @@ async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config) assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_modify_write(self, target, temp_rows, handler, cluster_config): + 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" @@ -446,7 +453,7 @@ async def test_read_modify_write(self, target, temp_rows, handler, cluster_confi row_key, value=0, family=family, qualifier=qualifier ) rule = IncrementRule(family, qualifier, 1) - await target.read_modify_write_row(row_key, rule) + await table.read_modify_write_row(row_key, rule) # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 @@ -475,7 +482,7 @@ async def test_read_modify_write(self, target, temp_rows, handler, cluster_confi assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_config): + 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 @@ -491,7 +498,7 @@ async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_co family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value ) predicate = ValueRangeFilter(0, 2) - await target.check_and_mutate_row( + await table.check_and_mutate_row( row_key, predicate, true_case_mutations=true_mutation, @@ -525,7 +532,7 @@ async def test_check_and_mutate_row(self, target, temp_rows, handler, cluster_co @CrossSync.pytest async def test_check_and_mutate_row_failure_grpc( - self, target, temp_rows, handler, error_injector + self, table, temp_rows, handler, error_injector ): """ Test failure in grpc layer by injecting an error into an interceptor @@ -544,7 +551,7 @@ async def test_check_and_mutate_row_failure_grpc( exc = RuntimeError("injected") error_injector.push(exc) with pytest.raises(RuntimeError): - await target.check_and_mutate_row( + await table.check_and_mutate_row( row_key, predicate=ValueRangeFilter(0,2), ) @@ -579,13 +586,13 @@ async def test_check_and_mutate_row_failure_grpc( @CrossSync.pytest - async def test_check_and_mutate_row_failure_invalid_argument( - self, target, temp_rows, handler + async def test_check_and_mutate_row_failure_timeout( + self, table, temp_rows, handler ): """ - Test failure on backend by passing invalid argument + Test failure in gapic layer by passing very low timeout - We expect a server-timing header, but no cluster/zone info + No grpc headers expected """ from google.cloud.bigtable.data.mutations import SetCell from google.cloud.bigtable.data.row_filters import ValueRangeFilter @@ -595,11 +602,16 @@ async def test_check_and_mutate_row_failure_invalid_argument( qualifier = b"test-qualifier" await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) - predicate = ValueRangeFilter(-1, -1) + 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 target.check_and_mutate_row( + await table.check_and_mutate_row( row_key, - predicate, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=true_mutation, + operation_timeout=0.001 ) # validate counts assert len(handler.completed_operations) == 1 @@ -608,41 +620,35 @@ async def test_check_and_mutate_row_failure_invalid_argument( # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "INVALID_ARGUMENT" + 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 >= 0 and attempt.gfe_latency_ns < operation.duration_ns - + assert attempt.gfe_latency_ns is None @CrossSync.pytest - async def test_check_and_mutate_row_failure_timeout( - self, target, temp_rows, handler, error_injector + async def test_check_and_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config ): """ - Test failure in gapic layer by passing very low timeout - - No grpc headers expected + 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" - 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 + mutation_value = b"true-mutation-value" + mutation = SetCell( + family="unauthorized", qualifier=qualifier, new_value=mutation_value ) with pytest.raises(GoogleAPICallError): - await target.check_and_mutate_row( + await authorized_view.check_and_mutate_row( row_key, predicate=ValueRangeFilter(0, 2), - true_case_mutations=true_mutation, - operation_timeout=0.001 + true_case_mutations=mutation, + false_case_mutations=mutation, ) # validate counts assert len(handler.completed_operations) == 1 @@ -651,9 +657,9 @@ async def test_check_and_mutate_row_failure_timeout( # 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" + 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 is None \ No newline at end of file + assert attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns \ No newline at end of file From eb8a229fd476d87ffc7fea718fc66cc3e3b370c6 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 12:19:14 -0700 Subject: [PATCH 31/78] added stubs --- tests/system/data/test_metrics_async.py | 309 +++++++++++++++++++++++- 1 file changed, 308 insertions(+), 1 deletion(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 49fe44895..d815b4770 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -180,6 +180,35 @@ async def test_read_rows(self, table, temp_rows, handler, cluster_config): 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_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + @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 + """ + pass + + @CrossSync.pytest + async def test_read_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + pass + @CrossSync.pytest async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): await temp_rows.add_row(b"row_key_1") @@ -216,6 +245,46 @@ async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config) 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_stream_failure_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + @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 + """ + pass + + @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 + """ + pass + + @CrossSync.pytest + async def test_read_rows_stream_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + pass + @CrossSync.pytest async def test_read_row(self, table, temp_rows, handler, cluster_config): await temp_rows.add_row(b"row_key_1") @@ -248,6 +317,35 @@ async def test_read_row(self, table, temp_rows, handler, cluster_config): 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_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + @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 + """ + pass + + @CrossSync.pytest + async def test_read_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + pass + @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 @@ -287,6 +385,46 @@ async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config 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_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + @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 + """ + pass + + @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 + """ + pass + + @CrossSync.pytest + async def test_read_rows_sharded_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + pass + @CrossSync.pytest async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): from google.cloud.bigtable.data.mutations import RowMutationEntry @@ -326,6 +464,37 @@ async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config) 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_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + @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 + """ + pass + + @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 + """ + pass + @CrossSync.pytest async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): from google.cloud.bigtable.data.mutations import RowMutationEntry @@ -376,9 +545,40 @@ async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_conf 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_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + @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 + """ + pass + + @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 + """ + pass + @CrossSync.pytest async def test_mutate_row(self, table, temp_rows, handler, cluster_config): - row_key = b"bulk_mutate" + 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 @@ -412,6 +612,38 @@ async def test_mutate_row(self, table, temp_rows, handler, cluster_config): assert attempt.application_blocking_time_ns == 0 assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + @CrossSync.pytest + async def test_mutate_row_failure_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + + @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 + """ + pass + + @CrossSync.pytest + async def test_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + pass + @CrossSync.pytest async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): await table.sample_row_keys() @@ -442,6 +674,48 @@ async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): assert attempt.application_blocking_time_ns == 0 assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + @CrossSync.pytest + async def test_sample_row_keys_failure_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + + @CrossSync.pytest + async def test_sample_row_keys_failure_timeout( + self, table, temp_rows, handler + ): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + pass + + @CrossSync.pytest + async def test_sample_row_keys_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + pass + + @CrossSync.pytest + async def test_sample_row_keys_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + pass + + @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 @@ -481,6 +755,39 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config assert attempt.application_blocking_time_ns == 0 assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm + @CrossSync.pytest + async def test_read_modify_write_row_failure_grpc( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + pass + + + @CrossSync.pytest + async def test_read_modify_write_row_failure_timeout( + self, table, temp_rows, handler + ): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + pass + + @CrossSync.pytest + async def test_read_modify_write_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + pass + + @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 761cf9cbcecd6171e8b4a863db3471b1f8cde850 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 14:56:36 -0700 Subject: [PATCH 32/78] added tests for sample_row_keys --- tests/system/data/test_metrics_async.py | 182 +++++++++++++++++++++--- 1 file changed, 165 insertions(+), 17 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index d815b4770..3fae8a152 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -19,7 +19,9 @@ from grpc.aio import AioRpcError from grpc.aio import Metadata +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_v2.types import ResponseParams from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric, ActiveOperationMetric, OperationState @@ -66,19 +68,21 @@ def __repr__(self): return f"{self.__class__}(completed_operations={len(self.completed_operations)}, cancelled_operations={len(self.cancelled_operations)}, completed_attempts={len(self.completed_attempts)}" -class _ErrorInjectorInterceptor(UnaryUnaryClientInterceptor): +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 async def intercept_unary_unary( self, continuation, client_call_details, request @@ -87,6 +91,41 @@ async def intercept_unary_unary( raise self._exc_list.pop(0) return await continuation(client_call_details, request) + 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 + + def __aiter__(self): + return self + + 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): @@ -123,6 +162,9 @@ 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 + ) yield client @CrossSync.convert @@ -679,41 +721,147 @@ async def test_sample_row_keys_failure_grpc( self, table, temp_rows, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor + Test failure in grpc layer by injecting errors into an interceptor + test with retryable errors, then a terminal one + + No headers expected + """ + exc = Aborted("injected") + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(exc) + error_injector.push(RuntimeError) + with pytest.raises(RuntimeError): + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable+1 + assert len(handler.cancelled_operations) == 0 + # 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_grpc_eventual_success( + self, table, temp_rows, handler, error_injector, cluster_config + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + test with retryable errors, then a success + No headers expected """ - pass + exc = Aborted("injected") + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(exc) + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable+1 + assert len(handler.cancelled_operations) == 0 + # 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, temp_rows, handler + self, table, handler ): """ Test failure in gapic layer by passing very low timeout No grpc headers expected """ - pass - - @CrossSync.pytest - async def test_sample_row_keys_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """ - Test failure in backend by accessing an unauthorized family - """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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 + self, table, temp_rows, handler, error_injector, cluster_config ): """ Test failure in grpc stream """ - pass + error_injector.fail_mid_stream = True + error_injector.push(Aborted("retryable")) + error_injector.push(PermissionDenied("terminal")) + 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 + assert len(handler.cancelled_operations) == 0 + # 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 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # 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 @@ -969,4 +1117,4 @@ async def test_check_and_mutate_row_failure_unauthorized( 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 \ No newline at end of file + assert attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns From c33b8135eea5030d992092c16388fb204415db3a Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 15:05:08 -0700 Subject: [PATCH 33/78] sped up test --- google/cloud/bigtable/data/_async/metrics_interceptor.py | 3 ++- tests/system/data/test_metrics_async.py | 7 +------ 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index b4b6da642..d9f6dfb3d 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -173,7 +173,8 @@ async def response_wrapper(call): raise finally: if call is not None: - _end_attempt(operation, encountered_exc, await _get_metadata(call)) + metadata = await _get_metadata(encountered_exc or call) + _end_attempt(operation, encountered_exc, metadata) try: return response_wrapper(await continuation(client_call_details, request)) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 3fae8a152..5200b4509 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -831,7 +831,7 @@ async def test_sample_row_keys_failure_timeout( @CrossSync.pytest async def test_sample_row_keys_failure_mid_stream( - self, table, temp_rows, handler, error_injector, cluster_config + self, table, temp_rows, handler, error_injector ): """ Test failure in grpc stream @@ -851,11 +851,6 @@ async def test_sample_row_keys_failure_mid_stream( assert operation.op_type.value == "SampleRowKeys" assert operation.is_streaming is False assert len(operation.completed_attempts) == 2 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) # validate retried attempt attempt = handler.completed_attempts[0] assert attempt.end_status.name == "ABORTED" From 661d6eb623a514069ccfcdaf24bc09ecd16008be Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 15:10:31 -0700 Subject: [PATCH 34/78] added read_modify_write tests --- tests/system/data/test_metrics_async.py | 99 +++++++++++++++++++++++-- 1 file changed, 93 insertions(+), 6 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 5200b4509..a6c205578 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -899,7 +899,7 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_modify_write_row_failure_grpc( + async def test_read_modify_write_failure_grpc( self, table, temp_rows, handler, error_injector ): """ @@ -907,11 +907,54 @@ async def test_read_modify_write_row_failure_grpc( No headers expected """ - pass + 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 = RuntimeError("injected") + error_injector.push(exc) + with pytest.raises(RuntimeError): + await table.read_modify_write_row(row_key, rule) + + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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_row_failure_timeout( + async def test_read_modify_write_failure_timeout( self, table, temp_rows, handler ): """ @@ -919,16 +962,60 @@ async def test_read_modify_write_row_failure_timeout( No grpc headers expected """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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_row_failure_unauthorized( + async def test_read_modify_write_failure_unauthorized( self, handler, authorized_view, cluster_config ): """ Test failure in backend by accessing an unauthorized family """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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 From 393574481db706b173367d3566389044fe392416 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 16:47:01 -0700 Subject: [PATCH 35/78] added operation logic to retry factory --- google/cloud/bigtable/data/_async/_read_rows.py | 5 ++--- google/cloud/bigtable/data/_helpers.py | 11 +++++++++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index b935798ec..1433b0920 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -18,6 +18,7 @@ from typing import Sequence, TYPE_CHECKING import time +from functools import partial from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB @@ -121,7 +122,7 @@ def start_operation(self) -> CrossSync.Iterable[Row]: self._predicate, self._operation_metric.backoff_generator, self.operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=partial(_retry_exception_factory, operation=self._operation_metric), ) def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: @@ -340,8 +341,6 @@ async def merge_rows( except CrossSync.StopIteration: raise InvalidChunk("premature end of stream") except Exception as generic_exception: - if not self._predicate(generic_exception): - self._operation_metric.end_attempt_with_status(generic_exception) raise generic_exception else: self._operation_metric.end_with_success() diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index e848ebc6f..6b25ab1b9 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -90,6 +90,7 @@ def _retry_exception_factory( exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None, + operation: "ActiveOperationMetric" | None=None, ) -> tuple[Exception, Exception | None]: """ Build retry error based on exceptions encountered during operation @@ -117,6 +118,16 @@ def _retry_exception_factory( # use the retry exception group as the cause of the exception cause_exc: Exception | None = RetryExceptionGroup(exc_list) if exc_list else None source_exc.__cause__ = cause_exc + if operation: + try: + if isinstance(source_exc, core_exceptions.GoogleAPICallError) and source_exc.errors: + rpc_error = source_exc.errors[-1] + metadata = list(rpc_error.trailing_metadata()) + list(rpc_error.initial_metadata()) + operation.add_response_metadata({k:v for k,v in metadata}) + except Exception: + # ignore errors in metadata collection + pass + operation.end_with_status(source_exc) return source_exc, cause_exc From 0a9af2d53b23db813763d2957c8b2f3396060b08 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 16:47:54 -0700 Subject: [PATCH 36/78] added read_rows tests --- tests/system/data/test_metrics_async.py | 193 ++++++++++++++++++++++-- 1 file changed, 184 insertions(+), 9 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index a6c205578..ef9a82823 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -25,6 +25,7 @@ from google.cloud.bigtable_v2.types import ResponseParams from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric, ActiveOperationMetric, OperationState +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data._cross_sync import CrossSync @@ -193,7 +194,7 @@ 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({}) + row_list = await table.read_rows(ReadRowsQuery()) assert len(row_list) == 2 # validate counts assert len(handler.completed_operations) == 1 @@ -231,7 +232,38 @@ async def test_read_rows_failure_grpc( No headers expected """ - pass + await temp_rows.add_row(b"row_key_1") + handler.clear() + exc = Aborted("injected") + num_retryable = 2 + for i in range(num_retryable): + error_injector.push(exc) + error_injector.push(PermissionDenied("terminal")) + 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 + assert len(handler.cancelled_operations) == 0 + # 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 == "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 == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None @CrossSync.pytest async def test_read_rows_failure_timeout(self, table, temp_rows, handler): @@ -240,7 +272,28 @@ async def test_read_rows_failure_timeout(self, table, temp_rows, handler): No grpc headers expected """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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( @@ -249,7 +302,29 @@ async def test_read_rows_failure_unauthorized( """ Test failure in backend by accessing an unauthorized family """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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): @@ -257,7 +332,7 @@ async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config) await temp_rows.add_row(b"row_key_2") handler.clear() # full table scan - generator = await table.read_rows_stream({}) + generator = await table.read_rows_stream(ReadRowsQuery()) row_list = [r async for r in generator] assert len(row_list) == 2 # validate counts @@ -296,7 +371,39 @@ async def test_read_rows_stream_failure_grpc( No headers expected """ - pass + await temp_rows.add_row(b"row_key_1") + handler.clear() + exc = Aborted("injected") + num_retryable = 2 + for i in range(num_retryable): + error_injector.push(exc) + error_injector.push(PermissionDenied("terminal")) + 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 + assert len(handler.cancelled_operations) == 0 + # 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 == "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 == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None @CrossSync.pytest async def test_read_rows_stream_failure_timeout( @@ -307,7 +414,29 @@ async def test_read_rows_stream_failure_timeout( No grpc headers expected """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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( @@ -316,7 +445,30 @@ async def test_read_rows_stream_failure_unauthorized( """ Test failure in backend by accessing an unauthorized family """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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_mid_stream( @@ -325,7 +477,30 @@ async def test_read_rows_stream_failure_mid_stream( """ Test failure in grpc stream """ - pass + await temp_rows.add_row(b"row_key_1") + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(Aborted("retryable")) + error_injector.push(PermissionDenied("terminal")) + 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 + assert len(handler.cancelled_operations) == 0 + # 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): From b52c871afaca66b6e5c69fa71e4332409b666692 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 16 Sep 2025 17:14:58 -0700 Subject: [PATCH 37/78] added read_row and read_rows_sharded tests --- google/cloud/bigtable/data/_async/client.py | 2 +- tests/system/data/test_metrics_async.py | 220 +++++++++++++++++++- 2 files changed, 214 insertions(+), 8 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 9cbf99a4f..91e98f0c1 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1391,7 +1391,7 @@ async def execute_rpc(): predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=partial(_retry_exception_factory, operation=operation_metric) ) @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"}) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index ef9a82823..886274ed1 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -543,7 +543,38 @@ async def test_read_row_failure_grpc( No headers expected """ - pass + await temp_rows.add_row(b"row_key_1") + handler.clear() + exc = Aborted("injected") + num_retryable = 2 + for i in range(num_retryable): + error_injector.push(exc) + error_injector.push(PermissionDenied("terminal")) + 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 + assert len(handler.cancelled_operations) == 0 + # 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 == "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 == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None @CrossSync.pytest async def test_read_row_failure_timeout(self, table, temp_rows, handler): @@ -552,7 +583,28 @@ async def test_read_row_failure_timeout(self, table, temp_rows, handler): No grpc headers expected """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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( @@ -561,7 +613,29 @@ async def test_read_row_failure_unauthorized( """ Test failure in backend by accessing an unauthorized family """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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): @@ -611,7 +685,46 @@ async def test_read_rows_sharded_failure_grpc( No headers expected """ - pass + 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.push(PermissionDenied("terminal")) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2]) + 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) == 2 + assert len(handler.cancelled_operations) == 0 + # 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 == "unspecified" + assert failed_op.zone == "global" + # validate failed attempt + failed_attempt = failed_op.completed_attempts[0] + assert failed_attempt.end_status.name == "PERMISSION_DENIED" + assert failed_attempt.gfe_latency_ns is None + # 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_timeout( @@ -622,7 +735,38 @@ async def test_read_rows_sharded_failure_timeout( No grpc headers expected """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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( @@ -631,7 +775,44 @@ async def test_read_rows_sharded_failure_unauthorized( """ Test failure in backend by accessing an unauthorized family """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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( @@ -640,7 +821,32 @@ async def test_read_rows_sharded_failure_mid_stream( """ Test failure in grpc stream """ - pass + 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(PermissionDenied("terminal")) + error_injector.push(PermissionDenied("terminal")) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2]) + assert len(e.value.exceptions) == 2 + assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + assert len(handler.cancelled_operations) == 0 + for operation in handler.completed_operations: + 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 + # validate attempt + attempt = operation.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" @CrossSync.pytest async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): From 759d1986608f6014b817929dd073f12a48e075aa Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Sep 2025 11:17:10 -0700 Subject: [PATCH 38/78] added bulk_mutate_row tests --- .../bigtable/data/_async/_mutate_rows.py | 6 +- .../bigtable/data/_metrics/data_model.py | 11 +- tests/system/data/test_metrics_async.py | 136 ++++++++++++++++-- 3 files changed, 133 insertions(+), 20 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 9c1fc30b3..d2fb089a6 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -16,6 +16,8 @@ from typing import Sequence, TYPE_CHECKING +from functools import partial + from google.api_core import exceptions as core_exceptions from google.api_core import retry as retries import google.cloud.bigtable_v2.types.bigtable as types_pb @@ -106,7 +108,9 @@ def __init__( self.is_retryable, metric.backoff_generator, operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=partial( + _retry_exception_factory, operation=metric + ), ) # initialize state self.timeout_generator = _attempt_timeout_generator( diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 6c4572d24..4b4c77604 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -418,8 +418,11 @@ def __exit__(self, exc_type, exc_val, exc_tb): The operation is automatically ended on exit, with the status determined by the exception type and value. + + If operation was already ended manually, do nothing. """ - if exc_val is None: - self.end_with_success() - else: - self.end_with_status(exc_val) + if not self.state == OperationState.COMPLETED: + if exc_val is None: + self.end_with_success() + else: + self.end_with_status(exc_val) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 886274ed1..e433d94b3 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -830,23 +830,35 @@ async def test_read_rows_sharded_failure_mid_stream( query2 = ReadRowsQuery(row_keys=[b"b"]) handler.clear() error_injector.fail_mid_stream = True - error_injector.push(PermissionDenied("terminal")) + error_injector.push(Aborted("retryable")) error_injector.push(PermissionDenied("terminal")) with pytest.raises(ShardedReadRowsExceptionGroup) as e: - await table.read_rows_sharded([query1, query2]) - assert len(e.value.exceptions) == 2 + 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) == 2 + assert len(handler.completed_attempts) == 3 assert len(handler.cancelled_operations) == 0 - for operation in handler.completed_operations: - 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 - # validate attempt - attempt = operation.completed_attempts[0] - assert attempt.end_status.name == "PERMISSION_DENIED" + # 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) == 2 + # validate retried attempt + attempt = failed_op.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + # validate final attempt + final_attempt = failed_op.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + # validate successful operation + assert success_op.final_status.name == "OK" + assert len(success_op.completed_attempts) == 1 + assert success_op.completed_attempts[0].end_status.name == "OK" @CrossSync.pytest async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): @@ -896,7 +908,45 @@ async def test_bulk_mutate_rows_failure_grpc( No headers expected """ - pass + 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() + exc = Aborted("injected") + num_retryable = 2 + for i in range(num_retryable): + error_injector.push(exc) + error_injector.push(PermissionDenied("terminal")) + with pytest.raises(MutationsExceptionGroup) as e: + 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 + assert len(handler.cancelled_operations) == 0 + # 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 == "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 == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None @CrossSync.pytest async def test_bulk_mutate_rows_failure_timeout( @@ -907,7 +957,35 @@ async def test_bulk_mutate_rows_failure_timeout( No grpc headers expected """ - pass + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.api_core.exceptions import DeadlineExceeded + + 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) as e: + await table.bulk_mutate_rows([entry], operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + assert len(handler.cancelled_operations) == 0 + # 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( @@ -916,7 +994,35 @@ async def test_bulk_mutate_rows_failure_unauthorized( """ Test failure in backend by accessing an unauthorized family """ - pass + 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]) + 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 + assert len(handler.cancelled_operations) == 0 + # 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_rows_batcher(self, table, temp_rows, handler, cluster_config): From 18fba8a8c1dbc3a9f18266f631b49aaa50fa457a Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Sep 2025 11:27:01 -0700 Subject: [PATCH 39/78] added batcher tests --- tests/system/data/test_metrics_async.py | 104 +++++++++++++++++++++++- 1 file changed, 101 insertions(+), 3 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index e433d94b3..6670f1b69 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -1083,7 +1083,43 @@ async def test_mutate_rows_batcher_failure_grpc( No headers expected """ - pass + 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() + + exc = Aborted("injected") + num_retryable = 2 + for i in range(num_retryable): + error_injector.push(exc) + error_injector.push(PermissionDenied("terminal")) + with pytest.raises(MutationsExceptionGroup) as e: + 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 + assert len(handler.cancelled_operations) == 1 # from batcher auto-closing + # 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 == "unspecified" + assert operation.zone == "global" + # 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( @@ -1094,7 +1130,33 @@ async def test_mutate_rows_batcher_failure_timeout( No grpc headers expected """ - pass + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.api_core.exceptions import DeadlineExceeded + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup) as e: + 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 + assert len(handler.cancelled_operations) == 1 # from batcher auto-closing + # 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( @@ -1103,7 +1165,43 @@ async def test_mutate_rows_batcher_failure_unauthorized( """ Test failure in backend by accessing an unauthorized family """ - pass + 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 + assert len(handler.cancelled_operations) == 1 # from batcher auto-closing + # 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): From c1071186ab7ca10e7c2838b936bf8d66d876255c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Sep 2025 12:16:17 -0700 Subject: [PATCH 40/78] added mutate_row tests --- tests/system/data/test_metrics_async.py | 89 ++++++++++++++++++++++++- 1 file changed, 86 insertions(+), 3 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 6670f1b69..bde7967f6 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -1248,7 +1248,41 @@ async def test_mutate_row_failure_grpc( No headers expected """ - pass + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + + exc = Aborted("injected") + num_retryable = 2 + for i in range(num_retryable): + error_injector.push(exc) + error_injector.push(PermissionDenied("terminal")) + 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 + assert len(handler.cancelled_operations) == 0 + # 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 == "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 == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None @CrossSync.pytest @@ -1260,7 +1294,31 @@ async def test_mutate_row_failure_timeout( No grpc headers expected """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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( @@ -1269,7 +1327,32 @@ async def test_mutate_row_failure_unauthorized( """ Test failure in backend by accessing an unauthorized family """ - pass + 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 + assert len(handler.cancelled_operations) == 0 + # 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_sample_row_keys(self, table, temp_rows, handler, cluster_config): From 2a42b9012fe0596d8967b2343f143c22bbfdd4ab Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Sep 2025 12:35:14 -0700 Subject: [PATCH 41/78] fixed read_rows_sharded test --- tests/system/data/test_metrics_async.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index bde7967f6..7d9608d18 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -848,17 +848,19 @@ async def test_read_rows_sharded_failure_mid_stream( 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) == 2 - # validate retried attempt - attempt = failed_op.completed_attempts[0] - assert attempt.end_status.name == "ABORTED" - # validate final attempt - final_attempt = failed_op.completed_attempts[-1] - assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert len(failed_op.completed_attempts) == 1 # validate successful operation assert success_op.final_status.name == "OK" - assert len(success_op.completed_attempts) == 1 - assert success_op.completed_attempts[0].end_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 fb84b4bf538aac78681b0f214a4b267ea015a028 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Sep 2025 16:14:02 -0700 Subject: [PATCH 42/78] refacotred tracked exception factory --- .../bigtable/data/_async/_mutate_rows.py | 6 ++-- .../cloud/bigtable/data/_async/_read_rows.py | 4 +-- google/cloud/bigtable/data/_async/client.py | 3 +- google/cloud/bigtable/data/_helpers.py | 36 ++++++++++++++++--- 4 files changed, 38 insertions(+), 11 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index d2fb089a6..da9037bdc 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -23,7 +23,7 @@ import google.cloud.bigtable_v2.types.bigtable as types_pb import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._helpers import _tracked_exception_factory # mutate_rows requests are limited to this number of mutations from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT @@ -108,9 +108,7 @@ def __init__( self.is_retryable, metric.backoff_generator, operation_timeout, - exception_factory=partial( - _retry_exception_factory, operation=metric - ), + exception_factory=_tracked_exception_factory(self._operation_metric), ) # initialize state self.timeout_generator = _attempt_timeout_generator( diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 1433b0920..36a0cc298 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -31,7 +31,7 @@ from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._helpers import _tracked_exception_factory from google.api_core import retry as retries @@ -122,7 +122,7 @@ def start_operation(self) -> CrossSync.Iterable[Row]: self._predicate, self._operation_metric.backoff_generator, self.operation_timeout, - exception_factory=partial(_retry_exception_factory, operation=self._operation_metric), + exception_factory=_tracked_exception_factory(self._operation_metric), ) def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 91e98f0c1..2e284373a 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -73,6 +73,7 @@ from google.cloud.bigtable.data._helpers import _WarmedInstanceKey from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._helpers import _tracked_exception_factory from google.cloud.bigtable.data._helpers import _validate_timeouts from google.cloud.bigtable.data._helpers import _get_error_type from google.cloud.bigtable.data._helpers import _get_retryable_errors @@ -1391,7 +1392,7 @@ async def execute_rpc(): predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=partial(_retry_exception_factory, operation=operation_metric) + exception_factory=_tracked_exception_factory(operation_metric), ) @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"}) diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 6b25ab1b9..b1ac9eb6b 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -16,7 +16,7 @@ """ from __future__ import annotations -from typing import Sequence, List, Tuple, TYPE_CHECKING, Union +from typing import Callable, Sequence, List, Tuple, TYPE_CHECKING, Union import time import enum from collections import namedtuple @@ -31,6 +31,7 @@ import grpc from google.cloud.bigtable.data._async.client import _DataApiTargetAsync from google.cloud.bigtable.data._sync_autogen.client import _DataApiTarget + from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric """ Helper functions used in various places in the library. @@ -90,7 +91,6 @@ def _retry_exception_factory( exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None, - operation: "ActiveOperationMetric" | None=None, ) -> tuple[Exception, Exception | None]: """ Build retry error based on exceptions encountered during operation @@ -118,8 +118,35 @@ def _retry_exception_factory( # use the retry exception group as the cause of the exception cause_exc: Exception | None = RetryExceptionGroup(exc_list) if exc_list else None source_exc.__cause__ = cause_exc - if operation: + return source_exc, cause_exc + + +def _tracked_exception_factory( + operation: "ActiveOperationMetric", +) -> Callable[[list[Exception], RetryFailureReason, float | None], tuple[Exception, Exception | None]]: + """ + wraps and extends _retry_exception_factory to add client-side metrics tracking. + + When the rpc raises a terminal error, record any discovered metadata and finalize + the associated operation + + Used by streaming rpcs, which can't always be perfectly captured by context managers + for operation termination. + + Args: + exc_list: list of exceptions encountered during operation + is_timeout: whether the operation failed due to timeout + timeout_val: the operation timeout value in seconds, for constructing + the error message + operation: the operation to finalize when an exception is built + Returns: + tuple[Exception, Exception|None]: + tuple of the exception to raise, and a cause exception if applicabl + """ + def wrapper(exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None) -> tuple[Exception, Exception | None]: + source_exc, cause_exc = _retry_exception_factory(exc_list, reason, timeout_val) try: + # record metadata from failed rpc if isinstance(source_exc, core_exceptions.GoogleAPICallError) and source_exc.errors: rpc_error = source_exc.errors[-1] metadata = list(rpc_error.trailing_metadata()) + list(rpc_error.initial_metadata()) @@ -128,7 +155,8 @@ def _retry_exception_factory( # ignore errors in metadata collection pass operation.end_with_status(source_exc) - return source_exc, cause_exc + return source_exc, cause_exc + return wrapper def _get_timeouts( From c560589afb7f419a1bc41801f8a5051fd599a858 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 17 Sep 2025 16:21:22 -0700 Subject: [PATCH 43/78] fixed lint --- .../bigtable/data/_async/_mutate_rows.py | 2 - .../cloud/bigtable/data/_async/_read_rows.py | 1 - .../data/_async/metrics_interceptor.py | 5 +- google/cloud/bigtable/data/_helpers.py | 22 +- tests/system/data/__init__.py | 4 +- tests/system/data/test_metrics_async.py | 436 ++++++++++++------ tests/system/data/test_system_async.py | 8 +- 7 files changed, 323 insertions(+), 155 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index da9037bdc..03d870e49 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -16,8 +16,6 @@ from typing import Sequence, TYPE_CHECKING -from functools import partial - from google.api_core import exceptions as core_exceptions from google.api_core import retry as retries import google.cloud.bigtable_v2.types.bigtable as types_pb diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 36a0cc298..d4ee5b2b1 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -18,7 +18,6 @@ from typing import Sequence, TYPE_CHECKING import time -from functools import partial from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index d9f6dfb3d..fe16b2d24 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -156,7 +156,10 @@ async def intercept_unary_stream( ): async def response_wrapper(call): # only track has_first response for READ_ROWS - has_first_response = operation.first_response_latency_ns is not None or operation.op_type != OperationType.READ_ROWS + has_first_response = ( + operation.first_response_latency_ns is not None + or operation.op_type != OperationType.READ_ROWS + ) encountered_exc = None try: async for response in call: diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index b1ac9eb6b..0db1de46f 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -123,7 +123,10 @@ def _retry_exception_factory( def _tracked_exception_factory( operation: "ActiveOperationMetric", -) -> Callable[[list[Exception], RetryFailureReason, float | None], tuple[Exception, Exception | None]]: +) -> Callable[ + [list[Exception], RetryFailureReason, float | None], + tuple[Exception, Exception | None], +]: """ wraps and extends _retry_exception_factory to add client-side metrics tracking. @@ -143,19 +146,28 @@ def _tracked_exception_factory( tuple[Exception, Exception|None]: tuple of the exception to raise, and a cause exception if applicabl """ - def wrapper(exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None) -> tuple[Exception, Exception | None]: + + def wrapper( + exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None + ) -> tuple[Exception, Exception | None]: source_exc, cause_exc = _retry_exception_factory(exc_list, reason, timeout_val) try: # record metadata from failed rpc - if isinstance(source_exc, core_exceptions.GoogleAPICallError) and source_exc.errors: + if ( + isinstance(source_exc, core_exceptions.GoogleAPICallError) + and source_exc.errors + ): rpc_error = source_exc.errors[-1] - metadata = list(rpc_error.trailing_metadata()) + list(rpc_error.initial_metadata()) - operation.add_response_metadata({k:v for k,v in metadata}) + metadata = list(rpc_error.trailing_metadata()) + list( + rpc_error.initial_metadata() + ) + operation.add_response_metadata({k: v for k, v in metadata}) except Exception: # ignore errors in metadata collection pass operation.end_with_status(source_exc) return source_exc, cause_exc + return wrapper diff --git a/tests/system/data/__init__.py b/tests/system/data/__init__.py index d123f43fe..dcd14c5f9 100644 --- a/tests/system/data/__init__.py +++ b/tests/system/data/__init__.py @@ -15,12 +15,12 @@ # import pytest import uuid -from google.cloud.bigtable.data._cross_sync import CrossSync TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" TEST_AGGREGATE_FAMILY = "test-aggregate-family" + class SystemTestRunner: """ configures a system test class with configuration for clusters/tables/etc @@ -67,4 +67,4 @@ def column_family_config(self): TEST_AGGREGATE_FAMILY: types.ColumnFamily( value_type=types.Type(aggregate_type=int_aggregate_type) ), - } \ No newline at end of file + } diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 7d9608d18..558781e75 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -15,16 +15,17 @@ import os import pytest import uuid -from grpc import RpcError -from grpc.aio import AioRpcError -from grpc.aio import Metadata 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_v2.types import ResponseParams from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler -from google.cloud.bigtable.data._metrics.data_model import CompletedOperationMetric, CompletedAttemptMetric, ActiveOperationMetric, OperationState +from google.cloud.bigtable.data._metrics.data_model import ( + CompletedOperationMetric, + CompletedAttemptMetric, + ActiveOperationMetric, + OperationState, +) from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data._cross_sync import CrossSync @@ -69,7 +70,9 @@ def __repr__(self): return f"{self.__class__}(completed_operations={len(self.completed_operations)}, cancelled_operations={len(self.cancelled_operations)}, completed_attempts={len(self.completed_attempts)}" -class _ErrorInjectorInterceptor(UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor): +class _ErrorInjectorInterceptor( + UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor +): """ Gprc interceptor used to inject errors into rpc calls, to test failures """ @@ -85,16 +88,12 @@ def clear(self): self._exc_list.clear() self.fail_mid_stream = False - async def intercept_unary_unary( - self, continuation, client_call_details, request - ): + 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) - async def intercept_unary_stream( - self, continuation, client_call_details, request - ): + 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) @@ -129,7 +128,6 @@ def __getattr__(self, name): @CrossSync.convert_class(sync_name="TestMetrics") class TestMetricsAsync(SystemTestRunner): - @CrossSync.drop @pytest.fixture(scope="session") def event_loop(self): @@ -162,7 +160,9 @@ async def _clear_state(self, handler, error_injector): 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_unary_interceptors.append( + error_injector + ) client.transport.grpc_channel._unary_stream_interceptors.append( error_injector ) @@ -184,8 +184,12 @@ async def table(self, client, table_id, instance_id, handler): @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: + 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 @@ -209,9 +213,15 @@ async def test_read_rows(self, table, temp_rows, handler, cluster_config): 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.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.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] @@ -219,8 +229,13 @@ async def test_read_rows(self, table, temp_rows, handler, cluster_config): 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.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 @@ -305,7 +320,9 @@ async def test_read_rows_failure_unauthorized( 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"))) + 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 @@ -319,12 +336,18 @@ async def test_read_rows_failure_unauthorized( 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] + 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 + 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): @@ -348,9 +371,15 @@ async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config) 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.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.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] @@ -358,8 +387,13 @@ async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config) 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.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 @@ -378,7 +412,9 @@ async def test_read_rows_stream_failure_grpc( for i in range(num_retryable): error_injector.push(exc) error_injector.push(PermissionDenied("terminal")) - generator = await table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) with pytest.raises(PermissionDenied): [_ async for _ in generator] # validate counts @@ -406,9 +442,7 @@ async def test_read_rows_stream_failure_grpc( assert final_attempt.gfe_latency_ns is None @CrossSync.pytest - async def test_read_rows_stream_failure_timeout( - self, table, temp_rows, handler - ): + async def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): """ Test failure in gapic layer by passing very low timeout @@ -416,7 +450,9 @@ async def test_read_rows_stream_failure_timeout( """ await temp_rows.add_row(b"row_key_1") handler.clear() - generator = await table.read_rows_stream(ReadRowsQuery(), operation_timeout=0.001) + generator = await table.read_rows_stream( + ReadRowsQuery(), operation_timeout=0.001 + ) with pytest.raises(GoogleAPICallError): [_ async for _ in generator] # validate counts @@ -448,7 +484,9 @@ async def test_read_rows_stream_failure_unauthorized( 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"))) + 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 @@ -463,12 +501,18 @@ async def test_read_rows_stream_failure_unauthorized( 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] + 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 + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) @CrossSync.pytest async def test_read_rows_stream_failure_mid_stream( @@ -482,7 +526,9 @@ async def test_read_rows_stream_failure_mid_stream( error_injector.fail_mid_stream = True error_injector.push(Aborted("retryable")) error_injector.push(PermissionDenied("terminal")) - generator = await table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) with pytest.raises(PermissionDenied): [_ async for _ in generator] # validate counts @@ -520,9 +566,15 @@ async def test_read_row(self, table, temp_rows, handler, cluster_config): 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.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.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] @@ -530,8 +582,13 @@ async def test_read_row(self, table, temp_rows, handler, cluster_config): 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.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 @@ -616,7 +673,9 @@ async def test_read_row_failure_unauthorized( 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")) + 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 @@ -630,16 +689,23 @@ async def test_read_row_failure_unauthorized( 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] + 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 + 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") @@ -663,17 +729,31 @@ async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config 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.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.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.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.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 @@ -704,8 +784,12 @@ async def test_read_rows_sharded_failure_grpc( assert len(handler.completed_attempts) == 2 assert len(handler.cancelled_operations) == 0 # 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") + 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" @@ -727,9 +811,7 @@ async def test_read_rows_sharded_failure_grpc( assert success_attempt.end_status.name == "OK" @CrossSync.pytest - async def test_read_rows_sharded_failure_timeout( - self, table, temp_rows, handler - ): + async def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): """ Test failure in gapic layer by passing very low timeout @@ -786,25 +868,37 @@ async def test_read_rows_sharded_failure_unauthorized( 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" + 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 assert len(handler.cancelled_operations) == 0 # 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") + 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] + 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 + 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" @@ -842,8 +936,12 @@ async def test_read_rows_sharded_failure_mid_stream( assert len(handler.completed_attempts) == 3 assert len(handler.cancelled_operations) == 0 # 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") + 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" @@ -887,9 +985,14 @@ async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config) 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.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.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] @@ -897,7 +1000,9 @@ async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config) 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.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 @@ -924,7 +1029,7 @@ async def test_bulk_mutate_rows_failure_grpc( for i in range(num_retryable): error_injector.push(exc) error_injector.push(PermissionDenied("terminal")) - with pytest.raises(MutationsExceptionGroup) as e: + with pytest.raises(MutationsExceptionGroup): await table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) # validate counts assert len(handler.completed_operations) == 1 @@ -951,9 +1056,7 @@ async def test_bulk_mutate_rows_failure_grpc( assert final_attempt.gfe_latency_ns is None @CrossSync.pytest - async def test_bulk_mutate_rows_failure_timeout( - self, table, temp_rows, handler - ): + async def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): """ Test failure in gapic layer by passing very low timeout @@ -961,14 +1064,13 @@ async def test_bulk_mutate_rows_failure_timeout( """ from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.api_core.exceptions import DeadlineExceeded 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) as e: + with pytest.raises(MutationsExceptionGroup): await table.bulk_mutate_rows([entry], operation_timeout=0.001) # validate counts assert len(handler.completed_operations) == 1 @@ -1008,7 +1110,9 @@ async def test_bulk_mutate_rows_failure_unauthorized( await authorized_view.bulk_mutate_rows([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 ( + 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 @@ -1020,11 +1124,17 @@ async def test_bulk_mutate_rows_failure_unauthorized( 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] + 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 + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) @CrossSync.pytest async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): @@ -1062,17 +1172,27 @@ async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_conf 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.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 + 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.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 @@ -1098,8 +1218,10 @@ async def test_mutate_rows_batcher_failure_grpc( for i in range(num_retryable): error_injector.push(exc) error_injector.push(PermissionDenied("terminal")) - with pytest.raises(MutationsExceptionGroup) as e: - async with table.mutations_batcher(batch_retryable_errors=[Aborted]) as batcher: + 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 @@ -1124,9 +1246,7 @@ async def test_mutate_rows_batcher_failure_grpc( assert final_attempt.gfe_latency_ns is None @CrossSync.pytest - async def test_mutate_rows_batcher_failure_timeout( - self, table, temp_rows, handler - ): + async def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): """ Test failure in gapic layer by passing very low timeout @@ -1134,14 +1254,15 @@ async def test_mutate_rows_batcher_failure_timeout( """ from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.api_core.exceptions import DeadlineExceeded row_key = b"row_key_1" mutation = SetCell(TEST_FAMILY, b"q", b"v") entry = RowMutationEntry(row_key, [mutation]) - with pytest.raises(MutationsExceptionGroup) as e: - async with table.mutations_batcher(batch_operation_timeout=0.001) as batcher: + 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 @@ -1227,9 +1348,14 @@ async def test_mutate_row(self, table, temp_rows, handler, cluster_config): 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.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.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] @@ -1237,7 +1363,9 @@ async def test_mutate_row(self, table, temp_rows, handler, cluster_config): 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.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 @@ -1286,11 +1414,8 @@ async def test_mutate_row_failure_grpc( 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 - ): + async def test_mutate_row_failure_timeout(self, table, temp_rows, handler): """ Test failure in gapic layer by passing very low timeout @@ -1349,12 +1474,18 @@ async def test_mutate_row_failure_unauthorized( 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] + 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 + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) @CrossSync.pytest async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): @@ -1372,9 +1503,14 @@ async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): 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.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.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] @@ -1382,7 +1518,9 @@ async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): 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.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 @@ -1393,7 +1531,7 @@ async def test_sample_row_keys_failure_grpc( """ Test failure in grpc layer by injecting errors into an interceptor test with retryable errors, then a terminal one - + No headers expected """ exc = Aborted("injected") @@ -1405,7 +1543,7 @@ async def test_sample_row_keys_failure_grpc( await table.sample_row_keys(retryable_errors=[Aborted]) # validate counts assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable+1 + assert len(handler.completed_attempts) == num_retryable + 1 assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] @@ -1413,7 +1551,7 @@ async def test_sample_row_keys_failure_grpc( 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 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" @@ -1435,7 +1573,7 @@ async def test_sample_row_keys_failure_grpc_eventual_success( """ Test failure in grpc layer by injecting errors into an interceptor test with retryable errors, then a success - + No headers expected """ exc = Aborted("injected") @@ -1445,7 +1583,7 @@ async def test_sample_row_keys_failure_grpc_eventual_success( await table.sample_row_keys(retryable_errors=[Aborted]) # validate counts assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable+1 + assert len(handler.completed_attempts) == num_retryable + 1 assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] @@ -1453,10 +1591,13 @@ async def test_sample_row_keys_failure_grpc_eventual_success( 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 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] + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) # validate attempts for i in range(num_retryable): attempt = handler.completed_attempts[i] @@ -1466,13 +1607,13 @@ async def test_sample_row_keys_failure_grpc_eventual_success( 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 - + 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 - ): + async def test_sample_row_keys_failure_timeout(self, table, handler): """ Test failure in gapic layer by passing very low timeout @@ -1528,7 +1669,6 @@ async def test_sample_row_keys_failure_mid_stream( 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 @@ -1536,9 +1676,7 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config 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 - ) + 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 @@ -1554,9 +1692,14 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config 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.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.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] @@ -1564,7 +1707,9 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config 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.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 @@ -1582,9 +1727,7 @@ async def test_read_modify_write_failure_grpc( 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 - ) + await temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) rule = IncrementRule(family, qualifier, 1) # trigger an exception @@ -1622,11 +1765,8 @@ async def test_read_modify_write_failure_grpc( 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 - ): + async def test_read_modify_write_failure_timeout(self, table, temp_rows, handler): """ Test failure in gapic layer by passing very low timeout @@ -1637,9 +1777,7 @@ async def test_read_modify_write_failure_timeout( 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 - ) + 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) @@ -1682,23 +1820,28 @@ async def test_read_modify_write_failure_unauthorized( 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] + 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 - + 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): + 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 - ) + await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) true_mutation_value = b"true-mutation-value" true_mutation = SetCell( @@ -1723,9 +1866,14 @@ async def test_check_and_mutate_row(self, table, temp_rows, handler, cluster_con 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.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.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] @@ -1733,7 +1881,9 @@ async def test_check_and_mutate_row(self, table, temp_rows, handler, cluster_con 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.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 @@ -1746,7 +1896,6 @@ async def test_check_and_mutate_row_failure_grpc( No 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" @@ -1760,7 +1909,7 @@ async def test_check_and_mutate_row_failure_grpc( with pytest.raises(RuntimeError): await table.check_and_mutate_row( row_key, - predicate=ValueRangeFilter(0,2), + predicate=ValueRangeFilter(0, 2), ) # validate counts assert len(handler.completed_operations) == 1 @@ -1791,7 +1940,6 @@ async def test_check_and_mutate_row_failure_grpc( 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 @@ -1818,7 +1966,7 @@ async def test_check_and_mutate_row_failure_timeout( row_key, predicate=ValueRangeFilter(0, 2), true_case_mutations=true_mutation, - operation_timeout=0.001 + operation_timeout=0.001, ) # validate counts assert len(handler.completed_operations) == 1 @@ -1866,7 +2014,13 @@ async def test_check_and_mutate_row_failure_unauthorized( 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] + 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 + 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 44eed7a42..beea316bb 100644 --- a/tests/system/data/test_system_async.py +++ b/tests/system/data/test_system_async.py @@ -132,7 +132,7 @@ async def retrieve_cell_value(self, target, row_key): @CrossSync.convert async def create_row_and_mutation( - self, table, *, start_value=b"start", new_value=b"new_value" + 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 @@ -151,9 +151,9 @@ async def create_row_and_mutation( mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) return row_key, mutation + @CrossSync.convert_class(sync_name="TestSystem") class TestSystemAsync(SystemTestRunner): - @CrossSync.drop @pytest.fixture(scope="session") def event_loop(self): @@ -570,7 +570,9 @@ async def test_mutations_batcher_no_flush(self, client, target, temp_rows): assert len(batcher._flush_jobs) == 0 # ensure cells were not updated assert (await temp_rows.retrieve_cell_value(target, row_key)) == start_value - assert (await temp_rows.retrieve_cell_value(target, row_key2)) == start_value + assert ( + await temp_rows.retrieve_cell_value(target, row_key2) + ) == start_value @pytest.mark.usefixtures("client") @pytest.mark.usefixtures("target") From 0f4ee8d9c3a9029cd4709a58469e084d110d27fe Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Sep 2025 12:32:18 -0700 Subject: [PATCH 44/78] use cancelled error to test non-retryable methods --- google/cloud/bigtable/data/_metrics/data_model.py | 10 +++++----- tests/system/data/test_metrics_async.py | 12 ++++++------ 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 4b4c77604..6cafa4938 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -280,7 +280,7 @@ def _parse_response_metadata_blob(blob: bytes) -> Tuple[str, str] | None: # failed to parse metadata return None - def end_attempt_with_status(self, status: StatusCode | Exception) -> None: + def end_attempt_with_status(self, status: StatusCode | BaseException) -> None: """ Called to mark the end of an attempt for the operation. @@ -297,7 +297,7 @@ def end_attempt_with_status(self, status: StatusCode | Exception) -> None: return self._handle_error( INVALID_STATE_ERROR.format("end_attempt_with_status", self.state) ) - if isinstance(status, Exception): + if isinstance(status, BaseException): status = self._exc_to_status(status) complete_attempt = CompletedAttemptMetric( duration_ns=time.monotonic_ns() - self.active_attempt.start_time_ns, @@ -312,7 +312,7 @@ def end_attempt_with_status(self, status: StatusCode | Exception) -> None: for handler in self.handlers: handler.on_attempt_complete(complete_attempt, self) - def end_with_status(self, status: StatusCode | Exception) -> None: + def end_with_status(self, status: StatusCode | BaseException) -> None: """ Called to mark the end of the operation. If there is an active attempt, end_attempt_with_status will be called with the same status. @@ -329,7 +329,7 @@ def end_with_status(self, status: StatusCode | Exception) -> None: INVALID_STATE_ERROR.format("end_with_status", self.state) ) final_status = ( - self._exc_to_status(status) if isinstance(status, Exception) else status + self._exc_to_status(status) if isinstance(status, BaseException) else status ) if self.state == OperationState.ACTIVE_ATTEMPT: self.end_attempt_with_status(final_status) @@ -367,7 +367,7 @@ def cancel(self): handler.on_operation_cancelled(self) @staticmethod - def _exc_to_status(exc: Exception) -> StatusCode: + def _exc_to_status(exc: BaseException) -> StatusCode: """ Extracts the grpc status code from an exception. diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 558781e75..95b8458ff 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -1538,8 +1538,8 @@ async def test_sample_row_keys_failure_grpc( num_retryable = 3 for i in range(num_retryable): error_injector.push(exc) - error_injector.push(RuntimeError) - with pytest.raises(RuntimeError): + 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 @@ -1731,9 +1731,9 @@ async def test_read_modify_write_failure_grpc( rule = IncrementRule(family, qualifier, 1) # trigger an exception - exc = RuntimeError("injected") + exc = asyncio.CancelledError("injected") error_injector.push(exc) - with pytest.raises(RuntimeError): + with pytest.raises(asyncio.CancelledError): await table.read_modify_write_row(row_key, rule) # validate counts @@ -1904,9 +1904,9 @@ async def test_check_and_mutate_row_failure_grpc( await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) # trigger an exception - exc = RuntimeError("injected") + exc = asyncio.CancelledError("injected") error_injector.push(exc) - with pytest.raises(RuntimeError): + with pytest.raises(asyncio.CancelledError): await table.check_and_mutate_row( row_key, predicate=ValueRangeFilter(0, 2), From 2683b50af4dc0684e3c01a3dadfbd684e60f78b8 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Sep 2025 12:33:53 -0700 Subject: [PATCH 45/78] added aclose test to read_rows_stream --- .../cloud/bigtable/data/_async/_read_rows.py | 7 ++++ tests/system/data/test_metrics_async.py | 33 +++++++++++++++++-- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index d4ee5b2b1..1f4314cc9 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -19,6 +19,8 @@ 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 @@ -339,7 +341,12 @@ async def merge_rows( 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 else: self._operation_metric.end_with_success() diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 95b8458ff..c05b9c7b4 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -397,13 +397,40 @@ async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config) assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_rows_stream_failure_grpc( + async def test_read_rows_stream_failure_closed( self, table, temp_rows, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor + 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 + assert len(handler.cancelled_operations) == 0 + # 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 + - No headers expected """ await temp_rows.add_row(b"row_key_1") handler.clear() From 685b62a072a2500cbcc97f52b112d78a886d454e Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Sep 2025 12:34:32 -0700 Subject: [PATCH 46/78] changed test names --- tests/system/data/test_metrics_async.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index c05b9c7b4..23afad5d7 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -1741,7 +1741,7 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_modify_write_failure_grpc( + async def test_read_modify_write_failure_cancelled( self, table, temp_rows, handler, error_injector ): """ @@ -1915,7 +1915,7 @@ async def test_check_and_mutate_row( assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_check_and_mutate_row_failure_grpc( + async def test_check_and_mutate_row_failure_cancelled( self, table, temp_rows, handler, error_injector ): """ From 2336e6451973746f79d70122aacc241fab930b00 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Sep 2025 14:45:45 -0700 Subject: [PATCH 47/78] fixed warnings --- .../data/_async/metrics_interceptor.py | 15 +- .../bigtable/data/_metrics/data_model.py | 3 + tests/system/data/test_metrics_async.py | 226 ++++++++++-------- 3 files changed, 137 insertions(+), 107 deletions(-) diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index fe16b2d24..7acdc705c 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -28,6 +28,7 @@ if CrossSync.is_async: from grpc.aio import UnaryUnaryClientInterceptor from grpc.aio import UnaryStreamClientInterceptor + from grpc.aio import AioRpcError else: from grpc import UnaryUnaryClientInterceptor from grpc import UnaryStreamClientInterceptor @@ -86,12 +87,18 @@ def _end_attempt(operation, exc, metadata): @CrossSync.convert -async def _get_metadata(source): +async def _get_metadata(source) -> dict[str, str|bytes] | None: """Helper to extract metadata from a call or RpcError""" try: - return (await source.trailing_metadata() or {}) + ( - await source.initial_metadata() or {} - ) + if isinstance(source, AioRpcError) and CrossSync.is_async: + return { + k:v for k,v + in list(source.trailing_metadata()) + list(source.initial_metadata()) + } + else: + return (await source.trailing_metadata() or {}) + ( + await source.initial_metadata() or {} + ) except Exception: # ignore errors while fetching metadata return None diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 6cafa4938..dce6f2af1 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -25,6 +25,7 @@ from dataclasses import dataclass from dataclasses import field from grpc import StatusCode +from grpc.aio import AioRpcError import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable_v2.types.response_params import ResponseParams @@ -389,6 +390,8 @@ def _exc_to_status(exc: BaseException) -> StatusCode: and exc.__cause__.grpc_status_code is not None ): return exc.__cause__.grpc_status_code + if isinstance(exc, AioRpcError): + return exc.code() return StatusCode.UNKNOWN @staticmethod diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 23afad5d7..cbb52df35 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -16,6 +16,8 @@ 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 @@ -27,6 +29,7 @@ OperationState, ) 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 @@ -35,6 +38,8 @@ 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 @@ -140,6 +145,17 @@ 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) + @pytest.fixture(scope="session") def handler(self): return _MetricsTestHandler() @@ -239,21 +255,21 @@ async def test_read_rows(self, table, temp_rows, handler, cluster_config): assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_rows_failure_grpc( + async def test_read_rows_failure_with_retries( self, table, temp_rows, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor - - No headers expected + 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() - exc = Aborted("injected") + expected_zone = "my_zone" + expected_cluster = "my_cluster" num_retryable = 2 for i in range(num_retryable): - error_injector.push(exc) - error_injector.push(PermissionDenied("terminal")) + 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 @@ -267,8 +283,8 @@ async def test_read_rows_failure_grpc( assert operation.op_type.value == "ReadRows" assert operation.is_streaming is True assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == "unspecified" - assert operation.zone == "global" + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone # validate attempts for i in range(num_retryable): attempt = handler.completed_attempts[i] @@ -431,14 +447,26 @@ async def test_read_rows_stream_failure_closed( 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() - exc = Aborted("injected") + expected_zone = "my_zone" + expected_cluster = "my_cluster" num_retryable = 2 for i in range(num_retryable): - error_injector.push(exc) - error_injector.push(PermissionDenied("terminal")) + 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] ) @@ -455,8 +483,8 @@ async def test_read_rows_stream_failure_closed( assert operation.op_type.value == "ReadRows" assert operation.is_streaming is True assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == "unspecified" - assert operation.zone == "global" + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone # validate attempts for i in range(num_retryable): attempt = handler.completed_attempts[i] @@ -551,8 +579,8 @@ async def test_read_rows_stream_failure_mid_stream( await temp_rows.add_row(b"row_key_1") handler.clear() error_injector.fail_mid_stream = True - error_injector.push(Aborted("retryable")) - error_injector.push(PermissionDenied("terminal")) + 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] ) @@ -619,21 +647,25 @@ async def test_read_row(self, table, temp_rows, handler, cluster_config): assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_row_failure_grpc( + async def test_read_row_failure_with_retries( self, table, temp_rows, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor - - No headers expected + 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() - exc = Aborted("injected") + expected_zone = "my_zone" + expected_cluster = "my_cluster" num_retryable = 2 for i in range(num_retryable): - error_injector.push(exc) - error_injector.push(PermissionDenied("terminal")) + 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 @@ -647,8 +679,8 @@ async def test_read_row_failure_grpc( assert operation.op_type.value == "ReadRows" assert operation.is_streaming is False assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == "unspecified" - assert operation.zone == "global" + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone # validate attempts for i in range(num_retryable): attempt = handler.completed_attempts[i] @@ -784,13 +816,12 @@ async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_read_rows_sharded_failure_grpc( + async def test_read_rows_sharded_failure_with_retries( self, table, temp_rows, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor - - No headers expected + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors """ from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup @@ -801,41 +832,20 @@ async def test_read_rows_sharded_failure_grpc( query2 = ReadRowsQuery(row_keys=[b"b"]) handler.clear() - error_injector.push(PermissionDenied("terminal")) - with pytest.raises(ShardedReadRowsExceptionGroup) as e: - await table.read_rows_sharded([query1, query2]) - assert len(e.value.exceptions) == 1 - assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + 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) == 2 + assert len(handler.completed_attempts) == 3 assert len(handler.cancelled_operations) == 0 - # 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 == "unspecified" - assert failed_op.zone == "global" - # validate failed attempt - failed_attempt = failed_op.completed_attempts[0] - assert failed_attempt.end_status.name == "PERMISSION_DENIED" - assert failed_attempt.gfe_latency_ns is None - # 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" + # 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): @@ -951,8 +961,8 @@ async def test_read_rows_sharded_failure_mid_stream( query2 = ReadRowsQuery(row_keys=[b"b"]) handler.clear() error_injector.fail_mid_stream = True - error_injector.push(Aborted("retryable")) - error_injector.push(PermissionDenied("terminal")) + 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 @@ -1034,13 +1044,12 @@ async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config) assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_bulk_mutate_rows_failure_grpc( + async def test_bulk_mutate_rows_failure_with_retries( self, table, temp_rows, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor - - No headers expected + 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 @@ -1051,11 +1060,16 @@ async def test_bulk_mutate_rows_failure_grpc( assert entry.is_idempotent() handler.clear() - exc = Aborted("injected") + expected_zone = "my_zone" + expected_cluster = "my_cluster" num_retryable = 2 for i in range(num_retryable): - error_injector.push(exc) - error_injector.push(PermissionDenied("terminal")) + 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 @@ -1069,8 +1083,8 @@ async def test_bulk_mutate_rows_failure_grpc( assert operation.op_type.value == "MutateRows" assert operation.is_streaming is False assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == "unspecified" - assert operation.zone == "global" + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone # validate attempts for i in range(num_retryable): attempt = handler.completed_attempts[i] @@ -1224,13 +1238,12 @@ async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_conf assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_mutate_rows_batcher_failure_grpc( - self, table, temp_rows, handler, error_injector + async def test_mutate_rows_batcher_failure_with_retries( + self, table, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor - - No headers expected + 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 @@ -1240,11 +1253,17 @@ async def test_mutate_rows_batcher_failure_grpc( entry = RowMutationEntry(row_key, [mutation]) assert entry.is_idempotent() - exc = Aborted("injected") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" num_retryable = 2 for i in range(num_retryable): - error_injector.push(exc) - error_injector.push(PermissionDenied("terminal")) + 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] @@ -1261,8 +1280,8 @@ async def test_mutate_rows_batcher_failure_grpc( assert operation.op_type.value == "MutateRows" assert operation.is_streaming is False assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == "unspecified" - assert operation.zone == "global" + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone # validate attempts for i in range(num_retryable): attempt = handler.completed_attempts[i] @@ -1397,24 +1416,29 @@ async def test_mutate_row(self, table, temp_rows, handler, cluster_config): assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_mutate_row_failure_grpc( - self, table, temp_rows, handler, error_injector + async def test_mutate_row_failure_with_retries( + self, table, handler, error_injector ): """ - Test failure in grpc layer by injecting an error into an interceptor - - No headers expected + 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") - exc = Aborted("injected") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" num_retryable = 2 for i in range(num_retryable): - error_injector.push(exc) - error_injector.push(PermissionDenied("terminal")) + 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 @@ -1428,8 +1452,8 @@ async def test_mutate_row_failure_grpc( assert operation.op_type.value == "MutateRow" assert operation.is_streaming is False assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == "unspecified" - assert operation.zone == "global" + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone # validate attempts for i in range(num_retryable): attempt = handler.completed_attempts[i] @@ -1552,7 +1576,7 @@ async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): assert attempt.grpc_throttling_time_ns == 0 # TODO: confirm @CrossSync.pytest - async def test_sample_row_keys_failure_grpc( + async def test_sample_row_keys_failure_cancelled( self, table, temp_rows, handler, error_injector ): """ @@ -1561,10 +1585,9 @@ async def test_sample_row_keys_failure_grpc( No headers expected """ - exc = Aborted("injected") num_retryable = 3 for i in range(num_retryable): - error_injector.push(exc) + 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]) @@ -1594,19 +1617,16 @@ async def test_sample_row_keys_failure_grpc( assert final_attempt.gfe_latency_ns is None @CrossSync.pytest - async def test_sample_row_keys_failure_grpc_eventual_success( + 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 - test with retryable errors, then a success - - No headers expected + with retryable errors, then a success """ - exc = Aborted("injected") num_retryable = 3 for i in range(num_retryable): - error_injector.push(exc) + error_injector.push(self._make_exception(StatusCode.ABORTED)) await table.sample_row_keys(retryable_errors=[Aborted]) # validate counts assert len(handler.completed_operations) == 1 @@ -1675,8 +1695,8 @@ async def test_sample_row_keys_failure_mid_stream( Test failure in grpc stream """ error_injector.fail_mid_stream = True - error_injector.push(Aborted("retryable")) - error_injector.push(PermissionDenied("terminal")) + 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 From e239d56dd026e5070ddb27084f13f7f9f53f6ae1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 18 Sep 2025 17:40:59 -0700 Subject: [PATCH 48/78] adding sync tests --- .../data/_async/metrics_interceptor.py | 20 ++++++++------ .../bigtable/data/_metrics/data_model.py | 3 ++- tests/system/data/test_metrics_async.py | 26 +++++++++++++++++++ 3 files changed, 40 insertions(+), 9 deletions(-) diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 7acdc705c..922d57f3e 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -13,6 +13,8 @@ # limitations under the License from __future__ import annotations +from typing import Sequence + import time from functools import wraps from google.cloud.bigtable.data._metrics.data_model import ( @@ -90,15 +92,17 @@ def _end_attempt(operation, exc, metadata): async def _get_metadata(source) -> dict[str, str|bytes] | None: """Helper to extract metadata from a call or RpcError""" try: - if isinstance(source, AioRpcError) and CrossSync.is_async: - return { - k:v for k,v - in list(source.trailing_metadata()) + list(source.initial_metadata()) - } + if CrossSync.is_async: + # grpc.aio returns metadata in Metadata objects + if isinstance(source, AioRpcError): + metadata = list(source.trailing_metadata()) + list(source.initial_metadata()) + else: + metadata = list(await source.trailing_metadata()) + list(await source.initial_metadata()) else: - return (await source.trailing_metadata() or {}) + ( - await source.initial_metadata() or {} - ) + # sync grpc returns metadata as a sequence of tuples + metadata: Sequence[tuple[str. str|bytes]] = source.trailing_metadata() + source.initial_metadata() + # convert metadata to dict format + return {k:v for k,v in metadata} except Exception: # ignore errors while fetching metadata return None diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index dce6f2af1..52aa32671 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -25,6 +25,7 @@ from dataclasses import dataclass from dataclasses import field from grpc import StatusCode +from grpc import RpcError from grpc.aio import AioRpcError import google.cloud.bigtable.data.exceptions as bt_exceptions @@ -390,7 +391,7 @@ def _exc_to_status(exc: BaseException) -> StatusCode: and exc.__cause__.grpc_status_code is not None ): return exc.__cause__.grpc_status_code - if isinstance(exc, AioRpcError): + if isinstance(exc, AioRpcError) or isinstance(exc, RpcError): return exc.code() return StatusCode.UNKNOWN diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index cbb52df35..236f697ca 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -43,6 +43,8 @@ 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" @@ -75,6 +77,7 @@ def __repr__(self): return f"{self.__class__}(completed_operations={len(self.completed_operations)}, cancelled_operations={len(self.cancelled_operations)}, completed_attempts={len(self.completed_attempts)}" +@CrossSync.convert_class class _ErrorInjectorInterceptor( UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor ): @@ -93,11 +96,13 @@ 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) @@ -113,9 +118,12 @@ def __init__(self, call, exc_to_raise): 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 @@ -155,6 +163,16 @@ def _make_exception(self, status, cluster_id=None, zone_id=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): @@ -182,6 +200,10 @@ async def client(self, 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 @@ -413,6 +435,7 @@ async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config) 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 ): @@ -1575,6 +1598,7 @@ async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): 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 @@ -1760,6 +1784,7 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config 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 @@ -1934,6 +1959,7 @@ async def test_check_and_mutate_row( 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 From 1a54a6941537c9081f7d04f19c0cf7f852c72a3f Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 19 Sep 2025 11:23:43 -0700 Subject: [PATCH 49/78] capture status for unary failed attempts --- .../data/_async/metrics_interceptor.py | 13 +++++-- tests/system/data/test_metrics_async.py | 39 +++++++++++++++++++ 2 files changed, 49 insertions(+), 3 deletions(-) diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 922d57f3e..80d9a8b4e 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -17,6 +17,8 @@ import time from functools import wraps +from grpc import StatusCode + from google.cloud.bigtable.data._metrics.data_model import ( OPERATION_INTERCEPTOR_METADATA_KEY, ) @@ -147,18 +149,23 @@ def on_operation_cancelled(self, op): async def intercept_unary_unary( self, operation, continuation, client_call_details, request ): - encountered_exc: Exception | None = None + encountered_status: Exception | StatusCode | None = None metadata = None try: call = await continuation(client_call_details, request) metadata = await _get_metadata(call) + if CrossSync.is_async: + encountered_status = await call.code() + elif isinstance(call, Exception): + # sync unary calls return exception objects without raising + encountered_status = call return call except Exception as rpc_error: metadata = await _get_metadata(rpc_error) - encountered_exc = rpc_error + encountered_status = rpc_error raise rpc_error finally: - _end_attempt(operation, encountered_exc, metadata) + _end_attempt(operation, encountered_status, metadata) @CrossSync.convert @_with_operation_from_metadata diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 236f697ca..ed8811158 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -1561,6 +1561,45 @@ async def test_mutate_row_failure_unauthorized( 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=30) + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + assert len(handler.cancelled_operations) == 0 + # 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 == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + @CrossSync.pytest async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): await table.sample_row_keys() From a23d92064a285fd470a84c4d2f6199255e8cc4eb Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 19 Sep 2025 12:11:20 -0700 Subject: [PATCH 50/78] added test for streaming retries --- tests/system/data/test_metrics_async.py | 44 ++++++++++++++++++++++++- 1 file changed, 43 insertions(+), 1 deletion(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index ed8811158..7ae3c45e4 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -592,6 +592,48 @@ async def test_read_rows_stream_failure_unauthorized( 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.1, + ) + [_ 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 + assert len(handler.cancelled_operations) == 0 + # 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 == "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_mid_stream( self, table, temp_rows, handler, error_injector @@ -1574,7 +1616,7 @@ async def test_mutate_row_failure_unauthorized_with_retries( 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=30) + await authorized_view.mutate_row(row_key, [mutation], retryable_errors=[PermissionDenied], operation_timeout=0.1) assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" # validate counts assert len(handler.completed_operations) == 1 From b33458fe98074fbcc84ec507aca6c31b437746ea Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 19 Sep 2025 17:41:36 -0700 Subject: [PATCH 51/78] added test --- tests/system/data/test_metrics_async.py | 44 ++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 5 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 7ae3c45e4..5467a9ef0 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -1214,11 +1214,6 @@ async def test_bulk_mutate_rows_failure_unauthorized( handler.clear() with pytest.raises(MutationsExceptionGroup) as e: await authorized_view.bulk_mutate_rows([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 @@ -1242,6 +1237,45 @@ async def test_bulk_mutate_rows_failure_unauthorized( 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""" + 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=25) + assert len(e.value.exceptions) == 1 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + assert len(handler.cancelled_operations) == 0 + # 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 == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + @CrossSync.pytest async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): from google.cloud.bigtable.data.mutations import RowMutationEntry From 3b146f5149fc4a0596c7351cc1de17fc506794ce Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 19 Sep 2025 17:41:57 -0700 Subject: [PATCH 52/78] fixed tracked flow control --- google/cloud/bigtable/data/_async/mutations_batcher.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index 42fbc4d3b..73c9a5022 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -181,6 +181,7 @@ 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], @@ -193,7 +194,7 @@ async def add_to_flow_with_metrics( flow_start_time = time.monotonic_ns() try: value = await inner_generator.__anext__() - except StopAsyncIteration: + except CrossSync.StopIteration: metric.cancel() return metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time From b91da1c17df7906d102b89aa6b8ed8e0948ebf7b Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 22 Sep 2025 14:35:06 -0700 Subject: [PATCH 53/78] updated mutate_rows test --- tests/system/data/test_metrics_async.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 5467a9ef0..7c96896e7 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -1241,7 +1241,13 @@ async def test_bulk_mutate_rows_failure_unauthorized( async def test_bulk_mutate_rows_failure_unauthorized_with_retries( self, handler, authorized_view, cluster_config ): - """retry unauthorized request multiple times before timing out""" + """ + 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 @@ -1270,7 +1276,7 @@ async def test_bulk_mutate_rows_failure_unauthorized_with_retries( ) # validate attempts for attempt in handler.completed_attempts: - assert attempt.end_status.name == "PERMISSION_DENIED" + assert attempt.end_status.name == "OK" assert ( attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns From 64ce0a4a561607075e8ddef1638289604eef0b9b Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 22 Sep 2025 15:52:39 -0700 Subject: [PATCH 54/78] improved retry instrumentation --- .../bigtable/data/_async/_mutate_rows.py | 5 +- .../cloud/bigtable/data/_async/_read_rows.py | 5 +- google/cloud/bigtable/data/_async/client.py | 7 +-- .../data/_async/metrics_interceptor.py | 29 +++-------- google/cloud/bigtable/data/_helpers.py | 51 ------------------- 5 files changed, 17 insertions(+), 80 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 03d870e49..3d603fc6c 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -21,7 +21,7 @@ import google.cloud.bigtable_v2.types.bigtable as types_pb import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _tracked_exception_factory +from google.cloud.bigtable.data._helpers import _retry_exception_factory # mutate_rows requests are limited to this number of mutations from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT @@ -106,7 +106,8 @@ def __init__( self.is_retryable, metric.backoff_generator, operation_timeout, - exception_factory=_tracked_exception_factory(self._operation_metric), + exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory), + on_error=self._operation_metric.track_retryable_error(), ) # initialize state self.timeout_generator = _attempt_timeout_generator( diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 1f4314cc9..c533492de 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -32,7 +32,7 @@ from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _tracked_exception_factory +from google.cloud.bigtable.data._helpers import _retry_exception_factory from google.api_core import retry as retries @@ -123,7 +123,8 @@ def start_operation(self) -> CrossSync.Iterable[Row]: self._predicate, self._operation_metric.backoff_generator, self.operation_timeout, - exception_factory=_tracked_exception_factory(self._operation_metric), + 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]: diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 2e284373a..df636849f 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -73,7 +73,6 @@ from google.cloud.bigtable.data._helpers import _WarmedInstanceKey from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT from google.cloud.bigtable.data._helpers import _retry_exception_factory -from google.cloud.bigtable.data._helpers import _tracked_exception_factory from google.cloud.bigtable.data._helpers import _validate_timeouts from google.cloud.bigtable.data._helpers import _get_error_type from google.cloud.bigtable.data._helpers import _get_retryable_errors @@ -1392,7 +1391,8 @@ async def execute_rpc(): predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=_tracked_exception_factory(operation_metric), + exception_factory=operation_metric.track_terminal_error(_retry_exception_factory), + on_error=operation_metric.track_retryable_error(), ) @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"}) @@ -1526,7 +1526,8 @@ async def mutate_row( predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=operation_metric.track_terminal_error(_retry_exception_factory), + on_error=operation_metric.track_retryable_error(), ) @CrossSync.convert diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 80d9a8b4e..b10e081ed 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -80,16 +80,6 @@ def wrapper(self, continuation, client_call_details, request): return wrapper - -def _end_attempt(operation, exc, metadata): - """Helper to add metadata and exception to an operation""" - if metadata is not None: - operation.add_response_metadata(metadata) - if exc is not None: - # end attempt. If it succeeded, let higher levels decide when to end operation - operation.end_attempt_with_status(exc) - - @CrossSync.convert async def _get_metadata(source) -> dict[str, str|bytes] | None: """Helper to extract metadata from a call or RpcError""" @@ -129,7 +119,6 @@ def register_operation(self, operation): When registered, the operation will receive metadata updates: - start_attempt if attempt not started when rpc is being sent - add_response_metadata after call is complete - - end_attempt_with_status if attempt receives an error The interceptor will register itself as a handeler for the operation, so it can unregister the operation when it is complete @@ -149,23 +138,17 @@ def on_operation_cancelled(self, op): async def intercept_unary_unary( self, operation, continuation, client_call_details, request ): - encountered_status: Exception | StatusCode | None = None metadata = None try: call = await continuation(client_call_details, request) metadata = await _get_metadata(call) - if CrossSync.is_async: - encountered_status = await call.code() - elif isinstance(call, Exception): - # sync unary calls return exception objects without raising - encountered_status = call return call except Exception as rpc_error: metadata = await _get_metadata(rpc_error) - encountered_status = rpc_error raise rpc_error finally: - _end_attempt(operation, encountered_status, metadata) + if metadata is not None: + operation.add_response_metadata(metadata) @CrossSync.convert @_with_operation_from_metadata @@ -195,11 +178,13 @@ async def response_wrapper(call): finally: if call is not None: metadata = await _get_metadata(encountered_exc or call) - _end_attempt(operation, encountered_exc, metadata) + if metadata is not None: + operation.add_response_metadata(metadata) try: return response_wrapper(await continuation(client_call_details, request)) except Exception as rpc_error: - # handle errors while intializing stream - _end_attempt(operation, rpc_error, await _get_metadata(rpc_error)) + metadata = await _get_metadata(rpc_error) + if metadata is not None: + operation.add_response_metadata(metadata) raise rpc_error diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 0db1de46f..592bea0d4 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -31,7 +31,6 @@ import grpc from google.cloud.bigtable.data._async.client import _DataApiTargetAsync from google.cloud.bigtable.data._sync_autogen.client import _DataApiTarget - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric """ Helper functions used in various places in the library. @@ -121,56 +120,6 @@ def _retry_exception_factory( return source_exc, cause_exc -def _tracked_exception_factory( - operation: "ActiveOperationMetric", -) -> Callable[ - [list[Exception], RetryFailureReason, float | None], - tuple[Exception, Exception | None], -]: - """ - wraps and extends _retry_exception_factory to add client-side metrics tracking. - - When the rpc raises a terminal error, record any discovered metadata and finalize - the associated operation - - Used by streaming rpcs, which can't always be perfectly captured by context managers - for operation termination. - - Args: - exc_list: list of exceptions encountered during operation - is_timeout: whether the operation failed due to timeout - timeout_val: the operation timeout value in seconds, for constructing - the error message - operation: the operation to finalize when an exception is built - Returns: - tuple[Exception, Exception|None]: - tuple of the exception to raise, and a cause exception if applicabl - """ - - def wrapper( - exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None - ) -> tuple[Exception, Exception | None]: - source_exc, cause_exc = _retry_exception_factory(exc_list, reason, timeout_val) - try: - # record metadata from failed rpc - if ( - isinstance(source_exc, core_exceptions.GoogleAPICallError) - and source_exc.errors - ): - rpc_error = source_exc.errors[-1] - metadata = list(rpc_error.trailing_metadata()) + list( - rpc_error.initial_metadata() - ) - operation.add_response_metadata({k: v for k, v in metadata}) - except Exception: - # ignore errors in metadata collection - pass - operation.end_with_status(source_exc) - return source_exc, cause_exc - - return wrapper - - def _get_timeouts( operation: float | TABLE_DEFAULT, attempt: float | None | TABLE_DEFAULT, From 7eb50d87dd9150c487d11698c1ea5bf7aeac65d9 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 22 Sep 2025 16:42:25 -0700 Subject: [PATCH 55/78] added trackers to data model --- .../bigtable/data/_metrics/data_model.py | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 52aa32671..b2349122b 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -28,13 +28,17 @@ from grpc import RpcError from grpc.aio import AioRpcError +from google.api_core.exceptions import GoogleAPICallError import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable_v2.types.response_params import ResponseParams from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator +from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete +from google.cloud.bigtable.data.exceptions import RetryExceptionGroup from google.protobuf.message import DecodeError if TYPE_CHECKING: from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler + from google.api_core.retry import RetryFailureReason LOGGER = logging.getLogger(__name__) @@ -395,6 +399,66 @@ def _exc_to_status(exc: BaseException) -> StatusCode: return exc.code() return StatusCode.UNKNOWN + def track_retryable_error(self) -> callable[[Exception], None]: + """ + Used as input to api_core.Retry classes, to track when retryable errors are encountered + + Should be passed as on_error callback + """ + + def wrapper(exc: Exception) -> None: + try: + # record metadata from failed rpc + if ( + isinstance(exc, GoogleAPICallError) + and exc.errors + ): + rpc_error = exc.errors[-1] + metadata = list(rpc_error.trailing_metadata()) + list( + rpc_error.initial_metadata() + ) + self.add_response_metadata({k: v for k, v in metadata}) + except Exception: + # ignore errors in metadata collection + pass + if isinstance(exc, _MutateRowsIncomplete): + # _MutateRowsIncomplete represents a successful rpc with some failed mutations + # mark the attempt as successful + self.end_attempt_with_status(StatusCode.OK) + else: + self.end_attempt_with_status(exc) + return wrapper + + def track_terminal_error(self, exception_factory:callable[ + [list[Exception], RetryFailureReason, float | None],tuple[Exception, Exception | None], + ]) -> callable[[list[Exception], RetryFailureReason, float | None], None]: + """ + Used as input to api_core.Retry classes, to track when terminal errors are encountered + + Should be used as a wrapper over an exception_factory callback + """ + def wrapper( + exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None + ) -> tuple[Exception, Exception | None]: + source_exc, cause_exc = exception_factory(exc_list, reason, timeout_val) + try: + # record metadata from failed rpc + if ( + isinstance(source_exc, GoogleAPICallError) + and source_exc.errors + ): + rpc_error = source_exc.errors[-1] + metadata = list(rpc_error.trailing_metadata()) + list( + rpc_error.initial_metadata() + ) + self.add_response_metadata({k: v for k, v in metadata}) + except Exception: + # ignore errors in metadata collection + pass + self.end_with_status(source_exc) + return source_exc, cause_exc + return wrapper + @staticmethod def _handle_error(message: str) -> None: """ From c1baf8151cddc6376822642555ecad1aeadeaf72 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 22 Sep 2025 17:27:02 -0700 Subject: [PATCH 56/78] fixed bug in application blocking time --- google/cloud/bigtable/data/_async/_read_rows.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index c533492de..5443ea96f 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -325,7 +325,7 @@ async def merge_rows( 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 - ) * 1000 + ) break c = await it.__anext__() except _ResetRow as e: From 9d6396336157530459744bfbde078ed11776b648 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 22 Sep 2025 17:31:02 -0700 Subject: [PATCH 57/78] record last attempt in exception factory --- google/cloud/bigtable/data/_helpers.py | 1 + google/cloud/bigtable/data/_metrics/data_model.py | 6 +++++- tests/system/data/test_metrics_async.py | 6 +++--- 3 files changed, 9 insertions(+), 4 deletions(-) diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 592bea0d4..6045483a3 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/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index b2349122b..6355be9d9 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -29,6 +29,7 @@ from grpc.aio import AioRpcError from google.api_core.exceptions import GoogleAPICallError +from google.api_core.retry import RetryFailureReason import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable_v2.types.response_params import ResponseParams from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator @@ -38,7 +39,6 @@ if TYPE_CHECKING: from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler - from google.api_core.retry import RetryFailureReason LOGGER = logging.getLogger(__name__) @@ -455,6 +455,10 @@ def wrapper( except Exception: # ignore errors in metadata collection pass + if reason == RetryFailureReason.TIMEOUT and self.state == OperationState.ACTIVE_ATTEMPT and exc_list: + # record ending attempt for timeout failures + attempt_exc = exc_list[-1] + self.track_retryable_error()(attempt_exc) self.end_with_status(source_exc) return source_exc, cause_exc return wrapper diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 7c96896e7..bc530b4af 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -605,7 +605,7 @@ async def test_read_rows_stream_failure_unauthorized_with_retries( generator = await authorized_view.read_rows_stream( ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), retryable_errors=[PermissionDenied], - operation_timeout=0.1, + operation_timeout=0.5, ) [_ async for _ in generator] assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" @@ -1257,7 +1257,7 @@ async def test_bulk_mutate_rows_failure_unauthorized_with_retries( handler.clear() with pytest.raises(MutationsExceptionGroup) as e: - await authorized_view.bulk_mutate_rows([entry], retryable_errors=[PermissionDenied], operation_timeout=25) + 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 @@ -1656,7 +1656,7 @@ async def test_mutate_row_failure_unauthorized_with_retries( 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.1) + 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 From d3f9d054b2d84c9fb3c33976d9e193d59bd4bc4f Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 22 Sep 2025 17:38:49 -0700 Subject: [PATCH 58/78] simplified helper --- .../bigtable/data/_async/_mutate_rows.py | 2 +- .../cloud/bigtable/data/_async/_read_rows.py | 2 +- google/cloud/bigtable/data/_async/client.py | 4 +- .../bigtable/data/_metrics/data_model.py | 47 +++++++++---------- 4 files changed, 26 insertions(+), 29 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 3d603fc6c..33ee2df72 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -107,7 +107,7 @@ def __init__( metric.backoff_generator, operation_timeout, exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory), - on_error=self._operation_metric.track_retryable_error(), + on_error=self._operation_metric.track_retryable_error, ) # initialize state self.timeout_generator = _attempt_timeout_generator( diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 5443ea96f..85dc5ceb2 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -124,7 +124,7 @@ def start_operation(self) -> CrossSync.Iterable[Row]: self._operation_metric.backoff_generator, self.operation_timeout, exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory), - on_error=self._operation_metric.track_retryable_error(), + on_error=self._operation_metric.track_retryable_error, ) def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index df636849f..ca53a03f6 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1392,7 +1392,7 @@ async def execute_rpc(): operation_metric.backoff_generator, operation_timeout, exception_factory=operation_metric.track_terminal_error(_retry_exception_factory), - on_error=operation_metric.track_retryable_error(), + on_error=operation_metric.track_retryable_error, ) @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"}) @@ -1527,7 +1527,7 @@ async def mutate_row( operation_metric.backoff_generator, operation_timeout, exception_factory=operation_metric.track_terminal_error(_retry_exception_factory), - on_error=operation_metric.track_retryable_error(), + on_error=operation_metric.track_retryable_error, ) @CrossSync.convert diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 6355be9d9..6ff51103a 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -399,35 +399,32 @@ def _exc_to_status(exc: BaseException) -> StatusCode: return exc.code() return StatusCode.UNKNOWN - def track_retryable_error(self) -> callable[[Exception], None]: + def track_retryable_error(self, exc: Exception) -> None: """ Used as input to api_core.Retry classes, to track when retryable errors are encountered Should be passed as on_error callback """ - - def wrapper(exc: Exception) -> None: - try: - # record metadata from failed rpc - if ( - isinstance(exc, GoogleAPICallError) - and exc.errors - ): - rpc_error = exc.errors[-1] - metadata = list(rpc_error.trailing_metadata()) + list( - rpc_error.initial_metadata() - ) - self.add_response_metadata({k: v for k, v in metadata}) - except Exception: - # ignore errors in metadata collection - pass - if isinstance(exc, _MutateRowsIncomplete): - # _MutateRowsIncomplete represents a successful rpc with some failed mutations - # mark the attempt as successful - self.end_attempt_with_status(StatusCode.OK) - else: - self.end_attempt_with_status(exc) - return wrapper + try: + # record metadata from failed rpc + if ( + isinstance(exc, GoogleAPICallError) + and exc.errors + ): + rpc_error = exc.errors[-1] + metadata = list(rpc_error.trailing_metadata()) + list( + rpc_error.initial_metadata() + ) + self.add_response_metadata({k: v for k, v in metadata}) + except Exception: + # ignore errors in metadata collection + pass + if isinstance(exc, _MutateRowsIncomplete): + # _MutateRowsIncomplete represents a successful rpc with some failed mutations + # mark the attempt as successful + self.end_attempt_with_status(StatusCode.OK) + else: + self.end_attempt_with_status(exc) def track_terminal_error(self, exception_factory:callable[ [list[Exception], RetryFailureReason, float | None],tuple[Exception, Exception | None], @@ -458,7 +455,7 @@ def wrapper( if reason == RetryFailureReason.TIMEOUT and self.state == OperationState.ACTIVE_ATTEMPT and exc_list: # record ending attempt for timeout failures attempt_exc = exc_list[-1] - self.track_retryable_error()(attempt_exc) + self.track_retryable_error(attempt_exc) self.end_with_status(source_exc) return source_exc, cause_exc return wrapper From d1b74fcdeaf09f94d710a4ac2d1daaee0c009463 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 25 Sep 2025 13:54:03 -0700 Subject: [PATCH 59/78] swapped out custom metadata with contextvar --- .../bigtable/data/_async/_mutate_rows.py | 1 - .../cloud/bigtable/data/_async/_read_rows.py | 1 - google/cloud/bigtable/data/_async/client.py | 4 ---- .../data/_async/metrics_interceptor.py | 21 ++--------------- .../bigtable/data/_metrics/data_model.py | 23 +++++++++++-------- 5 files changed, 15 insertions(+), 35 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 33ee2df72..66baef89a 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -188,7 +188,6 @@ async def _run_attempt(self): ), timeout=next(self.timeout_generator), retry=None, - metadata=[self._operation_metric.interceptor_metadata], ) async for result_list in result_generator: for result in result_list.entries: diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 85dc5ceb2..994cfe2da 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -159,7 +159,6 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: self.request, timeout=next(self.attempt_timeout_gen), retry=None, - metadata=[self._operation_metric.interceptor_metadata], ) chunked_stream = self.chunk_stream(gapic_stream) return self.merge_rows(chunked_stream) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index ca53a03f6..183d89975 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1382,7 +1382,6 @@ async def execute_rpc(): ), timeout=next(attempt_timeout_gen), retry=None, - metadata=[operation_metric.interceptor_metadata], ) return [(s.row_key, s.offset_bytes) async for s in results] @@ -1519,7 +1518,6 @@ async def mutate_row( ), timeout=attempt_timeout, retry=None, - metadata=[operation_metric.interceptor_metadata], ) return await CrossSync.retry_target( target, @@ -1657,7 +1655,6 @@ async def check_and_mutate_row( ), timeout=operation_timeout, retry=None, - metadata=[op.interceptor_metadata], ) return result.predicate_matched @@ -1712,7 +1709,6 @@ async def read_modify_write_row( ), timeout=operation_timeout, retry=None, - metadata=[op.interceptor_metadata], ) # construct Row from result return Row._from_pb(result.row) diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index b10e081ed..6812a8648 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -19,9 +19,6 @@ from functools import wraps from grpc import StatusCode -from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, -) from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric from google.cloud.bigtable.data._metrics.data_model import OperationState from google.cloud.bigtable.data._metrics.data_model import OperationType @@ -49,22 +46,8 @@ def _with_operation_from_metadata(func): @wraps(func) def wrapper(self, continuation, client_call_details, request): - found_operation_id: str | None = None - try: - new_metadata: list[tuple[str, str]] = [] - if client_call_details.metadata: - # find operation key from metadata - for k, v in client_call_details.metadata: - if k == OPERATION_INTERCEPTOR_METADATA_KEY: - found_operation_id = v - else: - new_metadata.append((k, v)) - # update client_call_details to drop the operation key metadata - client_call_details.metadata = new_metadata - except Exception: - pass - - operation: "ActiveOperationMetric" = self.operation_map.get(found_operation_id) + operation: "ActiveOperationMetric" | None = ActiveOperationMetric.get_active() + if operation: # start a new attempt if not started if ( diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 6ff51103a..1e6a2f133 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -13,12 +13,14 @@ # limitations under the License. from __future__ import annotations -from typing import Tuple, cast, TYPE_CHECKING +from typing import ClassVar, Tuple, cast, TYPE_CHECKING + import time import re import logging import uuid +import contextvars from enum import Enum from functools import lru_cache @@ -54,8 +56,6 @@ INVALID_STATE_ERROR = "Invalid state for {}: {}" -OPERATION_INTERCEPTOR_METADATA_KEY = "x-goog-operation-key" - class OperationType(Enum): """Enum for the type of operation being performed.""" @@ -169,14 +169,12 @@ class ActiveOperationMetric: # time waiting on flow control, in nanoseconds flow_throttling_time_ns: int = 0 - @property - def interceptor_metadata(self) -> tuple[str, str]: - """ - returns a tuple to attach to the grpc metadata. - This metadata field will be read by the BigtableMetricsInterceptor to associate a request with an operation - """ - return OPERATION_INTERCEPTOR_METADATA_KEY, self.uuid + _active_operation_context: ClassVar[contextvars.ContextVar] = contextvars.ContextVar("active_operation_context") + + @classmethod + def get_active(cls): + return cls._active_operation_context.get(None) @property def state(self) -> OperationState: @@ -190,6 +188,9 @@ def state(self) -> OperationState: else: return OperationState.ACTIVE_ATTEMPT + def __post_init__(self): + self._active_operation_context.set(self) + def start(self) -> None: """ Optionally called to mark the start of the operation. If not called, @@ -200,6 +201,7 @@ def start(self) -> None: if self.state != OperationState.CREATED: return self._handle_error(INVALID_STATE_ERROR.format("start", self.state)) self.start_time_ns = time.monotonic_ns() + self._active_operation_context.set(self) def start_attempt(self) -> ActiveAttemptMetric | None: """ @@ -214,6 +216,7 @@ def start_attempt(self) -> ActiveAttemptMetric | None: return self._handle_error( INVALID_STATE_ERROR.format("start_attempt", self.state) ) + self._active_operation_context.set(self) try: # find backoff value before this attempt From 8707d406e561ddb1758e7da9288b6be823e9e108 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 25 Sep 2025 14:33:48 -0700 Subject: [PATCH 60/78] simplified interceptor --- google/cloud/bigtable/data/_async/client.py | 1 - .../data/_async/metrics_interceptor.py | 25 ------------------- .../data/_metrics/metrics_controller.py | 9 +------ 3 files changed, 1 insertion(+), 34 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 183d89975..f2b910ce1 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -972,7 +972,6 @@ def __init__( ) self._metrics = BigtableClientSideMetricsController( - client._metrics_interceptor, handlers=[], project_id=self.client.project, instance_id=instance_id, diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 6812a8648..3a4670225 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -91,31 +91,6 @@ class AsyncBigtableMetricsInterceptor( An async gRPC interceptor to add client metadata and print server metadata. """ - def __init__(self): - super().__init__() - self.operation_map = {} - - def register_operation(self, operation): - """ - Register an operation object to be tracked my the interceptor - - When registered, the operation will receive metadata updates: - - start_attempt if attempt not started when rpc is being sent - - add_response_metadata after call is complete - - The interceptor will register itself as a handeler for the operation, - so it can unregister the operation when it is complete - """ - self.operation_map[operation.uuid] = operation - operation.handlers.append(self) - - def on_operation_complete(self, op): - if op.uuid in self.operation_map: - del self.operation_map[op.uuid] - - def on_operation_cancelled(self, op): - self.on_operation_complete(op) - @CrossSync.convert @_with_operation_from_metadata async def intercept_unary_unary( diff --git a/google/cloud/bigtable/data/_metrics/metrics_controller.py b/google/cloud/bigtable/data/_metrics/metrics_controller.py index f13590f7c..4ef39814a 100644 --- a/google/cloud/bigtable/data/_metrics/metrics_controller.py +++ b/google/cloud/bigtable/data/_metrics/metrics_controller.py @@ -38,7 +38,6 @@ class BigtableClientSideMetricsController: def __init__( self, - interceptor: AsyncBigtableMetricsInterceptor | BigtableMetricsInterceptor, handlers: list[MetricsHandler] | None = None, **kwargs, ): @@ -52,10 +51,6 @@ def __init__( """ self.interceptor = interceptor self.handlers: list[MetricsHandler] = handlers or [] - if handlers is None: - # handlers not given. Use default handlers. - # TODO: add default handlers - pass def add_handler(self, handler: MetricsHandler) -> None: """ @@ -72,6 +67,4 @@ def create_operation( """ Creates a new operation and registers it with the subscribed handlers. """ - new_op = ActiveOperationMetric(op_type, **kwargs, handlers=self.handlers) - self.interceptor.register_operation(new_op) - return new_op + return ActiveOperationMetric(op_type, **kwargs, handlers=self.handlers) From e9877fdb8499150785ecc560d90a94f17644c3b3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 25 Sep 2025 14:34:21 -0700 Subject: [PATCH 61/78] removed cancel from spec --- .../cloud/bigtable/data/_async/mutations_batcher.py | 1 - google/cloud/bigtable/data/_metrics/data_model.py | 7 ------- google/cloud/bigtable/data/_metrics/handlers/_base.py | 3 --- tests/unit/data/_metrics/test_data_model.py | 11 ----------- 4 files changed, 22 deletions(-) diff --git a/google/cloud/bigtable/data/_async/mutations_batcher.py b/google/cloud/bigtable/data/_async/mutations_batcher.py index 73c9a5022..8b8571b27 100644 --- a/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -195,7 +195,6 @@ async def add_to_flow_with_metrics( try: value = await inner_generator.__anext__() except CrossSync.StopIteration: - metric.cancel() return metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time yield value, metric diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 1e6a2f133..ba03b9d89 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -368,13 +368,6 @@ def end_with_success(self): """ return self.end_with_status(StatusCode.OK) - def cancel(self): - """ - Called to cancel an operation without processing emitting it. - """ - for handler in self.handlers: - handler.on_operation_cancelled(self) - @staticmethod def _exc_to_status(exc: BaseException) -> StatusCode: """ diff --git a/google/cloud/bigtable/data/_metrics/handlers/_base.py b/google/cloud/bigtable/data/_metrics/handlers/_base.py index 64cc89b05..72f5aa550 100644 --- a/google/cloud/bigtable/data/_metrics/handlers/_base.py +++ b/google/cloud/bigtable/data/_metrics/handlers/_base.py @@ -29,9 +29,6 @@ def __init__(self, **kwargs): def on_operation_complete(self, op: CompletedOperationMetric) -> None: pass - def on_operation_cancelled(self, op: ActiveOperationMetric) -> None: - pass - def on_attempt_complete( self, attempt: CompletedAttemptMetric, op: ActiveOperationMetric ) -> None: diff --git a/tests/unit/data/_metrics/test_data_model.py b/tests/unit/data/_metrics/test_data_model.py index 7d9b6671f..e77519629 100644 --- a/tests/unit/data/_metrics/test_data_model.py +++ b/tests/unit/data/_metrics/test_data_model.py @@ -545,17 +545,6 @@ def test_interceptor_metadata(self): assert key == OPERATION_INTERCEPTOR_METADATA_KEY assert value == metric.uuid - def test_cancel(self): - """ - cancel should call on_operation_cancelled on handlers - """ - handlers = [mock.Mock(), mock.Mock()] - metric = self._make_one(mock.Mock(), handlers=handlers) - metric.cancel() - for h in handlers: - assert h.on_operation_cancelled.call_count == 1 - assert h.on_operation_cancelled.call_args[0][0] == metric - def test_end_with_status_with_default_cluster_zone(self): """ ending the operation should use default cluster and zone if not set From f92d66b5abb4a86c7f251b282cb068338381dae1 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 25 Sep 2025 15:25:20 -0700 Subject: [PATCH 62/78] fixed tests --- tests/system/data/test_metrics_async.py | 58 +----- .../data/_async/test_metrics_interceptor.py | 187 ++---------------- 2 files changed, 22 insertions(+), 223 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index bc530b4af..a4d58c9a6 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -57,24 +57,19 @@ class _MetricsTestHandler(MetricsHandler): def __init__(self, **kwargs): self.completed_operations = [] self.completed_attempts = [] - self.cancelled_operations = [] def on_operation_complete(self, op): self.completed_operations.append(op) - def on_operation_cancelled(self, op): - self.cancelled_operations.append(op) - def on_attempt_complete(self, attempt, _): self.completed_attempts.append(attempt) def clear(self): - self.cancelled_operations.clear() self.completed_operations.clear() self.completed_attempts.clear() def __repr__(self): - return f"{self.__class__}(completed_operations={len(self.completed_operations)}, cancelled_operations={len(self.cancelled_operations)}, completed_attempts={len(self.completed_attempts)}" + return f"{self.__class__}(completed_operations={len(self.completed_operations)}, completed_attempts={len(self.completed_attempts)}" @CrossSync.convert_class @@ -241,7 +236,6 @@ async def test_read_rows(self, table, temp_rows, handler, cluster_config): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -297,7 +291,6 @@ async def test_read_rows_failure_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -332,7 +325,6 @@ async def test_read_rows_failure_timeout(self, table, temp_rows, handler): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -365,7 +357,6 @@ async def test_read_rows_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -399,7 +390,6 @@ async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config) # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -455,7 +445,6 @@ async def test_read_rows_stream_failure_closed( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert operation.final_status.name == "CANCELLED" @@ -498,7 +487,6 @@ async def test_read_rows_stream_failure_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -536,7 +524,6 @@ async def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler) # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -570,7 +557,6 @@ async def test_read_rows_stream_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -612,7 +598,6 @@ async def test_read_rows_stream_failure_unauthorized_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) > 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -654,7 +639,6 @@ async def test_read_rows_stream_failure_mid_stream( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" @@ -676,7 +660,6 @@ async def test_read_row(self, table, temp_rows, handler, cluster_config): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -736,7 +719,6 @@ async def test_read_row_failure_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -771,7 +753,6 @@ async def test_read_row_failure_timeout(self, table, temp_rows, handler): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -804,7 +785,6 @@ async def test_read_row_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -842,7 +822,6 @@ async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config # validate counts assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 # validate operations for operation in handler.completed_operations: assert isinstance(operation, CompletedOperationMetric) @@ -902,7 +881,6 @@ async def test_read_rows_sharded_failure_with_retries( assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 3 - assert len(handler.cancelled_operations) == 0 # validate operations for op in handler.completed_operations: assert op.final_status.name == "OK" @@ -936,7 +914,6 @@ async def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler # both shards should fail assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 # validate operations for operation in handler.completed_operations: assert isinstance(operation, CompletedOperationMetric) @@ -976,7 +953,6 @@ async def test_read_rows_sharded_failure_unauthorized( # one shard will fail, the other will succeed assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 # sort operations by status failed_op = next( op for op in handler.completed_operations if op.final_status.name != "OK" @@ -1036,7 +1012,6 @@ async def test_read_rows_sharded_failure_mid_stream( # the failing shard will have one retry assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 3 - assert len(handler.cancelled_operations) == 0 # sort operations by status failed_op = next( op for op in handler.completed_operations if op.final_status.name != "OK" @@ -1077,7 +1052,6 @@ async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config) # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1140,7 +1114,6 @@ async def test_bulk_mutate_rows_failure_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1181,7 +1154,6 @@ async def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler) # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1217,7 +1189,6 @@ async def test_bulk_mutate_rows_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" @@ -1262,7 +1233,6 @@ async def test_bulk_mutate_rows_failure_unauthorized_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) > 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1304,11 +1274,6 @@ async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_conf assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 # bacher expects to cancel staged operation on close - assert len(handler.cancelled_operations) == 1 - cancelled = handler.cancelled_operations[0] - assert isinstance(cancelled, ActiveOperationMetric) - assert cancelled.state == OperationState.CREATED - assert not cancelled.completed_attempts # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1377,7 +1342,6 @@ async def test_mutate_rows_batcher_failure_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 1 # from batcher auto-closing # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1418,7 +1382,6 @@ async def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handl # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 1 # from batcher auto-closing # validate operation operation = handler.completed_operations[0] assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1457,7 +1420,6 @@ async def test_mutate_rows_batcher_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 1 # from batcher auto-closing # validate operation operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" @@ -1489,7 +1451,6 @@ async def test_mutate_row(self, table, temp_rows, handler, cluster_config): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1549,7 +1510,6 @@ async def test_mutate_row_failure_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1587,7 +1547,6 @@ async def test_mutate_row_failure_timeout(self, table, temp_rows, handler): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1621,7 +1580,6 @@ async def test_mutate_row_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1661,7 +1619,6 @@ async def test_mutate_row_failure_unauthorized_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) > 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1688,7 +1645,6 @@ async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1739,7 +1695,6 @@ async def test_sample_row_keys_failure_cancelled( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1776,7 +1731,6 @@ async def test_sample_row_keys_failure_with_retries( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1816,7 +1770,6 @@ async def test_sample_row_keys_failure_timeout(self, table, handler): # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1847,7 +1800,6 @@ async def test_sample_row_keys_failure_mid_stream( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" @@ -1874,7 +1826,6 @@ async def test_read_modify_write(self, table, temp_rows, handler, cluster_config # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1932,7 +1883,6 @@ async def test_read_modify_write_failure_cancelled( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -1977,7 +1927,6 @@ async def test_read_modify_write_failure_timeout(self, table, temp_rows, handler # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -2006,7 +1955,6 @@ async def test_read_modify_write_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -2049,7 +1997,6 @@ async def test_check_and_mutate_row( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -2108,7 +2055,6 @@ async def test_check_and_mutate_row_failure_cancelled( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -2165,7 +2111,6 @@ async def test_check_and_mutate_row_failure_timeout( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) @@ -2202,7 +2147,6 @@ async def test_check_and_mutate_row_failure_unauthorized( # validate counts assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 # validate operation operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) diff --git a/tests/unit/data/_async/test_metrics_interceptor.py b/tests/unit/data/_async/test_metrics_interceptor.py index caa9bbb46..1f99473e3 100644 --- a/tests/unit/data/_async/test_metrics_interceptor.py +++ b/tests/unit/data/_async/test_metrics_interceptor.py @@ -16,6 +16,7 @@ from grpc import RpcError from grpc import ClientCallDetails +from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric from google.cloud.bigtable.data._metrics.data_model import OperationState from google.cloud.bigtable.data._cross_sync import CrossSync @@ -71,92 +72,10 @@ def _make_one(self, *args, **kwargs): def test_ctor(self): instance = self._make_one() - assert instance.operation_map == {} - - def test_register_operation(self): - """ - adding a new operation should register it in operation_map - """ - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - assert instance.operation_map[op.uuid] == op - assert instance in op.handlers - - def test_on_operation_comple_mock(self): - """ - completing or cancelling an operation should call on_operation_complete on interceptor - """ - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - instance.on_operation_complete = mock.Mock() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - op.end_with_success() - assert instance.on_operation_complete.call_count == 1 - op.cancel() - assert instance.on_operation_complete.call_count == 2 - - def test_on_operation_complete(self): - """ - completing an operation should remove it from the operation map - """ - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - op.end_with_success() - instance.on_operation_complete(op) - assert op.uuid not in instance.operation_map - - def test_on_operation_cancelled(self): - """ - completing an operation should remove it from the operation map - """ - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - op.cancel() - assert op.uuid not in instance.operation_map - - @CrossSync.pytest - async def test_strip_operation_id_metadata(self): - """ - After operation id is detected in metadata, the field should be stripped out before calling continuation - """ - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - - instance = self._make_one() - op = mock.Mock() - op.uuid = "test-uuid" - op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op - continuation = CrossSync.Mock() - details = ClientCallDetails() - details.metadata = [ - (OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid), - ("other_key", "other_value"), - ] - await instance.intercept_unary_unary(continuation, details, mock.Mock()) - assert details.metadata == [("other_key", "other_value")] - assert continuation.call_count == 1 - assert continuation.call_args[0][0].metadata == [("other_key", "other_value")] @CrossSync.pytest async def test_unary_unary_interceptor_op_not_found(self): - """Test that interceptor call cuntinuation if op is not found""" + """Test that interceptor call continuation if op is not found""" instance = self._make_one() continuation = CrossSync.Mock() details = ClientCallDetails() @@ -168,107 +87,84 @@ async def test_unary_unary_interceptor_op_not_found(self): @CrossSync.pytest async def test_unary_unary_interceptor_success(self): """Test that interceptor handles successful unary-unary calls""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync.Mock() call = continuation.return_value call.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() result = await instance.intercept_unary_unary(continuation, details, request) assert result == call continuation.assert_called_once_with(details, request) - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) op.end_attempt_with_status.assert_not_called() @CrossSync.pytest async def test_unary_unary_interceptor_failure(self): """test a failed RpcError with metadata""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") exc.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")]) exc.initial_metadata = CrossSync.Mock(return_value=[("c", "d")]) continuation = CrossSync.Mock(side_effect=exc) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: await instance.intercept_unary_unary(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) - op.end_attempt_with_status.assert_called_once_with(exc) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) @CrossSync.pytest async def test_unary_unary_interceptor_failure_no_metadata(self): """test with RpcError without without metadata attached""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") continuation = CrossSync.Mock(side_effect=exc) call = continuation.return_value call.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: await instance.intercept_unary_unary(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) @CrossSync.pytest async def test_unary_unary_interceptor_failure_generic(self): """test generic exception""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = ValueError("test") continuation = CrossSync.Mock(side_effect=exc) call = continuation.return_value call.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(ValueError) as e: await instance.intercept_unary_unary(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) @CrossSync.pytest async def test_unary_stream_interceptor_op_not_found(self): @@ -284,78 +180,61 @@ async def test_unary_stream_interceptor_op_not_found(self): @CrossSync.pytest async def test_unary_stream_interceptor_success(self): """Test that interceptor handles successful unary-stream calls""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync.Mock(return_value=_make_mock_stream_call([1, 2])) call = continuation.return_value call.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() wrapper = await instance.intercept_unary_stream(continuation, details, request) results = [val async for val in wrapper] assert results == [1, 2] continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) op.end_attempt_with_status.assert_not_called() @CrossSync.pytest async def test_unary_stream_interceptor_failure_mid_stream(self): """Test that interceptor handles failures mid-stream""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - + from grpc.aio import AioRpcError, Metadata instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op - exc = ValueError("test") + ActiveOperationMetric._active_operation_context.set(op) + exc = AioRpcError(0, Metadata(), Metadata(("a", "b"), ("c", "d"))) continuation = CrossSync.Mock(return_value=_make_mock_stream_call([1], exc=exc)) - call = continuation.return_value - call.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")]) - call.initial_metadata = CrossSync.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() wrapper = await instance.intercept_unary_stream(continuation, details, request) - with pytest.raises(ValueError) as e: + with pytest.raises(AioRpcError) as e: [val async for val in wrapper] assert e.value == exc continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) - op.end_attempt_with_status.assert_called_once_with(exc) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) @CrossSync.pytest async def test_unary_stream_interceptor_failure_start_stream(self): """Test that interceptor handles failures at start of stream with RpcError with metadata""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") exc.trailing_metadata = CrossSync.Mock(return_value=[("a", "b")]) exc.initial_metadata = CrossSync.Mock(return_value=[("c", "d")]) @@ -363,36 +242,29 @@ async def test_unary_stream_interceptor_failure_start_stream(self): continuation = CrossSync.Mock() continuation.side_effect = exc details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: await instance.intercept_unary_stream(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) - op.end_attempt_with_status.assert_called_once_with(exc) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) @CrossSync.pytest async def test_unary_stream_interceptor_failure_start_stream_no_metadata(self): """Test that interceptor handles failures at start of stream with RpcError with no metadata""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") continuation = CrossSync.Mock() continuation.side_effect = exc details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: await instance.intercept_unary_stream(continuation, details, request) @@ -400,28 +272,22 @@ async def test_unary_stream_interceptor_failure_start_stream_no_metadata(self): continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) @CrossSync.pytest async def test_unary_stream_interceptor_failure_start_stream_generic(self): """Test that interceptor handles failures at start of stream with generic exception""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = ValueError("test") continuation = CrossSync.Mock() continuation.side_effect = exc details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(ValueError) as e: await instance.intercept_unary_stream(continuation, details, request) @@ -429,7 +295,6 @@ async def test_unary_stream_interceptor_failure_start_stream_generic(self): continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) @CrossSync.pytest @pytest.mark.parametrize( @@ -437,21 +302,16 @@ async def test_unary_stream_interceptor_failure_start_stream_generic(self): ) async def test_unary_unary_interceptor_start_operation(self, initial_state): """if called with a newly created operation, it should be started""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = initial_state - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync.Mock() call = continuation.return_value call.trailing_metadata = CrossSync.Mock(return_value=[]) call.initial_metadata = CrossSync.Mock(return_value=[]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() await instance.intercept_unary_unary(continuation, details, request) op.start_attempt.assert_called_once() @@ -462,22 +322,17 @@ async def test_unary_unary_interceptor_start_operation(self, initial_state): ) async def test_unary_stream_interceptor_start_operation(self, initial_state): """if called with a newly created operation, it should be started""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = initial_state - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync.Mock(return_value=_make_mock_stream_call([1, 2])) call = continuation.return_value call.trailing_metadata = CrossSync.Mock(return_value=[]) call.initial_metadata = CrossSync.Mock(return_value=[]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() await instance.intercept_unary_stream(continuation, details, request) op.start_attempt.assert_called_once() From d9de44d59a0dd4425b4b74b243bc3f1241aaaad0 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 25 Sep 2025 15:25:29 -0700 Subject: [PATCH 63/78] removed dead pointer --- google/cloud/bigtable/data/_metrics/metrics_controller.py | 1 - 1 file changed, 1 deletion(-) diff --git a/google/cloud/bigtable/data/_metrics/metrics_controller.py b/google/cloud/bigtable/data/_metrics/metrics_controller.py index 4ef39814a..087c5058b 100644 --- a/google/cloud/bigtable/data/_metrics/metrics_controller.py +++ b/google/cloud/bigtable/data/_metrics/metrics_controller.py @@ -49,7 +49,6 @@ def __init__( - handlers: A list of MetricsHandler objects to subscribe to metrics events. - **kwargs: Optional arguments to pass to the metrics handlers. """ - self.interceptor = interceptor self.handlers: list[MetricsHandler] = handlers or [] def add_handler(self, handler: MetricsHandler) -> None: From 2e0e40276f6aac0a631015a9be229935faca3770 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 25 Sep 2025 16:32:00 -0700 Subject: [PATCH 64/78] updated sync files --- .../data/_sync_autogen/_mutate_rows.py | 5 +- .../bigtable/data/_sync_autogen/_read_rows.py | 13 +- .../bigtable/data/_sync_autogen/client.py | 20 +- .../data/_sync_autogen/metrics_interceptor.py | 76 ++------ .../data/_sync_autogen/mutations_batcher.py | 7 +- tests/system/data/test_metrics_autogen.py | 57 +----- tests/system/data/test_system_autogen.py | 165 +++++++---------- .../_sync_autogen/test_metrics_interceptor.py | 175 +++--------------- 8 files changed, 135 insertions(+), 383 deletions(-) diff --git a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index e7df85431..a490ab10b 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -84,7 +84,10 @@ def __init__( self.is_retryable, metric.backoff_generator, 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, ) self.timeout_generator = _attempt_timeout_generator( attempt_timeout, operation_timeout diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index 1a40b2fe5..75ad79931 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -19,6 +19,7 @@ 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 @@ -107,7 +108,10 @@ def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]: self._predicate, 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]: @@ -268,7 +272,7 @@ def merge_rows( if self._operation_metric.active_attempt is not None: self._operation_metric.active_attempt.application_blocking_time_ns += ( time.monotonic_ns() - block_time - ) * 1000 + ) break c = it.__next__() except _ResetRow as e: @@ -285,9 +289,10 @@ def merge_rows( 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: - if not self._predicate(generic_exception): - self._operation_metric.end_attempt_with_status(generic_exception) raise generic_exception else: self._operation_metric.end_with_success() diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index 9f84d8e14..ae1c88449 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -759,7 +759,6 @@ def __init__( default_retryable_errors or () ) self._metrics = BigtableClientSideMetricsController( - client._metrics_interceptor, handlers=[], project_id=self.client.project, instance_id=instance_id, @@ -941,8 +940,9 @@ def read_row( ) results_generator = row_merger.start_operation() try: - return results_generator.__next__() - except StopIteration: + results = [a for a in results_generator] + return results[0] + except IndexError: return None def read_rows_sharded( @@ -1142,7 +1142,10 @@ def execute_rpc(): predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory + ), + on_error=operation_metric.track_retryable_error, ) def mutations_batcher( @@ -1265,7 +1268,10 @@ def mutate_row( predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=_retry_exception_factory, + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory + ), + on_error=operation_metric.track_retryable_error, ) def bulk_mutate_rows( @@ -1371,7 +1377,7 @@ def check_and_mutate_row( ): false_case_mutations = [false_case_mutations] false_case_list = [m._to_pb() for m in false_case_mutations or []] - with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE): + with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op: result = self.client._gapic_client.check_and_mutate_row( request=CheckAndMutateRowRequest( true_mutations=true_case_list, @@ -1425,7 +1431,7 @@ def read_modify_write_row( rules = [rules] if not rules: raise ValueError("rules must contain at least one item") - with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE): + with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op: result = self.client._gapic_client.read_modify_write_row( request=ReadModifyWriteRowRequest( rules=[rule._to_pb() for rule in rules], diff --git a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py index 4cefed824..e592c6329 100644 --- a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py @@ -15,13 +15,12 @@ # This file is automatically generated by CrossSync. Do not edit manually. from __future__ import annotations +from typing import Sequence import time from functools import wraps -from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, -) from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric from google.cloud.bigtable.data._metrics.data_model import OperationState +from google.cloud.bigtable.data._metrics.data_model import OperationType from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from grpc import UnaryUnaryClientInterceptor from grpc import UnaryStreamClientInterceptor @@ -33,19 +32,7 @@ def _with_operation_from_metadata(func): @wraps(func) def wrapper(self, continuation, client_call_details, request): - found_operation_id: str | None = None - try: - new_metadata: list[tuple[str, str]] = [] - if client_call_details.metadata: - for k, v in client_call_details.metadata: - if k == OPERATION_INTERCEPTOR_METADATA_KEY: - found_operation_id = v - else: - new_metadata.append((k, v)) - client_call_details.metadata = new_metadata - except Exception: - pass - operation: "ActiveOperationMetric" = self.operation_map.get(found_operation_id) + operation: "ActiveOperationMetric" | None = ActiveOperationMetric.get_active() if operation: if ( operation.state == OperationState.CREATED @@ -59,18 +46,13 @@ def wrapper(self, continuation, client_call_details, request): return wrapper -def _end_attempt(operation, exc, metadata): - """Helper to add metadata and exception to an operation""" - if metadata is not None: - operation.add_response_metadata(metadata) - if exc is not None: - operation.end_attempt_with_status(exc) - - -def _get_metadata(source): +def _get_metadata(source) -> dict[str, str | bytes] | None: """Helper to extract metadata from a call or RpcError""" try: - return (source.trailing_metadata() or []) + (source.initial_metadata() or []) + metadata: Sequence[tuple[str.str | bytes]] = ( + source.trailing_metadata() + source.initial_metadata() + ) + return {k: v for (k, v) in metadata} except Exception: return None @@ -82,35 +64,10 @@ class BigtableMetricsInterceptor( An async gRPC interceptor to add client metadata and print server metadata. """ - def __init__(self): - super().__init__() - self.operation_map = {} - - def register_operation(self, operation): - """Register an operation object to be tracked my the interceptor - - When registered, the operation will receive metadata updates: - - start_attempt if attempt not started when rpc is being sent - - add_response_metadata after call is complete - - end_attempt_with_status if attempt receives an error - - The interceptor will register itself as a handeler for the operation, - so it can unregister the operation when it is complete""" - self.operation_map[operation.uuid] = operation - operation.handlers.append(self) - - def on_operation_complete(self, op): - if op.uuid in self.operation_map: - del self.operation_map[op.uuid] - - def on_operation_cancelled(self, op): - self.on_operation_complete(op) - @_with_operation_from_metadata def intercept_unary_unary( self, operation, continuation, client_call_details, request ): - encountered_exc: Exception | None = None metadata = None try: call = continuation(client_call_details, request) @@ -118,17 +75,20 @@ def intercept_unary_unary( return call except Exception as rpc_error: metadata = _get_metadata(rpc_error) - encountered_exc = rpc_error raise rpc_error finally: - _end_attempt(operation, encountered_exc, metadata) + if metadata is not None: + operation.add_response_metadata(metadata) @_with_operation_from_metadata def intercept_unary_stream( self, operation, continuation, client_call_details, request ): def response_wrapper(call): - has_first_response = operation.first_response_latency is not None + has_first_response = ( + operation.first_response_latency_ns is not None + or operation.op_type != OperationType.READ_ROWS + ) encountered_exc = None try: for response in call: @@ -143,10 +103,14 @@ def response_wrapper(call): raise finally: if call is not None: - _end_attempt(operation, encountered_exc, _get_metadata(call)) + metadata = _get_metadata(encountered_exc or call) + if metadata is not None: + operation.add_response_metadata(metadata) try: return response_wrapper(continuation(client_call_details, request)) except Exception as rpc_error: - _end_attempt(operation, rpc_error, _get_metadata(rpc_error)) + metadata = _get_metadata(rpc_error) + if metadata is not None: + operation.add_response_metadata(metadata) raise rpc_error diff --git a/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py b/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py index 3f4c61bc3..9f47eb39a 100644 --- a/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py +++ b/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py @@ -151,7 +151,7 @@ def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]): ) yield mutations[start_idx:end_idx] - async def add_to_flow_with_metrics( + def add_to_flow_with_metrics( self, mutations: RowMutationEntry | list[RowMutationEntry], metrics_controller: BigtableClientSideMetricsController, @@ -161,9 +161,8 @@ async def add_to_flow_with_metrics( metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS) flow_start_time = time.monotonic_ns() try: - value = await inner_generator.__anext__() - except StopAsyncIteration: - metric.cancel() + value = inner_generator.__next__() + except CrossSync._Sync_Impl.StopIteration: return metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time yield (value, metric) diff --git a/tests/system/data/test_metrics_autogen.py b/tests/system/data/test_metrics_autogen.py index 4528c708b..bcd74de8d 100644 --- a/tests/system/data/test_metrics_autogen.py +++ b/tests/system/data/test_metrics_autogen.py @@ -25,8 +25,6 @@ from google.cloud.bigtable.data._metrics.data_model import ( CompletedOperationMetric, CompletedAttemptMetric, - ActiveOperationMetric, - OperationState, ) from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable_v2.types import ResponseParams @@ -46,24 +44,19 @@ class _MetricsTestHandler(MetricsHandler): def __init__(self, **kwargs): self.completed_operations = [] self.completed_attempts = [] - self.cancelled_operations = [] def on_operation_complete(self, op): self.completed_operations.append(op) - def on_operation_cancelled(self, op): - self.cancelled_operations.append(op) - def on_attempt_complete(self, attempt, _): self.completed_attempts.append(attempt) def clear(self): - self.cancelled_operations.clear() self.completed_operations.clear() self.completed_attempts.clear() def __repr__(self): - return f"{self.__class__}(completed_operations={len(self.completed_operations)}, cancelled_operations={len(self.cancelled_operations)}, completed_attempts={len(self.completed_attempts)}" + return f"{self.__class__}(completed_operations={len(self.completed_operations)}, completed_attempts={len(self.completed_attempts)}" class _ErrorInjectorInterceptor( @@ -199,7 +192,6 @@ def test_read_rows(self, table, temp_rows, handler, cluster_config): assert len(row_list) == 2 assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -253,7 +245,6 @@ def test_read_rows_failure_with_retries( table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -282,7 +273,6 @@ def test_read_rows_failure_timeout(self, table, temp_rows, handler): table.read_rows(ReadRowsQuery(), operation_timeout=0.001) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -309,7 +299,6 @@ def test_read_rows_failure_unauthorized( assert e.value.grpc_status_code.name == "PERMISSION_DENIED" assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -338,7 +327,6 @@ def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): assert len(row_list) == 2 assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -385,7 +373,6 @@ def test_read_rows_stream_failure_closed( generator.__next__() assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert operation.final_status.name == "CANCELLED" assert operation.op_type.value == "ReadRows" @@ -419,7 +406,6 @@ def test_read_rows_stream_failure_with_retries( [_ for _ in generator] assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -449,7 +435,6 @@ def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): [_ for _ in generator] assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -477,7 +462,6 @@ def test_read_rows_stream_failure_unauthorized( assert e.value.grpc_status_code.name == "PERMISSION_DENIED" assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -513,7 +497,6 @@ def test_read_rows_stream_failure_unauthorized_with_retries( assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) > 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -547,7 +530,6 @@ def test_read_rows_stream_failure_mid_stream( [_ for _ in generator] assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" assert operation.op_type.value == "ReadRows" @@ -564,7 +546,6 @@ def test_read_row(self, table, temp_rows, handler, cluster_config): table.read_row(b"row_key_1") assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -618,7 +599,6 @@ def test_read_row_failure_with_retries( table.read_row(b"row_key_1", retryable_errors=[Aborted]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -647,7 +627,6 @@ def test_read_row_failure_timeout(self, table, temp_rows, handler): table.read_row(b"row_key_1", operation_timeout=0.001) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -674,7 +653,6 @@ def test_read_row_failure_unauthorized( assert e.value.grpc_status_code.name == "PERMISSION_DENIED" assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -708,7 +686,6 @@ def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): assert len(row_list) == 4 assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 for operation in handler.completed_operations: assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -760,7 +737,6 @@ def test_read_rows_sharded_failure_with_retries( table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 3 - assert len(handler.cancelled_operations) == 0 for op in handler.completed_operations: assert op.final_status.name == "OK" assert op.op_type.value == "ReadRows" @@ -800,7 +776,6 @@ def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): assert isinstance(sub_exc.__cause__, DeadlineExceeded) assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 for operation in handler.completed_operations: assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -834,7 +809,6 @@ def test_read_rows_sharded_failure_unauthorized( ) assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 failed_op = next( (op for op in handler.completed_operations if op.final_status.name != "OK") ) @@ -884,7 +858,6 @@ def test_read_rows_sharded_failure_mid_stream( assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) assert len(handler.completed_operations) == 2 assert len(handler.completed_attempts) == 3 - assert len(handler.cancelled_operations) == 0 failed_op = next( (op for op in handler.completed_operations if op.final_status.name != "OK") ) @@ -916,7 +889,6 @@ def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): table.bulk_mutate_rows([bulk_mutation]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -970,7 +942,6 @@ def test_bulk_mutate_rows_failure_with_retries( table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -1004,7 +975,6 @@ def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): table.bulk_mutate_rows([entry], operation_timeout=0.001) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1033,7 +1003,6 @@ def test_bulk_mutate_rows_failure_unauthorized( authorized_view.bulk_mutate_rows([entry]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" assert operation.op_type.value == "MutateRows" @@ -1073,7 +1042,6 @@ def test_bulk_mutate_rows_failure_unauthorized_with_retries( assert len(e.value.exceptions) == 1 assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) > 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert operation.final_status.name == "DEADLINE_EXCEEDED" assert operation.op_type.value == "MutateRows" @@ -1109,11 +1077,6 @@ def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): batcher.append(bulk_mutation2) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 1 - cancelled = handler.cancelled_operations[0] - assert isinstance(cancelled, ActiveOperationMetric) - assert cancelled.state == OperationState.CREATED - assert not cancelled.completed_attempts operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -1171,7 +1134,6 @@ def test_mutate_rows_batcher_failure_with_retries( batcher.append(entry) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 1 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -1203,7 +1165,6 @@ def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): batcher.append(entry) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 1 operation = handler.completed_operations[0] assert operation.final_status.name == "DEADLINE_EXCEEDED" assert operation.op_type.value == "MutateRows" @@ -1235,7 +1196,6 @@ def test_mutate_rows_batcher_failure_unauthorized( ) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 1 operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" assert operation.op_type.value == "MutateRows" @@ -1263,7 +1223,6 @@ def test_mutate_row(self, table, temp_rows, handler, cluster_config): table.mutate_row(row_key, mutation) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -1312,7 +1271,6 @@ def test_mutate_row_failure_with_retries(self, table, handler, error_injector): table.mutate_row(row_key, [mutation], retryable_errors=[Aborted]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -1343,7 +1301,6 @@ def test_mutate_row_failure_timeout(self, table, temp_rows, handler): table.mutate_row(row_key, [mutation], operation_timeout=0.001) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1370,7 +1327,6 @@ def test_mutate_row_failure_unauthorized( assert e.value.grpc_status_code.name == "PERMISSION_DENIED" assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -1408,7 +1364,6 @@ def test_mutate_row_failure_unauthorized_with_retries( assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) > 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1431,7 +1386,6 @@ 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 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -1469,7 +1423,6 @@ def test_sample_row_keys_failure_with_retries( table.sample_row_keys(retryable_errors=[Aborted]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == num_retryable + 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "OK" @@ -1503,7 +1456,6 @@ def test_sample_row_keys_failure_timeout(self, table, handler): table.sample_row_keys(operation_timeout=0.001) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1528,7 +1480,6 @@ def test_sample_row_keys_failure_mid_stream( table.sample_row_keys(retryable_errors=[Aborted]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 2 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert operation.final_status.name == "PERMISSION_DENIED" assert operation.op_type.value == "SampleRowKeys" @@ -1550,7 +1501,6 @@ def test_read_modify_write(self, table, temp_rows, handler, cluster_config): table.read_modify_write_row(row_key, rule) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -1592,7 +1542,6 @@ def test_read_modify_write_failure_timeout(self, table, temp_rows, handler): table.read_modify_write_row(row_key, rule, operation_timeout=0.001) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1615,7 +1564,6 @@ def test_read_modify_write_failure_unauthorized( authorized_view.read_modify_write_row(row_key, rule) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" @@ -1649,7 +1597,6 @@ def test_check_and_mutate_row(self, table, temp_rows, handler, cluster_config): ) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.value[0] == 0 @@ -1700,7 +1647,6 @@ def test_check_and_mutate_row_failure_timeout(self, table, temp_rows, handler): ) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "DEADLINE_EXCEEDED" @@ -1731,7 +1677,6 @@ def test_check_and_mutate_row_failure_unauthorized( ) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 - assert len(handler.cancelled_operations) == 0 operation = handler.completed_operations[0] assert isinstance(operation, CompletedOperationMetric) assert operation.final_status.name == "PERMISSION_DENIED" 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/_sync_autogen/test_metrics_interceptor.py b/tests/unit/data/_sync_autogen/test_metrics_interceptor.py index 283814e27..ec3e51583 100644 --- a/tests/unit/data/_sync_autogen/test_metrics_interceptor.py +++ b/tests/unit/data/_sync_autogen/test_metrics_interceptor.py @@ -18,6 +18,7 @@ import pytest from grpc import RpcError from grpc import ClientCallDetails +from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric from google.cloud.bigtable.data._metrics.data_model import OperationState from google.cloud.bigtable.data._cross_sync import CrossSync @@ -54,80 +55,9 @@ def _make_one(self, *args, **kwargs): def test_ctor(self): instance = self._make_one() - assert instance.operation_map == {} - - def test_register_operation(self): - """adding a new operation should register it in operation_map""" - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - assert instance.operation_map[op.uuid] == op - assert instance in op.handlers - - def test_on_operation_comple_mock(self): - """completing or cancelling an operation should call on_operation_complete on interceptor""" - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - instance.on_operation_complete = mock.Mock() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - op.end_with_success() - assert instance.on_operation_complete.call_count == 1 - op.cancel() - assert instance.on_operation_complete.call_count == 2 - - def test_on_operation_complete(self): - """completing an operation should remove it from the operation map""" - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - op.end_with_success() - instance.on_operation_complete(op) - assert op.uuid not in instance.operation_map - - def test_on_operation_cancelled(self): - """completing an operation should remove it from the operation map""" - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric - from google.cloud.bigtable.data._metrics.data_model import OperationType - - instance = self._make_one() - op = ActiveOperationMetric(OperationType.READ_ROWS) - instance.register_operation(op) - op.cancel() - assert op.uuid not in instance.operation_map - - def test_strip_operation_id_metadata(self): - """After operation id is detected in metadata, the field should be stripped out before calling continuation""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - - instance = self._make_one() - op = mock.Mock() - op.uuid = "test-uuid" - op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op - continuation = CrossSync._Sync_Impl.Mock() - details = ClientCallDetails() - details.metadata = [ - (OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid), - ("other_key", "other_value"), - ] - instance.intercept_unary_unary(continuation, details, mock.Mock()) - assert details.metadata == [("other_key", "other_value")] - assert continuation.call_count == 1 - assert continuation.call_args[0][0].metadata == [("other_key", "other_value")] def test_unary_unary_interceptor_op_not_found(self): - """Test that interceptor call cuntinuation if op is not found""" + """Test that interceptor call continuation if op is not found""" instance = self._make_one() continuation = CrossSync._Sync_Impl.Mock() details = ClientCallDetails() @@ -138,104 +68,81 @@ def test_unary_unary_interceptor_op_not_found(self): def test_unary_unary_interceptor_success(self): """Test that interceptor handles successful unary-unary calls""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync._Sync_Impl.Mock() call = continuation.return_value call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() result = instance.intercept_unary_unary(continuation, details, request) assert result == call continuation.assert_called_once_with(details, request) - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) op.end_attempt_with_status.assert_not_called() def test_unary_unary_interceptor_failure(self): """test a failed RpcError with metadata""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") exc.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")]) exc.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")]) continuation = CrossSync._Sync_Impl.Mock(side_effect=exc) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: instance.intercept_unary_unary(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) - op.end_attempt_with_status.assert_called_once_with(exc) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) def test_unary_unary_interceptor_failure_no_metadata(self): """test with RpcError without without metadata attached""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") continuation = CrossSync._Sync_Impl.Mock(side_effect=exc) call = continuation.return_value call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: instance.intercept_unary_unary(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) def test_unary_unary_interceptor_failure_generic(self): """test generic exception""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = ValueError("test") continuation = CrossSync._Sync_Impl.Mock(side_effect=exc) call = continuation.return_value call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(ValueError) as e: instance.intercept_unary_unary(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) def test_unary_stream_interceptor_op_not_found(self): """Test that interceptor calls continuation if op is not found""" @@ -249,17 +156,13 @@ def test_unary_stream_interceptor_op_not_found(self): def test_unary_stream_interceptor_success(self): """Test that interceptor handles successful unary-stream calls""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync._Sync_Impl.Mock( return_value=_make_mock_stream_call([1, 2]) ) @@ -267,21 +170,18 @@ def test_unary_stream_interceptor_success(self): call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")]) call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() wrapper = instance.intercept_unary_stream(continuation, details, request) results = [val for val in wrapper] assert results == [1, 2] continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) op.end_attempt_with_status.assert_not_called() def test_unary_stream_interceptor_failure_mid_stream(self): """Test that interceptor handles failures mid-stream""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) + from grpc.aio import AioRpcError, Metadata instance = self._make_one() op = mock.Mock() @@ -289,73 +189,57 @@ def test_unary_stream_interceptor_failure_mid_stream(self): op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op - exc = ValueError("test") + ActiveOperationMetric._active_operation_context.set(op) + exc = AioRpcError(0, Metadata(), Metadata(("a", "b"), ("c", "d"))) continuation = CrossSync._Sync_Impl.Mock( return_value=_make_mock_stream_call([1], exc=exc) ) - call = continuation.return_value - call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")]) - call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() wrapper = instance.intercept_unary_stream(continuation, details, request) - with pytest.raises(ValueError) as e: + with pytest.raises(AioRpcError) as e: [val for val in wrapper] assert e.value == exc continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) - op.end_attempt_with_status.assert_called_once_with(exc) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) def test_unary_stream_interceptor_failure_start_stream(self): """Test that interceptor handles failures at start of stream with RpcError with metadata""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") exc.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[("a", "b")]) exc.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[("c", "d")]) continuation = CrossSync._Sync_Impl.Mock() continuation.side_effect = exc details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: instance.intercept_unary_stream(continuation, details, request) assert e.value == exc continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None - op.add_response_metadata.assert_called_once_with([("a", "b"), ("c", "d")]) - op.end_attempt_with_status.assert_called_once_with(exc) + op.add_response_metadata.assert_called_once_with({"a": "b", "c": "d"}) def test_unary_stream_interceptor_failure_start_stream_no_metadata(self): """Test that interceptor handles failures at start of stream with RpcError with no metadata""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = RpcError("test") continuation = CrossSync._Sync_Impl.Mock() continuation.side_effect = exc details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(RpcError) as e: instance.intercept_unary_stream(continuation, details, request) @@ -363,26 +247,20 @@ def test_unary_stream_interceptor_failure_start_stream_no_metadata(self): continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) def test_unary_stream_interceptor_failure_start_stream_generic(self): """Test that interceptor handles failures at start of stream with generic exception""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = OperationState.ACTIVE_ATTEMPT op.start_time_ns = 0 op.first_response_latency = None - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) exc = ValueError("test") continuation = CrossSync._Sync_Impl.Mock() continuation.side_effect = exc details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() with pytest.raises(ValueError) as e: instance.intercept_unary_stream(continuation, details, request) @@ -390,28 +268,22 @@ def test_unary_stream_interceptor_failure_start_stream_generic(self): continuation.assert_called_once_with(details, request) assert op.first_response_latency_ns is not None op.add_response_metadata.assert_not_called() - op.end_attempt_with_status.assert_called_once_with(exc) @pytest.mark.parametrize( "initial_state", [OperationState.CREATED, OperationState.BETWEEN_ATTEMPTS] ) def test_unary_unary_interceptor_start_operation(self, initial_state): """if called with a newly created operation, it should be started""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = initial_state - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync._Sync_Impl.Mock() call = continuation.return_value call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[]) call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() instance.intercept_unary_unary(continuation, details, request) op.start_attempt.assert_called_once() @@ -421,15 +293,11 @@ def test_unary_unary_interceptor_start_operation(self, initial_state): ) def test_unary_stream_interceptor_start_operation(self, initial_state): """if called with a newly created operation, it should be started""" - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" op.state = initial_state - instance.operation_map[op.uuid] = op + ActiveOperationMetric._active_operation_context.set(op) continuation = CrossSync._Sync_Impl.Mock( return_value=_make_mock_stream_call([1, 2]) ) @@ -437,7 +305,6 @@ def test_unary_stream_interceptor_start_operation(self, initial_state): call.trailing_metadata = CrossSync._Sync_Impl.Mock(return_value=[]) call.initial_metadata = CrossSync._Sync_Impl.Mock(return_value=[]) details = ClientCallDetails() - details.metadata = [(OPERATION_INTERCEPTOR_METADATA_KEY, op.uuid)] request = mock.Mock() instance.intercept_unary_stream(continuation, details, request) op.start_attempt.assert_called_once() From 88644b2e540d3fe9d8c16d20ddaf13a71c32f532 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 26 Sep 2025 14:39:56 -0700 Subject: [PATCH 65/78] fixed lint --- .../bigtable/data/_async/_mutate_rows.py | 4 +- .../cloud/bigtable/data/_async/_read_rows.py | 4 +- google/cloud/bigtable/data/_async/client.py | 12 ++- .../data/_async/metrics_interceptor.py | 91 ++++++++++++------- google/cloud/bigtable/data/_helpers.py | 2 +- .../bigtable/data/_metrics/data_model.py | 39 ++++---- .../data/_metrics/metrics_controller.py | 10 -- .../bigtable/data/_sync_autogen/client.py | 4 +- .../data/_sync_autogen/metrics_interceptor.py | 60 +++++++----- tests/system/data/test_metrics_async.py | 65 +++++++++---- tests/system/data/test_metrics_autogen.py | 2 +- .../data/_async/test_metrics_interceptor.py | 4 +- .../_sync_autogen/test_metrics_interceptor.py | 3 - 13 files changed, 180 insertions(+), 120 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 66baef89a..9884929ba 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -106,7 +106,9 @@ def __init__( self.is_retryable, metric.backoff_generator, operation_timeout, - exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory), + exception_factory=self._operation_metric.track_terminal_error( + _retry_exception_factory + ), on_error=self._operation_metric.track_retryable_error, ) # initialize state diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 994cfe2da..aa5014a56 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -123,7 +123,9 @@ def start_operation(self) -> CrossSync.Iterable[Row]: self._predicate, self._operation_metric.backoff_generator, self.operation_timeout, - exception_factory=self._operation_metric.track_terminal_error(_retry_exception_factory), + exception_factory=self._operation_metric.track_terminal_error( + _retry_exception_factory + ), on_error=self._operation_metric.track_retryable_error, ) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index f2b910ce1..844cbc390 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -1389,7 +1389,9 @@ async def execute_rpc(): predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=operation_metric.track_terminal_error(_retry_exception_factory), + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory + ), on_error=operation_metric.track_retryable_error, ) @@ -1523,7 +1525,9 @@ async def mutate_row( predicate, operation_metric.backoff_generator, operation_timeout, - exception_factory=operation_metric.track_terminal_error(_retry_exception_factory), + exception_factory=operation_metric.track_terminal_error( + _retry_exception_factory + ), on_error=operation_metric.track_retryable_error, ) @@ -1638,7 +1642,7 @@ 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 []] - with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op: + 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, @@ -1696,7 +1700,7 @@ async def read_modify_write_row( if not rules: raise ValueError("rules must contain at least one item") - with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op: + 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], diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 3a4670225..3f1383bd2 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -17,7 +17,6 @@ import time from functools import wraps -from grpc import StatusCode from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric from google.cloud.bigtable.data._metrics.data_model import OperationState @@ -63,21 +62,28 @@ def wrapper(self, continuation, client_call_details, request): return wrapper + @CrossSync.convert -async def _get_metadata(source) -> dict[str, str|bytes] | None: +async def _get_metadata(source) -> dict[str, str | bytes] | None: """Helper to extract metadata from a call or RpcError""" try: if CrossSync.is_async: # grpc.aio returns metadata in Metadata objects if isinstance(source, AioRpcError): - metadata = list(source.trailing_metadata()) + list(source.initial_metadata()) + metadata = list(source.trailing_metadata()) + list( + source.initial_metadata() + ) else: - metadata = list(await source.trailing_metadata()) + list(await source.initial_metadata()) + metadata = list(await source.trailing_metadata()) + list( + await source.initial_metadata() + ) else: # sync grpc returns metadata as a sequence of tuples - metadata: Sequence[tuple[str. str|bytes]] = source.trailing_metadata() + source.initial_metadata() + metadata: Sequence[tuple[str.str | bytes]] = ( + source.trailing_metadata() + source.initial_metadata() + ) # convert metadata to dict format - return {k:v for k,v in metadata} + return {k: v for k, v in metadata} except Exception: # ignore errors while fetching metadata return None @@ -96,6 +102,12 @@ class AsyncBigtableMetricsInterceptor( async def intercept_unary_unary( self, operation, continuation, client_call_details, request ): + """ + Interceptor for unary rpcs: + - MutateRow + - CheckAndMutateRow + - ReadModifyWriteRow + """ metadata = None try: call = await continuation(client_call_details, request) @@ -113,36 +125,49 @@ async def intercept_unary_unary( async def intercept_unary_stream( self, operation, continuation, client_call_details, request ): - async def response_wrapper(call): - # only track has_first response for READ_ROWS - has_first_response = ( - operation.first_response_latency_ns is not None - or operation.op_type != OperationType.READ_ROWS - ) - encountered_exc = None - try: - async for response in call: - # record time to first response. Currently only used for READ_ROWs - if not has_first_response: - operation.first_response_latency_ns = ( - time.monotonic_ns() - operation.start_time_ns - ) - has_first_response = True - yield response - except Exception as e: - # handle errors while processing stream - encountered_exc = e - raise - finally: - if call is not None: - metadata = await _get_metadata(encountered_exc or call) - if metadata is not None: - operation.add_response_metadata(metadata) - + """ + Interceptor for streaming rpcs: + - ReadRows + - MutateRows + - SampleRowKeys + """ try: - return response_wrapper(await continuation(client_call_details, request)) + return self._streaming_generator_wrapper( + operation, await continuation(client_call_details, request) + ) except Exception as rpc_error: metadata = await _get_metadata(rpc_error) if metadata is not None: operation.add_response_metadata(metadata) raise rpc_error + + @staticmethod + @CrossSync.convert + async def _streaming_generator_wrapper(operation, call): + """ + Wrapped generator to be returned by intercept_unary_stream + """ + # only track has_first response for READ_ROWS + has_first_response = ( + operation.first_response_latency_ns is not None + or operation.op_type != OperationType.READ_ROWS + ) + encountered_exc = None + try: + async for response in call: + # record time to first response. Currently only used for READ_ROWs + if not has_first_response: + operation.first_response_latency_ns = ( + time.monotonic_ns() - operation.start_time_ns + ) + has_first_response = True + yield response + except Exception as e: + # handle errors while processing stream + encountered_exc = e + raise + finally: + if call is not None: + metadata = await _get_metadata(encountered_exc or call) + if metadata is not None: + operation.add_response_metadata(metadata) diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 6045483a3..13bcfcc29 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -16,7 +16,7 @@ """ from __future__ import annotations -from typing import Callable, Sequence, List, Tuple, TYPE_CHECKING, Union +from typing import Sequence, List, Tuple, TYPE_CHECKING, Union import time import enum from collections import namedtuple diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index ba03b9d89..70c589a04 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -15,7 +15,6 @@ from typing import ClassVar, Tuple, cast, TYPE_CHECKING - import time import re import logging @@ -36,7 +35,6 @@ from google.cloud.bigtable_v2.types.response_params import ResponseParams from google.cloud.bigtable.data._helpers import TrackedBackoffGenerator from google.cloud.bigtable.data.exceptions import _MutateRowsIncomplete -from google.cloud.bigtable.data.exceptions import RetryExceptionGroup from google.protobuf.message import DecodeError if TYPE_CHECKING: @@ -169,8 +167,9 @@ class ActiveOperationMetric: # time waiting on flow control, in nanoseconds flow_throttling_time_ns: int = 0 - - _active_operation_context: ClassVar[contextvars.ContextVar] = contextvars.ContextVar("active_operation_context") + _active_operation_context: ClassVar[ + contextvars.ContextVar + ] = contextvars.ContextVar("active_operation_context") @classmethod def get_active(cls): @@ -403,10 +402,7 @@ def track_retryable_error(self, exc: Exception) -> None: """ try: # record metadata from failed rpc - if ( - isinstance(exc, GoogleAPICallError) - and exc.errors - ): + if isinstance(exc, GoogleAPICallError) and exc.errors: rpc_error = exc.errors[-1] metadata = list(rpc_error.trailing_metadata()) + list( rpc_error.initial_metadata() @@ -422,24 +418,28 @@ def track_retryable_error(self, exc: Exception) -> None: else: self.end_attempt_with_status(exc) - def track_terminal_error(self, exception_factory:callable[ - [list[Exception], RetryFailureReason, float | None],tuple[Exception, Exception | None], - ]) -> callable[[list[Exception], RetryFailureReason, float | None], None]: + def track_terminal_error( + self, + exception_factory: callable[ + [list[Exception], RetryFailureReason, float | None], + tuple[Exception, Exception | None], + ], + ) -> callable[[list[Exception], RetryFailureReason, float | None], None]: """ Used as input to api_core.Retry classes, to track when terminal errors are encountered Should be used as a wrapper over an exception_factory callback """ + def wrapper( - exc_list: list[Exception], reason: RetryFailureReason, timeout_val: float | None + exc_list: list[Exception], + reason: RetryFailureReason, + timeout_val: float | None, ) -> tuple[Exception, Exception | None]: source_exc, cause_exc = exception_factory(exc_list, reason, timeout_val) try: # record metadata from failed rpc - if ( - isinstance(source_exc, GoogleAPICallError) - and source_exc.errors - ): + if isinstance(source_exc, GoogleAPICallError) and source_exc.errors: rpc_error = source_exc.errors[-1] metadata = list(rpc_error.trailing_metadata()) + list( rpc_error.initial_metadata() @@ -448,12 +448,17 @@ def wrapper( except Exception: # ignore errors in metadata collection pass - if reason == RetryFailureReason.TIMEOUT and self.state == OperationState.ACTIVE_ATTEMPT and exc_list: + if ( + reason == RetryFailureReason.TIMEOUT + and self.state == OperationState.ACTIVE_ATTEMPT + and exc_list + ): # record ending attempt for timeout failures attempt_exc = exc_list[-1] self.track_retryable_error(attempt_exc) self.end_with_status(source_exc) return source_exc, cause_exc + return wrapper @staticmethod diff --git a/google/cloud/bigtable/data/_metrics/metrics_controller.py b/google/cloud/bigtable/data/_metrics/metrics_controller.py index 087c5058b..6d3c3ef81 100644 --- a/google/cloud/bigtable/data/_metrics/metrics_controller.py +++ b/google/cloud/bigtable/data/_metrics/metrics_controller.py @@ -13,20 +13,10 @@ # limitations under the License. from __future__ import annotations -from typing import TYPE_CHECKING - from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from google.cloud.bigtable.data._metrics.data_model import OperationType -if TYPE_CHECKING: - from google.cloud.bigtable.data._async.metrics_interceptor import ( - AsyncBigtableMetricsInterceptor, - ) - from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import ( - BigtableMetricsInterceptor, - ) - class BigtableClientSideMetricsController: """ diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index ae1c88449..070636a8b 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -1377,7 +1377,7 @@ def check_and_mutate_row( ): false_case_mutations = [false_case_mutations] false_case_list = [m._to_pb() for m in false_case_mutations or []] - with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op: + 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, @@ -1431,7 +1431,7 @@ def read_modify_write_row( rules = [rules] if not rules: raise ValueError("rules must contain at least one item") - with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op: + 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], diff --git a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py index e592c6329..2be8afc6c 100644 --- a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py @@ -68,6 +68,10 @@ class BigtableMetricsInterceptor( def intercept_unary_unary( self, operation, continuation, client_call_details, request ): + """Interceptor for unary rpcs: + - MutateRow + - CheckAndMutateRow + - ReadModifyWriteRow""" metadata = None try: call = continuation(client_call_details, request) @@ -84,33 +88,41 @@ def intercept_unary_unary( def intercept_unary_stream( self, operation, continuation, client_call_details, request ): - def response_wrapper(call): - has_first_response = ( - operation.first_response_latency_ns is not None - or operation.op_type != OperationType.READ_ROWS - ) - encountered_exc = None - try: - for response in call: - if not has_first_response: - operation.first_response_latency_ns = ( - time.monotonic_ns() - operation.start_time_ns - ) - has_first_response = True - yield response - except Exception as e: - encountered_exc = e - raise - finally: - if call is not None: - metadata = _get_metadata(encountered_exc or call) - if metadata is not None: - operation.add_response_metadata(metadata) - + """Interceptor for streaming rpcs: + - ReadRows + - MutateRows + - SampleRowKeys""" try: - return response_wrapper(continuation(client_call_details, request)) + return self._streaming_generator_wrapper( + operation, continuation(client_call_details, request) + ) except Exception as rpc_error: metadata = _get_metadata(rpc_error) if metadata is not None: operation.add_response_metadata(metadata) raise rpc_error + + @staticmethod + def _streaming_generator_wrapper(operation, call): + """Wrapped generator to be returned by intercept_unary_stream""" + has_first_response = ( + operation.first_response_latency_ns is not None + or operation.op_type != OperationType.READ_ROWS + ) + encountered_exc = None + try: + for response in call: + if not has_first_response: + operation.first_response_latency_ns = ( + time.monotonic_ns() - operation.start_time_ns + ) + has_first_response = True + yield response + except Exception as e: + encountered_exc = e + raise + finally: + if call is not None: + metadata = _get_metadata(encountered_exc or call) + if metadata is not None: + operation.add_response_metadata(metadata) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index a4d58c9a6..c8ac0ff1c 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -25,8 +25,6 @@ from google.cloud.bigtable.data._metrics.data_model import ( CompletedOperationMetric, CompletedAttemptMetric, - ActiveOperationMetric, - OperationState, ) from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable_v2.types import ResponseParams @@ -117,8 +115,9 @@ def __init__(self, call, exc_to_raise): def __aiter__(self): return self - - @CrossSync.convert(sync_name="__next__", replace_symbols={"__anext__": "__next__"}) + @CrossSync.convert( + sync_name="__next__", replace_symbols={"__anext__": "__next__"} + ) async def __anext__(self): if not self._raised: self._raised = True @@ -150,9 +149,12 @@ def _make_client(self): 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) - )) + metadata = ( + "x-goog-ext-425905942-bin", + ResponseParams.serialize( + ResponseParams(cluster_id=cluster_id, zone_id=zone_id) + ), + ) else: metadata = None if CrossSync.is_async: @@ -164,8 +166,10 @@ def _make_exception(self, status, cluster_id=None, zone_id=None): exc.initial_metadata = lambda: [] exc.code = lambda: status exc.details = lambda: None + def _result(): raise exc + exc.result = _result return exc @@ -198,7 +202,9 @@ async def client(self, 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) + client.transport._grpc_channel._channel._channel = intercept_channel( + metrics_channel, error_injector + ) yield client @CrossSync.convert @@ -284,8 +290,12 @@ async def test_read_rows_failure_with_retries( 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)) + 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 @@ -435,9 +445,7 @@ async def test_read_rows_stream_failure_closed( 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() - ) + generator = await table.read_rows_stream(ReadRowsQuery()) await generator.__anext__() await generator.aclose() with pytest.raises(CrossSync.StopIteration): @@ -458,7 +466,6 @@ async def test_read_rows_stream_failure_closed( 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 @@ -868,7 +875,6 @@ async def test_read_rows_sharded_failure_with_retries( with retryable errors """ 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") @@ -887,8 +893,20 @@ async def test_read_rows_sharded_failure_with_retries( 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 + 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): @@ -1184,7 +1202,7 @@ async def test_bulk_mutate_rows_failure_unauthorized( entry = RowMutationEntry(row_key, [mutation]) handler.clear() - with pytest.raises(MutationsExceptionGroup) as e: + with pytest.raises(MutationsExceptionGroup): await authorized_view.bulk_mutate_rows([entry]) # validate counts assert len(handler.completed_operations) == 1 @@ -1228,7 +1246,9 @@ async def test_bulk_mutate_rows_failure_unauthorized_with_retries( handler.clear() with pytest.raises(MutationsExceptionGroup) as e: - await authorized_view.bulk_mutate_rows([entry], retryable_errors=[PermissionDenied], operation_timeout=0.5) + 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 @@ -1614,7 +1634,12 @@ async def test_mutate_row_failure_unauthorized_with_retries( 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) + 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 diff --git a/tests/system/data/test_metrics_autogen.py b/tests/system/data/test_metrics_autogen.py index bcd74de8d..2ec714351 100644 --- a/tests/system/data/test_metrics_autogen.py +++ b/tests/system/data/test_metrics_autogen.py @@ -999,7 +999,7 @@ def test_bulk_mutate_rows_failure_unauthorized( mutation = SetCell("unauthorized", b"q", b"v") entry = RowMutationEntry(row_key, [mutation]) handler.clear() - with pytest.raises(MutationsExceptionGroup) as e: + with pytest.raises(MutationsExceptionGroup): authorized_view.bulk_mutate_rows([entry]) assert len(handler.completed_operations) == 1 assert len(handler.completed_attempts) == 1 diff --git a/tests/unit/data/_async/test_metrics_interceptor.py b/tests/unit/data/_async/test_metrics_interceptor.py index 1f99473e3..1593b8c99 100644 --- a/tests/unit/data/_async/test_metrics_interceptor.py +++ b/tests/unit/data/_async/test_metrics_interceptor.py @@ -70,9 +70,6 @@ def _get_target_class(): def _make_one(self, *args, **kwargs): return self._get_target_class()(*args, **kwargs) - def test_ctor(self): - instance = self._make_one() - @CrossSync.pytest async def test_unary_unary_interceptor_op_not_found(self): """Test that interceptor call continuation if op is not found""" @@ -206,6 +203,7 @@ async def test_unary_stream_interceptor_success(self): async def test_unary_stream_interceptor_failure_mid_stream(self): """Test that interceptor handles failures mid-stream""" from grpc.aio import AioRpcError, Metadata + instance = self._make_one() op = mock.Mock() op.uuid = "test-uuid" diff --git a/tests/unit/data/_sync_autogen/test_metrics_interceptor.py b/tests/unit/data/_sync_autogen/test_metrics_interceptor.py index ec3e51583..c4efcc5b9 100644 --- a/tests/unit/data/_sync_autogen/test_metrics_interceptor.py +++ b/tests/unit/data/_sync_autogen/test_metrics_interceptor.py @@ -53,9 +53,6 @@ def _get_target_class(): def _make_one(self, *args, **kwargs): return self._get_target_class()(*args, **kwargs) - def test_ctor(self): - instance = self._make_one() - def test_unary_unary_interceptor_op_not_found(self): """Test that interceptor call continuation if op is not found""" instance = self._make_one() From 0340fcebd754208eee74b8005e695566925ef683 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 26 Sep 2025 16:36:37 -0700 Subject: [PATCH 66/78] fixed tests --- tests/system/data/test_metrics_async.py | 6 +-- tests/system/data/test_metrics_autogen.py | 6 +-- .../data/_async/test_read_rows_acceptance.py | 38 +++++++++--------- .../test_read_rows_acceptance.py | 39 ++++++++----------- 4 files changed, 41 insertions(+), 48 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index c8ac0ff1c..21a5400ec 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -620,7 +620,7 @@ async def test_read_rows_stream_failure_unauthorized_with_retries( # validate attempts for attempt in handler.completed_attempts: assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] assert ( attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns @@ -1266,7 +1266,7 @@ async def test_bulk_mutate_rows_failure_unauthorized_with_retries( ) # validate attempts for attempt in handler.completed_attempts: - assert attempt.end_status.name == "OK" + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] assert ( attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns @@ -1658,7 +1658,7 @@ async def test_mutate_row_failure_unauthorized_with_retries( ) # validate attempts for attempt in handler.completed_attempts: - assert attempt.end_status.name == "PERMISSION_DENIED" + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] 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 index 2ec714351..fa633db65 100644 --- a/tests/system/data/test_metrics_autogen.py +++ b/tests/system/data/test_metrics_autogen.py @@ -510,7 +510,7 @@ def test_read_rows_stream_failure_unauthorized_with_retries( ) for attempt in handler.completed_attempts: assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] assert ( attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns @@ -1053,7 +1053,7 @@ def test_bulk_mutate_rows_failure_unauthorized_with_retries( == cluster_config[operation.cluster_id].location.split("/")[-1] ) for attempt in handler.completed_attempts: - assert attempt.end_status.name == "OK" + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] assert ( attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns @@ -1376,7 +1376,7 @@ def test_mutate_row_failure_unauthorized_with_retries( == cluster_config[operation.cluster_id].location.split("/")[-1] ) for attempt in handler.completed_attempts: - assert attempt.end_status.name == "PERMISSION_DENIED" + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] assert ( attempt.gfe_latency_ns >= 0 and attempt.gfe_latency_ns < operation.duration_ns diff --git a/tests/unit/data/_async/test_read_rows_acceptance.py b/tests/unit/data/_async/test_read_rows_acceptance.py index ab9502223..87fbd91d7 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,12 @@ 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 +88,10 @@ 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()) + chunker = self._make_operation().chunk_stream( + self._coro_wrapper(_row_stream()) ) - merger = self._get_operation_class().merge_rows(chunker) + merger = self._make_operation().merge_rows(chunker) results = [] async for row in merger: results.append(row) @@ -106,13 +108,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 +198,12 @@ 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()) + op = self._make_operation() + op._last_yielded_row_key = b"b" + chunker = op.chunk_stream( + self._coro_wrapper(_row_stream()) ) - merger = self._get_operation_class().merge_rows(chunker) + 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_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 From 6798ea2bbb59c17e97f03e680e446d78ed96f766 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Sep 2025 17:23:17 -0700 Subject: [PATCH 67/78] fixed more tests --- tests/unit/data/_async/test__mutate_rows.py | 20 +++- tests/unit/data/_async/test__read_rows.py | 10 +- tests/unit/data/_async/test_client.py | 99 +++++++++++++------ .../data/_async/test_mutations_batcher.py | 17 ++-- .../data/_sync_autogen/test__mutate_rows.py | 25 ++++- .../data/_sync_autogen/test__read_rows.py | 16 ++- .../_sync_autogen/test_mutations_batcher.py | 22 +++-- 7 files changed, 150 insertions(+), 59 deletions(-) diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index f14fa6dee..827247a3d 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,7 @@ 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 +92,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 +100,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 +120,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 +144,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 +158,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 +166,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 +180,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 +189,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 +211,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 +224,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 +248,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 +265,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 +282,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 +294,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..7cead5521 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,7 @@ 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 +312,7 @@ 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 +338,7 @@ 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 48ba94ba0..783df5897 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1205,7 +1205,6 @@ async def test_ctor(self): assert instance_key in client._active_instances assert client._instance_owners[instance_key] == {id(table)} assert isinstance(table._metrics, BigtableClientSideMetricsController) - assert table._metrics.interceptor == client._metrics_interceptor assert table.default_operation_timeout == expected_operation_timeout assert table.default_attempt_timeout == expected_attempt_timeout assert ( @@ -1547,7 +1546,6 @@ async def test_ctor(self): assert instance_key in client._active_instances assert client._instance_owners[instance_key] == {id(view)} assert isinstance(view._metrics, BigtableClientSideMetricsController) - assert view._metrics.interceptor == client._metrics_interceptor assert view.default_operation_timeout == expected_operation_timeout assert view.default_attempt_timeout == expected_attempt_timeout assert ( @@ -1959,9 +1957,19 @@ 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( @@ -1970,16 +1978,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): @@ -1987,14 +1996,22 @@ 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, @@ -2002,11 +2019,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] @@ -2020,9 +2037,19 @@ 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( @@ -2031,8 +2058,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) @@ -2055,22 +2082,34 @@ 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": [ @@ -2079,10 +2118,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..5f6d492d9 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -307,6 +307,7 @@ 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 +319,7 @@ def _make_one(self, table=None, **kwargs): DeadlineExceeded, ServiceUnavailable, ) + table._metrics = BigtableClientSideMetricsController([]) return self._get_target_class()(table, **kwargs) @@ -935,14 +937,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 +967,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 +1097,7 @@ 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 +1195,7 @@ 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 +1211,7 @@ 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/_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_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 ) From c96dd2547ae6ad242ee25d11dc4987f739a0a8be Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Sep 2025 17:25:22 -0700 Subject: [PATCH 68/78] removed unneeded kwargs --- google/cloud/bigtable/data/_async/client.py | 4 -- .../data/_metrics/metrics_controller.py | 1 - .../bigtable/data/_sync_autogen/client.py | 8 +-- tests/unit/data/_sync_autogen/test_client.py | 72 +++++++++++-------- 4 files changed, 43 insertions(+), 42 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 844cbc390..d90e710f6 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -973,10 +973,6 @@ def __init__( self._metrics = BigtableClientSideMetricsController( handlers=[], - project_id=self.client.project, - instance_id=instance_id, - table_id=table_id, - app_profile_id=app_profile_id, ) try: diff --git a/google/cloud/bigtable/data/_metrics/metrics_controller.py b/google/cloud/bigtable/data/_metrics/metrics_controller.py index 6d3c3ef81..fd88a13dc 100644 --- a/google/cloud/bigtable/data/_metrics/metrics_controller.py +++ b/google/cloud/bigtable/data/_metrics/metrics_controller.py @@ -29,7 +29,6 @@ class BigtableClientSideMetricsController: def __init__( self, handlers: list[MetricsHandler] | None = None, - **kwargs, ): """ Initializes the metrics controller. diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index 070636a8b..acfcec465 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -758,13 +758,7 @@ def __init__( self.default_retryable_errors: Sequence[type[Exception]] = ( default_retryable_errors or () ) - self._metrics = BigtableClientSideMetricsController( - handlers=[], - project_id=self.client.project, - instance_id=instance_id, - table_id=table_id, - app_profile_id=app_profile_id, - ) + self._metrics = BigtableClientSideMetricsController(handlers=[]) try: self._register_instance_future = CrossSync._Sync_Impl.create_task( self.client._register_instance, diff --git a/tests/unit/data/_sync_autogen/test_client.py b/tests/unit/data/_sync_autogen/test_client.py index afe741e57..df4ca3675 100644 --- a/tests/unit/data/_sync_autogen/test_client.py +++ b/tests/unit/data/_sync_autogen/test_client.py @@ -982,7 +982,6 @@ def test_ctor(self): assert instance_key in client._active_instances assert client._instance_owners[instance_key] == {id(table)} assert isinstance(table._metrics, BigtableClientSideMetricsController) - assert table._metrics.interceptor == client._metrics_interceptor assert table.default_operation_timeout == expected_operation_timeout assert table.default_attempt_timeout == expected_attempt_timeout assert ( @@ -1251,7 +1250,6 @@ def test_ctor(self): assert instance_key in client._active_instances assert client._instance_owners[instance_key] == {id(view)} assert isinstance(view._metrics, BigtableClientSideMetricsController) - assert view._metrics.interceptor == client._metrics_interceptor assert view.default_operation_timeout == expected_operation_timeout assert view.default_attempt_timeout == expected_attempt_timeout assert ( @@ -1621,9 +1619,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( @@ -1632,30 +1634,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, @@ -1663,11 +1668,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] @@ -1680,8 +1685,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( @@ -1690,8 +1699,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) @@ -1709,21 +1718,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": [ @@ -1732,10 +1748,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 From 7eb83a8981ea56830aaa6393239595a5e988cd40 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Sep 2025 20:51:44 -0700 Subject: [PATCH 69/78] fixed unit tests --- tests/unit/data/_metrics/test_data_model.py | 10 ------- .../data/_metrics/test_metrics_controller.py | 28 ++----------------- 2 files changed, 2 insertions(+), 36 deletions(-) diff --git a/tests/unit/data/_metrics/test_data_model.py b/tests/unit/data/_metrics/test_data_model.py index e77519629..42aa96093 100644 --- a/tests/unit/data/_metrics/test_data_model.py +++ b/tests/unit/data/_metrics/test_data_model.py @@ -535,16 +535,6 @@ def test_end_with_status_w_exception(self): final_op = handlers[0].on_operation_complete.call_args[0][0] assert final_op.final_status == expected_status - def test_interceptor_metadata(self): - from google.cloud.bigtable.data._metrics.data_model import ( - OPERATION_INTERCEPTOR_METADATA_KEY, - ) - - metric = self._make_one(mock.Mock()) - key, value = metric.interceptor_metadata - assert key == OPERATION_INTERCEPTOR_METADATA_KEY - assert value == metric.uuid - def test_end_with_status_with_default_cluster_zone(self): """ ending the operation should use default cluster and zone if not set diff --git a/tests/unit/data/_metrics/test_metrics_controller.py b/tests/unit/data/_metrics/test_metrics_controller.py index 701af737b..66ebe56f6 100644 --- a/tests/unit/data/_metrics/test_metrics_controller.py +++ b/tests/unit/data/_metrics/test_metrics_controller.py @@ -21,19 +21,13 @@ def _make_one(self, *args, **kwargs): BigtableClientSideMetricsController, ) - # add mock interceptor if called with no arguments - if not args and "interceptor" not in kwargs: - args = [mock.Mock()] - return BigtableClientSideMetricsController(*args, **kwargs) def test_ctor_defaults(self): """ should create instance with GCP Exporter handler by default """ - expected_interceptor = object() - instance = self._make_one(expected_interceptor) - assert instance.interceptor == expected_interceptor + instance = self._make_one() assert len(instance.handlers) == 0 def ctor_custom_handlers(self): @@ -92,22 +86,4 @@ def test_create_operation(self): assert op.is_streaming is expected_is_streaming assert op.zone is expected_zone assert len(op.handlers) == 1 - assert op.handlers[0] is handler - - def test_create_operation_registers_interceptor(self): - """ - creating an operation should link the operation with the controller's interceptor, - and add the interceptor as a handler to the operation - """ - from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import ( - BigtableMetricsInterceptor, - ) - - custom_handler = object() - controller = self._make_one( - BigtableMetricsInterceptor(), handlers=[custom_handler] - ) - op = controller.create_operation(object()) - assert custom_handler in op.handlers - assert op.uuid in controller.interceptor.operation_map - assert controller.interceptor.operation_map[op.uuid] == op + assert op.handlers[0] is handler \ No newline at end of file From 61f8b853f62d0b4a36595af97222f6dfe6827cfe Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Sep 2025 20:59:18 -0700 Subject: [PATCH 70/78] ran blacken --- tests/unit/data/_async/test__mutate_rows.py | 4 +++- tests/unit/data/_async/test__read_rows.py | 12 +++++++++--- tests/unit/data/_async/test_client.py | 16 ++++++++++++---- tests/unit/data/_async/test_mutations_batcher.py | 13 ++++++++++--- .../data/_async/test_read_rows_acceptance.py | 9 ++------- .../data/_metrics/test_metrics_controller.py | 2 +- 6 files changed, 37 insertions(+), 19 deletions(-) diff --git a/tests/unit/data/_async/test__mutate_rows.py b/tests/unit/data/_async/test__mutate_rows.py index 827247a3d..a43b0a35f 100644 --- a/tests/unit/data/_async/test__mutate_rows.py +++ b/tests/unit/data/_async/test__mutate_rows.py @@ -49,7 +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")) + kwargs["metric"] = kwargs.pop( + "metric", ActiveOperationMetric("MUTATE_ROWS") + ) return self._target_class()(*args, **kwargs) def _make_mutation(self, count=1, size=1): diff --git a/tests/unit/data/_async/test__read_rows.py b/tests/unit/data/_async/test__read_rows.py index 7cead5521..4bf0f933e 100644 --- a/tests/unit/data/_async/test__read_rows.py +++ b/tests/unit/data/_async/test__read_rows.py @@ -273,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, ActiveOperationMetric("READ_ROWS")) + 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()): @@ -312,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, ActiveOperationMetric("READ_ROWS")) + 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 @@ -338,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, ActiveOperationMetric("READ_ROWS")) + 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 783df5897..c4b9923db 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1957,7 +1957,9 @@ 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(CrossSync, "_ReadRowsOperation") as mock_op_constructor: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: mock_op = mock.Mock() expected_result = object() @@ -1996,7 +1998,9 @@ 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(CrossSync, "_ReadRowsOperation") as mock_op_constructor: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: mock_op = mock.Mock() expected_result = object() @@ -2037,7 +2041,9 @@ 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(CrossSync, "_ReadRowsOperation") as mock_op_constructor: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: mock_op = mock.Mock() if CrossSync.is_async: @@ -2082,7 +2088,9 @@ 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(CrossSync, "_ReadRowsOperation") as mock_op_constructor: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: mock_op = mock.Mock() if CrossSync.is_async: diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index 5f6d492d9..d94b1e98c 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -307,7 +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 + from google.cloud.bigtable.data._metrics import ( + BigtableClientSideMetricsController, + ) if table is None: table = mock.Mock() @@ -1097,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()], mock.Mock()) + 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 @@ -1196,6 +1200,7 @@ async def test_customizable_retryable_errors( 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: @@ -1211,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], ActiveOperationMetric("MUTATE_ROWS")) + 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 87fbd91d7..1fbbd7c82 100644 --- a/tests/unit/data/_async/test_read_rows_acceptance.py +++ b/tests/unit/data/_async/test_read_rows_acceptance.py @@ -46,7 +46,6 @@ def _make_operation(): op._remaining_count = None return op - @staticmethod @CrossSync.convert def _get_client_class(): @@ -88,9 +87,7 @@ async def _process_chunks(self, *chunks): async def _row_stream(): yield ReadRowsResponse(chunks=chunks) - chunker = self._make_operation().chunk_stream( - self._coro_wrapper(_row_stream()) - ) + chunker = self._make_operation().chunk_stream(self._coro_wrapper(_row_stream())) merger = self._make_operation().merge_rows(chunker) results = [] async for row in merger: @@ -200,9 +197,7 @@ async def _row_stream(): op = self._make_operation() op._last_yielded_row_key = b"b" - chunker = op.chunk_stream( - self._coro_wrapper(_row_stream()) - ) + chunker = op.chunk_stream(self._coro_wrapper(_row_stream())) merger = self._make_operation().merge_rows(chunker) with pytest.raises(InvalidChunk): async for _ in merger: diff --git a/tests/unit/data/_metrics/test_metrics_controller.py b/tests/unit/data/_metrics/test_metrics_controller.py index 66ebe56f6..7fdbaef07 100644 --- a/tests/unit/data/_metrics/test_metrics_controller.py +++ b/tests/unit/data/_metrics/test_metrics_controller.py @@ -86,4 +86,4 @@ def test_create_operation(self): assert op.is_streaming is expected_is_streaming assert op.zone is expected_zone assert len(op.handlers) == 1 - assert op.handlers[0] is handler \ No newline at end of file + assert op.handlers[0] is handler From dd7453b5e3008e2d9a194d45c9d6235f52fe9e67 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Sep 2025 22:39:55 -0700 Subject: [PATCH 71/78] removed unneeded kwargs --- google/cloud/bigtable/data/_async/client.py | 4 ---- google/cloud/bigtable/data/_metrics/metrics_controller.py | 1 - google/cloud/bigtable/data/_sync_autogen/client.py | 8 +------- 3 files changed, 1 insertion(+), 12 deletions(-) diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 844cbc390..d90e710f6 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -973,10 +973,6 @@ def __init__( self._metrics = BigtableClientSideMetricsController( handlers=[], - project_id=self.client.project, - instance_id=instance_id, - table_id=table_id, - app_profile_id=app_profile_id, ) try: diff --git a/google/cloud/bigtable/data/_metrics/metrics_controller.py b/google/cloud/bigtable/data/_metrics/metrics_controller.py index 6d3c3ef81..fd88a13dc 100644 --- a/google/cloud/bigtable/data/_metrics/metrics_controller.py +++ b/google/cloud/bigtable/data/_metrics/metrics_controller.py @@ -29,7 +29,6 @@ class BigtableClientSideMetricsController: def __init__( self, handlers: list[MetricsHandler] | None = None, - **kwargs, ): """ Initializes the metrics controller. diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index 070636a8b..acfcec465 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -758,13 +758,7 @@ def __init__( self.default_retryable_errors: Sequence[type[Exception]] = ( default_retryable_errors or () ) - self._metrics = BigtableClientSideMetricsController( - handlers=[], - project_id=self.client.project, - instance_id=instance_id, - table_id=table_id, - app_profile_id=app_profile_id, - ) + self._metrics = BigtableClientSideMetricsController(handlers=[]) try: self._register_instance_future = CrossSync._Sync_Impl.create_task( self.client._register_instance, From 911a299a6d00bb224824a87598f2a5a78718da4b Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 29 Sep 2025 22:40:40 -0700 Subject: [PATCH 72/78] fixed lint --- .../cloud/bigtable/data/_async/_mutate_rows.py | 6 ++---- google/cloud/bigtable/data/_async/_read_rows.py | 2 -- .../bigtable/data/_async/metrics_interceptor.py | 7 +++---- .../cloud/bigtable/data/_metrics/data_model.py | 17 +++++++++-------- .../bigtable/data/_sync_autogen/_mutate_rows.py | 6 ++---- .../bigtable/data/_sync_autogen/_read_rows.py | 2 -- .../data/_sync_autogen/metrics_interceptor.py | 5 ++--- 7 files changed, 18 insertions(+), 27 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index 9884929ba..a4eb93bfd 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -106,10 +106,8 @@ def __init__( self.is_retryable, metric.backoff_generator, operation_timeout, - exception_factory=self._operation_metric.track_terminal_error( - _retry_exception_factory - ), - on_error=self._operation_metric.track_retryable_error, + exception_factory=metric.track_terminal_error(_retry_exception_factory), + on_error=metric.track_retryable_error, ) # initialize state self.timeout_generator = _attempt_timeout_generator( diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index aa5014a56..35b2e44e9 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -350,8 +350,6 @@ async def merge_rows( except Exception as generic_exception: # handle exceptions in retry wrapper raise generic_exception - else: - self._operation_metric.end_with_success() @staticmethod def _revise_request_rowset( diff --git a/google/cloud/bigtable/data/_async/metrics_interceptor.py b/google/cloud/bigtable/data/_async/metrics_interceptor.py index 3f1383bd2..0bd401a78 100644 --- a/google/cloud/bigtable/data/_async/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_async/metrics_interceptor.py @@ -67,6 +67,7 @@ def wrapper(self, continuation, client_call_details, request): async def _get_metadata(source) -> dict[str, str | bytes] | None: """Helper to extract metadata from a call or RpcError""" try: + metadata: Sequence[tuple[str, str | bytes]] if CrossSync.is_async: # grpc.aio returns metadata in Metadata objects if isinstance(source, AioRpcError): @@ -79,11 +80,9 @@ async def _get_metadata(source) -> dict[str, str | bytes] | None: ) else: # sync grpc returns metadata as a sequence of tuples - metadata: Sequence[tuple[str.str | bytes]] = ( - source.trailing_metadata() + source.initial_metadata() - ) + metadata = source.trailing_metadata() + source.initial_metadata() # convert metadata to dict format - return {k: v for k, v in metadata} + return {k: v for (k, v) in metadata} except Exception: # ignore errors while fetching metadata return None diff --git a/google/cloud/bigtable/data/_metrics/data_model.py b/google/cloud/bigtable/data/_metrics/data_model.py index 70c589a04..d0d9b5f52 100644 --- a/google/cloud/bigtable/data/_metrics/data_model.py +++ b/google/cloud/bigtable/data/_metrics/data_model.py @@ -13,7 +13,7 @@ # limitations under the License. from __future__ import annotations -from typing import ClassVar, Tuple, cast, TYPE_CHECKING +from typing import Callable, ClassVar, List, Tuple, Optional, cast, TYPE_CHECKING import time import re @@ -54,6 +54,11 @@ INVALID_STATE_ERROR = "Invalid state for {}: {}" +ExceptionFactoryType = Callable[ + [List[Exception], RetryFailureReason, Optional[float]], + Tuple[Exception, Optional[Exception]], +] + class OperationType(Enum): """Enum for the type of operation being performed.""" @@ -168,7 +173,7 @@ class ActiveOperationMetric: flow_throttling_time_ns: int = 0 _active_operation_context: ClassVar[ - contextvars.ContextVar + contextvars.ContextVar[ActiveOperationMetric] ] = contextvars.ContextVar("active_operation_context") @classmethod @@ -419,12 +424,8 @@ def track_retryable_error(self, exc: Exception) -> None: self.end_attempt_with_status(exc) def track_terminal_error( - self, - exception_factory: callable[ - [list[Exception], RetryFailureReason, float | None], - tuple[Exception, Exception | None], - ], - ) -> callable[[list[Exception], RetryFailureReason, float | None], None]: + self, exception_factory: ExceptionFactoryType + ) -> ExceptionFactoryType: """ Used as input to api_core.Retry classes, to track when terminal errors are encountered diff --git a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index a490ab10b..81a5be4cf 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -84,10 +84,8 @@ def __init__( self.is_retryable, metric.backoff_generator, operation_timeout, - exception_factory=self._operation_metric.track_terminal_error( - _retry_exception_factory - ), - on_error=self._operation_metric.track_retryable_error, + 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 diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index 75ad79931..adbe819eb 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -294,8 +294,6 @@ def merge_rows( raise close_exception except Exception as generic_exception: raise generic_exception - else: - self._operation_metric.end_with_success() @staticmethod def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSetPB: diff --git a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py index 2be8afc6c..dcc17e591 100644 --- a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py @@ -49,9 +49,8 @@ def wrapper(self, continuation, client_call_details, request): def _get_metadata(source) -> dict[str, str | bytes] | None: """Helper to extract metadata from a call or RpcError""" try: - metadata: Sequence[tuple[str.str | bytes]] = ( - source.trailing_metadata() + source.initial_metadata() - ) + metadata: Sequence[tuple[str, str | bytes]] + metadata = source.trailing_metadata() + source.initial_metadata() return {k: v for (k, v) in metadata} except Exception: return None From abaf5b2111b9d816fe42ddaa2fec0bd9d71fbc76 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Thu, 2 Oct 2025 13:12:14 -0700 Subject: [PATCH 73/78] fixed flakes --- tests/system/data/test_metrics_async.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index 21a5400ec..ee339965d 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -621,10 +621,6 @@ async def test_read_rows_stream_failure_unauthorized_with_retries( for attempt in handler.completed_attempts: assert isinstance(attempt, CompletedAttemptMetric) assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) @CrossSync.pytest async def test_read_rows_stream_failure_mid_stream( @@ -1267,10 +1263,6 @@ async def test_bulk_mutate_rows_failure_unauthorized_with_retries( # validate attempts for attempt in handler.completed_attempts: assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) @CrossSync.pytest async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): @@ -1659,10 +1651,6 @@ async def test_mutate_row_failure_unauthorized_with_retries( # validate attempts for attempt in handler.completed_attempts: assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) @CrossSync.pytest async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): From 183282455d97026bdaa57c8e8293810588dbb259 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 3 Oct 2025 13:07:29 -0700 Subject: [PATCH 74/78] generated sync --- tests/system/data/test_metrics_autogen.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/tests/system/data/test_metrics_autogen.py b/tests/system/data/test_metrics_autogen.py index fa633db65..54a9f2256 100644 --- a/tests/system/data/test_metrics_autogen.py +++ b/tests/system/data/test_metrics_autogen.py @@ -511,10 +511,6 @@ def test_read_rows_stream_failure_unauthorized_with_retries( for attempt in handler.completed_attempts: assert isinstance(attempt, CompletedAttemptMetric) assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) def test_read_rows_stream_failure_mid_stream( self, table, temp_rows, handler, error_injector @@ -1054,10 +1050,6 @@ def test_bulk_mutate_rows_failure_unauthorized_with_retries( ) for attempt in handler.completed_attempts: assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): from google.cloud.bigtable.data.mutations import RowMutationEntry @@ -1377,10 +1369,6 @@ def test_mutate_row_failure_unauthorized_with_retries( ) for attempt in handler.completed_attempts: assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): table.sample_row_keys() From 6bbac87f01495d3bc93945013b500aee111ebcf4 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 3 Oct 2025 13:21:56 -0700 Subject: [PATCH 75/78] moved all fixtures into one place --- tests/system/conftest.py | 4 - tests/system/data/__init__.py | 184 +++++++++++++++ tests/system/data/setup_fixtures.py | 210 ------------------ .../data/_metrics/test_metrics_controller.py | 2 +- 4 files changed, 185 insertions(+), 215 deletions(-) delete mode 100644 tests/system/data/setup_fixtures.py 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 dcd14c5f9..6f836fb96 100644 --- a/tests/system/data/__init__.py +++ b/tests/system/data/__init__.py @@ -14,12 +14,17 @@ # 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: """ @@ -68,3 +73,182 @@ def column_family_config(self): 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/unit/data/_metrics/test_metrics_controller.py b/tests/unit/data/_metrics/test_metrics_controller.py index 66ebe56f6..7fdbaef07 100644 --- a/tests/unit/data/_metrics/test_metrics_controller.py +++ b/tests/unit/data/_metrics/test_metrics_controller.py @@ -86,4 +86,4 @@ def test_create_operation(self): assert op.is_streaming is expected_is_streaming assert op.zone is expected_zone assert len(op.handlers) == 1 - assert op.handlers[0] is handler \ No newline at end of file + assert op.handlers[0] is handler From 2bde30d5c87a55b879d95b0b226f2c02c94fa4bb Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 3 Oct 2025 15:04:47 -0700 Subject: [PATCH 76/78] fixed event loop error --- tests/system/data/test_metrics_async.py | 8 -------- tests/system/data/test_system_async.py | 8 -------- 2 files changed, 16 deletions(-) diff --git a/tests/system/data/test_metrics_async.py b/tests/system/data/test_metrics_async.py index ee339965d..af4735d97 100644 --- a/tests/system/data/test_metrics_async.py +++ b/tests/system/data/test_metrics_async.py @@ -135,14 +135,6 @@ def __getattr__(self, name): @CrossSync.convert_class(sync_name="TestMetrics") class TestMetricsAsync(SystemTestRunner): - @CrossSync.drop - @pytest.fixture(scope="session") - def event_loop(self): - loop = asyncio.get_event_loop() - yield loop - loop.stop() - loop.close() - def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None return CrossSync.DataClient(project=project) diff --git a/tests/system/data/test_system_async.py b/tests/system/data/test_system_async.py index beea316bb..4f9c1ae4e 100644 --- a/tests/system/data/test_system_async.py +++ b/tests/system/data/test_system_async.py @@ -154,14 +154,6 @@ async def create_row_and_mutation( @CrossSync.convert_class(sync_name="TestSystem") class TestSystemAsync(SystemTestRunner): - @CrossSync.drop - @pytest.fixture(scope="session") - def event_loop(self): - loop = asyncio.get_event_loop() - yield loop - loop.stop() - loop.close() - def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None return CrossSync.DataClient(project=project) From 08f5ce4ebf58c54a81942949b0bff9d3717634d3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Mon, 6 Oct 2025 15:38:45 -0700 Subject: [PATCH 77/78] remvoed docstring --- google/cloud/bigtable/data/_async/metrics_interceptor.py | 4 ---- 1 file changed, 4 deletions(-) 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( From 9fe5d4eb264b85ff1d9154107f51cb7f08b35c13 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Wed, 26 Nov 2025 16:41:12 -0800 Subject: [PATCH 78/78] use new tracked_retry function --- .../bigtable/data/_async/_mutate_rows.py | 15 +++++---- .../cloud/bigtable/data/_async/_read_rows.py | 17 +++++----- google/cloud/bigtable/data/_async/client.py | 31 ++++++++----------- .../data/_sync_autogen/_mutate_rows.py | 15 +++++---- .../bigtable/data/_sync_autogen/_read_rows.py | 17 +++++----- .../bigtable/data/_sync_autogen/client.py | 31 ++++++++----------- .../data/_sync_autogen/metrics_interceptor.py | 4 --- tests/unit/data/_async/test_client.py | 4 +-- .../data/_async/test_mutations_batcher.py | 4 +-- tests/unit/data/_sync_autogen/test_client.py | 4 +-- .../_sync_autogen/test_mutations_batcher.py | 4 +-- 11 files changed, 62 insertions(+), 84 deletions(-) diff --git a/google/cloud/bigtable/data/_async/_mutate_rows.py b/google/cloud/bigtable/data/_async/_mutate_rows.py index a4eb93bfd..5502e3a5e 100644 --- a/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -21,7 +21,7 @@ import google.cloud.bigtable_v2.types.bigtable as types_pb import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._metrics import tracked_retry # mutate_rows requests are limited to this number of mutations from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT @@ -101,13 +101,12 @@ def __init__( # Entry level errors bt_exceptions._MutateRowsIncomplete, ) - self._operation = lambda: CrossSync.retry_target( - self._run_attempt, - self.is_retryable, - metric.backoff_generator, - operation_timeout, - exception_factory=metric.track_terminal_error(_retry_exception_factory), - on_error=metric.track_retryable_error, + self._operation = lambda: tracked_retry( + retry_fn=CrossSync.retry_target, + operation=metric, + target=self._run_attempt, + predicate=self.is_retryable, + timeout=operation_timeout, ) # initialize state self.timeout_generator = _attempt_timeout_generator( diff --git a/google/cloud/bigtable/data/_async/_read_rows.py b/google/cloud/bigtable/data/_async/_read_rows.py index 35b2e44e9..50b478544 100644 --- a/google/cloud/bigtable/data/_async/_read_rows.py +++ b/google/cloud/bigtable/data/_async/_read_rows.py @@ -32,7 +32,7 @@ from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._metrics import tracked_retry from google.api_core import retry as retries @@ -118,15 +118,12 @@ def start_operation(self) -> CrossSync.Iterable[Row]: Yields: Row: The next row in the stream """ - return CrossSync.retry_target_stream( - self._read_rows_attempt, - self._predicate, - self._operation_metric.backoff_generator, - self.operation_timeout, - exception_factory=self._operation_metric.track_terminal_error( - _retry_exception_factory - ), - on_error=self._operation_metric.track_retryable_error, + return tracked_retry( + retry_fn=CrossSync.retry_target_stream, + operation=self._operation_metric, + target=self._read_rows_attempt, + predicate=self._predicate, + timeout=self.operation_timeout, ) def _read_rows_attempt(self) -> CrossSync.Iterable[Row]: diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 15e52796e..9b995b66f 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -89,6 +89,7 @@ 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._metrics import tracked_retry from google.cloud.bigtable.data._cross_sync import CrossSync @@ -1448,15 +1449,12 @@ async def execute_rpc(): ) 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 - ), - on_error=operation_metric.track_retryable_error, + return await tracked_retry( + retry_fn=CrossSync.retry_target, + operation=operation_metric, + target=execute_rpc, + predicate=predicate, + timeout=operation_timeout, ) @CrossSync.convert(replace_symbols={"MutationsBatcherAsync": "MutationsBatcher"}) @@ -1584,15 +1582,12 @@ async def mutate_row( 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, + return await tracked_retry( + retry_fn=CrossSync.retry_target, + operation=operation_metric, + target=target, + predicate=predicate, + timeout=operation_timeout, ) @CrossSync.convert diff --git a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index 81a5be4cf..5daea0045 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -22,7 +22,7 @@ import google.cloud.bigtable_v2.types.bigtable as types_pb import google.cloud.bigtable.data.exceptions as bt_exceptions from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._metrics import tracked_retry from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT from google.cloud.bigtable.data.mutations import _EntryWithProto from google.cloud.bigtable.data._cross_sync import CrossSync @@ -79,13 +79,12 @@ def __init__( self.is_retryable = retries.if_exception_type( *retryable_exceptions, bt_exceptions._MutateRowsIncomplete ) - self._operation = lambda: CrossSync._Sync_Impl.retry_target( - self._run_attempt, - self.is_retryable, - metric.backoff_generator, - operation_timeout, - exception_factory=metric.track_terminal_error(_retry_exception_factory), - on_error=metric.track_retryable_error, + self._operation = lambda: tracked_retry( + retry_fn=CrossSync._Sync_Impl.retry_target, + operation=metric, + target=self._run_attempt, + predicate=self.is_retryable, + timeout=operation_timeout, ) self.timeout_generator = _attempt_timeout_generator( attempt_timeout, operation_timeout diff --git a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index adbe819eb..542a64856 100644 --- a/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -30,7 +30,7 @@ from google.cloud.bigtable.data.exceptions import _RowSetComplete from google.cloud.bigtable.data.exceptions import _ResetRow from google.cloud.bigtable.data._helpers import _attempt_timeout_generator -from google.cloud.bigtable.data._helpers import _retry_exception_factory +from google.cloud.bigtable.data._metrics import tracked_retry from google.api_core import retry as retries from google.cloud.bigtable.data._cross_sync import CrossSync @@ -103,15 +103,12 @@ def start_operation(self) -> CrossSync._Sync_Impl.Iterable[Row]: Yields: Row: The next row in the stream""" - return CrossSync._Sync_Impl.retry_target_stream( - self._read_rows_attempt, - self._predicate, - self._operation_metric.backoff_generator, - self.operation_timeout, - exception_factory=self._operation_metric.track_terminal_error( - _retry_exception_factory - ), - on_error=self._operation_metric.track_retryable_error, + return tracked_retry( + retry_fn=CrossSync._Sync_Impl.retry_target_stream, + operation=self._operation_metric, + target=self._read_rows_attempt, + predicate=self._predicate, + timeout=self.operation_timeout, ) def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index a3fbddb60..ddf3e09f4 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -76,6 +76,7 @@ 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._metrics import tracked_retry from google.cloud.bigtable.data._cross_sync import CrossSync from typing import Iterable from grpc import insecure_channel @@ -1196,15 +1197,12 @@ def execute_rpc(): ) 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 - ), - on_error=operation_metric.track_retryable_error, + return tracked_retry( + retry_fn=CrossSync._Sync_Impl.retry_target, + operation=operation_metric, + target=execute_rpc, + predicate=predicate, + timeout=operation_timeout, ) def mutations_batcher( @@ -1322,15 +1320,12 @@ def mutate_row( 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, + return tracked_retry( + retry_fn=CrossSync._Sync_Impl.retry_target, + operation=operation_metric, + target=target, + predicate=predicate, + timeout=operation_timeout, ) def bulk_mutate_rows( diff --git a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py index c5a59787c..c82f62139 100644 --- a/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py +++ b/google/cloud/bigtable/data/_sync_autogen/metrics_interceptor.py @@ -58,10 +58,6 @@ def _get_metadata(source) -> dict[str, str | bytes] | None: class BigtableMetricsInterceptor( UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor ): - """ - An async gRPC interceptor to add client metadata and print server metadata. - """ - @_with_active_operation def intercept_unary_unary( self, operation, continuation, client_call_details, request diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index 32bcab5c2..6b3a5510d 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -1400,9 +1400,9 @@ async def test_customizable_retryable_errors( predicate_builder_mock.assert_called_once_with( *expected_retryables, *extra_retryables ) - retry_call_args = retry_fn_mock.call_args_list[0].args + retry_call_kwargs = retry_fn_mock.call_args_list[0].kwargs # output of if_exception_type should be sent in to retry constructor - assert retry_call_args[1] is expected_predicate + assert retry_call_kwargs["predicate"] is expected_predicate @pytest.mark.parametrize( "fn_name,fn_args,gapic_fn", diff --git a/tests/unit/data/_async/test_mutations_batcher.py b/tests/unit/data/_async/test_mutations_batcher.py index d94b1e98c..42f62a936 100644 --- a/tests/unit/data/_async/test_mutations_batcher.py +++ b/tests/unit/data/_async/test_mutations_batcher.py @@ -1223,9 +1223,9 @@ async def test_customizable_retryable_errors( predicate_builder_mock.assert_called_once_with( *expected_retryables, _MutateRowsIncomplete ) - retry_call_args = retry_fn_mock.call_args_list[0].args + retry_call_kwargs = retry_fn_mock.call_args_list[0].kwargs # output of if_exception_type should be sent in to retry constructor - assert retry_call_args[1] is expected_predicate + assert retry_call_kwargs["predicate"] is expected_predicate @CrossSync.pytest async def test_large_batch_write(self): diff --git a/tests/unit/data/_sync_autogen/test_client.py b/tests/unit/data/_sync_autogen/test_client.py index 49a561b0e..61fa3a7d1 100644 --- a/tests/unit/data/_sync_autogen/test_client.py +++ b/tests/unit/data/_sync_autogen/test_client.py @@ -1120,8 +1120,8 @@ def test_customizable_retryable_errors( predicate_builder_mock.assert_called_once_with( *expected_retryables, *extra_retryables ) - retry_call_args = retry_fn_mock.call_args_list[0].args - assert retry_call_args[1] is expected_predicate + retry_call_kwargs = retry_fn_mock.call_args_list[0].kwargs + assert retry_call_kwargs["predicate"] is expected_predicate @pytest.mark.parametrize( "fn_name,fn_args,gapic_fn", diff --git a/tests/unit/data/_sync_autogen/test_mutations_batcher.py b/tests/unit/data/_sync_autogen/test_mutations_batcher.py index cb16ecfff..7b08df1bc 100644 --- a/tests/unit/data/_sync_autogen/test_mutations_batcher.py +++ b/tests/unit/data/_sync_autogen/test_mutations_batcher.py @@ -1069,8 +1069,8 @@ def test_customizable_retryable_errors(self, input_retryables, expected_retryabl predicate_builder_mock.assert_called_once_with( *expected_retryables, _MutateRowsIncomplete ) - retry_call_args = retry_fn_mock.call_args_list[0].args - assert retry_call_args[1] is expected_predicate + retry_call_kwargs = retry_fn_mock.call_args_list[0].kwargs + assert retry_call_kwargs["predicate"] is expected_predicate def test_large_batch_write(self): """Test that a large batch of mutations can be written"""