From a78cb52ef6f684f1751bf8da27b14fffe42ce958 Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 13:36:50 -0700 Subject: [PATCH 1/4] 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 0ab1b05187850f6cd186ab5fbee4519af7b9125c Mon Sep 17 00:00:00 2001 From: Daniel Sanche Date: Tue, 21 Apr 2026 15:39:39 -0700 Subject: [PATCH 2/4] 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 3/4] 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 4/4] 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()