From 8a3af451311a79d4aa6a803736698530d5071b96 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 17 Apr 2026 12:38:50 -0700 Subject: [PATCH 01/11] moved over instrumentation code --- .../cloud/bigtable/data/_async/client.py | 155 ++++++++++-------- .../google/cloud/bigtable/data/_helpers.py | 1 + .../bigtable/data/_sync_autogen/client.py | 148 +++++++++-------- 3 files changed, 165 insertions(+), 139 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 62d233bed3ba..b2c13521240f 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -63,7 +63,11 @@ _validate_timeouts, _WarmedInstanceKey, ) -from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController +from google.cloud.bigtable.data._metrics import ( + BigtableClientSideMetricsController, + OperationType, + tracked_retry, +) from google.cloud.bigtable.data.exceptions import ( FailedQueryShardError, ShardedReadRowsExceptionGroup, @@ -1431,26 +1435,28 @@ async def sample_row_keys( retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - - @CrossSync.convert - async def execute_rpc(): - results = await self.client._gapic_client.sample_row_keys( - request=SampleRowKeysRequest( - app_profile_id=self.app_profile_id, **self._request_path - ), - timeout=next(attempt_timeout_gen), - retry=None, + with self._metrics.create_operation( + OperationType.SAMPLE_ROW_KEYS + ) as operation_metric: + + @CrossSync.convert + async def execute_rpc(): + results = await self.client._gapic_client.sample_row_keys( + request=SampleRowKeysRequest( + app_profile_id=self.app_profile_id, **self._request_path + ), + timeout=next(attempt_timeout_gen), + retry=None, + ) + return [(s.row_key, s.offset_bytes) async for s in results] + + return await tracked_retry( + retry_fn=CrossSync.retry_target, + operation=operation_metric, + target=execute_rpc, + predicate=predicate, + timeout=operation_timeout, ) - 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( @@ -1561,28 +1567,29 @@ async def mutate_row( # mutations should not be retried predicate = retries.if_exception_type() - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - - target = partial( - self.client._gapic_client.mutate_row, - request=MutateRowRequest( - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - mutations=[mutation._to_pb() for mutation in mutations_list], - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=attempt_timeout, - retry=None, - ) - return await CrossSync.retry_target( - target, - predicate, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, - ) + with self._metrics.create_operation( + OperationType.MUTATE_ROW + ) as operation_metric: + target = partial( + self.client._gapic_client.mutate_row, + request=MutateRowRequest( + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + mutations=[mutation._to_pb() for mutation in mutations_list], + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=attempt_timeout, + retry=None, + ) + return await tracked_retry( + retry_fn=CrossSync.retry_target, + operation=operation_metric, + target=target, + predicate=predicate, + timeout=operation_timeout, + ) @CrossSync.convert async def bulk_mutate_rows( @@ -1693,21 +1700,25 @@ async def check_and_mutate_row( ): false_case_mutations = [false_case_mutations] false_case_list = [m._to_pb() for m in false_case_mutations or []] - result = await self.client._gapic_client.check_and_mutate_row( - request=CheckAndMutateRowRequest( - true_mutations=true_case_list, - false_mutations=false_case_list, - predicate_filter=predicate._to_pb() if predicate is not None else None, - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=operation_timeout, - retry=None, - ) - return result.predicate_matched + + with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE): + result = await self.client._gapic_client.check_and_mutate_row( + request=CheckAndMutateRowRequest( + true_mutations=true_case_list, + false_mutations=false_case_list, + predicate_filter=predicate._to_pb() + if predicate is not None + else None, + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=operation_timeout, + retry=None, + ) + return result.predicate_matched @CrossSync.convert async def read_modify_write_row( @@ -1747,20 +1758,22 @@ async def read_modify_write_row( rules = [rules] if not rules: raise ValueError("rules must contain at least one item") - result = await self.client._gapic_client.read_modify_write_row( - request=ReadModifyWriteRowRequest( - rules=[rule._to_pb() for rule in rules], - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=operation_timeout, - retry=None, - ) - # construct Row from result - return Row._from_pb(result.row) + + with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE): + result = await self.client._gapic_client.read_modify_write_row( + request=ReadModifyWriteRowRequest( + rules=[rule._to_pb() for rule in rules], + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=operation_timeout, + retry=None, + ) + # construct Row from result + return Row._from_pb(result.row) @CrossSync.convert async def close(self): diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py index 01eda4ec7591..71549b53ceb5 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_helpers.py @@ -105,6 +105,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/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index a5873ecc0931..ff28f9e3cf7c 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -54,7 +54,11 @@ _validate_timeouts, _WarmedInstanceKey, ) -from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController +from google.cloud.bigtable.data._metrics import ( + BigtableClientSideMetricsController, + OperationType, + tracked_retry, +) from google.cloud.bigtable.data.exceptions import ( FailedQueryShardError, ShardedReadRowsExceptionGroup, @@ -1182,25 +1186,27 @@ def sample_row_keys( ) retryable_excs = _get_retryable_errors(retryable_errors, self) predicate = retries.if_exception_type(*retryable_excs) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - - def execute_rpc(): - results = self.client._gapic_client.sample_row_keys( - request=SampleRowKeysRequest( - app_profile_id=self.app_profile_id, **self._request_path - ), - timeout=next(attempt_timeout_gen), - retry=None, + with self._metrics.create_operation( + OperationType.SAMPLE_ROW_KEYS + ) as operation_metric: + + def execute_rpc(): + results = self.client._gapic_client.sample_row_keys( + request=SampleRowKeysRequest( + app_profile_id=self.app_profile_id, **self._request_path + ), + timeout=next(attempt_timeout_gen), + retry=None, + ) + return [(s.row_key, s.offset_bytes) for s in results] + + return tracked_retry( + retry_fn=CrossSync._Sync_Impl.retry_target, + operation=operation_metric, + target=execute_rpc, + predicate=predicate, + timeout=operation_timeout, ) - 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, @@ -1301,27 +1307,29 @@ def mutate_row( ) else: predicate = retries.if_exception_type() - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - target = partial( - self.client._gapic_client.mutate_row, - request=MutateRowRequest( - row_key=row_key.encode("utf-8") - if isinstance(row_key, str) - else row_key, - mutations=[mutation._to_pb() for mutation in mutations_list], - app_profile_id=self.app_profile_id, - **self._request_path, - ), - timeout=attempt_timeout, - retry=None, - ) - return CrossSync._Sync_Impl.retry_target( - target, - predicate, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, - ) + with self._metrics.create_operation( + OperationType.MUTATE_ROW + ) as operation_metric: + target = partial( + self.client._gapic_client.mutate_row, + request=MutateRowRequest( + row_key=row_key.encode("utf-8") + if isinstance(row_key, str) + else row_key, + mutations=[mutation._to_pb() for mutation in mutations_list], + app_profile_id=self.app_profile_id, + **self._request_path, + ), + timeout=attempt_timeout, + retry=None, + ) + return tracked_retry( + retry_fn=CrossSync._Sync_Impl.retry_target, + operation=operation_metric, + target=target, + predicate=predicate, + timeout=operation_timeout, + ) def bulk_mutate_rows( self, @@ -1425,21 +1433,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, @@ -1476,19 +1487,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.""" From cf58c575d9d1291095bde4f8a8bf742bc593207c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 17 Apr 2026 13:15:58 -0700 Subject: [PATCH 02/11] refactored system tests --- .../tests/system/conftest.py | 4 - .../tests/system/data/__init__.py | 235 ++++++++++++++++++ .../tests/system/data/setup_fixtures.py | 210 ---------------- .../tests/system/data/test_system_async.py | 190 ++++++-------- .../tests/system/data/test_system_autogen.py | 171 +++++-------- 5 files changed, 376 insertions(+), 434 deletions(-) delete mode 100644 packages/google-cloud-bigtable/tests/system/data/setup_fixtures.py diff --git a/packages/google-cloud-bigtable/tests/system/conftest.py b/packages/google-cloud-bigtable/tests/system/conftest.py index 5557803fffc9..e9cb2fcc91e8 100644 --- a/packages/google-cloud-bigtable/tests/system/conftest.py +++ b/packages/google-cloud-bigtable/tests/system/conftest.py @@ -24,10 +24,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/packages/google-cloud-bigtable/tests/system/data/__init__.py b/packages/google-cloud-bigtable/tests/system/data/__init__.py index 2b35cea8f778..6f836fb9678a 100644 --- a/packages/google-cloud-bigtable/tests/system/data/__init__.py +++ b/packages/google-cloud-bigtable/tests/system/data/__init__.py @@ -13,7 +13,242 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import pytest +import os +import uuid TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" TEST_AGGREGATE_FAMILY = "test-aggregate-family" + +# authorized view subset to allow all qualifiers +ALLOW_ALL = "" +ALL_QUALIFIERS = {"qualifier_prefixes": [ALLOW_ALL]} + + +class SystemTestRunner: + """ + configures a system test class with configuration for clusters/tables/etc + + used by standard system tests, and metrics tests + """ + + @pytest.fixture(scope="session") + def init_table_id(self): + """ + The table_id to use when creating a new test table + """ + return f"test-table-{uuid.uuid4().hex}" + + @pytest.fixture(scope="session") + def cluster_config(self, project_id): + """ + Configuration for the clusters to use when creating a new instance + """ + from google.cloud.bigtable_admin_v2 import types + + cluster = { + "test-cluster": types.Cluster( + location=f"projects/{project_id}/locations/us-central1-b", + serve_nodes=1, + ) + } + return cluster + + @pytest.fixture(scope="session") + def column_family_config(self): + """ + specify column families to create when creating a new test table + """ + from google.cloud.bigtable_admin_v2 import types + + int_aggregate_type = types.Type.Aggregate( + input_type=types.Type(int64_type={"encoding": {"big_endian_bytes": {}}}), + sum={}, + ) + return { + TEST_FAMILY: types.ColumnFamily(), + TEST_FAMILY_2: types.ColumnFamily(), + TEST_AGGREGATE_FAMILY: types.ColumnFamily( + value_type=types.Type(aggregate_type=int_aggregate_type) + ), + } + + @pytest.fixture(scope="session") + def admin_client(self): + """ + Client for interacting with Table and Instance admin APIs + """ + from google.cloud.bigtable.client import Client + + client = Client(admin=True) + yield client + + @pytest.fixture(scope="session") + def instance_id(self, admin_client, project_id, cluster_config): + """ + Returns BIGTABLE_TEST_INSTANCE if set, otherwise creates a new temporary instance for the test session + """ + from google.cloud.bigtable_admin_v2 import types + from google.api_core import exceptions + from google.cloud.environment_vars import BIGTABLE_EMULATOR + + # use user-specified instance if available + user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") + if user_specified_instance: + print("Using user-specified instance: {}".format(user_specified_instance)) + yield user_specified_instance + return + + # create a new temporary test instance + instance_id = f"python-bigtable-tests-{uuid.uuid4().hex[:6]}" + if os.getenv(BIGTABLE_EMULATOR): + # don't create instance if in emulator mode + yield instance_id + else: + try: + operation = admin_client.instance_admin_client.create_instance( + parent=f"projects/{project_id}", + instance_id=instance_id, + instance=types.Instance( + display_name="Test Instance", + # labels={"python-system-test": "true"}, + ), + clusters=cluster_config, + ) + operation.result(timeout=240) + except exceptions.AlreadyExists: + pass + yield instance_id + admin_client.instance_admin_client.delete_instance( + name=f"projects/{project_id}/instances/{instance_id}" + ) + + @pytest.fixture(scope="session") + def column_split_config(self): + """ + specify initial splits to create when creating a new test table + """ + return [(num * 1000).to_bytes(8, "big") for num in range(1, 10)] + + @pytest.fixture(scope="session") + def table_id( + self, + admin_client, + project_id, + instance_id, + column_family_config, + init_table_id, + column_split_config, + ): + """ + Returns BIGTABLE_TEST_TABLE if set, otherwise creates a new temporary table for the test session + + Args: + - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture. + - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. + - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. + - init_column_families: A list of column families to initialize the table with, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. + Supplied by the init_column_families fixture. + - init_table_id: The table ID to give to the test table, if pre-initialized table is not given with BIGTABLE_TEST_TABLE. + Supplied by the init_table_id fixture. + - column_split_config: A list of row keys to use as initial splits when creating the test table. + """ + from google.api_core import exceptions + from google.api_core import retry + + # use user-specified instance if available + user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") + if user_specified_table: + print("Using user-specified table: {}".format(user_specified_table)) + yield user_specified_table + return + + retry = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + try: + parent_path = f"projects/{project_id}/instances/{instance_id}" + print(f"Creating table: {parent_path}/tables/{init_table_id}") + admin_client.table_admin_client.create_table( + request={ + "parent": parent_path, + "table_id": init_table_id, + "table": {"column_families": column_family_config}, + "initial_splits": [{"key": key} for key in column_split_config], + }, + retry=retry, + ) + except exceptions.AlreadyExists: + pass + yield init_table_id + print(f"Deleting table: {parent_path}/tables/{init_table_id}") + try: + admin_client.table_admin_client.delete_table( + name=f"{parent_path}/tables/{init_table_id}" + ) + except exceptions.NotFound: + print(f"Table {init_table_id} not found, skipping deletion") + + @pytest.fixture(scope="session") + def authorized_view_id( + self, + admin_client, + project_id, + instance_id, + table_id, + ): + """ + Creates and returns a new temporary authorized view for the test session + + Args: + - admin_client: Client for interacting with the Table Admin API. Supplied by the admin_client fixture. + - project_id: The project ID of the GCP project to test against. Supplied by the project_id fixture. + - instance_id: The ID of the Bigtable instance to test against. Supplied by the instance_id fixture. + - table_id: The ID of the table to create the authorized view for. Supplied by the table_id fixture. + """ + from google.api_core import exceptions + from google.api_core import retry + + retry = retry.Retry( + predicate=retry.if_exception_type(exceptions.FailedPrecondition) + ) + new_view_id = uuid.uuid4().hex[:8] + parent_path = f"projects/{project_id}/instances/{instance_id}/tables/{table_id}" + new_path = f"{parent_path}/authorizedViews/{new_view_id}" + try: + print(f"Creating view: {new_path}") + admin_client.table_admin_client.create_authorized_view( + request={ + "parent": parent_path, + "authorized_view_id": new_view_id, + "authorized_view": { + "subset_view": { + "row_prefixes": [ALLOW_ALL], + "family_subsets": { + TEST_FAMILY: ALL_QUALIFIERS, + TEST_FAMILY_2: ALL_QUALIFIERS, + TEST_AGGREGATE_FAMILY: ALL_QUALIFIERS, + }, + }, + }, + }, + retry=retry, + ) + except exceptions.AlreadyExists: + pass + except exceptions.MethodNotImplemented: + # will occur when run in emulator. Pass empty id + new_view_id = None + yield new_view_id + if new_view_id: + print(f"Deleting view: {new_path}") + try: + admin_client.table_admin_client.delete_authorized_view(name=new_path) + except exceptions.NotFound: + print(f"View {new_view_id} not found, skipping deletion") + + @pytest.fixture(scope="session") + def project_id(self, client): + """Returns the project ID from the client.""" + yield client.project diff --git a/packages/google-cloud-bigtable/tests/system/data/setup_fixtures.py b/packages/google-cloud-bigtable/tests/system/data/setup_fixtures.py deleted file mode 100644 index 1416d6b7ab90..000000000000 --- a/packages/google-cloud-bigtable/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 os -import uuid - -import pytest - -from . import TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 - -# 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.api_core import exceptions - from google.cloud.environment_vars import BIGTABLE_EMULATOR - - from google.cloud.bigtable_admin_v2 import types - - # 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, 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, 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/packages/google-cloud-bigtable/tests/system/data/test_system_async.py b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py index b23b8f0aeff8..bd8f3e1f67cc 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py @@ -26,7 +26,7 @@ from google.cloud.bigtable.data.execute_query.metadata import SqlType from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE -from . import TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 +from . import SystemTestRunner, TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 if CrossSync.is_async: from google.cloud.bigtable_v2.services.bigtable.transports.grpc_asyncio import ( @@ -116,9 +116,43 @@ async def delete_rows(self): } await self.target.client._gapic_client.mutate_rows(request) + @CrossSync.convert + async def retrieve_cell_value(self, target, row_key): + """ + Helper to read an individual row + """ + from google.cloud.bigtable.data import ReadRowsQuery + + row_list = await target.read_rows(ReadRowsQuery(row_keys=row_key)) + assert len(row_list) == 1 + row = row_list[0] + cell = row.cells[0] + return cell.value + + @CrossSync.convert + async def create_row_and_mutation( + self, table, *, start_value=b"start", new_value=b"new_value" + ): + """ + Helper to create a new row, and a sample set_cell mutation to change its value + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = uuid.uuid4().hex.encode() + family = TEST_FAMILY + qualifier = b"test-qualifier" + await self.add_row( + row_key, family=family, qualifier=qualifier, value=start_value + ) + # ensure cell is initialized + assert await self.retrieve_cell_value(table, row_key) == start_value + + mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) + return row_key, mutation + @CrossSync.convert_class(sync_name="TestSystem") -class TestSystemAsync: +class TestSystemAsync(SystemTestRunner): def _make_client(self): project = os.getenv("GOOGLE_CLOUD_PROJECT") or None return CrossSync.DataClient(project=project) @@ -148,82 +182,6 @@ async def target(self, client, table_id, authorized_view_id, instance_id, reques else: raise ValueError(f"unknown target type: {request.param}") - @pytest.fixture(scope="session") - def column_family_config(self): - """ - specify column families to create when creating a new test table - """ - from google.cloud.bigtable_admin_v2 import types - - int_aggregate_type = types.Type.Aggregate( - input_type=types.Type(int64_type={"encoding": {"big_endian_bytes": {}}}), - sum={}, - ) - return { - TEST_FAMILY: types.ColumnFamily(), - TEST_FAMILY_2: types.ColumnFamily(), - TEST_AGGREGATE_FAMILY: types.ColumnFamily( - value_type=types.Type(aggregate_type=int_aggregate_type) - ), - } - - @pytest.fixture(scope="session") - def init_table_id(self): - """ - The table_id to use when creating a new test table - """ - return f"test-table-{uuid.uuid4().hex}" - - @pytest.fixture(scope="session") - def cluster_config(self, project_id): - """ - Configuration for the clusters to use when creating a new instance - """ - from google.cloud.bigtable_admin_v2 import types - - cluster = { - "test-cluster": types.Cluster( - location=f"projects/{project_id}/locations/us-central1-b", - serve_nodes=1, - ) - } - return cluster - - @CrossSync.convert - @pytest.mark.usefixtures("target") - async def _retrieve_cell_value(self, target, row_key): - """ - Helper to read an individual row - """ - from google.cloud.bigtable.data import ReadRowsQuery - - row_list = await target.read_rows(ReadRowsQuery(row_keys=row_key)) - assert len(row_list) == 1 - row = row_list[0] - cell = row.cells[0] - return cell.value - - @CrossSync.convert - async def _create_row_and_mutation( - self, table, temp_rows, *, start_value=b"start", new_value=b"new_value" - ): - """ - Helper to create a new row, and a sample set_cell mutation to change its value - """ - from google.cloud.bigtable.data.mutations import SetCell - - row_key = uuid.uuid4().hex.encode() - family = TEST_FAMILY - qualifier = b"test-qualifier" - await temp_rows.add_row( - row_key, family=family, qualifier=qualifier, value=start_value - ) - # ensure cell is initialized - assert await self._retrieve_cell_value(table, row_key) == start_value - - mutation = SetCell(family=TEST_FAMILY, qualifier=qualifier, new_value=new_value) - return row_key, mutation - @CrossSync.convert @CrossSync.pytest_fixture(scope="function") async def temp_rows(self, target): @@ -321,13 +279,13 @@ async def test_mutation_set_cell(self, target, temp_rows): """ row_key = b"bulk_mutate" new_value = uuid.uuid4().hex.encode() - row_key, mutation = await self._create_row_and_mutation( - target, temp_rows, new_value=new_value + row_key, mutation = await temp_rows.create_row_and_mutation( + target, new_value=new_value ) await target.mutate_row(row_key, mutation) # ensure cell is updated - assert (await self._retrieve_cell_value(target, row_key)) == new_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == new_value @CrossSync.pytest @pytest.mark.usefixtures("target") @@ -349,14 +307,14 @@ async def test_mutation_add_to_cell(self, target, temp_rows): await target.mutate_row( row_key, AddToCell(family, qualifier, 1, timestamp_micros=0) ) - encoded_result = await self._retrieve_cell_value(target, row_key) + encoded_result = await temp_rows.retrieve_cell_value(target, row_key) int_result = int.from_bytes(encoded_result, byteorder="big") assert int_result == 1 # update again await target.mutate_row( row_key, AddToCell(family, qualifier, 9, timestamp_micros=0) ) - encoded_result = await self._retrieve_cell_value(target, row_key) + encoded_result = await temp_rows.retrieve_cell_value(target, row_key) int_result = int.from_bytes(encoded_result, byteorder="big") assert int_result == 10 @@ -398,15 +356,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): @@ -446,11 +404,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]) @@ -459,7 +417,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") @@ -475,8 +433,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 @@ -487,7 +445,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") @@ -502,12 +460,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]) @@ -527,8 +485,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") @@ -543,12 +501,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]) @@ -568,8 +526,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") @@ -582,12 +540,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]) @@ -604,8 +562,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") @@ -676,7 +634,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") @@ -716,7 +674,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") @@ -755,7 +713,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") @@ -804,7 +762,7 @@ async def test_check_and_mutate( expected_value = ( true_mutation_value if expected_result else false_mutation_value ) - assert (await self._retrieve_cell_value(target, row_key)) == expected_value + assert (await temp_rows.retrieve_cell_value(target, row_key)) == expected_value @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py index 396f7b1a0605..c24372a8ed02 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py @@ -26,7 +26,7 @@ from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data.execute_query.metadata import SqlType from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE -from . import TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 +from . import SystemTestRunner, TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 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) @@ -257,11 +220,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( @@ -276,11 +239,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 @@ -311,12 +274,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""" @@ -349,18 +312,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") @@ -373,8 +336,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 @@ -384,7 +347,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") @@ -396,12 +359,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: @@ -415,8 +378,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") @@ -428,12 +391,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 @@ -447,8 +410,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") @@ -458,12 +421,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 @@ -477,8 +440,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") @@ -536,7 +499,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") @@ -569,7 +532,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") @@ -603,7 +566,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") @@ -641,7 +604,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)), @@ -933,9 +896,9 @@ def test_literal_value_filter( temp_rows.add_row(b"row_key_1", value=cell_value) query = ReadRowsQuery(row_filter=f) row_list = target.read_rows(query) - assert len(row_list) == bool(expect_match), ( - f"row {type(cell_value)}({cell_value}) not found with {type(filter_input)}({filter_input}) filter" - ) + assert len(row_list) == bool( + expect_match + ), f"row {type(cell_value)}({cell_value}) not found with {type(filter_input)}({filter_input}) filter" @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't support SQL" From 00618f88f663559df60e98b97c4fd84f806ffc52 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 17 Apr 2026 13:18:12 -0700 Subject: [PATCH 03/11] updated tests --- .../tests/system/data/test_metrics_async.py | 999 ++++++++++++++++++ .../tests/system/data/test_metrics_autogen.py | 710 +++++++++++++ .../tests/unit/data/_async/test_client.py | 10 +- .../unit/data/_sync_autogen/test_client.py | 45 +- 4 files changed, 1740 insertions(+), 24 deletions(-) create mode 100644 packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py create mode 100644 packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py new file mode 100644 index 000000000000..83635b29cef8 --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py @@ -0,0 +1,999 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +import asyncio +import os +import pytest +import uuid + +from grpc import StatusCode + +from google.api_core.exceptions import Aborted +from google.api_core.exceptions import GoogleAPICallError +from google.api_core.exceptions import PermissionDenied +from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data._metrics.data_model import ( + CompletedOperationMetric, + CompletedAttemptMetric, +) +from google.cloud.bigtable_v2.types import ResponseParams +from google.cloud.environment_vars import BIGTABLE_EMULATOR + +from google.cloud.bigtable.data._cross_sync import CrossSync + +from . import TEST_FAMILY, SystemTestRunner + +if CrossSync.is_async: + from grpc.aio import UnaryUnaryClientInterceptor + from grpc.aio import UnaryStreamClientInterceptor + from grpc.aio import AioRpcError + from grpc.aio import Metadata +else: + from grpc import UnaryUnaryClientInterceptor + from grpc import UnaryStreamClientInterceptor + from grpc import RpcError + from grpc import intercept_channel + +__CROSS_SYNC_OUTPUT__ = "tests.system.data.test_metrics_autogen" + + +class _MetricsTestHandler(MetricsHandler): + """ + Store completed metrics events in internal lists for testing + """ + + def __init__(self, **kwargs): + self.completed_operations = [] + self.completed_attempts = [] + + def on_operation_complete(self, op): + self.completed_operations.append(op) + + def on_attempt_complete(self, attempt, _): + self.completed_attempts.append(attempt) + + def clear(self): + self.completed_operations.clear() + self.completed_attempts.clear() + + def __repr__(self): + return f"{self.__class__}(completed_operations={len(self.completed_operations)}, completed_attempts={len(self.completed_attempts)}" + + +@CrossSync.convert_class +class _ErrorInjectorInterceptor( + UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor +): + """ + Gprc interceptor used to inject errors into rpc calls, to test failures + """ + + def __init__(self): + self._exc_list = [] + self.fail_mid_stream = False + + def push(self, exc: Exception): + self._exc_list.append(exc) + + def clear(self): + self._exc_list.clear() + self.fail_mid_stream = False + + @CrossSync.convert + async def intercept_unary_unary(self, continuation, client_call_details, request): + if self._exc_list: + raise self._exc_list.pop(0) + return await continuation(client_call_details, request) + + @CrossSync.convert + async def intercept_unary_stream(self, continuation, client_call_details, request): + if not self.fail_mid_stream and self._exc_list: + raise self._exc_list.pop(0) + + response = await continuation(client_call_details, request) + + if self.fail_mid_stream and self._exc_list: + exc = self._exc_list.pop(0) + + class CallWrapper: + def __init__(self, call, exc_to_raise): + self._call = call + self._exc = exc_to_raise + self._raised = False + + @CrossSync.convert(sync_name="__iter__") + def __aiter__(self): + return self + + @CrossSync.convert( + sync_name="__next__", replace_symbols={"__anext__": "__next__"} + ) + async def __anext__(self): + if not self._raised: + self._raised = True + if self._exc: + raise self._exc + return await self._call.__anext__() + + def __getattr__(self, name): + return getattr(self._call, name) + + return CallWrapper(response, exc) + + return response + + +@CrossSync.convert_class(sync_name="TestMetrics") +class TestMetricsAsync(SystemTestRunner): + def _make_client(self): + project = os.getenv("GOOGLE_CLOUD_PROJECT") or None + return CrossSync.DataClient(project=project) + + def _make_exception(self, status, cluster_id=None, zone_id=None): + if cluster_id or zone_id: + metadata = ( + "x-goog-ext-425905942-bin", + ResponseParams.serialize( + ResponseParams(cluster_id=cluster_id, zone_id=zone_id) + ), + ) + else: + metadata = None + if CrossSync.is_async: + metadata = Metadata(metadata) if metadata else Metadata() + return AioRpcError(status, Metadata(), metadata) + else: + exc = RpcError(status) + exc.trailing_metadata = lambda: [metadata] if metadata else [] + exc.initial_metadata = lambda: [] + exc.code = lambda: status + exc.details = lambda: None + + def _result(): + raise exc + + exc.result = _result + return exc + + @pytest.fixture(scope="session") + def handler(self): + return _MetricsTestHandler() + + @pytest.fixture(scope="session") + def error_injector(self): + return _ErrorInjectorInterceptor() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function", autouse=True) + async def _clear_state(self, handler, error_injector): + """Clear handler and interceptor between each test""" + handler.clear() + error_injector.clear() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def client(self, error_injector): + async with self._make_client() as client: + if CrossSync.is_async: + client.transport.grpc_channel._unary_unary_interceptors.append( + error_injector + ) + client.transport.grpc_channel._unary_stream_interceptors.append( + error_injector + ) + else: + # inject interceptor after bigtable metrics interceptors + metrics_channel = client.transport._grpc_channel._channel._channel + client.transport._grpc_channel._channel._channel = intercept_channel( + metrics_channel, error_injector + ) + yield client + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="function") + async def temp_rows(self, table): + builder = CrossSync.TempRowBuilder(table) + yield builder + await builder.delete_rows() + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def table(self, client, table_id, instance_id, handler): + async with client.get_table(instance_id, table_id) as table: + table._metrics.add_handler(handler) + yield table + + @CrossSync.convert + @CrossSync.pytest_fixture(scope="session") + async def authorized_view( + self, client, table_id, instance_id, authorized_view_id, handler + ): + async with client.get_authorized_view( + instance_id, table_id, authorized_view_id + ) as table: + table._metrics.add_handler(handler) + yield table + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_mutate_row(self, table, temp_rows, handler, cluster_config): + row_key = b"mutate" + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + handler.clear() + await table.mutate_row(row_key, mutation) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_mutate_row_failure_with_retries( + self, table, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.mutate_row(row_key, [mutation], retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + @CrossSync.pytest + async def test_mutate_row_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + + with pytest.raises(GoogleAPICallError): + await table.mutate_row(row_key, [mutation], operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + 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 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.mutate_row(row_key, [mutation]) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_mutate_row_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + """ + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.mutate_row( + row_key, + [mutation], + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): + await table.sample_row_keys() + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "SampleRowKeys" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @CrossSync.drop + @CrossSync.pytest + async def test_sample_row_keys_failure_cancelled( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + test with retryable errors, then a terminal one + + No headers expected + """ + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(asyncio.CancelledError) + with pytest.raises(asyncio.CancelledError): + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "UNKNOWN" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == "" + 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 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_sample_row_keys_failure_with_retries( + self, table, temp_rows, handler, error_injector, cluster_config + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a success + """ + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(self._make_exception(StatusCode.ABORTED)) + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "OK" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "OK" + assert ( + final_attempt.gfe_latency_ns > 0 + and final_attempt.gfe_latency_ns < operation.duration_ns + ) + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + @CrossSync.pytest + async def test_sample_row_keys_failure_timeout(self, table, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + with pytest.raises(GoogleAPICallError): + await table.sample_row_keys(operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + 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 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_sample_row_keys_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(PermissionDenied): + await table.sample_row_keys(retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 2 + # validate retried attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + # validate final attempt + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_read_modify_write(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + await table.read_modify_write_row(row_key, rule) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadModifyWriteRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @CrossSync.drop + @CrossSync.pytest + async def test_read_modify_write_failure_cancelled( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + + # trigger an exception + exc = asyncio.CancelledError("injected") + error_injector.push(exc) + with pytest.raises(asyncio.CancelledError): + await table.read_modify_write_row(row_key, rule) + + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "UNKNOWN" + assert operation.is_streaming is False + assert operation.op_type.value == "ReadModifyWriteRow" + assert len(operation.completed_attempts) == len(handler.completed_attempts) + assert operation.completed_attempts == handler.completed_attempts + assert operation.cluster_id == "" + 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 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + @CrossSync.pytest + async def test_read_modify_write_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + with pytest.raises(GoogleAPICallError): + await table.read_modify_write_row(row_key, rule, operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + @CrossSync.pytest + async def test_read_modify_write_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + rule = IncrementRule("unauthorized", qualifier, 1) + with pytest.raises(GoogleAPICallError): + await authorized_view.read_modify_write_row(row_key, rule) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_check_and_mutate_row( + self, table, temp_rows, handler, cluster_config + ): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + predicate = ValueRangeFilter(0, 2) + await table.check_and_mutate_row( + row_key, + predicate, + true_case_mutations=true_mutation, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "CheckAndMutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @CrossSync.drop + @CrossSync.pytest + async def test_check_and_mutate_row_failure_cancelled( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting an error into an interceptor + + No headers expected + """ + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + + # trigger an exception + exc = asyncio.CancelledError("injected") + error_injector.push(exc) + with pytest.raises(asyncio.CancelledError): + await table.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "UNKNOWN" + assert operation.is_streaming is False + assert operation.op_type.value == "CheckAndMutateRow" + assert len(operation.completed_attempts) == len(handler.completed_attempts) + assert operation.completed_attempts == handler.completed_attempts + assert operation.cluster_id == "" + 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 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + @CrossSync.pytest + async def test_check_and_mutate_row_failure_timeout( + self, table, temp_rows, handler + ): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + await temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + with pytest.raises(GoogleAPICallError): + await table.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=true_mutation, + operation_timeout=0.001, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + @CrossSync.pytest + async def test_check_and_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + mutation_value = b"true-mutation-value" + mutation = SetCell( + family="unauthorized", qualifier=qualifier, new_value=mutation_value + ) + with pytest.raises(GoogleAPICallError): + await authorized_view.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=mutation, + false_case_mutations=mutation, + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py new file mode 100644 index 000000000000..2fec721e6c6d --- /dev/null +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py @@ -0,0 +1,710 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This file is automatically generated by CrossSync. Do not edit manually. + +import os +import pytest +import uuid +from grpc import StatusCode +from google.api_core.exceptions import Aborted +from google.api_core.exceptions import GoogleAPICallError +from google.api_core.exceptions import PermissionDenied +from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data._metrics.data_model import ( + CompletedOperationMetric, + CompletedAttemptMetric, +) +from google.cloud.bigtable_v2.types import ResponseParams +from google.cloud.environment_vars import BIGTABLE_EMULATOR +from google.cloud.bigtable.data._cross_sync import CrossSync +from . import TEST_FAMILY, SystemTestRunner +from grpc import UnaryUnaryClientInterceptor +from grpc import UnaryStreamClientInterceptor +from grpc import RpcError +from grpc import intercept_channel + + +class _MetricsTestHandler(MetricsHandler): + """ + Store completed metrics events in internal lists for testing + """ + + def __init__(self, **kwargs): + self.completed_operations = [] + self.completed_attempts = [] + + def on_operation_complete(self, op): + self.completed_operations.append(op) + + def on_attempt_complete(self, attempt, _): + self.completed_attempts.append(attempt) + + def clear(self): + self.completed_operations.clear() + self.completed_attempts.clear() + + def __repr__(self): + return f"{self.__class__}(completed_operations={len(self.completed_operations)}, completed_attempts={len(self.completed_attempts)}" + + +class _ErrorInjectorInterceptor( + UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor +): + """ + Gprc interceptor used to inject errors into rpc calls, to test failures + """ + + def __init__(self): + self._exc_list = [] + self.fail_mid_stream = False + + def push(self, exc: Exception): + self._exc_list.append(exc) + + def clear(self): + self._exc_list.clear() + self.fail_mid_stream = False + + def intercept_unary_unary(self, continuation, client_call_details, request): + if self._exc_list: + raise self._exc_list.pop(0) + return continuation(client_call_details, request) + + def intercept_unary_stream(self, continuation, client_call_details, request): + if not self.fail_mid_stream and self._exc_list: + raise self._exc_list.pop(0) + response = continuation(client_call_details, request) + if self.fail_mid_stream and self._exc_list: + exc = self._exc_list.pop(0) + + class CallWrapper: + def __init__(self, call, exc_to_raise): + self._call = call + self._exc = exc_to_raise + self._raised = False + + def __iter__(self): + return self + + def __next__(self): + if not self._raised: + self._raised = True + if self._exc: + raise self._exc + return self._call.__next__() + + def __getattr__(self, name): + return getattr(self._call, name) + + return CallWrapper(response, exc) + return response + + +class TestMetrics(SystemTestRunner): + def _make_client(self): + project = os.getenv("GOOGLE_CLOUD_PROJECT") or None + return CrossSync._Sync_Impl.DataClient(project=project) + + def _make_exception(self, status, cluster_id=None, zone_id=None): + if cluster_id or zone_id: + metadata = ( + "x-goog-ext-425905942-bin", + ResponseParams.serialize( + ResponseParams(cluster_id=cluster_id, zone_id=zone_id) + ), + ) + else: + metadata = None + exc = RpcError(status) + exc.trailing_metadata = lambda: [metadata] if metadata else [] + exc.initial_metadata = lambda: [] + exc.code = lambda: status + exc.details = lambda: None + + def _result(): + raise exc + + exc.result = _result + return exc + + @pytest.fixture(scope="session") + def handler(self): + return _MetricsTestHandler() + + @pytest.fixture(scope="session") + def error_injector(self): + return _ErrorInjectorInterceptor() + + @pytest.fixture(scope="function", autouse=True) + def _clear_state(self, handler, error_injector): + """Clear handler and interceptor between each test""" + handler.clear() + error_injector.clear() + + @pytest.fixture(scope="session") + def client(self, error_injector): + with self._make_client() as client: + metrics_channel = client.transport._grpc_channel._channel._channel + client.transport._grpc_channel._channel._channel = intercept_channel( + metrics_channel, error_injector + ) + yield client + + @pytest.fixture(scope="function") + def temp_rows(self, table): + builder = CrossSync._Sync_Impl.TempRowBuilder(table) + yield builder + builder.delete_rows() + + @pytest.fixture(scope="session") + def table(self, client, table_id, instance_id, handler): + with client.get_table(instance_id, table_id) as table: + table._metrics.add_handler(handler) + yield table + + @pytest.fixture(scope="session") + def authorized_view( + self, client, table_id, instance_id, authorized_view_id, handler + ): + with client.get_authorized_view( + instance_id, table_id, authorized_view_id + ) as table: + table._metrics.add_handler(handler) + yield table + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_mutate_row(self, table, temp_rows, handler, cluster_config): + row_key = b"mutate" + new_value = uuid.uuid4().hex.encode() + row_key, mutation = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + handler.clear() + table.mutate_row(row_key, mutation) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_mutate_row_failure_with_retries(self, table, handler, error_injector): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.mutate_row(row_key, [mutation], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + def test_mutate_row_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + with pytest.raises(GoogleAPICallError): + table.mutate_row(row_key, [mutation], operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + with pytest.raises(GoogleAPICallError) as e: + authorized_view.mutate_row(row_key, [mutation]) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_mutate_row_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out""" + from google.cloud.bigtable.data.mutations import SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + with pytest.raises(GoogleAPICallError) as e: + authorized_view.mutate_row( + row_key, + [mutation], + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRow" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_sample_row_keys(self, table, temp_rows, handler, cluster_config): + table.sample_row_keys() + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "SampleRowKeys" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_sample_row_keys_failure_with_retries( + self, table, temp_rows, handler, error_injector, cluster_config + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a success""" + num_retryable = 3 + for i in range(num_retryable): + error_injector.push(self._make_exception(StatusCode.ABORTED)) + table.sample_row_keys(retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "OK" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "OK" + assert ( + final_attempt.gfe_latency_ns > 0 + and final_attempt.gfe_latency_ns < operation.duration_ns + ) + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + def test_sample_row_keys_failure_timeout(self, table, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + with pytest.raises(GoogleAPICallError): + table.sample_row_keys(operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_sample_row_keys_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(PermissionDenied): + table.sample_row_keys(retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "SampleRowKeys" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 2 + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_read_modify_write(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + table.read_modify_write_row(row_key, rule) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadModifyWriteRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + def test_read_modify_write_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=0, family=family, qualifier=qualifier) + rule = IncrementRule(family, qualifier, 1) + with pytest.raises(GoogleAPICallError): + table.read_modify_write_row(row_key, rule, operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + def test_read_modify_write_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + rule = IncrementRule("unauthorized", qualifier, 1) + with pytest.raises(GoogleAPICallError): + authorized_view.read_modify_write_row(row_key, rule) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadModifyWriteRow" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_check_and_mutate_row(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + predicate = ValueRangeFilter(0, 2) + table.check_and_mutate_row( + row_key, predicate, true_case_mutations=true_mutation + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "CheckAndMutateRow" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" + ) + def test_check_and_mutate_row_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + family = TEST_FAMILY + qualifier = b"test-qualifier" + temp_rows.add_row(row_key, value=1, family=family, qualifier=qualifier) + true_mutation_value = b"true-mutation-value" + true_mutation = SetCell( + family=TEST_FAMILY, qualifier=qualifier, new_value=true_mutation_value + ) + with pytest.raises(GoogleAPICallError): + table.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=true_mutation, + operation_timeout=0.001, + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.gfe_latency_ns is None + + @pytest.mark.skipif( + bool(os.environ.get(BIGTABLE_EMULATOR)), + reason="emulator doesn't suport cluster_config", + ) + def test_check_and_mutate_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.mutations import SetCell + from google.cloud.bigtable.data.row_filters import ValueRangeFilter + + row_key = b"test-row-key" + qualifier = b"test-qualifier" + mutation_value = b"true-mutation-value" + mutation = SetCell( + family="unauthorized", qualifier=qualifier, new_value=mutation_value + ) + with pytest.raises(GoogleAPICallError): + authorized_view.check_and_mutate_row( + row_key, + predicate=ValueRangeFilter(0, 2), + true_case_mutations=mutation, + false_case_mutations=mutation, + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index f40059823773..ad2ae9c8bb42 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -1402,9 +1402,15 @@ 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 # output of if_exception_type should be sent in to retry constructor - assert retry_call_args[1] is expected_predicate + retry_call_kwargs = retry_fn_mock.call_args_list[0].kwargs + # check for predicate passed as kwarg + if "predicate" in retry_call_kwargs: + assert retry_call_kwargs["predicate"] is expected_predicate + else: + # check for predicate passed as arg + retry_call_args = retry_fn_mock.call_args_list[0].args + assert retry_call_args[1] is expected_predicate @pytest.mark.parametrize( "fn_name,fn_args,gapic_fn", diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index 0d5e47b1072b..fbcd8c598801 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -172,9 +172,9 @@ def test_veneer_grpc_headers(self): wrapped_user_agent_sorted = " ".join( sorted(client_info.to_user_agent().split(" ")) ) - assert VENEER_HEADER_REGEX.match(wrapped_user_agent_sorted), ( - f"'{wrapped_user_agent_sorted}' does not match {VENEER_HEADER_REGEX}" - ) + assert VENEER_HEADER_REGEX.match( + wrapped_user_agent_sorted + ), f"'{wrapped_user_agent_sorted}' does not match {VENEER_HEADER_REGEX}" client.close() def test__start_background_channel_refresh_task_exists(self): @@ -301,9 +301,9 @@ def test__manage_channel_first_sleep( pass sleep.assert_called_once() call_time = sleep.call_args[0][1] - assert abs(call_time - expected_sleep) < 0.1, ( - f"refresh_interval: {refresh_interval}, wait_time: {wait_time}, expected_sleep: {expected_sleep}" - ) + assert ( + abs(call_time - expected_sleep) < 0.1 + ), f"refresh_interval: {refresh_interval}, wait_time: {wait_time}, expected_sleep: {expected_sleep}" client.close() def test__manage_channel_ping_and_warm(self): @@ -319,9 +319,9 @@ def test__manage_channel_ping_and_warm(self): ) with mock.patch.object(*sleep_tuple) as sleep_mock: sleep_mock.side_effect = [None, asyncio.CancelledError] - ping_and_warm = client._ping_and_warm_instances = ( - CrossSync._Sync_Impl.Mock() - ) + ping_and_warm = ( + client._ping_and_warm_instances + ) = CrossSync._Sync_Impl.Mock() try: client._manage_channel(10) except asyncio.CancelledError: @@ -363,9 +363,9 @@ def test__manage_channel_sleeps(self, refresh_interval, num_cycles, expected_sle pass assert sleep.call_count == num_cycles total_sleep = sum([call[0][1] for call in sleep.call_args_list]) - assert abs(total_sleep - expected_sleep) < 0.5, ( - f"refresh_interval={refresh_interval}, num_cycles={num_cycles}, expected_sleep={expected_sleep}" - ) + assert ( + abs(total_sleep - expected_sleep) < 0.5 + ), f"refresh_interval={refresh_interval}, num_cycles={num_cycles}, expected_sleep={expected_sleep}" client.close() def test__manage_channel_random(self): @@ -1120,8 +1120,12 @@ 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 + if "predicate" in retry_call_kwargs: + assert retry_call_kwargs["predicate"] is expected_predicate + else: + retry_call_args = retry_fn_mock.call_args_list[0].args + assert retry_call_args[1] is expected_predicate @pytest.mark.parametrize( "fn_name,fn_args,gapic_fn", @@ -1772,14 +1776,11 @@ def test_read_rows_sharded_multiple_queries(self): with mock.patch.object( table.client._gapic_client, "read_rows" ) as read_rows: - read_rows.side_effect = ( - lambda *args, - **kwargs: CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( - [ - CrossSync._Sync_Impl.TestReadRows._make_chunk(row_key=k) - for k in args[0].rows.row_keys - ] - ) + read_rows.side_effect = lambda *args, **kwargs: CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( + [ + CrossSync._Sync_Impl.TestReadRows._make_chunk(row_key=k) + for k in args[0].rows.row_keys + ] ) query_1 = ReadRowsQuery(b"test_1") query_2 = ReadRowsQuery(b"test_2") From ead3bf1b5bbb139895497aa16f92593bf182fe0b Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Fri, 17 Apr 2026 15:35:05 -0700 Subject: [PATCH 04/11] fixed lint --- .../tests/system/data/test_metrics_async.py | 6 +-- .../tests/system/data/test_system_async.py | 4 +- .../tests/system/data/test_system_autogen.py | 6 +-- .../unit/data/_sync_autogen/test_client.py | 37 ++++++++++--------- 4 files changed, 29 insertions(+), 24 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py index 83635b29cef8..5b335857d2fa 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py @@ -265,7 +265,7 @@ async def test_mutate_row(self, table, temp_rows, handler, cluster_config): 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 + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns ) assert attempt.application_blocking_time_ns == 0 @@ -315,12 +315,12 @@ async def test_mutate_row_failure_with_retries( for i in range(num_retryable): attempt = handler.completed_attempts[i] assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" + 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 + assert final_attempt.gfe_latency_ns is None @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="not supported by emulator" diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py index bd8f3e1f67cc..3636e6993e55 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py @@ -563,7 +563,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") diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py index c24372a8ed02..d916aae14f25 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py @@ -896,9 +896,9 @@ def test_literal_value_filter( temp_rows.add_row(b"row_key_1", value=cell_value) query = ReadRowsQuery(row_filter=f) row_list = target.read_rows(query) - assert len(row_list) == bool( - expect_match - ), f"row {type(cell_value)}({cell_value}) not found with {type(filter_input)}({filter_input}) filter" + assert len(row_list) == bool(expect_match), ( + f"row {type(cell_value)}({cell_value}) not found with {type(filter_input)}({filter_input}) filter" + ) @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't support SQL" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index fbcd8c598801..f314b8c5bd9c 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -172,9 +172,9 @@ def test_veneer_grpc_headers(self): wrapped_user_agent_sorted = " ".join( sorted(client_info.to_user_agent().split(" ")) ) - assert VENEER_HEADER_REGEX.match( - wrapped_user_agent_sorted - ), f"'{wrapped_user_agent_sorted}' does not match {VENEER_HEADER_REGEX}" + assert VENEER_HEADER_REGEX.match(wrapped_user_agent_sorted), ( + f"'{wrapped_user_agent_sorted}' does not match {VENEER_HEADER_REGEX}" + ) client.close() def test__start_background_channel_refresh_task_exists(self): @@ -301,9 +301,9 @@ def test__manage_channel_first_sleep( pass sleep.assert_called_once() call_time = sleep.call_args[0][1] - assert ( - abs(call_time - expected_sleep) < 0.1 - ), f"refresh_interval: {refresh_interval}, wait_time: {wait_time}, expected_sleep: {expected_sleep}" + assert abs(call_time - expected_sleep) < 0.1, ( + f"refresh_interval: {refresh_interval}, wait_time: {wait_time}, expected_sleep: {expected_sleep}" + ) client.close() def test__manage_channel_ping_and_warm(self): @@ -319,9 +319,9 @@ def test__manage_channel_ping_and_warm(self): ) with mock.patch.object(*sleep_tuple) as sleep_mock: sleep_mock.side_effect = [None, asyncio.CancelledError] - ping_and_warm = ( - client._ping_and_warm_instances - ) = CrossSync._Sync_Impl.Mock() + ping_and_warm = client._ping_and_warm_instances = ( + CrossSync._Sync_Impl.Mock() + ) try: client._manage_channel(10) except asyncio.CancelledError: @@ -363,9 +363,9 @@ def test__manage_channel_sleeps(self, refresh_interval, num_cycles, expected_sle pass assert sleep.call_count == num_cycles total_sleep = sum([call[0][1] for call in sleep.call_args_list]) - assert ( - abs(total_sleep - expected_sleep) < 0.5 - ), f"refresh_interval={refresh_interval}, num_cycles={num_cycles}, expected_sleep={expected_sleep}" + assert abs(total_sleep - expected_sleep) < 0.5, ( + f"refresh_interval={refresh_interval}, num_cycles={num_cycles}, expected_sleep={expected_sleep}" + ) client.close() def test__manage_channel_random(self): @@ -1776,11 +1776,14 @@ def test_read_rows_sharded_multiple_queries(self): with mock.patch.object( table.client._gapic_client, "read_rows" ) as read_rows: - read_rows.side_effect = lambda *args, **kwargs: CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( - [ - CrossSync._Sync_Impl.TestReadRows._make_chunk(row_key=k) - for k in args[0].rows.row_keys - ] + read_rows.side_effect = ( + lambda *args, + **kwargs: CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( + [ + CrossSync._Sync_Impl.TestReadRows._make_chunk(row_key=k) + for k in args[0].rows.row_keys + ] + ) ) query_1 = ReadRowsQuery(b"test_1") query_2 = ReadRowsQuery(b"test_2") From d0c03d07f91f6976d7356990b1e6f0e09813ab51 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 12:40:19 -0700 Subject: [PATCH 05/11] regenerated files --- .../bigtable/data/_sync_autogen/client.py | 7 +++++ .../tests/system/data/test_metrics_autogen.py | 29 ++++++++++--------- .../tests/system/data/test_system_autogen.py | 3 +- 3 files changed, 24 insertions(+), 15 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index 1ee10324dba9..9dc118de0289 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -62,6 +62,13 @@ OperationType, tracked_retry, ) +from google.cloud.bigtable.data._sync_autogen._swappable_channel import ( + SwappableChannel as SwappableChannelType, +) +from google.cloud.bigtable.data._sync_autogen.metrics_interceptor import ( + BigtableMetricsInterceptor as MetricsInterceptorType, +) +from google.cloud.bigtable.data._sync_autogen.mutations_batcher import _MB_SIZE from google.cloud.bigtable.data.exceptions import ( FailedQueryShardError, ShardedReadRowsExceptionGroup, diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py index 2fec721e6c6d..bb14d565c267 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py @@ -15,25 +15,28 @@ # This file is automatically generated by CrossSync. Do not edit manually. import os -import pytest import uuid -from grpc import StatusCode -from google.api_core.exceptions import Aborted -from google.api_core.exceptions import GoogleAPICallError -from google.api_core.exceptions import PermissionDenied -from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler + +import pytest +from google.api_core.exceptions import Aborted, GoogleAPICallError, PermissionDenied +from google.cloud.environment_vars import BIGTABLE_EMULATOR +from grpc import ( + RpcError, + StatusCode, + UnaryStreamClientInterceptor, + UnaryUnaryClientInterceptor, + intercept_channel, +) + +from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._metrics.data_model import ( - CompletedOperationMetric, CompletedAttemptMetric, + CompletedOperationMetric, ) +from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from google.cloud.bigtable_v2.types import ResponseParams -from google.cloud.environment_vars import BIGTABLE_EMULATOR -from google.cloud.bigtable.data._cross_sync import CrossSync + from . import TEST_FAMILY, SystemTestRunner -from grpc import UnaryUnaryClientInterceptor -from grpc import UnaryStreamClientInterceptor -from grpc import RpcError -from grpc import intercept_channel class _MetricsTestHandler(MetricsHandler): diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py index 6747948bca9e..c31b2c20a4b8 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_autogen.py @@ -28,12 +28,11 @@ from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data.execute_query.metadata import SqlType from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE -from . import SystemTestRunner, TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 from google.cloud.bigtable_v2.services.bigtable.transports.grpc import ( _LoggingClientInterceptor as GapicInterceptor, ) -from . import TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 +from . import TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2, SystemTestRunner TARGETS = ["table"] if not os.environ.get(BIGTABLE_EMULATOR): From a78cb52ef6f684f1751bf8da27b14fffe42ce958 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 13:36:50 -0700 Subject: [PATCH 06/11] added new metrics tests --- .../tests/system/data/test_metrics_async.py | 1213 +++++++++++++++++ .../tests/system/data/test_metrics_autogen.py | 1016 ++++++++++++++ 2 files changed, 2229 insertions(+) diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py index 5b335857d2fa..49cd9ba2b744 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py @@ -224,6 +224,1219 @@ async def authorized_view( table._metrics.add_handler(handler) yield table + @CrossSync.pytest + async def test_read_rows(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + row_list = await table.read_rows(ReadRowsQuery()) + assert len(row_list) == 2 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + await table.read_rows(ReadRowsQuery(), operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.read_rows( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + # full table scan + generator = await table.read_rows_stream(ReadRowsQuery()) + row_list = [r async for r in generator] + assert len(row_list) == 2 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + @CrossSync.convert(replace_symbols={"__anext__": "__next__", "aclose": "close"}) + async def test_read_rows_stream_failure_closed( + self, table, temp_rows, handler, error_injector + ): + """ + Test how metrics collection handles closed generator + """ + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + generator = await table.read_rows_stream(ReadRowsQuery()) + await generator.__anext__() + await generator.aclose() + with pytest.raises(CrossSync.StopIteration): + await generator.__anext__() + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "CANCELLED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "CANCELLED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) + with pytest.raises(PermissionDenied): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + generator = await table.read_rows_stream( + ReadRowsQuery(), operation_timeout=0.001 + ) + with pytest.raises(GoogleAPICallError): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = await authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + [_ async for _ in generator] + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_stream_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = await authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + [_ async for _ in generator] + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + @CrossSync.pytest + async def test_read_rows_stream_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) + with pytest.raises(PermissionDenied): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 2 + # validate retried attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + # validate final attempt + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + @CrossSync.pytest + async def test_read_row(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + handler.clear() + await table.read_row(b"row_key_1") + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns > 0 + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_row_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.read_row(b"row_key_1", retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_row_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + await table.read_row(b"row_key_1", operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.read_row( + b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) + handler.clear() + row_list = await table.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + # validate counts + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # validate operations + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + attempt = operation.completed_attempts[0] + assert attempt in handler.completed_attempts + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + assert isinstance(attempt, CompletedAttemptMetric) + assert ( + attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + ) + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 + and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_sharded_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + + error_injector.push(self._make_exception(StatusCode.ABORTED)) + await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + # validate operations + for op in handler.completed_operations: + assert op.final_status.name == "OK" + assert op.op_type.value == "ReadRows" + assert op.is_streaming is True + # validate attempts + assert ( + len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) + == 2 + ) + assert ( + len( + [ + a + for a in handler.completed_attempts + if a.end_status.name == "ABORTED" + ] + ) + == 1 + ) + + @CrossSync.pytest + async def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.api_core.exceptions import DeadlineExceeded + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2], operation_timeout=0.005) + assert len(e.value.exceptions) == 2 + for sub_exc in e.value.exceptions: + assert isinstance(sub_exc.__cause__, DeadlineExceeded) + # both shards should fail + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # validate operations + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = operation.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_sharded_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await authorized_view.read_rows_sharded([query1, query2]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + # one shard will fail, the other will succeed + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # sort operations by status + failed_op = next( + op for op in handler.completed_operations if op.final_status.name != "OK" + ) + success_op = next( + op for op in handler.completed_operations if op.final_status.name == "OK" + ) + # validate failed operation + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert failed_op.cluster_id == next(iter(cluster_config.keys())) + assert ( + failed_op.zone + == cluster_config[failed_op.cluster_id].location.split("/")[-1] + ) + # validate failed attempt + failed_attempt = failed_op.completed_attempts[0] + assert failed_attempt.end_status.name == "PERMISSION_DENIED" + assert ( + failed_attempt.gfe_latency_ns >= 0 + and failed_attempt.gfe_latency_ns < failed_op.duration_ns + ) + # validate successful operation + assert success_op.final_status.name == "OK" + assert success_op.op_type.value == "ReadRows" + assert success_op.is_streaming is True + assert len(success_op.completed_attempts) == 1 + # validate successful attempt + success_attempt = success_op.completed_attempts[0] + assert success_attempt.end_status.name == "OK" + + @CrossSync.pytest + async def test_read_rows_sharded_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + # one shard will fail, the other will succeed + # the failing shard will have one retry + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + # sort operations by status + failed_op = next( + op for op in handler.completed_operations if op.final_status.name != "OK" + ) + success_op = next( + op for op in handler.completed_operations if op.final_status.name == "OK" + ) + # validate failed operation + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + # validate successful operation + assert success_op.final_status.name == "OK" + assert len(success_op.completed_attempts) == 2 + # validate failed attempt + attempt = failed_op.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + # validate retried attempt + retried_attempt = success_op.completed_attempts[0] + assert retried_attempt.end_status.name == "ABORTED" + # validate successful attempt + success_attempt = success_op.completed_attempts[-1] + assert success_attempt.end_status.name == "OK" + + @CrossSync.pytest + async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + + handler.clear() + await table.bulk_mutate_rows([bulk_mutation]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + await table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup): + await table.bulk_mutate_rows([entry], operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup): + await authorized_view.bulk_mutate_rows([entry]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + + For bulk_mutate, the rpc returns success, with failures returned in the response. + For this reason, We expect the attempts to be marked as successful, even though + the underlying mutation is retried + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup) as e: + await authorized_view.bulk_mutate_rows( + [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 + ) + assert len(e.value.exceptions) == 1 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] + + @CrossSync.pytest + async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + row_key2, mutation2 = await temp_rows.create_row_and_mutation( + table, new_value=new_value2 + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + + handler.clear() + async with table.mutations_batcher() as batcher: + await batcher.append(bulk_mutation) + await batcher.append(bulk_mutation2) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # bacher expects to cancel staged operation on close + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert ( + operation.flow_throttling_time_ns > 0 + and operation.flow_throttling_time_ns < operation.duration_ns + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_with_retries( + self, table, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + async with table.mutations_batcher( + batch_retryable_errors=[Aborted] + ) as batcher: + await batcher.append(entry) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup): + async with table.mutations_batcher( + batch_operation_timeout=0.001 + ) as batcher: + await batcher.append(entry) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup) as e: + async with authorized_view.mutations_batcher() as batcher: + await batcher.append(entry) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't suport cluster_config", diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py index bb14d565c267..4f315652b15c 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py @@ -187,6 +187,1022 @@ def authorized_view( table._metrics.add_handler(handler) yield table + def test_read_rows(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + row_list = table.read_rows(ReadRowsQuery()) + assert len(row_list) == 2 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_rows_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + table.read_rows(ReadRowsQuery(), operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + authorized_view.read_rows( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery()) + row_list = [r for r in generator] + assert len(row_list) == 2 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_rows_stream_failure_closed( + self, table, temp_rows, handler, error_injector + ): + """Test how metrics collection handles closed generator""" + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery()) + generator.__next__() + generator.close() + with pytest.raises(CrossSync._Sync_Impl.StopIteration): + generator.__next__() + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "CANCELLED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "CANCELLED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + with pytest.raises(PermissionDenied): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery(), operation_timeout=0.001) + with pytest.raises(GoogleAPICallError): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + [_ for _ in generator] + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_stream_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + [_ for _ in generator] + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + def test_read_rows_stream_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + temp_rows.add_row(b"row_key_1") + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + with pytest.raises(PermissionDenied): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 2 + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + def test_read_row(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + handler.clear() + table.read_row(b"row_key_1") + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns > 0 + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_row_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.read_row(b"row_key_1", retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_row_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + table.read_row(b"row_key_1", operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + authorized_view.read_row( + b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + temp_rows.add_row(b"c") + temp_rows.add_row(b"d") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) + handler.clear() + row_list = table.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + attempt = operation.completed_attempts[0] + assert attempt in handler.completed_attempts + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + assert isinstance(attempt, CompletedAttemptMetric) + assert ( + attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + ) + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 + and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_rows_sharded_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors""" + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.push(self._make_exception(StatusCode.ABORTED)) + table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + for op in handler.completed_operations: + assert op.final_status.name == "OK" + assert op.op_type.value == "ReadRows" + assert op.is_streaming is True + assert ( + len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) + == 2 + ) + assert ( + len( + [ + a + for a in handler.completed_attempts + if a.end_status.name == "ABORTED" + ] + ) + == 1 + ) + + def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.api_core.exceptions import DeadlineExceeded + + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + table.read_rows_sharded([query1, query2], operation_timeout=0.005) + assert len(e.value.exceptions) == 2 + for sub_exc in e.value.exceptions: + assert isinstance(sub_exc.__cause__, DeadlineExceeded) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = operation.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_sharded_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + authorized_view.read_rows_sharded([query1, query2]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + failed_op = next( + (op for op in handler.completed_operations if op.final_status.name != "OK") + ) + success_op = next( + (op for op in handler.completed_operations if op.final_status.name == "OK") + ) + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert failed_op.cluster_id == next(iter(cluster_config.keys())) + assert ( + failed_op.zone + == cluster_config[failed_op.cluster_id].location.split("/")[-1] + ) + failed_attempt = failed_op.completed_attempts[0] + assert failed_attempt.end_status.name == "PERMISSION_DENIED" + assert ( + failed_attempt.gfe_latency_ns >= 0 + and failed_attempt.gfe_latency_ns < failed_op.duration_ns + ) + assert success_op.final_status.name == "OK" + assert success_op.op_type.value == "ReadRows" + assert success_op.is_streaming is True + assert len(success_op.completed_attempts) == 1 + success_attempt = success_op.completed_attempts[0] + assert success_attempt.end_status.name == "OK" + + def test_read_rows_sharded_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + failed_op = next( + (op for op in handler.completed_operations if op.final_status.name != "OK") + ) + success_op = next( + (op for op in handler.completed_operations if op.final_status.name == "OK") + ) + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert success_op.final_status.name == "OK" + assert len(success_op.completed_attempts) == 2 + attempt = failed_op.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + retried_attempt = success_op.completed_attempts[0] + assert retried_attempt.end_status.name == "ABORTED" + success_attempt = success_op.completed_attempts[-1] + assert success_attempt.end_status.name == "OK" + + def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + handler.clear() + table.bulk_mutate_rows([bulk_mutation]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + def test_bulk_mutate_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup): + table.bulk_mutate_rows([entry], operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_bulk_mutate_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup): + authorized_view.bulk_mutate_rows([entry]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_bulk_mutate_rows_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out + + For bulk_mutate, the rpc returns success, with failures returned in the response. + For this reason, We expect the attempts to be marked as successful, even though + the underlying mutation is retried""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup) as e: + authorized_view.bulk_mutate_rows( + [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 + ) + assert len(e.value.exceptions) == 1 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] + + def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + row_key2, mutation2 = temp_rows.create_row_and_mutation( + table, new_value=new_value2 + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + handler.clear() + with table.mutations_batcher() as batcher: + batcher.append(bulk_mutation) + batcher.append(bulk_mutation2) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert ( + operation.flow_throttling_time_ns > 0 + and operation.flow_throttling_time_ns < operation.duration_ns + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + def test_mutate_rows_batcher_failure_with_retries( + self, table, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + with table.mutations_batcher(batch_retryable_errors=[Aborted]) as batcher: + batcher.append(entry) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + with pytest.raises(MutationsExceptionGroup): + with table.mutations_batcher(batch_operation_timeout=0.001) as batcher: + batcher.append(entry) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_mutate_rows_batcher_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + with pytest.raises(MutationsExceptionGroup) as e: + with authorized_view.mutations_batcher() as batcher: + batcher.append(entry) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + @pytest.mark.skipif( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't suport cluster_config", From 01f1739a0742098a4dfb0eeda088884243cfdaa2 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 15:31:51 -0700 Subject: [PATCH 07/11] fixed mtls tests --- .../tests/unit/data/_async/test_client.py | 24 +++++--- .../unit/data/_sync_autogen/test_client.py | 55 +++++++++++-------- 2 files changed, 47 insertions(+), 32 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index ad2ae9c8bb42..e116f4a4624b 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -90,7 +90,7 @@ def _get_target_class(): return CrossSync.DataClient @classmethod - def _make_client(cls, *args, use_emulator=True, **kwargs): + def _make_client(cls, *args, use_emulator=True, use_mtls="auto", **kwargs): import os env_mask = {} @@ -105,6 +105,8 @@ def _make_client(cls, *args, use_emulator=True, **kwargs): # set some default values kwargs["credentials"] = kwargs.get("credentials", AnonymousCredentials()) kwargs["project"] = kwargs.get("project", "project-id") + if use_mtls is not None: + env_mask["GOOGLE_API_USE_MTLS_ENDPOINT"] = use_mtls with mock.patch.dict(os.environ, env_mask): return cls._get_target_class()(*args, **kwargs) @@ -1046,13 +1048,17 @@ def test_client_ctor_sync(self): assert client._channel_refresh_task is None @CrossSync.pytest - async def test_default_universe_domain(self): + @pytest.mark.paramtrize( + "use_mtls, expected_domain", + [("never", "googleapis.com"), ("always", "mtls.googleapis.com")] + ) + async def test_default_universe_domain(self, use_mtls, expected_domain): """ When not passed, universe_domain should default to googleapis.com """ - async with self._make_client(project="project-id", credentials=None) as client: - assert client.universe_domain == "googleapis.com" - assert client.api_endpoint == "bigtable.googleapis.com" + async with self._make_client(project="project-id", credentials=None, use_mtls=expected_domain) as client: + assert client.universe_domain == expected_domain + assert client.api_endpoint == f"bigtable.{expected_domain}" @CrossSync.pytest async def test_custom_universe_domain(self): @@ -1064,6 +1070,7 @@ async def test_custom_universe_domain(self): client_options=options, use_emulator=True, credentials=None, + use_mtls="never", ) as client: assert client.universe_domain == universe_domain assert client.api_endpoint == f"bigtable.{universe_domain}" @@ -1077,7 +1084,6 @@ async def test_configured_universe_domain_matches_GDU(self): project="project_id", client_options=options, credentials=None ) as client: assert client.universe_domain == "googleapis.com" - assert client.api_endpoint == "bigtable.googleapis.com" @CrossSync.pytest async def test_credential_universe_domain_matches_GDU(self): @@ -1086,13 +1092,12 @@ async def test_credential_universe_domain_matches_GDU(self): creds._universe_domain = "googleapis.com" async with self._make_client(project="project_id", credentials=creds) as client: assert client.universe_domain == "googleapis.com" - assert client.api_endpoint == "bigtable.googleapis.com" @CrossSync.pytest async def test_anomynous_credential_universe_domain(self): """Anomynopus credentials should use default universe domain""" creds = AnonymousCredentials() - async with self._make_client(project="project_id", credentials=creds) as client: + async with self._make_client(project="project_id", credentials=creds, use_mtls="never") as client: assert client.universe_domain == "googleapis.com" assert client.api_endpoint == "bigtable.googleapis.com" @@ -1111,6 +1116,7 @@ async def test_configured_universe_domain_mismatched_credentials(self): client_options=options, use_emulator=False, credentials=creds, + use_mtls="never", ) err_msg = ( f"The configured universe domain ({universe_domain}) does " @@ -1131,7 +1137,7 @@ async def test_configured_universe_domain_matches_credentials(self): creds = AnonymousCredentials() creds._universe_domain = universe_domain async with self._make_client( - project="project_id", credentials=creds, client_options=options + project="project_id", credentials=creds, client_options=options, use_mtls="never" ) as client: assert client.universe_domain == universe_domain assert client.api_endpoint == f"bigtable.{universe_domain}" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index 343507cfc854..9d8bbd7f6834 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -65,7 +65,7 @@ def _get_target_class(): return CrossSync._Sync_Impl.DataClient @classmethod - def _make_client(cls, *args, use_emulator=True, **kwargs): + def _make_client(cls, *args, use_emulator=True, use_mtls="auto", **kwargs): import os env_mask = {} @@ -77,6 +77,8 @@ def _make_client(cls, *args, use_emulator=True, **kwargs): else: kwargs["credentials"] = kwargs.get("credentials", AnonymousCredentials()) kwargs["project"] = kwargs.get("project", "project-id") + if use_mtls is not None: + env_mask["GOOGLE_API_USE_MTLS_ENDPOINT"] = use_mtls with mock.patch.dict(os.environ, env_mask): return cls._get_target_class()(*args, **kwargs) @@ -206,10 +208,8 @@ def test__start_background_channel_refresh(self): def test__ping_and_warm_instances(self): """test ping and warm with mocked asyncio.gather""" client_mock = mock.Mock() - client_mock._execute_ping_and_warms = ( - lambda *args: self._get_target_class()._execute_ping_and_warms( - client_mock, *args - ) + client_mock._execute_ping_and_warms = lambda *args: ( + self._get_target_class()._execute_ping_and_warms(client_mock, *args) ) with mock.patch.object( CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock() @@ -252,10 +252,8 @@ def test__ping_and_warm_instances(self): def test__ping_and_warm_single_instance(self): """should be able to call ping and warm with single instance""" client_mock = mock.Mock() - client_mock._execute_ping_and_warms = ( - lambda *args: self._get_target_class()._execute_ping_and_warms( - client_mock, *args - ) + client_mock._execute_ping_and_warms = lambda *args: ( + self._get_target_class()._execute_ping_and_warms(client_mock, *args) ) with mock.patch.object( CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock() @@ -844,11 +842,17 @@ def test_context_manager(self): close_mock.assert_called_once() true_close() - def test_default_universe_domain(self): + @pytest.mark.parametrize( + "use_mtls, expected_domain", + [("never", "googleapis.com"), ("always", "mtls.googleapis.com")], + ) + def test_default_universe_domain(self, use_mtls, expected_domain): """When not passed, universe_domain should default to googleapis.com""" - with self._make_client(project="project-id", credentials=None) as client: + with self._make_client( + project="project-id", credentials=None, use_mtls=use_mtls + ) as client: assert client.universe_domain == "googleapis.com" - assert client.api_endpoint == "bigtable.googleapis.com" + assert client.api_endpoint == f"bigtable.{expected_domain}" def test_custom_universe_domain(self): """test with a customized universe domain value and emulator enabled""" @@ -859,6 +863,7 @@ def test_custom_universe_domain(self): client_options=options, use_emulator=True, credentials=None, + use_mtls="never", ) as client: assert client.universe_domain == universe_domain assert client.api_endpoint == f"bigtable.{universe_domain}" @@ -871,7 +876,6 @@ def test_configured_universe_domain_matches_GDU(self): project="project_id", client_options=options, credentials=None ) as client: assert client.universe_domain == "googleapis.com" - assert client.api_endpoint == "bigtable.googleapis.com" def test_credential_universe_domain_matches_GDU(self): """Test with credentials""" @@ -879,12 +883,13 @@ def test_credential_universe_domain_matches_GDU(self): creds._universe_domain = "googleapis.com" with self._make_client(project="project_id", credentials=creds) as client: assert client.universe_domain == "googleapis.com" - assert client.api_endpoint == "bigtable.googleapis.com" def test_anomynous_credential_universe_domain(self): """Anomynopus credentials should use default universe domain""" creds = AnonymousCredentials() - with self._make_client(project="project_id", credentials=creds) as client: + with self._make_client( + project="project_id", credentials=creds, use_mtls="never" + ) as client: assert client.universe_domain == "googleapis.com" assert client.api_endpoint == "bigtable.googleapis.com" @@ -901,6 +906,7 @@ def test_configured_universe_domain_mismatched_credentials(self): client_options=options, use_emulator=False, credentials=creds, + use_mtls="never", ) err_msg = f"The configured universe domain ({universe_domain}) does not match the universe domain found in the credentials ({creds.universe_domain}). If you haven't configured the universe domain explicitly, `googleapis.com` is the default." assert exc.value.args[0] == err_msg @@ -913,7 +919,10 @@ def test_configured_universe_domain_matches_credentials(self): creds = AnonymousCredentials() creds._universe_domain = universe_domain with self._make_client( - project="project_id", credentials=creds, client_options=options + project="project_id", + credentials=creds, + client_options=options, + use_mtls="never", ) as client: assert client.universe_domain == universe_domain assert client.api_endpoint == f"bigtable.{universe_domain}" @@ -1313,11 +1322,11 @@ def _make_client(self, *args, **kwargs): def _make_table(self, *args, **kwargs): client_mock = mock.Mock() - client_mock._register_instance.side_effect = ( - lambda *args, **kwargs: CrossSync._Sync_Impl.yield_to_event_loop() + client_mock._register_instance.side_effect = lambda *args, **kwargs: ( + CrossSync._Sync_Impl.yield_to_event_loop() ) - client_mock._remove_instance_registration.side_effect = ( - lambda *args, **kwargs: CrossSync._Sync_Impl.yield_to_event_loop() + client_mock._remove_instance_registration.side_effect = lambda *args, **kwargs: ( + CrossSync._Sync_Impl.yield_to_event_loop() ) kwargs["instance_id"] = kwargs.get( "instance_id", args[0] if args else "instance" @@ -1779,9 +1788,8 @@ def test_read_rows_sharded_multiple_queries(self): with mock.patch.object( table.client._gapic_client, "read_rows" ) as read_rows: - read_rows.side_effect = ( - lambda *args, - **kwargs: CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( + read_rows.side_effect = lambda *args, **kwargs: ( + CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( [ CrossSync._Sync_Impl.TestReadRows._make_chunk(row_key=k) for k in args[0].rows.row_keys @@ -2870,6 +2878,7 @@ def prepare_mock(self, client): yield prepare_mock def _make_gapic_stream(self, sample_list: list["ExecuteQueryResponse" | Exception]): + class MockStream: def __init__(self, sample_list): self.sample_list = sample_list From 0ab1b05187850f6cd186ab5fbee4519af7b9125c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 15:39:39 -0700 Subject: [PATCH 08/11] copied over changes --- .../bigtable/data/_async/_mutate_rows.py | 82 +- .../cloud/bigtable/data/_async/_read_rows.py | 244 ++-- .../cloud/bigtable/data/_async/client.py | 36 +- .../bigtable/data/_async/mutations_batcher.py | 34 +- .../tests/system/data/test_metrics_async.py | 1214 +++++++++++++++++ .../unit/data/_async/test__mutate_rows.py | 22 +- .../tests/unit/data/_async/test__read_rows.py | 16 +- .../tests/unit/data/_async/test_client.py | 115 +- .../data/_async/test_mutations_batcher.py | 28 +- .../data/_async/test_read_rows_acceptance.py | 37 +- 10 files changed, 1591 insertions(+), 237 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py index 6efb9e5f25be..3c93b01bdf72 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py @@ -22,10 +22,8 @@ import google.cloud.bigtable.data.exceptions as bt_exceptions import google.cloud.bigtable_v2.types.bigtable as types_pb from google.cloud.bigtable.data._cross_sync import CrossSync -from google.cloud.bigtable.data._helpers import ( - _attempt_timeout_generator, - _retry_exception_factory, -) +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +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 ( @@ -35,6 +33,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.data._async.client import ( # type: ignore @@ -72,6 +71,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 @@ -82,6 +83,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 @@ -101,13 +103,12 @@ def __init__( # Entry level errors bt_exceptions._MutateRowsIncomplete, ) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - self._operation = lambda: CrossSync.retry_target( - self._run_attempt, - self.is_retryable, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, + 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( @@ -116,6 +117,8 @@ def __init__( self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries] self.remaining_indices = list(range(len(self.mutations))) self.errors: dict[int, list[Exception]] = {} + # set up metrics + self._operation_metric = metric @CrossSync.convert async def start(self): @@ -125,34 +128,35 @@ async def start(self): Raises: MutationsExceptionGroup: if any mutations failed """ - try: - # trigger mutate_rows - await self._operation() - except Exception as exc: - # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations - incomplete_indices = self.remaining_indices.copy() - for idx in incomplete_indices: - self._handle_entry_error(idx, exc) - finally: - # raise exception detailing incomplete mutations - all_errors: list[Exception] = [] - for idx, exc_list in self.errors.items(): - if len(exc_list) == 0: - raise core_exceptions.ClientError( - f"Mutation {idx} failed with no associated errors" + with self._operation_metric: + try: + # trigger mutate_rows + await self._operation() + except Exception as exc: + # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations + incomplete_indices = self.remaining_indices.copy() + for idx in incomplete_indices: + self._handle_entry_error(idx, exc) + finally: + # raise exception detailing incomplete mutations + all_errors: list[Exception] = [] + for idx, exc_list in self.errors.items(): + if len(exc_list) == 0: + raise core_exceptions.ClientError( + f"Mutation {idx} failed with no associated errors" + ) + elif len(exc_list) == 1: + cause_exc = exc_list[0] + else: + cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) + entry = self.mutations[idx].entry + all_errors.append( + bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) + ) + if all_errors: + raise bt_exceptions.MutationsExceptionGroup( + all_errors, len(self.mutations) ) - elif len(exc_list) == 1: - cause_exc = exc_list[0] - else: - cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) - entry = self.mutations[idx].entry - all_errors.append( - bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) - ) - if all_errors: - raise bt_exceptions.MutationsExceptionGroup( - all_errors, len(self.mutations) - ) @CrossSync.convert async def _run_attempt(self): @@ -164,6 +168,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/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py index f8e203bc10b3..5e03e6afe714 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_read_rows.py @@ -16,15 +16,16 @@ from __future__ import annotations from typing import TYPE_CHECKING, Sequence +import time +from grpc import StatusCode 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 from google.cloud.bigtable.data._helpers import ( _attempt_timeout_generator, - _retry_exception_factory, ) +from google.cloud.bigtable.data._metrics import tracked_retry from google.cloud.bigtable.data.exceptions import ( InvalidChunk, _ResetRow, @@ -38,6 +39,8 @@ from google.cloud.bigtable_v2.types import RowSet as RowSetPB 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, @@ -68,6 +71,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 """ @@ -79,6 +83,7 @@ class _ReadRowsOperationAsync: "_predicate", "_last_yielded_row_key", "_remaining_count", + "_operation_metric", ) def __init__( @@ -87,6 +92,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( @@ -105,6 +111,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]: """ @@ -113,12 +120,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, - exponential_sleep_generator(0.01, 60, multiplier=2), - self.operation_timeout, - exception_factory=_retry_exception_factory, + 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]: @@ -131,6 +138,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 @@ -208,12 +216,11 @@ async def chunk_stream( raise InvalidChunk("emit count exceeds row limit") current_key = None - @staticmethod @CrossSync.convert( replace_symbols={"__aiter__": "__iter__", "__anext__": "__next__"}, ) async def merge_rows( - chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None, + self, chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None ) -> CrossSync.Iterable[Row]: """ Merge chunks into rows @@ -223,108 +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: + self._operation_metric.end_with_success() return - row_key = c.row_key - - if not row_key: - raise InvalidChunk("first row chunk is missing key") - - cells = [] - - # shared per cell storage - family: str | None = None - qualifier: bytes | None = None - - try: - # for each cell - while True: - if c.reset_row: - raise _ResetRow(c) - k = c.row_key - f = c.family_name.value - q = c.qualifier.value if c.HasField("qualifier") else None - if k and k != row_key: - raise InvalidChunk("unexpected new row key") - if f: - family = f - if q is not None: - qualifier = q - else: - raise InvalidChunk("new family without qualifier") - elif family is None: - raise InvalidChunk("missing family") - elif q is not None: - if family is None: - raise InvalidChunk("new qualifier without family") - qualifier = q - elif qualifier is None: - raise InvalidChunk("missing qualifier") - - ts = c.timestamp_micros - labels = c.labels if c.labels else [] - value = c.value - - # merge split cells - if c.value_size > 0: - buffer = [value] - while c.value_size > 0: - # throws when premature end - c = await it.__anext__() - - t = c.timestamp_micros - cl = c.labels - k = c.row_key - if ( - c.HasField("family_name") - and c.family_name.value != family - ): - raise InvalidChunk("family changed mid cell") - if ( - c.HasField("qualifier") - and c.qualifier.value != qualifier - ): - raise InvalidChunk("qualifier changed mid cell") - if t and t != ts: - raise InvalidChunk("timestamp changed mid cell") - if cl and cl != labels: - raise InvalidChunk("labels changed mid cell") - if k and k != row_key: - raise InvalidChunk("row key changed mid cell") - - if c.reset_row: - raise _ResetRow(c) - buffer.append(c.value) - value = b"".join(buffer) - cells.append( - Cell(value, row_key, family, qualifier, ts, list(labels)) - ) - if c.commit_row: - yield Row(row_key, cells) - break + it = chunks.__aiter__() + # For each row + while True: + try: c = await it.__anext__() - except _ResetRow as e: - c = e.chunk - if ( - c.row_key - or c.HasField("family_name") - or c.HasField("qualifier") - or c.timestamp_micros - or c.labels - or c.value - ): - raise InvalidChunk("reset row with data") - continue - except CrossSync.StopIteration: - raise InvalidChunk("premature end of stream") + except CrossSync.StopIteration: + # stream complete + self._operation_metric.end_with_success() + return + row_key = c.row_key + + if not row_key: + raise InvalidChunk("first row chunk is missing key") + + cells = [] + + # shared per cell storage + family: str | None = None + qualifier: bytes | None = None + + try: + # for each cell + while True: + if c.reset_row: + raise _ResetRow(c) + k = c.row_key + f = c.family_name.value + q = c.qualifier.value if c.HasField("qualifier") else None + if k and k != row_key: + raise InvalidChunk("unexpected new row key") + if f: + family = f + if q is not None: + qualifier = q + else: + raise InvalidChunk("new family without qualifier") + elif family is None: + raise InvalidChunk("missing family") + elif q is not None: + if family is None: + raise InvalidChunk("new qualifier without family") + qualifier = q + elif qualifier is None: + raise InvalidChunk("missing qualifier") + + ts = c.timestamp_micros + labels = c.labels if c.labels else [] + value = c.value + + # merge split cells + if c.value_size > 0: + buffer = [value] + while c.value_size > 0: + # throws when premature end + c = await it.__anext__() + + t = c.timestamp_micros + cl = c.labels + k = c.row_key + if ( + c.HasField("family_name") + and c.family_name.value != family + ): + raise InvalidChunk("family changed mid cell") + if ( + c.HasField("qualifier") + and c.qualifier.value != qualifier + ): + raise InvalidChunk("qualifier changed mid cell") + if t and t != ts: + raise InvalidChunk("timestamp changed mid cell") + if cl and cl != labels: + raise InvalidChunk("labels changed mid cell") + if k and k != row_key: + raise InvalidChunk("row key changed mid cell") + + if c.reset_row: + raise _ResetRow(c) + buffer.append(c.value) + value = b"".join(buffer) + cells.append( + Cell(value, row_key, family, qualifier, ts, list(labels)) + ) + if c.commit_row: + block_time = time.monotonic_ns() + yield Row(row_key, cells) + # most metric operations use setters, but this one updates + # the value directly to avoid extra overhead + if self._operation_metric.active_attempt is not None: + self._operation_metric.active_attempt.application_blocking_time_ns += ( # type: ignore + time.monotonic_ns() - block_time + ) + break + c = await it.__anext__() + except _ResetRow as e: + c = e.chunk + if ( + c.row_key + or c.HasField("family_name") + or c.HasField("qualifier") + or c.timestamp_micros + or c.labels + or c.value + ): + raise InvalidChunk("reset row with data") + continue + except CrossSync.StopIteration: + raise InvalidChunk("premature end of stream") + except GeneratorExit as close_exception: + # handle aclose() + self._operation_metric.end_with_status(StatusCode.CANCELLED) + raise close_exception + except Exception as generic_exception: + # handle exceptions in retry wrapper + raise generic_exception @staticmethod def _revise_request_rowset( diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index b2c13521240f..1e3e6e3d202f 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -1131,6 +1131,9 @@ async def read_rows_stream( self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=True + ), retryable_exceptions=retryable_excs, ) return row_merger.start_operation() @@ -1223,15 +1226,28 @@ async def read_row( if row_key is None: raise ValueError("row_key must be string or bytes") query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1) - results = await self.read_rows( + + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self + ) + retryable_excs = _get_retryable_errors(retryable_errors, self) + + row_merger = CrossSync._ReadRowsOperation( query, + self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, - retryable_errors=retryable_errors, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=False + ), + retryable_exceptions=retryable_excs, ) - if len(results) == 0: + results_generator = row_merger.start_operation() + try: + results = [a async for a in results_generator] + return results[0] + except IndexError: return None - return results[0] @CrossSync.convert async def read_rows_sharded( @@ -1370,20 +1386,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( @@ -1643,6 +1656,7 @@ async def bulk_mutate_rows( mutation_entries, operation_timeout, attempt_timeout, + metric=self._metrics.create_operation(OperationType.BULK_MUTATE_ROWS), retryable_exceptions=retryable_excs, ) await operation.start() diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py index 405983393ee7..2466f87234fc 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py @@ -16,6 +16,7 @@ import atexit import concurrent.futures +import time import warnings from collections import deque from typing import TYPE_CHECKING, Sequence, cast @@ -30,6 +31,8 @@ FailedMutationEntryError, MutationsExceptionGroup, ) +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, Mutation, @@ -37,6 +40,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 ( @@ -181,6 +185,24 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry] ) yield mutations[start_idx:end_idx] + @CrossSync.convert(replace_symbols={"__anext__": "__next__"}) + async def add_to_flow_with_metrics( + self, + mutations: RowMutationEntry | list[RowMutationEntry], + metrics_controller: BigtableClientSideMetricsController, + ): + inner_generator = self.add_to_flow(mutations) + while True: + # start a new metric + metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS) + flow_start_time = time.monotonic_ns() + try: + value = await inner_generator.__anext__() + except CrossSync.StopIteration: + return + metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time + yield value, metric + @CrossSync.convert_class(sync_name="MutationsBatcher") class MutationsBatcherAsync: @@ -357,9 +379,14 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]): """ # flush new entries in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = [] - async for batch in self._flow_control.add_to_flow(new_entries): + async for batch, metric in self._flow_control.add_to_flow_with_metrics( + new_entries, self._target._metrics + ): batch_task = CrossSync.create_task( - self._execute_mutate_rows, batch, sync_executor=self._sync_rpc_executor + self._execute_mutate_rows, + batch, + metric, + sync_executor=self._sync_rpc_executor, ) in_process_requests.append(batch_task) # wait for all inflight requests to complete @@ -370,7 +397,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 @@ -391,6 +418,7 @@ async def _execute_mutate_rows( batch, operation_timeout=self._operation_timeout, attempt_timeout=self._attempt_timeout, + metric=metric, retryable_exceptions=self._retryable_errors, ) await operation.start() diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py index 49cd9ba2b744..c6969461f76f 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py @@ -26,6 +26,7 @@ CompletedOperationMetric, CompletedAttemptMetric, ) +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable_v2.types import ResponseParams from google.cloud.environment_vars import BIGTABLE_EMULATOR @@ -1441,6 +1442,1219 @@ async def test_mutate_rows_batcher_failure_unauthorized( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't suport cluster_config", ) + @CrossSync.pytest + async def test_read_rows(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + row_list = await table.read_rows(ReadRowsQuery()) + assert len(row_list) == 2 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + await table.read_rows(ReadRowsQuery(), operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.read_rows( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + # full table scan + generator = await table.read_rows_stream(ReadRowsQuery()) + row_list = [r async for r in generator] + assert len(row_list) == 2 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + @CrossSync.convert(replace_symbols={"__anext__": "__next__", "aclose": "close"}) + async def test_read_rows_stream_failure_closed( + self, table, temp_rows, handler, error_injector + ): + """ + Test how metrics collection handles closed generator + """ + await temp_rows.add_row(b"row_key_1") + await temp_rows.add_row(b"row_key_2") + handler.clear() + generator = await table.read_rows_stream(ReadRowsQuery()) + await generator.__anext__() + await generator.aclose() + with pytest.raises(CrossSync.StopIteration): + await generator.__anext__() + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "CANCELLED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "CANCELLED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) + with pytest.raises(PermissionDenied): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + generator = await table.read_rows_stream( + ReadRowsQuery(), operation_timeout=0.001 + ) + with pytest.raises(GoogleAPICallError): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_stream_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = await authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + [_ async for _ in generator] + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_stream_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = await authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + [_ async for _ in generator] + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + @CrossSync.pytest + async def test_read_rows_stream_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + generator = await table.read_rows_stream( + ReadRowsQuery(), retryable_errors=[Aborted] + ) + with pytest.raises(PermissionDenied): + [_ async for _ in generator] + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 2 + # validate retried attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + # validate final attempt + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + @CrossSync.pytest + async def test_read_row(self, table, temp_rows, handler, cluster_config): + await temp_rows.add_row(b"row_key_1") + handler.clear() + await table.read_row(b"row_key_1") + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns > 0 + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_row_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + await table.read_row(b"row_key_1", retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_row_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + await temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + await table.read_row(b"row_key_1", operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + await authorized_view.read_row( + b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + await temp_rows.add_row(b"c") + await temp_rows.add_row(b"d") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) + handler.clear() + row_list = await table.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + # validate counts + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # validate operations + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + attempt = operation.completed_attempts[0] + assert attempt in handler.completed_attempts + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + # validate attempt + assert isinstance(attempt, CompletedAttemptMetric) + assert ( + attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + ) + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 + and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_read_rows_sharded_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + + error_injector.push(self._make_exception(StatusCode.ABORTED)) + await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + # validate operations + for op in handler.completed_operations: + assert op.final_status.name == "OK" + assert op.op_type.value == "ReadRows" + assert op.is_streaming is True + # validate attempts + assert ( + len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) + == 2 + ) + assert ( + len( + [ + a + for a in handler.completed_attempts + if a.end_status.name == "ABORTED" + ] + ) + == 1 + ) + + @CrossSync.pytest + async def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.api_core.exceptions import DeadlineExceeded + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2], operation_timeout=0.005) + assert len(e.value.exceptions) == 2 + for sub_exc in e.value.exceptions: + assert isinstance(sub_exc.__cause__, DeadlineExceeded) + # both shards should fail + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # validate operations + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = operation.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_read_rows_sharded_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await authorized_view.read_rows_sharded([query1, query2]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + # one shard will fail, the other will succeed + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + # sort operations by status + failed_op = next( + op for op in handler.completed_operations if op.final_status.name != "OK" + ) + success_op = next( + op for op in handler.completed_operations if op.final_status.name == "OK" + ) + # validate failed operation + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert failed_op.cluster_id == next(iter(cluster_config.keys())) + assert ( + failed_op.zone + == cluster_config[failed_op.cluster_id].location.split("/")[-1] + ) + # validate failed attempt + failed_attempt = failed_op.completed_attempts[0] + assert failed_attempt.end_status.name == "PERMISSION_DENIED" + assert ( + failed_attempt.gfe_latency_ns >= 0 + and failed_attempt.gfe_latency_ns < failed_op.duration_ns + ) + # validate successful operation + assert success_op.final_status.name == "OK" + assert success_op.op_type.value == "ReadRows" + assert success_op.is_streaming is True + assert len(success_op.completed_attempts) == 1 + # validate successful attempt + success_attempt = success_op.completed_attempts[0] + assert success_attempt.end_status.name == "OK" + + @CrossSync.pytest + async def test_read_rows_sharded_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc stream + """ + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + + await temp_rows.add_row(b"a") + await temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + # one shard will fail, the other will succeed + # the failing shard will have one retry + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + # sort operations by status + failed_op = next( + op for op in handler.completed_operations if op.final_status.name != "OK" + ) + success_op = next( + op for op in handler.completed_operations if op.final_status.name == "OK" + ) + # validate failed operation + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + # validate successful operation + assert success_op.final_status.name == "OK" + assert len(success_op.completed_attempts) == 2 + # validate failed attempt + attempt = failed_op.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + # validate retried attempt + retried_attempt = success_op.completed_attempts[0] + assert retried_attempt.end_status.name == "ABORTED" + # validate successful attempt + success_attempt = success_op.completed_attempts[-1] + assert success_attempt.end_status.name == "OK" + + @CrossSync.pytest + async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + + handler.clear() + await table.bulk_mutate_rows([bulk_mutation]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert operation.flow_throttling_time_ns == 0 + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + await table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup): + await table.bulk_mutate_rows([entry], operation_timeout=0.001) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup): + await authorized_view.bulk_mutate_rows([entry]) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + @CrossSync.pytest + async def test_bulk_mutate_rows_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """ + retry unauthorized request multiple times before timing out + + For bulk_mutate, the rpc returns success, with failures returned in the response. + For this reason, We expect the attempts to be marked as successful, even though + the underlying mutation is retried + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + handler.clear() + with pytest.raises(MutationsExceptionGroup) as e: + await authorized_view.bulk_mutate_rows( + [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 + ) + assert len(e.value.exceptions) == 1 + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempts + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] + + @CrossSync.pytest + async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = await temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + row_key2, mutation2 = await temp_rows.create_row_and_mutation( + table, new_value=new_value2 + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + + handler.clear() + async with table.mutations_batcher() as batcher: + await batcher.append(bulk_mutation) + await batcher.append(bulk_mutation2) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # bacher expects to cancel staged operation on close + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1e9 + assert ( + operation.first_response_latency_ns is None + ) # populated for read_rows only + assert ( + operation.flow_throttling_time_ns > 0 + and operation.flow_throttling_time_ns < operation.duration_ns + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_with_retries( + self, table, handler, error_injector + ): + """ + Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + async with table.mutations_batcher( + batch_retryable_errors=[Aborted] + ) as batcher: + await batcher.append(entry) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + # validate operation + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + # validate attempts + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): + """ + Test failure in gapic layer by passing very low timeout + + No grpc headers expected + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup): + async with table.mutations_batcher( + batch_operation_timeout=0.001 + ) as batcher: + await batcher.append(entry) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + @CrossSync.pytest + async def test_mutate_rows_batcher_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """ + Test failure in backend by accessing an unauthorized family + """ + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + + with pytest.raises(MutationsExceptionGroup) as e: + async with authorized_view.mutations_batcher() as batcher: + await batcher.append(entry) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + # validate counts + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + # validate operation + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + # validate attempt + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + @CrossSync.pytest async def test_mutate_row(self, table, temp_rows, handler, cluster_config): row_key = b"mutate" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py index 82f234350a8c..f2d4f8eef42f 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test__mutate_rows.py @@ -14,6 +14,7 @@ import pytest from google.api_core.exceptions import DeadlineExceeded, Forbidden +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.rpc import status_pb2 from google.cloud.bigtable.data._cross_sync import CrossSync @@ -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): @@ -87,6 +91,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, @@ -94,6 +99,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 @@ -113,6 +119,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): """ @@ -136,6 +143,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 @@ -149,6 +157,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() @@ -156,7 +165,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 @@ -170,6 +179,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 @@ -178,7 +188,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: @@ -202,6 +212,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") @@ -214,7 +225,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: @@ -238,6 +249,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") @@ -254,6 +266,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() @@ -273,6 +286,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( @@ -284,7 +298,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/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py index 7fad973c43a3..29e2d3e84c0d 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test__read_rows.py @@ -15,6 +15,7 @@ import pytest from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data._metrics import ActiveOperationMetric # try/except added for compatibility with python < 3.8 try: @@ -59,6 +60,7 @@ def test_ctor(self): expected_operation_timeout = 42 expected_request_timeout = 44 time_gen_mock = mock.Mock() + expected_metric = mock.Mock() subpath = "_async" if CrossSync.is_async else "_sync_autogen" with mock.patch( f"google.cloud.bigtable.data.{subpath}._read_rows._attempt_timeout_generator", @@ -69,6 +71,7 @@ def test_ctor(self): table, operation_timeout=expected_operation_timeout, attempt_timeout=expected_request_timeout, + metric=expected_metric, ) assert time_gen_mock.call_count == 1 time_gen_mock.assert_called_once_with( @@ -81,6 +84,7 @@ def test_ctor(self): assert instance.request.table_name == "test_table" assert instance.request.app_profile_id == table.app_profile_id assert instance.request.rows_limit == row_limit + assert instance._operation_metric == expected_metric @pytest.mark.parametrize( "in_keys,last_key,expected", @@ -269,7 +273,9 @@ async def mock_stream(): table = mock.Mock() table._request_path = {"table_name": "table_name"} table.app_profile_id = "app_profile_id" - instance = self._make_one(query, table, 10, 10) + instance = self._make_one( + query, table, 10, 10, ActiveOperationMetric("READ_ROWS") + ) assert instance._remaining_count == start_limit # read emit_num rows async for val in instance.chunk_stream(awaitable_stream()): @@ -308,7 +314,9 @@ async def mock_stream(): table = mock.Mock() table._request_path = {"table_name": "table_name"} table.app_profile_id = "app_profile_id" - instance = self._make_one(query, table, 10, 10) + instance = self._make_one( + query, table, 10, 10, ActiveOperationMetric("READ_ROWS") + ) assert instance._remaining_count == start_limit with pytest.raises(InvalidChunk) as e: # read emit_num rows @@ -334,7 +342,9 @@ async def mock_stream(): with mock.patch.object( self._get_target_class(), "_read_rows_attempt" ) as mock_attempt: - instance = self._make_one(mock.Mock(), mock.Mock(), 1, 1) + instance = self._make_one( + mock.Mock(), mock.Mock(), 1, 1, ActiveOperationMetric("READ_ROWS") + ) wrapped_gen = mock_stream() mock_attempt.return_value = wrapped_gen gen = instance.start_operation() diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index ad2ae9c8bb42..83ad3b2b4a02 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -1402,15 +1402,9 @@ async def test_customizable_retryable_errors( predicate_builder_mock.assert_called_once_with( *expected_retryables, *extra_retryables ) - # output of if_exception_type should be sent in to retry constructor retry_call_kwargs = retry_fn_mock.call_args_list[0].kwargs - # check for predicate passed as kwarg - if "predicate" in retry_call_kwargs: - assert retry_call_kwargs["predicate"] is expected_predicate - else: - # check for predicate passed as arg - retry_call_args = retry_fn_mock.call_args_list[0].args - assert retry_call_args[1] is expected_predicate + # output of if_exception_type should be sent in to retry constructor + assert retry_call_kwargs["predicate"] is expected_predicate @pytest.mark.parametrize( "fn_name,fn_args,gapic_fn", @@ -1983,9 +1977,21 @@ async def test_read_row(self): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() expected_result = object() - read_rows.side_effect = lambda *args, **kwargs: [expected_result] + + if CrossSync.is_async: + + async def mock_generator(): + yield expected_result + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = [expected_result] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 row = await table.read_row( @@ -1994,16 +2000,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): @@ -2011,14 +2018,24 @@ async def test_read_row_w_filter(self): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() expected_result = object() - read_rows.side_effect = lambda *args, **kwargs: [expected_result] + + if CrossSync.is_async: + + async def mock_generator(): + yield expected_result + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = [expected_result] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 - mock_filter = mock.Mock() - expected_filter = {"filter": "mock filter"} - mock_filter._to_dict.return_value = expected_filter + expected_filter = mock.Mock() row = await table.read_row( row_key, operation_timeout=expected_op_timeout, @@ -2026,11 +2043,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] @@ -2044,9 +2061,21 @@ async def test_read_row_no_response(self): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: - # return no rows - read_rows.side_effect = lambda *args, **kwargs: [] + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() + + if CrossSync.is_async: + + async def mock_generator(): + if False: + yield + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = [] + mock_op_constructor.return_value = mock_op expected_op_timeout = 8 expected_req_timeout = 4 result = await table.read_row( @@ -2055,8 +2084,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) @@ -2079,22 +2108,36 @@ async def test_row_exists(self, return_value, expected_result): async with self._make_client() as client: table = client.get_table("instance", "table") row_key = b"test_1" - with mock.patch.object(table, "read_rows") as read_rows: - # return no rows - read_rows.side_effect = lambda *args, **kwargs: return_value - expected_op_timeout = 1 - expected_req_timeout = 2 + with mock.patch.object( + CrossSync, "_ReadRowsOperation" + ) as mock_op_constructor: + mock_op = mock.Mock() + if CrossSync.is_async: + + async def mock_generator(): + for item in return_value: + yield item + + mock_op.start_operation.return_value = mock_generator() + else: + mock_op.start_operation.return_value = return_value + mock_op_constructor.return_value = mock_op + expected_op_timeout = 2 + expected_req_timeout = 1 result = await table.row_exists( row_key, operation_timeout=expected_op_timeout, attempt_timeout=expected_req_timeout, ) assert expected_result == result - assert read_rows.call_count == 1 - args, kwargs = read_rows.call_args_list[0] + assert mock_op_constructor.call_count == 1 + args, kwargs = mock_op_constructor.call_args_list[0] assert kwargs["operation_timeout"] == expected_op_timeout assert kwargs["attempt_timeout"] == expected_req_timeout - assert isinstance(args[0], ReadRowsQuery) + query = args[0] + assert isinstance(query, ReadRowsQuery) + assert query.row_keys == [row_key] + assert query.limit == 1 expected_filter = { "chain": { "filters": [ @@ -2103,10 +2146,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/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py index 75de7c281332..e067febfc0fe 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py @@ -306,6 +306,9 @@ def _get_target_class(self): def _make_one(self, table=None, **kwargs): from google.api_core.exceptions import DeadlineExceeded, ServiceUnavailable + from google.cloud.bigtable.data._metrics import ( + BigtableClientSideMetricsController, + ) if table is None: table = mock.Mock() @@ -317,6 +320,7 @@ def _make_one(self, table=None, **kwargs): DeadlineExceeded, ServiceUnavailable, ) + table._metrics = BigtableClientSideMetricsController([]) return self._get_target_class()(table, **kwargs) @@ -935,14 +939,16 @@ async def test__execute_mutate_rows(self): table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: batch = [self._make_mutation()] - result = await instance._execute_mutate_rows(batch) + expected_metric = mock.Mock() + result = await instance._execute_mutate_rows(batch, expected_metric) assert start_operation.call_count == 1 args, kwargs = mutate_rows.call_args assert args[0] == table.client._gapic_client assert args[1] == table assert args[2] == batch - kwargs["operation_timeout"] == 17 - kwargs["attempt_timeout"] == 13 + assert kwargs["operation_timeout"] == 17 + assert kwargs["attempt_timeout"] == 13 + assert kwargs["metric"] == expected_metric assert result == [] @CrossSync.pytest @@ -963,7 +969,7 @@ async def test__execute_mutate_rows_returns_errors(self): table.default_mutate_rows_retryable_errors = () async with self._make_one(table) as instance: batch = [self._make_mutation()] - result = await instance._execute_mutate_rows(batch) + result = await instance._execute_mutate_rows(batch, mock.Mock()) assert len(result) == 2 assert result[0] == err1 assert result[1] == err2 @@ -1093,7 +1099,9 @@ async def test_timeout_args_passed(self): assert instance._operation_timeout == expected_operation_timeout assert instance._attempt_timeout == expected_attempt_timeout # make simulated gapic call - await instance._execute_mutate_rows([self._make_mutation()]) + await instance._execute_mutate_rows( + [self._make_mutation()], mock.Mock() + ) assert mutate_rows.call_count == 1 kwargs = mutate_rows.call_args[1] assert kwargs["operation_timeout"] == expected_operation_timeout @@ -1192,6 +1200,8 @@ async def test_customizable_retryable_errors( Test that retryable functions support user-configurable arguments, and that the configured retryables are passed down to the gapic layer. """ + from google.cloud.bigtable.data._metrics import ActiveOperationMetric + with mock.patch.object( google.api_core.retry, "if_exception_type" ) as predicate_builder_mock: @@ -1207,14 +1217,16 @@ 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 ) - 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/packages/google-cloud-bigtable/tests/unit/data/_async/test_read_rows_acceptance.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_read_rows_acceptance.py index d69b776bfe42..5e8dbcba123a 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_read_rows_acceptance.py +++ b/packages/google-cloud-bigtable/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_v2 import ReadRowsResponse +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from ...v2_client.test_row_merger import ReadRowsTest, TestFile @@ -36,8 +37,11 @@ class TestReadRowsAcceptanceAsync: @staticmethod @CrossSync.convert - def _get_operation_class(): - return CrossSync._ReadRowsOperation + def _make_operation(): + metric = ActiveOperationMetric("READ_ROWS") + op = CrossSync._ReadRowsOperation(mock.Mock(), mock.Mock(), 5, 5, metric) + op._remaining_count = None + return op @staticmethod @CrossSync.convert @@ -80,13 +84,8 @@ async def _process_chunks(self, *chunks): async def _row_stream(): yield ReadRowsResponse(chunks=chunks) - instance = mock.Mock() - instance._remaining_count = None - instance._last_yielded_row_key = None - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_row_stream()) - ) - merger = self._get_operation_class().merge_rows(chunker) + chunker = self._make_operation().chunk_stream(self._coro_wrapper(_row_stream())) + merger = self._make_operation().merge_rows(chunker) results = [] async for row in merger: results.append(row) @@ -103,13 +102,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( @@ -196,13 +192,10 @@ async def test_out_of_order_rows(self): async def _row_stream(): yield ReadRowsResponse(last_scanned_row_key=b"a") - instance = mock.Mock() - instance._remaining_count = None - instance._last_yielded_row_key = b"b" - chunker = self._get_operation_class().chunk_stream( - instance, self._coro_wrapper(_row_stream()) - ) - merger = self._get_operation_class().merge_rows(chunker) + op = self._make_operation() + op._last_yielded_row_key = b"b" + chunker = op.chunk_stream(self._coro_wrapper(_row_stream())) + merger = self._make_operation().merge_rows(chunker) with pytest.raises(InvalidChunk): async for _ in merger: pass From 83cbcf0e0a3ee4ff442fc3ba44977dc44da3dfd3 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 15:40:51 -0700 Subject: [PATCH 09/11] regenerated sync files --- .../data/_sync_autogen/_mutate_rows.py | 74 +- .../bigtable/data/_sync_autogen/_read_rows.py | 215 ++-- .../bigtable/data/_sync_autogen/client.py | 33 +- .../data/_sync_autogen/mutations_batcher.py | 31 +- .../tests/system/data/test_metrics_autogen.py | 1017 +++++++++++++++++ .../data/_sync_autogen/test__mutate_rows.py | 25 +- .../data/_sync_autogen/test__read_rows.py | 16 +- .../unit/data/_sync_autogen/test_client.py | 76 +- .../_sync_autogen/test_mutations_batcher.py | 27 +- .../test_read_rows_acceptance.py | 39 +- 10 files changed, 1335 insertions(+), 218 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py index c1c508a526f2..40e19dd85847 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py @@ -25,16 +25,15 @@ import google.cloud.bigtable.data.exceptions as bt_exceptions import google.cloud.bigtable_v2.types.bigtable as types_pb from google.cloud.bigtable.data._cross_sync import CrossSync -from google.cloud.bigtable.data._helpers import ( - _attempt_timeout_generator, - _retry_exception_factory, -) +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data._metrics import tracked_retry from google.cloud.bigtable.data.mutations import ( _MUTATE_ROWS_REQUEST_MUTATION_LIMIT, _EntryWithProto, ) if TYPE_CHECKING: + from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data._sync_autogen.client import ( _DataApiTarget as TargetType, ) @@ -61,6 +60,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__( @@ -70,6 +71,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)) @@ -82,13 +84,12 @@ def __init__( self.is_retryable = retries.if_exception_type( *retryable_exceptions, bt_exceptions._MutateRowsIncomplete ) - sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60) - self._operation = lambda: CrossSync._Sync_Impl.retry_target( - self._run_attempt, - self.is_retryable, - sleep_generator, - operation_timeout, - exception_factory=_retry_exception_factory, + 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 @@ -96,37 +97,39 @@ def __init__( self.mutations = [_EntryWithProto(m, m._to_pb()) for m in mutation_entries] self.remaining_indices = list(range(len(self.mutations))) self.errors: dict[int, list[Exception]] = {} + self._operation_metric = metric def start(self): """Start the operation, and run until completion Raises: MutationsExceptionGroup: if any mutations failed""" - try: - self._operation() - except Exception as exc: - incomplete_indices = self.remaining_indices.copy() - for idx in incomplete_indices: - self._handle_entry_error(idx, exc) - finally: - all_errors: list[Exception] = [] - for idx, exc_list in self.errors.items(): - if len(exc_list) == 0: - raise core_exceptions.ClientError( - f"Mutation {idx} failed with no associated errors" + with self._operation_metric: + try: + self._operation() + except Exception as exc: + incomplete_indices = self.remaining_indices.copy() + for idx in incomplete_indices: + self._handle_entry_error(idx, exc) + finally: + all_errors: list[Exception] = [] + for idx, exc_list in self.errors.items(): + if len(exc_list) == 0: + raise core_exceptions.ClientError( + f"Mutation {idx} failed with no associated errors" + ) + elif len(exc_list) == 1: + cause_exc = exc_list[0] + else: + cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) + entry = self.mutations[idx].entry + all_errors.append( + bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) + ) + if all_errors: + raise bt_exceptions.MutationsExceptionGroup( + all_errors, len(self.mutations) ) - elif len(exc_list) == 1: - cause_exc = exc_list[0] - else: - cause_exc = bt_exceptions.RetryExceptionGroup(exc_list) - entry = self.mutations[idx].entry - all_errors.append( - bt_exceptions.FailedMutationEntryError(idx, entry, cause_exc) - ) - if all_errors: - raise bt_exceptions.MutationsExceptionGroup( - all_errors, len(self.mutations) - ) def _run_attempt(self): """Run a single attempt of the mutate_rows rpc. @@ -135,6 +138,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 for req_idx, orig_idx in enumerate(self.remaining_indices) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py index a74374988161..b9c2a4bf8cb6 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_read_rows.py @@ -18,16 +18,15 @@ from __future__ import annotations +import time from typing import TYPE_CHECKING, Sequence from google.api_core import retry as retries -from google.api_core.retry import exponential_sleep_generator +from grpc import StatusCode from google.cloud.bigtable.data._cross_sync import CrossSync -from google.cloud.bigtable.data._helpers import ( - _attempt_timeout_generator, - _retry_exception_factory, -) +from google.cloud.bigtable.data._helpers import _attempt_timeout_generator +from google.cloud.bigtable.data._metrics import tracked_retry from google.cloud.bigtable.data.exceptions import ( InvalidChunk, _ResetRow, @@ -41,6 +40,7 @@ from google.cloud.bigtable_v2.types import RowSet as RowSetPB if TYPE_CHECKING: + from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data._sync_autogen.client import ( _DataApiTarget as TargetType, ) @@ -63,6 +63,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 """ @@ -74,6 +75,7 @@ class _ReadRowsOperation: "_predicate", "_last_yielded_row_key", "_remaining_count", + "_operation_metric", ) def __init__( @@ -82,6 +84,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( @@ -98,18 +101,19 @@ 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""" - return CrossSync._Sync_Impl.retry_target_stream( - self._read_rows_attempt, - self._predicate, - exponential_sleep_generator(0.01, 60, multiplier=2), - self.operation_timeout, - exception_factory=_retry_exception_factory, + 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]: @@ -120,6 +124,7 @@ def _read_rows_attempt(self) -> CrossSync._Sync_Impl.Iterable[Row]: Yields: Row: The next row in the stream""" + self._operation_metric.start_attempt() if self._last_yielded_row_key is not None: try: self.request.rows = self._revise_request_rowset( @@ -181,9 +186,8 @@ def chunk_stream( raise InvalidChunk("emit count exceeds row limit") current_key = None - @staticmethod def merge_rows( - chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None, + self, chunks: CrossSync._Sync_Impl.Iterable[ReadRowsResponsePB.CellChunk] | None ) -> CrossSync._Sync_Impl.Iterable[Row]: """Merge chunks into rows @@ -191,94 +195,107 @@ def merge_rows( chunks: the chunk stream to merge Yields: Row: the next row in the stream""" - if chunks is None: - return - it = chunks.__iter__() - while True: - try: - c = it.__next__() - except CrossSync._Sync_Impl.StopIteration: + try: + if chunks is None: + self._operation_metric.end_with_success() return - row_key = c.row_key - if not row_key: - raise InvalidChunk("first row chunk is missing key") - cells = [] - family: str | None = None - qualifier: bytes | None = None - try: - while True: - if c.reset_row: - raise _ResetRow(c) - k = c.row_key - f = c.family_name.value - q = c.qualifier.value if c.HasField("qualifier") else None - if k and k != row_key: - raise InvalidChunk("unexpected new row key") - if f: - family = f - if q is not None: - qualifier = q - else: - raise InvalidChunk("new family without qualifier") - elif family is None: - raise InvalidChunk("missing family") - elif q is not None: - if family is None: - raise InvalidChunk("new qualifier without family") - qualifier = q - elif qualifier is None: - raise InvalidChunk("missing qualifier") - ts = c.timestamp_micros - labels = c.labels if c.labels else [] - value = c.value - if c.value_size > 0: - buffer = [value] - while c.value_size > 0: - c = it.__next__() - t = c.timestamp_micros - cl = c.labels - k = c.row_key - if ( - c.HasField("family_name") - and c.family_name.value != family - ): - raise InvalidChunk("family changed mid cell") - if ( - c.HasField("qualifier") - and c.qualifier.value != qualifier - ): - raise InvalidChunk("qualifier changed mid cell") - if t and t != ts: - raise InvalidChunk("timestamp changed mid cell") - if cl and cl != labels: - raise InvalidChunk("labels changed mid cell") - if k and k != row_key: - raise InvalidChunk("row key changed mid cell") - if c.reset_row: - raise _ResetRow(c) - buffer.append(c.value) - value = b"".join(buffer) - cells.append( - Cell(value, row_key, family, qualifier, ts, list(labels)) - ) - if c.commit_row: - yield Row(row_key, cells) - break + it = chunks.__iter__() + while True: + try: c = it.__next__() - except _ResetRow as e: - c = e.chunk - if ( - c.row_key - or c.HasField("family_name") - or c.HasField("qualifier") - or c.timestamp_micros - or c.labels - or c.value - ): - raise InvalidChunk("reset row with data") - continue - except CrossSync._Sync_Impl.StopIteration: - raise InvalidChunk("premature end of stream") + except CrossSync._Sync_Impl.StopIteration: + self._operation_metric.end_with_success() + return + row_key = c.row_key + if not row_key: + raise InvalidChunk("first row chunk is missing key") + cells = [] + family: str | None = None + qualifier: bytes | None = None + try: + while True: + if c.reset_row: + raise _ResetRow(c) + k = c.row_key + f = c.family_name.value + q = c.qualifier.value if c.HasField("qualifier") else None + if k and k != row_key: + raise InvalidChunk("unexpected new row key") + if f: + family = f + if q is not None: + qualifier = q + else: + raise InvalidChunk("new family without qualifier") + elif family is None: + raise InvalidChunk("missing family") + elif q is not None: + if family is None: + raise InvalidChunk("new qualifier without family") + qualifier = q + elif qualifier is None: + raise InvalidChunk("missing qualifier") + ts = c.timestamp_micros + labels = c.labels if c.labels else [] + value = c.value + if c.value_size > 0: + buffer = [value] + while c.value_size > 0: + c = it.__next__() + t = c.timestamp_micros + cl = c.labels + k = c.row_key + if ( + c.HasField("family_name") + and c.family_name.value != family + ): + raise InvalidChunk("family changed mid cell") + if ( + c.HasField("qualifier") + and c.qualifier.value != qualifier + ): + raise InvalidChunk("qualifier changed mid cell") + if t and t != ts: + raise InvalidChunk("timestamp changed mid cell") + if cl and cl != labels: + raise InvalidChunk("labels changed mid cell") + if k and k != row_key: + raise InvalidChunk("row key changed mid cell") + if c.reset_row: + raise _ResetRow(c) + buffer.append(c.value) + value = b"".join(buffer) + cells.append( + Cell(value, row_key, family, qualifier, ts, list(labels)) + ) + if c.commit_row: + block_time = time.monotonic_ns() + yield Row(row_key, cells) + if self._operation_metric.active_attempt is not None: + self._operation_metric.active_attempt.application_blocking_time_ns += ( + time.monotonic_ns() - block_time + ) + break + c = it.__next__() + except _ResetRow as e: + c = e.chunk + if ( + c.row_key + or c.HasField("family_name") + or c.HasField("qualifier") + or c.timestamp_micros + or c.labels + or c.value + ): + raise InvalidChunk("reset row with data") + continue + except CrossSync._Sync_Impl.StopIteration: + raise InvalidChunk("premature end of stream") + except GeneratorExit as close_exception: + self._operation_metric.end_with_status(StatusCode.CANCELLED) + raise close_exception + except Exception as generic_exception: + raise generic_exception @staticmethod def _revise_request_rowset(row_set: RowSetPB, last_seen_row_key: bytes) -> RowSetPB: diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index 9dc118de0289..2d90a7b990c2 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -906,6 +906,9 @@ def read_rows_stream( self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=True + ), retryable_exceptions=retryable_excs, ) return row_merger.start_operation() @@ -992,15 +995,26 @@ def read_row( if row_key is None: raise ValueError("row_key must be string or bytes") query = ReadRowsQuery(row_keys=row_key, row_filter=row_filter, limit=1) - results = self.read_rows( + operation_timeout, attempt_timeout = _get_timeouts( + operation_timeout, attempt_timeout, self + ) + retryable_excs = _get_retryable_errors(retryable_errors, self) + row_merger = CrossSync._Sync_Impl._ReadRowsOperation( query, + self, operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, - retryable_errors=retryable_errors, + metric=self._metrics.create_operation( + OperationType.READ_ROWS, is_streaming=False + ), + retryable_exceptions=retryable_excs, ) - if len(results) == 0: + results_generator = row_merger.start_operation() + try: + results = [a for a in results_generator] + return results[0] + except IndexError: return None - return results[0] def read_rows_sharded( self, @@ -1122,19 +1136,17 @@ def row_exists( will be chained with a RetryExceptionGroup containing GoogleAPIError exceptions 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, @@ -1372,6 +1384,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() diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py index 5be449a49d4a..107c2cbf591b 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py @@ -19,6 +19,7 @@ import atexit import concurrent.futures +import time import warnings from collections import deque from typing import TYPE_CHECKING, Sequence, cast @@ -29,6 +30,7 @@ _get_retryable_errors, _get_timeouts, ) +from google.cloud.bigtable.data._metrics import ActiveOperationMetric, OperationType from google.cloud.bigtable.data.exceptions import ( FailedMutationEntryError, MutationsExceptionGroup, @@ -39,6 +41,7 @@ ) if TYPE_CHECKING: + from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController from google.cloud.bigtable.data._sync_autogen.client import ( _DataApiTarget as TargetType, ) @@ -154,6 +157,22 @@ def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]): ) yield mutations[start_idx:end_idx] + def add_to_flow_with_metrics( + self, + mutations: RowMutationEntry | list[RowMutationEntry], + metrics_controller: BigtableClientSideMetricsController, + ): + inner_generator = self.add_to_flow(mutations) + while True: + metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS) + flow_start_time = time.monotonic_ns() + try: + value = inner_generator.__next__() + except CrossSync._Sync_Impl.StopIteration: + return + metric.flow_throttling_time_ns = time.monotonic_ns() - flow_start_time + yield (value, metric) + class MutationsBatcher: """ @@ -309,9 +328,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) @@ -319,7 +343,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 @@ -338,6 +362,7 @@ def _execute_mutate_rows( batch, operation_timeout=self._operation_timeout, attempt_timeout=self._attempt_timeout, + metric=metric, retryable_exceptions=self._retryable_errors, ) operation.start() diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py index 4f315652b15c..6b28328eaab9 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py @@ -34,6 +34,7 @@ CompletedOperationMetric, ) from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable_v2.types import ResponseParams from . import TEST_FAMILY, SystemTestRunner @@ -1207,6 +1208,1022 @@ def test_mutate_rows_batcher_failure_unauthorized( bool(os.environ.get(BIGTABLE_EMULATOR)), reason="emulator doesn't suport cluster_config", ) + def test_read_rows(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + row_list = table.read_rows(ReadRowsQuery()) + assert len(row_list) == 2 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_rows_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + table.read_rows(ReadRowsQuery(), operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + authorized_view.read_rows( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery()) + row_list = [r for r in generator] + assert len(row_list) == 2 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_rows_stream_failure_closed( + self, table, temp_rows, handler, error_injector + ): + """Test how metrics collection handles closed generator""" + temp_rows.add_row(b"row_key_1") + temp_rows.add_row(b"row_key_2") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery()) + generator.__next__() + generator.close() + with pytest.raises(CrossSync._Sync_Impl.StopIteration): + generator.__next__() + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "CANCELLED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "CANCELLED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + with pytest.raises(PermissionDenied): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + generator = table.read_rows_stream(ReadRowsQuery(), operation_timeout=0.001) + with pytest.raises(GoogleAPICallError): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_stream_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + ) + [_ for _ in generator] + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_stream_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + generator = authorized_view.read_rows_stream( + ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), + retryable_errors=[PermissionDenied], + operation_timeout=0.5, + ) + [_ for _ in generator] + assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] + + def test_read_rows_stream_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + temp_rows.add_row(b"row_key_1") + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) + with pytest.raises(PermissionDenied): + [_ for _ in generator] + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 2 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 2 + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "ABORTED" + final_attempt = handler.completed_attempts[-1] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + + def test_read_row(self, table, temp_rows, handler, cluster_config): + temp_rows.add_row(b"row_key_1") + handler.clear() + table.read_row(b"row_key_1") + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns > 0 + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_row_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + temp_rows.add_row(b"row_key_1") + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(PermissionDenied): + table.read_row(b"row_key_1", retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_read_row_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + temp_rows.add_row(b"row_key_1") + handler.clear() + with pytest.raises(GoogleAPICallError): + table.read_row(b"row_key_1", operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_row_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + with pytest.raises(GoogleAPICallError) as e: + authorized_view.read_row( + b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") + ) + assert e.value.grpc_status_code.name == "PERMISSION_DENIED" + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + temp_rows.add_row(b"c") + temp_rows.add_row(b"d") + query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) + query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) + handler.clear() + row_list = table.read_rows_sharded([query1, query2]) + assert len(row_list) == 4 + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is True + assert operation.op_type.value == "ReadRows" + assert len(operation.completed_attempts) == 1 + attempt = operation.completed_attempts[0] + assert attempt in handler.completed_attempts + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert ( + operation.first_response_latency_ns is not None + and operation.first_response_latency_ns < operation.duration_ns + ) + assert operation.flow_throttling_time_ns == 0 + assert isinstance(attempt, CompletedAttemptMetric) + assert ( + attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + ) + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 + and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert ( + attempt.application_blocking_time_ns > 0 + and attempt.application_blocking_time_ns < operation.duration_ns + ) + + def test_read_rows_sharded_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors""" + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.push(self._make_exception(StatusCode.ABORTED)) + table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + for op in handler.completed_operations: + assert op.final_status.name == "OK" + assert op.op_type.value == "ReadRows" + assert op.is_streaming is True + assert ( + len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) + == 2 + ) + assert ( + len( + [ + a + for a in handler.completed_attempts + if a.end_status.name == "ABORTED" + ] + ) + == 1 + ) + + def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.api_core.exceptions import DeadlineExceeded + + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + table.read_rows_sharded([query1, query2], operation_timeout=0.005) + assert len(e.value.exceptions) == 2 + for sub_exc in e.value.exceptions: + assert isinstance(sub_exc.__cause__, DeadlineExceeded) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + for operation in handler.completed_operations: + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "ReadRows" + assert operation.is_streaming is True + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = operation.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_read_rows_sharded_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter + + query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) + query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) + handler.clear() + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + authorized_view.read_rows_sharded([query1, query2]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 2 + failed_op = next( + (op for op in handler.completed_operations if op.final_status.name != "OK") + ) + success_op = next( + (op for op in handler.completed_operations if op.final_status.name == "OK") + ) + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert failed_op.cluster_id == next(iter(cluster_config.keys())) + assert ( + failed_op.zone + == cluster_config[failed_op.cluster_id].location.split("/")[-1] + ) + failed_attempt = failed_op.completed_attempts[0] + assert failed_attempt.end_status.name == "PERMISSION_DENIED" + assert ( + failed_attempt.gfe_latency_ns >= 0 + and failed_attempt.gfe_latency_ns < failed_op.duration_ns + ) + assert success_op.final_status.name == "OK" + assert success_op.op_type.value == "ReadRows" + assert success_op.is_streaming is True + assert len(success_op.completed_attempts) == 1 + success_attempt = success_op.completed_attempts[0] + assert success_attempt.end_status.name == "OK" + + def test_read_rows_sharded_failure_mid_stream( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc stream""" + from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup + from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery + + temp_rows.add_row(b"a") + temp_rows.add_row(b"b") + query1 = ReadRowsQuery(row_keys=[b"a"]) + query2 = ReadRowsQuery(row_keys=[b"b"]) + handler.clear() + error_injector.fail_mid_stream = True + error_injector.push(self._make_exception(StatusCode.ABORTED)) + error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) + with pytest.raises(ShardedReadRowsExceptionGroup) as e: + table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) + assert len(handler.completed_operations) == 2 + assert len(handler.completed_attempts) == 3 + failed_op = next( + (op for op in handler.completed_operations if op.final_status.name != "OK") + ) + success_op = next( + (op for op in handler.completed_operations if op.final_status.name == "OK") + ) + assert failed_op.final_status.name == "PERMISSION_DENIED" + assert failed_op.op_type.value == "ReadRows" + assert failed_op.is_streaming is True + assert len(failed_op.completed_attempts) == 1 + assert success_op.final_status.name == "OK" + assert len(success_op.completed_attempts) == 2 + attempt = failed_op.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + retried_attempt = success_op.completed_attempts[0] + assert retried_attempt.end_status.name == "ABORTED" + success_attempt = success_op.completed_attempts[-1] + assert success_attempt.end_status.name == "OK" + + def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value = uuid.uuid4().hex.encode() + row_key, mutation = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + handler.clear() + table.bulk_mutate_rows([bulk_mutation]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert operation.flow_throttling_time_ns == 0 + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + def test_bulk_mutate_rows_failure_with_retries( + self, table, temp_rows, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert isinstance(final_attempt, CompletedAttemptMetric) + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup): + table.bulk_mutate_rows([entry], operation_timeout=0.001) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_bulk_mutate_rows_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup): + authorized_view.bulk_mutate_rows([entry]) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + + def test_bulk_mutate_rows_failure_unauthorized_with_retries( + self, handler, authorized_view, cluster_config + ): + """retry unauthorized request multiple times before timing out + + For bulk_mutate, the rpc returns success, with failures returned in the response. + For this reason, We expect the attempts to be marked as successful, even though + the underlying mutation is retried""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + handler.clear() + with pytest.raises(MutationsExceptionGroup) as e: + authorized_view.bulk_mutate_rows( + [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 + ) + assert len(e.value.exceptions) == 1 + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) > 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) > 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + for attempt in handler.completed_attempts: + assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] + + def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): + from google.cloud.bigtable.data.mutations import RowMutationEntry + + new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] + row_key, mutation = temp_rows.create_row_and_mutation( + table, new_value=new_value + ) + row_key2, mutation2 = temp_rows.create_row_and_mutation( + table, new_value=new_value2 + ) + bulk_mutation = RowMutationEntry(row_key, [mutation]) + bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) + handler.clear() + with table.mutations_batcher() as batcher: + batcher.append(bulk_mutation) + batcher.append(bulk_mutation2) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.value[0] == 0 + assert operation.is_streaming is False + assert operation.op_type.value == "MutateRows" + assert len(operation.completed_attempts) == 1 + assert operation.completed_attempts[0] == handler.completed_attempts[0] + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 + assert operation.first_response_latency_ns is None + assert ( + operation.flow_throttling_time_ns > 0 + and operation.flow_throttling_time_ns < operation.duration_ns + ) + attempt = handler.completed_attempts[0] + assert isinstance(attempt, CompletedAttemptMetric) + assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns + assert attempt.end_status.value[0] == 0 + assert attempt.backoff_before_attempt_ns == 0 + assert ( + attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns + ) + assert attempt.application_blocking_time_ns == 0 + + def test_mutate_rows_batcher_failure_with_retries( + self, table, handler, error_injector + ): + """Test failure in grpc layer by injecting errors into an interceptor + with retryable errors, then a terminal one""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + assert entry.is_idempotent() + handler.clear() + expected_zone = "my_zone" + expected_cluster = "my_cluster" + num_retryable = 2 + for i in range(num_retryable): + error_injector.push( + self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) + ) + error_injector.push( + self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) + ) + with pytest.raises(MutationsExceptionGroup): + with table.mutations_batcher(batch_retryable_errors=[Aborted]) as batcher: + batcher.append(entry) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == num_retryable + 1 + operation = handler.completed_operations[0] + assert isinstance(operation, CompletedOperationMetric) + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == num_retryable + 1 + assert operation.cluster_id == expected_cluster + assert operation.zone == expected_zone + for i in range(num_retryable): + attempt = handler.completed_attempts[i] + assert attempt.end_status.name == "ABORTED" + assert attempt.gfe_latency_ns is None + final_attempt = handler.completed_attempts[num_retryable] + assert final_attempt.end_status.name == "PERMISSION_DENIED" + assert final_attempt.gfe_latency_ns is None + + def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): + """Test failure in gapic layer by passing very low timeout + + No grpc headers expected""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell(TEST_FAMILY, b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + with pytest.raises(MutationsExceptionGroup): + with table.mutations_batcher(batch_operation_timeout=0.001) as batcher: + batcher.append(entry) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "DEADLINE_EXCEEDED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == "" + assert operation.zone == "global" + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "DEADLINE_EXCEEDED" + assert attempt.gfe_latency_ns is None + + def test_mutate_rows_batcher_failure_unauthorized( + self, handler, authorized_view, cluster_config + ): + """Test failure in backend by accessing an unauthorized family""" + from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup + from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell + + row_key = b"row_key_1" + mutation = SetCell("unauthorized", b"q", b"v") + entry = RowMutationEntry(row_key, [mutation]) + with pytest.raises(MutationsExceptionGroup) as e: + with authorized_view.mutations_batcher() as batcher: + batcher.append(entry) + assert len(e.value.exceptions) == 1 + assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) + assert ( + e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" + ) + assert len(handler.completed_operations) == 1 + assert len(handler.completed_attempts) == 1 + operation = handler.completed_operations[0] + assert operation.final_status.name == "PERMISSION_DENIED" + assert operation.op_type.value == "MutateRows" + assert operation.is_streaming is False + assert len(operation.completed_attempts) == 1 + assert operation.cluster_id == next(iter(cluster_config.keys())) + assert ( + operation.zone + == cluster_config[operation.cluster_id].location.split("/")[-1] + ) + attempt = handler.completed_attempts[0] + assert attempt.end_status.name == "PERMISSION_DENIED" + assert ( + attempt.gfe_latency_ns >= 0 + and attempt.gfe_latency_ns < operation.duration_ns + ) + def test_mutate_row(self, table, temp_rows, handler, cluster_config): row_key = b"mutate" new_value = uuid.uuid4().hex.encode() diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py index 3bce82f1e50f..1c0aa52ca0b8 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__mutate_rows.py @@ -20,6 +20,7 @@ from google.rpc import status_pb2 from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data.mutations import DeleteAllFromRow, RowMutationEntry from google.cloud.bigtable_v2.types import MutateRowsResponse @@ -44,6 +45,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): @@ -83,6 +87,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, @@ -90,6 +95,7 @@ def test_ctor(self): entries, operation_timeout, attempt_timeout, + metric, retryable_exceptions, ) assert client.mutate_rows.call_count == 0 @@ -105,6 +111,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""" @@ -119,7 +126,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 ) @@ -129,6 +138,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() @@ -136,7 +146,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 @@ -147,6 +157,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 @@ -155,7 +166,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: @@ -176,6 +187,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") @@ -186,7 +198,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: @@ -203,6 +215,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") @@ -217,6 +230,7 @@ def test_mutate_rows_exception_retryable_eventually_pass(self, exc_type): entries, operation_timeout, operation_timeout, + metric, retryable_exceptions=(exc_type,), ) instance.start() @@ -233,6 +247,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( @@ -242,7 +257,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/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py index f16a523e862d..9acc44209be4 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test__read_rows.py @@ -18,6 +18,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 @@ -54,6 +55,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", @@ -64,6 +66,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( @@ -76,6 +79,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/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index 343507cfc854..77c1fd8fc6ce 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -1124,11 +1124,7 @@ def test_customizable_retryable_errors( *expected_retryables, *extra_retryables ) retry_call_kwargs = retry_fn_mock.call_args_list[0].kwargs - if "predicate" in retry_call_kwargs: - assert retry_call_kwargs["predicate"] is expected_predicate - else: - retry_call_args = retry_fn_mock.call_args_list[0].args - 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", @@ -1643,9 +1639,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( @@ -1654,30 +1654,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, @@ -1685,11 +1688,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] @@ -1702,8 +1705,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( @@ -1712,8 +1719,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) @@ -1731,21 +1738,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": [ @@ -1754,10 +1768,6 @@ def test_row_exists(self, return_value, expected_result): ] } } - query = args[0] - assert query.row_keys == [row_key] - assert query.row_ranges == [] - assert query.limit == 1 assert query.filter._to_dict() == expected_filter diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py index f6568448ff8c..bf54a44ad35b 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py @@ -258,6 +258,10 @@ def _get_target_class(self): def _make_one(self, table=None, **kwargs): from google.api_core.exceptions import DeadlineExceeded, ServiceUnavailable + from google.cloud.bigtable.data._metrics import ( + BigtableClientSideMetricsController, + ) + if table is None: table = mock.Mock() table._request_path = {"table_name": "table"} @@ -268,6 +272,7 @@ def _make_one(self, table=None, **kwargs): DeadlineExceeded, ServiceUnavailable, ) + table._metrics = BigtableClientSideMetricsController([]) return self._get_target_class()(table, **kwargs) @staticmethod @@ -816,14 +821,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): @@ -845,7 +852,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 @@ -953,7 +960,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 @@ -1039,6 +1046,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: @@ -1056,12 +1065,14 @@ 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 ) - 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""" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py index 29332e712d35..77c55ce0183b 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_read_rows_acceptance.py @@ -24,6 +24,7 @@ import pytest from google.cloud.bigtable.data._cross_sync import CrossSync +from google.cloud.bigtable.data._metrics import ActiveOperationMetric from google.cloud.bigtable.data.exceptions import InvalidChunk from google.cloud.bigtable.data.row import Row from google.cloud.bigtable_v2 import ReadRowsResponse @@ -33,8 +34,13 @@ 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(): @@ -72,13 +78,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) @@ -94,13 +95,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( @@ -183,13 +181,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 7ca6e11a67fb1b108a6c51a6942ddc2b7dc14b95 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 15:48:04 -0700 Subject: [PATCH 10/11] removed duplicate files --- .../tests/system/data/test_metrics_async.py | 1217 ----------------- .../tests/system/data/test_metrics_autogen.py | 1020 -------------- 2 files changed, 2237 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py index c6969461f76f..e23336691830 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py @@ -1438,1223 +1438,6 @@ async def test_mutate_rows_batcher_failure_unauthorized( and attempt.gfe_latency_ns < operation.duration_ns ) - @pytest.mark.skipif( - bool(os.environ.get(BIGTABLE_EMULATOR)), - reason="emulator doesn't suport cluster_config", - ) - @CrossSync.pytest - async def test_read_rows(self, table, temp_rows, handler, cluster_config): - await temp_rows.add_row(b"row_key_1") - await temp_rows.add_row(b"row_key_2") - handler.clear() - row_list = await table.read_rows(ReadRowsQuery()) - assert len(row_list) == 2 - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is True - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1e9 - assert ( - operation.first_response_latency_ns is not None - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - @CrossSync.pytest - async def test_read_rows_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """ - Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one - """ - await temp_rows.add_row(b"row_key_1") - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(PermissionDenied): - await table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - # validate attempts - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_rows_failure_timeout(self, table, temp_rows, handler): - """ - Test failure in gapic layer by passing very low timeout - - No grpc headers expected - """ - await temp_rows.add_row(b"row_key_1") - handler.clear() - with pytest.raises(GoogleAPICallError): - await table.read_rows(ReadRowsQuery(), operation_timeout=0.001) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_rows_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """ - Test failure in backend by accessing an unauthorized family - """ - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - await authorized_view.read_rows( - ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) - ) - assert e.value.grpc_status_code.name == "PERMISSION_DENIED" - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - @CrossSync.pytest - async def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): - await temp_rows.add_row(b"row_key_1") - await temp_rows.add_row(b"row_key_2") - handler.clear() - # full table scan - generator = await table.read_rows_stream(ReadRowsQuery()) - row_list = [r async for r in generator] - assert len(row_list) == 2 - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is True - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1e9 - assert ( - operation.first_response_latency_ns is not None - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - @CrossSync.pytest - @CrossSync.convert(replace_symbols={"__anext__": "__next__", "aclose": "close"}) - async def test_read_rows_stream_failure_closed( - self, table, temp_rows, handler, error_injector - ): - """ - Test how metrics collection handles closed generator - """ - await temp_rows.add_row(b"row_key_1") - await temp_rows.add_row(b"row_key_2") - handler.clear() - generator = await table.read_rows_stream(ReadRowsQuery()) - await generator.__anext__() - await generator.aclose() - with pytest.raises(CrossSync.StopIteration): - await generator.__anext__() - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert operation.final_status.name == "CANCELLED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - # validate attempt - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "CANCELLED" - assert attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_rows_stream_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """ - Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one - """ - await temp_rows.add_row(b"row_key_1") - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - generator = await table.read_rows_stream( - ReadRowsQuery(), retryable_errors=[Aborted] - ) - with pytest.raises(PermissionDenied): - [_ async for _ in generator] - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - # validate attempts - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): - """ - Test failure in gapic layer by passing very low timeout - - No grpc headers expected - """ - await temp_rows.add_row(b"row_key_1") - handler.clear() - generator = await table.read_rows_stream( - ReadRowsQuery(), operation_timeout=0.001 - ) - with pytest.raises(GoogleAPICallError): - [_ async for _ in generator] - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_rows_stream_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """ - Test failure in backend by accessing an unauthorized family - """ - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - generator = await authorized_view.read_rows_stream( - ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) - ) - [_ async for _ in generator] - assert e.value.grpc_status_code.name == "PERMISSION_DENIED" - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - @CrossSync.pytest - async def test_read_rows_stream_failure_unauthorized_with_retries( - self, handler, authorized_view, cluster_config - ): - """ - retry unauthorized request multiple times before timing out - """ - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - generator = await authorized_view.read_rows_stream( - ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), - retryable_errors=[PermissionDenied], - operation_timeout=0.5, - ) - [_ async for _ in generator] - assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) > 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) > 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - # validate attempts - for attempt in handler.completed_attempts: - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] - - @CrossSync.pytest - async def test_read_rows_stream_failure_mid_stream( - self, table, temp_rows, handler, error_injector - ): - """ - Test failure in grpc stream - """ - await temp_rows.add_row(b"row_key_1") - handler.clear() - error_injector.fail_mid_stream = True - error_injector.push(self._make_exception(StatusCode.ABORTED)) - error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) - generator = await table.read_rows_stream( - ReadRowsQuery(), retryable_errors=[Aborted] - ) - with pytest.raises(PermissionDenied): - [_ async for _ in generator] - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 2 - # validate operation - operation = handler.completed_operations[0] - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 2 - # validate retried attempt - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "ABORTED" - # validate final attempt - final_attempt = handler.completed_attempts[-1] - assert final_attempt.end_status.name == "PERMISSION_DENIED" - - @CrossSync.pytest - async def test_read_row(self, table, temp_rows, handler, cluster_config): - await temp_rows.add_row(b"row_key_1") - handler.clear() - await table.read_row(b"row_key_1") - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is False - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1e9 - assert ( - operation.first_response_latency_ns > 0 - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - @CrossSync.pytest - async def test_read_row_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """ - Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one - """ - await temp_rows.add_row(b"row_key_1") - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(PermissionDenied): - await table.read_row(b"row_key_1", retryable_errors=[Aborted]) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - # validate attempts - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_row_failure_timeout(self, table, temp_rows, handler): - """ - Test failure in gapic layer by passing very low timeout - - No grpc headers expected - """ - await temp_rows.add_row(b"row_key_1") - handler.clear() - with pytest.raises(GoogleAPICallError): - await table.read_row(b"row_key_1", operation_timeout=0.001) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_row_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """ - Test failure in backend by accessing an unauthorized family - """ - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - await authorized_view.read_row( - b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") - ) - assert e.value.grpc_status_code.name == "PERMISSION_DENIED" - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - @CrossSync.pytest - async def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - - await temp_rows.add_row(b"a") - await temp_rows.add_row(b"b") - await temp_rows.add_row(b"c") - await temp_rows.add_row(b"d") - query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) - query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) - handler.clear() - row_list = await table.read_rows_sharded([query1, query2]) - assert len(row_list) == 4 - # validate counts - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 2 - # validate operations - for operation in handler.completed_operations: - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is True - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - attempt = operation.completed_attempts[0] - assert attempt in handler.completed_attempts - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1e9 - assert ( - operation.first_response_latency_ns is not None - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - # validate attempt - assert isinstance(attempt, CompletedAttemptMetric) - assert ( - attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - ) - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 - and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - @CrossSync.pytest - async def test_read_rows_sharded_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """ - Test failure in grpc layer by injecting errors into an interceptor - with retryable errors - """ - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - - await temp_rows.add_row(b"a") - await temp_rows.add_row(b"b") - query1 = ReadRowsQuery(row_keys=[b"a"]) - query2 = ReadRowsQuery(row_keys=[b"b"]) - handler.clear() - - error_injector.push(self._make_exception(StatusCode.ABORTED)) - await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) - - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 3 - # validate operations - for op in handler.completed_operations: - assert op.final_status.name == "OK" - assert op.op_type.value == "ReadRows" - assert op.is_streaming is True - # validate attempts - assert ( - len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) - == 2 - ) - assert ( - len( - [ - a - for a in handler.completed_attempts - if a.end_status.name == "ABORTED" - ] - ) - == 1 - ) - - @CrossSync.pytest - async def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): - """ - Test failure in gapic layer by passing very low timeout - - No grpc headers expected - """ - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup - from google.api_core.exceptions import DeadlineExceeded - - await temp_rows.add_row(b"a") - await temp_rows.add_row(b"b") - query1 = ReadRowsQuery(row_keys=[b"a"]) - query2 = ReadRowsQuery(row_keys=[b"b"]) - handler.clear() - with pytest.raises(ShardedReadRowsExceptionGroup) as e: - await table.read_rows_sharded([query1, query2], operation_timeout=0.005) - assert len(e.value.exceptions) == 2 - for sub_exc in e.value.exceptions: - assert isinstance(sub_exc.__cause__, DeadlineExceeded) - # both shards should fail - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 2 - # validate operations - for operation in handler.completed_operations: - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - # validate attempt - attempt = operation.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_read_rows_sharded_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """ - Test failure in backend by accessing an unauthorized family - """ - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup - - query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) - query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) - handler.clear() - with pytest.raises(ShardedReadRowsExceptionGroup) as e: - await authorized_view.read_rows_sharded([query1, query2]) - assert len(e.value.exceptions) == 1 - assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) - assert ( - e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" - ) - # one shard will fail, the other will succeed - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 2 - # sort operations by status - failed_op = next( - op for op in handler.completed_operations if op.final_status.name != "OK" - ) - success_op = next( - op for op in handler.completed_operations if op.final_status.name == "OK" - ) - # validate failed operation - assert failed_op.final_status.name == "PERMISSION_DENIED" - assert failed_op.op_type.value == "ReadRows" - assert failed_op.is_streaming is True - assert len(failed_op.completed_attempts) == 1 - assert failed_op.cluster_id == next(iter(cluster_config.keys())) - assert ( - failed_op.zone - == cluster_config[failed_op.cluster_id].location.split("/")[-1] - ) - # validate failed attempt - failed_attempt = failed_op.completed_attempts[0] - assert failed_attempt.end_status.name == "PERMISSION_DENIED" - assert ( - failed_attempt.gfe_latency_ns >= 0 - and failed_attempt.gfe_latency_ns < failed_op.duration_ns - ) - # validate successful operation - assert success_op.final_status.name == "OK" - assert success_op.op_type.value == "ReadRows" - assert success_op.is_streaming is True - assert len(success_op.completed_attempts) == 1 - # validate successful attempt - success_attempt = success_op.completed_attempts[0] - assert success_attempt.end_status.name == "OK" - - @CrossSync.pytest - async def test_read_rows_sharded_failure_mid_stream( - self, table, temp_rows, handler, error_injector - ): - """ - Test failure in grpc stream - """ - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup - - await temp_rows.add_row(b"a") - await temp_rows.add_row(b"b") - query1 = ReadRowsQuery(row_keys=[b"a"]) - query2 = ReadRowsQuery(row_keys=[b"b"]) - handler.clear() - error_injector.fail_mid_stream = True - error_injector.push(self._make_exception(StatusCode.ABORTED)) - error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) - with pytest.raises(ShardedReadRowsExceptionGroup) as e: - await table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) - assert len(e.value.exceptions) == 1 - assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) - # one shard will fail, the other will succeed - # the failing shard will have one retry - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 3 - # sort operations by status - failed_op = next( - op for op in handler.completed_operations if op.final_status.name != "OK" - ) - success_op = next( - op for op in handler.completed_operations if op.final_status.name == "OK" - ) - # validate failed operation - assert failed_op.final_status.name == "PERMISSION_DENIED" - assert failed_op.op_type.value == "ReadRows" - assert failed_op.is_streaming is True - assert len(failed_op.completed_attempts) == 1 - # validate successful operation - assert success_op.final_status.name == "OK" - assert len(success_op.completed_attempts) == 2 - # validate failed attempt - attempt = failed_op.completed_attempts[0] - assert attempt.end_status.name == "PERMISSION_DENIED" - # validate retried attempt - retried_attempt = success_op.completed_attempts[0] - assert retried_attempt.end_status.name == "ABORTED" - # validate successful attempt - success_attempt = success_op.completed_attempts[-1] - assert success_attempt.end_status.name == "OK" - - @CrossSync.pytest - async def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): - from google.cloud.bigtable.data.mutations import RowMutationEntry - - new_value = uuid.uuid4().hex.encode() - row_key, mutation = await temp_rows.create_row_and_mutation( - table, new_value=new_value - ) - bulk_mutation = RowMutationEntry(row_key, [mutation]) - - handler.clear() - await table.bulk_mutate_rows([bulk_mutation]) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is False - assert operation.op_type.value == "MutateRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1e9 - assert ( - operation.first_response_latency_ns is None - ) # populated for read_rows only - assert operation.flow_throttling_time_ns == 0 - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert attempt.application_blocking_time_ns == 0 - - @CrossSync.pytest - async def test_bulk_mutate_rows_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """ - Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one - """ - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - assert entry.is_idempotent() - - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(MutationsExceptionGroup): - await table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - # validate attempts - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): - """ - Test failure in gapic layer by passing very low timeout - - No grpc headers expected - """ - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - - handler.clear() - with pytest.raises(MutationsExceptionGroup): - await table.bulk_mutate_rows([entry], operation_timeout=0.001) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_bulk_mutate_rows_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """ - Test failure in backend by accessing an unauthorized family - """ - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - - row_key = b"row_key_1" - mutation = SetCell("unauthorized", b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - - handler.clear() - with pytest.raises(MutationsExceptionGroup): - await authorized_view.bulk_mutate_rows([entry]) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - # validate attempt - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - @CrossSync.pytest - async def test_bulk_mutate_rows_failure_unauthorized_with_retries( - self, handler, authorized_view, cluster_config - ): - """ - retry unauthorized request multiple times before timing out - - For bulk_mutate, the rpc returns success, with failures returned in the response. - For this reason, We expect the attempts to be marked as successful, even though - the underlying mutation is retried - """ - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - - row_key = b"row_key_1" - mutation = SetCell("unauthorized", b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - - handler.clear() - with pytest.raises(MutationsExceptionGroup) as e: - await authorized_view.bulk_mutate_rows( - [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 - ) - assert len(e.value.exceptions) == 1 - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) > 1 - # validate operation - operation = handler.completed_operations[0] - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) > 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - # validate attempts - for attempt in handler.completed_attempts: - assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] - - @CrossSync.pytest - async def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): - from google.cloud.bigtable.data.mutations import RowMutationEntry - - new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] - row_key, mutation = await temp_rows.create_row_and_mutation( - table, new_value=new_value - ) - row_key2, mutation2 = await temp_rows.create_row_and_mutation( - table, new_value=new_value2 - ) - bulk_mutation = RowMutationEntry(row_key, [mutation]) - bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) - - handler.clear() - async with table.mutations_batcher() as batcher: - await batcher.append(bulk_mutation) - await batcher.append(bulk_mutation2) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # bacher expects to cancel staged operation on close - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is False - assert operation.op_type.value == "MutateRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1e9 - assert ( - operation.first_response_latency_ns is None - ) # populated for read_rows only - assert ( - operation.flow_throttling_time_ns > 0 - and operation.flow_throttling_time_ns < operation.duration_ns - ) - # validate attempt - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert attempt.application_blocking_time_ns == 0 - - @CrossSync.pytest - async def test_mutate_rows_batcher_failure_with_retries( - self, table, handler, error_injector - ): - """ - Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one - """ - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - assert entry.is_idempotent() - - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(MutationsExceptionGroup): - async with table.mutations_batcher( - batch_retryable_errors=[Aborted] - ) as batcher: - await batcher.append(entry) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - # validate operation - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - # validate attempts - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): - """ - Test failure in gapic layer by passing very low timeout - - No grpc headers expected - """ - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - - with pytest.raises(MutationsExceptionGroup): - async with table.mutations_batcher( - batch_operation_timeout=0.001 - ) as batcher: - await batcher.append(entry) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - # validate attempt - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - @CrossSync.pytest - async def test_mutate_rows_batcher_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """ - Test failure in backend by accessing an unauthorized family - """ - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - - row_key = b"row_key_1" - mutation = SetCell("unauthorized", b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - - with pytest.raises(MutationsExceptionGroup) as e: - async with authorized_view.mutations_batcher() as batcher: - await batcher.append(entry) - assert len(e.value.exceptions) == 1 - assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) - assert ( - e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" - ) - # validate counts - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - # validate operation - operation = handler.completed_operations[0] - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - # validate attempt - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - @CrossSync.pytest async def test_mutate_row(self, table, temp_rows, handler, cluster_config): row_key = b"mutate" diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py index 6b28328eaab9..f3a50f35215e 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_autogen.py @@ -1204,1026 +1204,6 @@ def test_mutate_rows_batcher_failure_unauthorized( and attempt.gfe_latency_ns < operation.duration_ns ) - @pytest.mark.skipif( - bool(os.environ.get(BIGTABLE_EMULATOR)), - reason="emulator doesn't suport cluster_config", - ) - def test_read_rows(self, table, temp_rows, handler, cluster_config): - temp_rows.add_row(b"row_key_1") - temp_rows.add_row(b"row_key_2") - handler.clear() - row_list = table.read_rows(ReadRowsQuery()) - assert len(row_list) == 2 - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is True - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 - assert ( - operation.first_response_latency_ns is not None - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - def test_read_rows_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one""" - temp_rows.add_row(b"row_key_1") - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(PermissionDenied): - table.read_rows(ReadRowsQuery(), retryable_errors=[Aborted]) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - def test_read_rows_failure_timeout(self, table, temp_rows, handler): - """Test failure in gapic layer by passing very low timeout - - No grpc headers expected""" - temp_rows.add_row(b"row_key_1") - handler.clear() - with pytest.raises(GoogleAPICallError): - table.read_rows(ReadRowsQuery(), operation_timeout=0.001) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - def test_read_rows_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """Test failure in backend by accessing an unauthorized family""" - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - authorized_view.read_rows( - ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) - ) - assert e.value.grpc_status_code.name == "PERMISSION_DENIED" - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - def test_read_rows_stream(self, table, temp_rows, handler, cluster_config): - temp_rows.add_row(b"row_key_1") - temp_rows.add_row(b"row_key_2") - handler.clear() - generator = table.read_rows_stream(ReadRowsQuery()) - row_list = [r for r in generator] - assert len(row_list) == 2 - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is True - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 - assert ( - operation.first_response_latency_ns is not None - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - def test_read_rows_stream_failure_closed( - self, table, temp_rows, handler, error_injector - ): - """Test how metrics collection handles closed generator""" - temp_rows.add_row(b"row_key_1") - temp_rows.add_row(b"row_key_2") - handler.clear() - generator = table.read_rows_stream(ReadRowsQuery()) - generator.__next__() - generator.close() - with pytest.raises(CrossSync._Sync_Impl.StopIteration): - generator.__next__() - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert operation.final_status.name == "CANCELLED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "CANCELLED" - assert attempt.gfe_latency_ns is None - - def test_read_rows_stream_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one""" - temp_rows.add_row(b"row_key_1") - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) - with pytest.raises(PermissionDenied): - [_ for _ in generator] - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - def test_read_rows_stream_failure_timeout(self, table, temp_rows, handler): - """Test failure in gapic layer by passing very low timeout - - No grpc headers expected""" - temp_rows.add_row(b"row_key_1") - handler.clear() - generator = table.read_rows_stream(ReadRowsQuery(), operation_timeout=0.001) - with pytest.raises(GoogleAPICallError): - [_ for _ in generator] - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - def test_read_rows_stream_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """Test failure in backend by accessing an unauthorized family""" - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - generator = authorized_view.read_rows_stream( - ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) - ) - [_ for _ in generator] - assert e.value.grpc_status_code.name == "PERMISSION_DENIED" - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - def test_read_rows_stream_failure_unauthorized_with_retries( - self, handler, authorized_view, cluster_config - ): - """retry unauthorized request multiple times before timing out""" - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - generator = authorized_view.read_rows_stream( - ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")), - retryable_errors=[PermissionDenied], - operation_timeout=0.5, - ) - [_ for _ in generator] - assert e.value.grpc_status_code.name == "DEADLINE_EXCEEDED" - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) > 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) > 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - for attempt in handler.completed_attempts: - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name in ["PERMISSION_DENIED", "DEADLINE_EXCEEDED"] - - def test_read_rows_stream_failure_mid_stream( - self, table, temp_rows, handler, error_injector - ): - """Test failure in grpc stream""" - temp_rows.add_row(b"row_key_1") - handler.clear() - error_injector.fail_mid_stream = True - error_injector.push(self._make_exception(StatusCode.ABORTED)) - error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) - generator = table.read_rows_stream(ReadRowsQuery(), retryable_errors=[Aborted]) - with pytest.raises(PermissionDenied): - [_ for _ in generator] - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 2 - operation = handler.completed_operations[0] - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 2 - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "ABORTED" - final_attempt = handler.completed_attempts[-1] - assert final_attempt.end_status.name == "PERMISSION_DENIED" - - def test_read_row(self, table, temp_rows, handler, cluster_config): - temp_rows.add_row(b"row_key_1") - handler.clear() - table.read_row(b"row_key_1") - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is False - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 - assert ( - operation.first_response_latency_ns > 0 - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - def test_read_row_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one""" - temp_rows.add_row(b"row_key_1") - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(PermissionDenied): - table.read_row(b"row_key_1", retryable_errors=[Aborted]) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - def test_read_row_failure_timeout(self, table, temp_rows, handler): - """Test failure in gapic layer by passing very low timeout - - No grpc headers expected""" - temp_rows.add_row(b"row_key_1") - handler.clear() - with pytest.raises(GoogleAPICallError): - table.read_row(b"row_key_1", operation_timeout=0.001) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - def test_read_row_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """Test failure in backend by accessing an unauthorized family""" - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - with pytest.raises(GoogleAPICallError) as e: - authorized_view.read_row( - b"any_row", row_filter=FamilyNameRegexFilter("unauthorized") - ) - assert e.value.grpc_status_code.name == "PERMISSION_DENIED" - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - def test_read_rows_sharded(self, table, temp_rows, handler, cluster_config): - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - - temp_rows.add_row(b"a") - temp_rows.add_row(b"b") - temp_rows.add_row(b"c") - temp_rows.add_row(b"d") - query1 = ReadRowsQuery(row_keys=[b"a", b"c"]) - query2 = ReadRowsQuery(row_keys=[b"b", b"d"]) - handler.clear() - row_list = table.read_rows_sharded([query1, query2]) - assert len(row_list) == 4 - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 2 - for operation in handler.completed_operations: - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is True - assert operation.op_type.value == "ReadRows" - assert len(operation.completed_attempts) == 1 - attempt = operation.completed_attempts[0] - assert attempt in handler.completed_attempts - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 - assert ( - operation.first_response_latency_ns is not None - and operation.first_response_latency_ns < operation.duration_ns - ) - assert operation.flow_throttling_time_ns == 0 - assert isinstance(attempt, CompletedAttemptMetric) - assert ( - attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - ) - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 - and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert ( - attempt.application_blocking_time_ns > 0 - and attempt.application_blocking_time_ns < operation.duration_ns - ) - - def test_read_rows_sharded_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """Test failure in grpc layer by injecting errors into an interceptor - with retryable errors""" - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - - temp_rows.add_row(b"a") - temp_rows.add_row(b"b") - query1 = ReadRowsQuery(row_keys=[b"a"]) - query2 = ReadRowsQuery(row_keys=[b"b"]) - handler.clear() - error_injector.push(self._make_exception(StatusCode.ABORTED)) - table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 3 - for op in handler.completed_operations: - assert op.final_status.name == "OK" - assert op.op_type.value == "ReadRows" - assert op.is_streaming is True - assert ( - len([a for a in handler.completed_attempts if a.end_status.name == "OK"]) - == 2 - ) - assert ( - len( - [ - a - for a in handler.completed_attempts - if a.end_status.name == "ABORTED" - ] - ) - == 1 - ) - - def test_read_rows_sharded_failure_timeout(self, table, temp_rows, handler): - """Test failure in gapic layer by passing very low timeout - - No grpc headers expected""" - from google.api_core.exceptions import DeadlineExceeded - - from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - - temp_rows.add_row(b"a") - temp_rows.add_row(b"b") - query1 = ReadRowsQuery(row_keys=[b"a"]) - query2 = ReadRowsQuery(row_keys=[b"b"]) - handler.clear() - with pytest.raises(ShardedReadRowsExceptionGroup) as e: - table.read_rows_sharded([query1, query2], operation_timeout=0.005) - assert len(e.value.exceptions) == 2 - for sub_exc in e.value.exceptions: - assert isinstance(sub_exc.__cause__, DeadlineExceeded) - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 2 - for operation in handler.completed_operations: - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "ReadRows" - assert operation.is_streaming is True - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - attempt = operation.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - def test_read_rows_sharded_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """Test failure in backend by accessing an unauthorized family""" - from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - from google.cloud.bigtable.data.row_filters import FamilyNameRegexFilter - - query1 = ReadRowsQuery(row_filter=FamilyNameRegexFilter("unauthorized")) - query2 = ReadRowsQuery(row_filter=FamilyNameRegexFilter(TEST_FAMILY)) - handler.clear() - with pytest.raises(ShardedReadRowsExceptionGroup) as e: - authorized_view.read_rows_sharded([query1, query2]) - assert len(e.value.exceptions) == 1 - assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) - assert ( - e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" - ) - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 2 - failed_op = next( - (op for op in handler.completed_operations if op.final_status.name != "OK") - ) - success_op = next( - (op for op in handler.completed_operations if op.final_status.name == "OK") - ) - assert failed_op.final_status.name == "PERMISSION_DENIED" - assert failed_op.op_type.value == "ReadRows" - assert failed_op.is_streaming is True - assert len(failed_op.completed_attempts) == 1 - assert failed_op.cluster_id == next(iter(cluster_config.keys())) - assert ( - failed_op.zone - == cluster_config[failed_op.cluster_id].location.split("/")[-1] - ) - failed_attempt = failed_op.completed_attempts[0] - assert failed_attempt.end_status.name == "PERMISSION_DENIED" - assert ( - failed_attempt.gfe_latency_ns >= 0 - and failed_attempt.gfe_latency_ns < failed_op.duration_ns - ) - assert success_op.final_status.name == "OK" - assert success_op.op_type.value == "ReadRows" - assert success_op.is_streaming is True - assert len(success_op.completed_attempts) == 1 - success_attempt = success_op.completed_attempts[0] - assert success_attempt.end_status.name == "OK" - - def test_read_rows_sharded_failure_mid_stream( - self, table, temp_rows, handler, error_injector - ): - """Test failure in grpc stream""" - from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup - from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery - - temp_rows.add_row(b"a") - temp_rows.add_row(b"b") - query1 = ReadRowsQuery(row_keys=[b"a"]) - query2 = ReadRowsQuery(row_keys=[b"b"]) - handler.clear() - error_injector.fail_mid_stream = True - error_injector.push(self._make_exception(StatusCode.ABORTED)) - error_injector.push(self._make_exception(StatusCode.PERMISSION_DENIED)) - with pytest.raises(ShardedReadRowsExceptionGroup) as e: - table.read_rows_sharded([query1, query2], retryable_errors=[Aborted]) - assert len(e.value.exceptions) == 1 - assert isinstance(e.value.exceptions[0].__cause__, PermissionDenied) - assert len(handler.completed_operations) == 2 - assert len(handler.completed_attempts) == 3 - failed_op = next( - (op for op in handler.completed_operations if op.final_status.name != "OK") - ) - success_op = next( - (op for op in handler.completed_operations if op.final_status.name == "OK") - ) - assert failed_op.final_status.name == "PERMISSION_DENIED" - assert failed_op.op_type.value == "ReadRows" - assert failed_op.is_streaming is True - assert len(failed_op.completed_attempts) == 1 - assert success_op.final_status.name == "OK" - assert len(success_op.completed_attempts) == 2 - attempt = failed_op.completed_attempts[0] - assert attempt.end_status.name == "PERMISSION_DENIED" - retried_attempt = success_op.completed_attempts[0] - assert retried_attempt.end_status.name == "ABORTED" - success_attempt = success_op.completed_attempts[-1] - assert success_attempt.end_status.name == "OK" - - def test_bulk_mutate_rows(self, table, temp_rows, handler, cluster_config): - from google.cloud.bigtable.data.mutations import RowMutationEntry - - new_value = uuid.uuid4().hex.encode() - row_key, mutation = temp_rows.create_row_and_mutation( - table, new_value=new_value - ) - bulk_mutation = RowMutationEntry(row_key, [mutation]) - handler.clear() - table.bulk_mutate_rows([bulk_mutation]) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is False - assert operation.op_type.value == "MutateRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 - assert operation.first_response_latency_ns is None - assert operation.flow_throttling_time_ns == 0 - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert attempt.application_blocking_time_ns == 0 - - def test_bulk_mutate_rows_failure_with_retries( - self, table, temp_rows, handler, error_injector - ): - """Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one""" - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - assert entry.is_idempotent() - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(MutationsExceptionGroup): - table.bulk_mutate_rows([entry], retryable_errors=[Aborted]) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert isinstance(final_attempt, CompletedAttemptMetric) - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - def test_bulk_mutate_rows_failure_timeout(self, table, temp_rows, handler): - """Test failure in gapic layer by passing very low timeout - - No grpc headers expected""" - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - handler.clear() - with pytest.raises(MutationsExceptionGroup): - table.bulk_mutate_rows([entry], operation_timeout=0.001) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - def test_bulk_mutate_rows_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """Test failure in backend by accessing an unauthorized family""" - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - - row_key = b"row_key_1" - mutation = SetCell("unauthorized", b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - handler.clear() - with pytest.raises(MutationsExceptionGroup): - authorized_view.bulk_mutate_rows([entry]) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - - def test_bulk_mutate_rows_failure_unauthorized_with_retries( - self, handler, authorized_view, cluster_config - ): - """retry unauthorized request multiple times before timing out - - For bulk_mutate, the rpc returns success, with failures returned in the response. - For this reason, We expect the attempts to be marked as successful, even though - the underlying mutation is retried""" - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - - row_key = b"row_key_1" - mutation = SetCell("unauthorized", b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - handler.clear() - with pytest.raises(MutationsExceptionGroup) as e: - authorized_view.bulk_mutate_rows( - [entry], retryable_errors=[PermissionDenied], operation_timeout=0.5 - ) - assert len(e.value.exceptions) == 1 - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) > 1 - operation = handler.completed_operations[0] - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) > 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - for attempt in handler.completed_attempts: - assert attempt.end_status.name in ["OK", "DEADLINE_EXCEEDED"] - - def test_mutate_rows_batcher(self, table, temp_rows, handler, cluster_config): - from google.cloud.bigtable.data.mutations import RowMutationEntry - - new_value, new_value2 = [uuid.uuid4().hex.encode() for _ in range(2)] - row_key, mutation = temp_rows.create_row_and_mutation( - table, new_value=new_value - ) - row_key2, mutation2 = temp_rows.create_row_and_mutation( - table, new_value=new_value2 - ) - bulk_mutation = RowMutationEntry(row_key, [mutation]) - bulk_mutation2 = RowMutationEntry(row_key2, [mutation2]) - handler.clear() - with table.mutations_batcher() as batcher: - batcher.append(bulk_mutation) - batcher.append(bulk_mutation2) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.value[0] == 0 - assert operation.is_streaming is False - assert operation.op_type.value == "MutateRows" - assert len(operation.completed_attempts) == 1 - assert operation.completed_attempts[0] == handler.completed_attempts[0] - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - assert operation.duration_ns > 0 and operation.duration_ns < 1000000000.0 - assert operation.first_response_latency_ns is None - assert ( - operation.flow_throttling_time_ns > 0 - and operation.flow_throttling_time_ns < operation.duration_ns - ) - attempt = handler.completed_attempts[0] - assert isinstance(attempt, CompletedAttemptMetric) - assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns - assert attempt.end_status.value[0] == 0 - assert attempt.backoff_before_attempt_ns == 0 - assert ( - attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns - ) - assert attempt.application_blocking_time_ns == 0 - - def test_mutate_rows_batcher_failure_with_retries( - self, table, handler, error_injector - ): - """Test failure in grpc layer by injecting errors into an interceptor - with retryable errors, then a terminal one""" - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - assert entry.is_idempotent() - handler.clear() - expected_zone = "my_zone" - expected_cluster = "my_cluster" - num_retryable = 2 - for i in range(num_retryable): - error_injector.push( - self._make_exception(StatusCode.ABORTED, cluster_id=expected_cluster) - ) - error_injector.push( - self._make_exception(StatusCode.PERMISSION_DENIED, zone_id=expected_zone) - ) - with pytest.raises(MutationsExceptionGroup): - with table.mutations_batcher(batch_retryable_errors=[Aborted]) as batcher: - batcher.append(entry) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == num_retryable + 1 - operation = handler.completed_operations[0] - assert isinstance(operation, CompletedOperationMetric) - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == num_retryable + 1 - assert operation.cluster_id == expected_cluster - assert operation.zone == expected_zone - for i in range(num_retryable): - attempt = handler.completed_attempts[i] - assert attempt.end_status.name == "ABORTED" - assert attempt.gfe_latency_ns is None - final_attempt = handler.completed_attempts[num_retryable] - assert final_attempt.end_status.name == "PERMISSION_DENIED" - assert final_attempt.gfe_latency_ns is None - - def test_mutate_rows_batcher_failure_timeout(self, table, temp_rows, handler): - """Test failure in gapic layer by passing very low timeout - - No grpc headers expected""" - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - - row_key = b"row_key_1" - mutation = SetCell(TEST_FAMILY, b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - with pytest.raises(MutationsExceptionGroup): - with table.mutations_batcher(batch_operation_timeout=0.001) as batcher: - batcher.append(entry) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert operation.final_status.name == "DEADLINE_EXCEEDED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == "" - assert operation.zone == "global" - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "DEADLINE_EXCEEDED" - assert attempt.gfe_latency_ns is None - - def test_mutate_rows_batcher_failure_unauthorized( - self, handler, authorized_view, cluster_config - ): - """Test failure in backend by accessing an unauthorized family""" - from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup - from google.cloud.bigtable.data.mutations import RowMutationEntry, SetCell - - row_key = b"row_key_1" - mutation = SetCell("unauthorized", b"q", b"v") - entry = RowMutationEntry(row_key, [mutation]) - with pytest.raises(MutationsExceptionGroup) as e: - with authorized_view.mutations_batcher() as batcher: - batcher.append(entry) - assert len(e.value.exceptions) == 1 - assert isinstance(e.value.exceptions[0].__cause__, GoogleAPICallError) - assert ( - e.value.exceptions[0].__cause__.grpc_status_code.name == "PERMISSION_DENIED" - ) - assert len(handler.completed_operations) == 1 - assert len(handler.completed_attempts) == 1 - operation = handler.completed_operations[0] - assert operation.final_status.name == "PERMISSION_DENIED" - assert operation.op_type.value == "MutateRows" - assert operation.is_streaming is False - assert len(operation.completed_attempts) == 1 - assert operation.cluster_id == next(iter(cluster_config.keys())) - assert ( - operation.zone - == cluster_config[operation.cluster_id].location.split("/")[-1] - ) - attempt = handler.completed_attempts[0] - assert attempt.end_status.name == "PERMISSION_DENIED" - assert ( - attempt.gfe_latency_ns >= 0 - and attempt.gfe_latency_ns < operation.duration_ns - ) - def test_mutate_row(self, table, temp_rows, handler, cluster_config): row_key = b"mutate" new_value = uuid.uuid4().hex.encode() From d40fa0ac659ca565af9d760e826d0122932659f0 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 15:58:36 -0700 Subject: [PATCH 11/11] fixed format --- .../tests/system/data/__init__.py | 12 +++---- .../tests/system/data/test_metrics_async.py | 35 ++++++++++--------- .../tests/system/data/test_system_async.py | 2 +- .../tests/unit/data/_async/test_client.py | 19 ++++++---- .../unit/data/_sync_autogen/test_client.py | 26 ++++++++------ 5 files changed, 53 insertions(+), 41 deletions(-) diff --git a/packages/google-cloud-bigtable/tests/system/data/__init__.py b/packages/google-cloud-bigtable/tests/system/data/__init__.py index 6f836fb9678a..9e8a9c8ab9f3 100644 --- a/packages/google-cloud-bigtable/tests/system/data/__init__.py +++ b/packages/google-cloud-bigtable/tests/system/data/__init__.py @@ -13,10 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. # -import pytest import os import uuid +import pytest + TEST_FAMILY = "test-family" TEST_FAMILY_2 = "test-family-2" TEST_AGGREGATE_FAMILY = "test-aggregate-family" @@ -89,10 +90,11 @@ 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 + from google.cloud.bigtable_admin_v2 import types + # use user-specified instance if available user_specified_instance = os.getenv("BIGTABLE_TEST_INSTANCE") if user_specified_instance: @@ -154,8 +156,7 @@ def table_id( 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 + from google.api_core import exceptions, retry # use user-specified instance if available user_specified_table = os.getenv("BIGTABLE_TEST_TABLE") @@ -207,8 +208,7 @@ def authorized_view_id( - 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 + from google.api_core import exceptions, retry retry = retry.Retry( predicate=retry.if_exception_type(exceptions.FailedPrecondition) diff --git a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py index 5b335857d2fa..49f8bdfee8eb 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_metrics_async.py @@ -13,36 +13,37 @@ # limitations under the License. import asyncio import os -import pytest import uuid +import pytest +from google.api_core.exceptions import Aborted, GoogleAPICallError, PermissionDenied +from google.cloud.environment_vars import BIGTABLE_EMULATOR from grpc import StatusCode -from google.api_core.exceptions import Aborted -from google.api_core.exceptions import GoogleAPICallError -from google.api_core.exceptions import PermissionDenied -from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler +from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._metrics.data_model import ( - CompletedOperationMetric, CompletedAttemptMetric, + CompletedOperationMetric, ) +from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler from google.cloud.bigtable_v2.types import ResponseParams -from google.cloud.environment_vars import BIGTABLE_EMULATOR - -from google.cloud.bigtable.data._cross_sync import CrossSync from . import TEST_FAMILY, SystemTestRunner if CrossSync.is_async: - from grpc.aio import UnaryUnaryClientInterceptor - from grpc.aio import UnaryStreamClientInterceptor - from grpc.aio import AioRpcError - from grpc.aio import Metadata + from grpc.aio import ( + AioRpcError, + Metadata, + UnaryStreamClientInterceptor, + UnaryUnaryClientInterceptor, + ) else: - from grpc import UnaryUnaryClientInterceptor - from grpc import UnaryStreamClientInterceptor - from grpc import RpcError - from grpc import intercept_channel + from grpc import ( + RpcError, + UnaryStreamClientInterceptor, + UnaryUnaryClientInterceptor, + intercept_channel, + ) __CROSS_SYNC_OUTPUT__ = "tests.system.data.test_metrics_autogen" diff --git a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py index 3636e6993e55..b65f05e4bd17 100644 --- a/packages/google-cloud-bigtable/tests/system/data/test_system_async.py +++ b/packages/google-cloud-bigtable/tests/system/data/test_system_async.py @@ -26,7 +26,7 @@ from google.cloud.bigtable.data.execute_query.metadata import SqlType from google.cloud.bigtable.data.read_modify_write_rules import _MAX_INCREMENT_VALUE -from . import SystemTestRunner, TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2 +from . import TEST_AGGREGATE_FAMILY, TEST_FAMILY, TEST_FAMILY_2, SystemTestRunner if CrossSync.is_async: from google.cloud.bigtable_v2.services.bigtable.transports.grpc_asyncio import ( diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index e116f4a4624b..6c6719615c40 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -1048,16 +1048,18 @@ def test_client_ctor_sync(self): assert client._channel_refresh_task is None @CrossSync.pytest - @pytest.mark.paramtrize( + @pytest.mark.parametrize( "use_mtls, expected_domain", - [("never", "googleapis.com"), ("always", "mtls.googleapis.com")] + [("never", "googleapis.com"), ("always", "mtls.googleapis.com")], ) async def test_default_universe_domain(self, use_mtls, expected_domain): """ When not passed, universe_domain should default to googleapis.com """ - async with self._make_client(project="project-id", credentials=None, use_mtls=expected_domain) as client: - assert client.universe_domain == expected_domain + async with self._make_client( + project="project-id", credentials=None, use_mtls=use_mtls + ) as client: + assert client.universe_domain == "googleapis.com" assert client.api_endpoint == f"bigtable.{expected_domain}" @CrossSync.pytest @@ -1097,7 +1099,9 @@ async def test_credential_universe_domain_matches_GDU(self): async def test_anomynous_credential_universe_domain(self): """Anomynopus credentials should use default universe domain""" creds = AnonymousCredentials() - async with self._make_client(project="project_id", credentials=creds, use_mtls="never") as client: + async with self._make_client( + project="project_id", credentials=creds, use_mtls="never" + ) as client: assert client.universe_domain == "googleapis.com" assert client.api_endpoint == "bigtable.googleapis.com" @@ -1137,7 +1141,10 @@ async def test_configured_universe_domain_matches_credentials(self): creds = AnonymousCredentials() creds._universe_domain = universe_domain async with self._make_client( - project="project_id", credentials=creds, client_options=options, use_mtls="never" + project="project_id", + credentials=creds, + client_options=options, + use_mtls="never", ) as client: assert client.universe_domain == universe_domain assert client.api_endpoint == f"bigtable.{universe_domain}" diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index 9d8bbd7f6834..79ad903b6191 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -208,8 +208,10 @@ def test__start_background_channel_refresh(self): def test__ping_and_warm_instances(self): """test ping and warm with mocked asyncio.gather""" client_mock = mock.Mock() - client_mock._execute_ping_and_warms = lambda *args: ( - self._get_target_class()._execute_ping_and_warms(client_mock, *args) + client_mock._execute_ping_and_warms = ( + lambda *args: self._get_target_class()._execute_ping_and_warms( + client_mock, *args + ) ) with mock.patch.object( CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock() @@ -252,8 +254,10 @@ def test__ping_and_warm_instances(self): def test__ping_and_warm_single_instance(self): """should be able to call ping and warm with single instance""" client_mock = mock.Mock() - client_mock._execute_ping_and_warms = lambda *args: ( - self._get_target_class()._execute_ping_and_warms(client_mock, *args) + client_mock._execute_ping_and_warms = ( + lambda *args: self._get_target_class()._execute_ping_and_warms( + client_mock, *args + ) ) with mock.patch.object( CrossSync._Sync_Impl, "gather_partials", CrossSync._Sync_Impl.Mock() @@ -1322,11 +1326,11 @@ def _make_client(self, *args, **kwargs): def _make_table(self, *args, **kwargs): client_mock = mock.Mock() - client_mock._register_instance.side_effect = lambda *args, **kwargs: ( - CrossSync._Sync_Impl.yield_to_event_loop() + client_mock._register_instance.side_effect = ( + lambda *args, **kwargs: CrossSync._Sync_Impl.yield_to_event_loop() ) - client_mock._remove_instance_registration.side_effect = lambda *args, **kwargs: ( - CrossSync._Sync_Impl.yield_to_event_loop() + client_mock._remove_instance_registration.side_effect = ( + lambda *args, **kwargs: CrossSync._Sync_Impl.yield_to_event_loop() ) kwargs["instance_id"] = kwargs.get( "instance_id", args[0] if args else "instance" @@ -1788,8 +1792,9 @@ def test_read_rows_sharded_multiple_queries(self): with mock.patch.object( table.client._gapic_client, "read_rows" ) as read_rows: - read_rows.side_effect = lambda *args, **kwargs: ( - CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( + read_rows.side_effect = ( + lambda *args, + **kwargs: CrossSync._Sync_Impl.TestReadRows._make_gapic_stream( [ CrossSync._Sync_Impl.TestReadRows._make_chunk(row_key=k) for k in args[0].rows.row_keys @@ -2878,7 +2883,6 @@ def prepare_mock(self, client): yield prepare_mock def _make_gapic_stream(self, sample_list: list["ExecuteQueryResponse" | Exception]): - class MockStream: def __init__(self, sample_list): self.sample_list = sample_list