added test for sample_row_keys

Author: daniel-sanche

google/cloud/bigtable/data/_async/client.py

@@ -1353,6 +1353,7 @@ async def execute_rpc():
                     ),
                     timeout=next(attempt_timeout_gen),
                     retry=None,
+                    metadata=[operation_metric.interceptor_metadata],
                 )
                 return [(s.row_key, s.offset_bytes) async for s in results]
 

google/cloud/bigtable/data/_async/metrics_interceptor.py

@@ -20,6 +20,7 @@
 )
 from google.cloud.bigtable.data._metrics.data_model import ActiveOperationMetric
 from google.cloud.bigtable.data._metrics.data_model import OperationState
+from google.cloud.bigtable.data._metrics.data_model import OperationType
 from google.cloud.bigtable.data._metrics.handlers._base import MetricsHandler
 
 from google.cloud.bigtable.data._cross_sync import CrossSync
@@ -129,7 +130,8 @@ async def intercept_unary_stream(
     ):
         # TODO: benchmark
         async def response_wrapper(call):
-            has_first_response = operation.first_response_latency is not None
+            # only track has_first response for READ_ROWS
+            has_first_response = operation.first_response_latency_ns is not None or operation.op_type != OperationType.READ_ROWS
             encountered_exc = None
             try:
                 async for response in call:

tests/system/data/test_metrics_async.py

@@ -97,6 +97,35 @@ async def target(self, client, table_id, instance_id, handler):
             table._metrics.add_handler(handler)
             yield table
 
+    @CrossSync.pytest
+    async def test_sample_row_keys(self, target, temp_rows, handler, cluster_config):
+        await target.sample_row_keys()
+        # validate counts
+        assert len(handler.completed_operations) == 1
+        assert len(handler.completed_attempts) == 1
+        assert len(handler.cancelled_operations) == 0
+        # validate operation
+        operation = handler.completed_operations[0]
+        assert isinstance(operation, CompletedOperationMetric)
+        assert operation.final_status.value[0] == 0
+        assert operation.is_streaming is False
+        assert operation.op_type.value == "SampleRowKeys"
+        assert len(operation.completed_attempts) == 1
+        assert operation.completed_attempts[0] == handler.completed_attempts[0]
+        assert operation.cluster_id == next(iter(cluster_config.keys()))
+        assert operation.zone == cluster_config[operation.cluster_id].location.split("/")[-1]
+        assert operation.duration_ns > 0 and operation.duration_ns < 1e9
+        assert operation.first_response_latency_ns is None  # populated for read_rows only
+        assert operation.flow_throttling_time_ns == 0
+        # validate attempt
+        attempt = handler.completed_attempts[0]
+        assert attempt.duration_ns > 0 and attempt.duration_ns < operation.duration_ns
+        assert attempt.end_status.value[0] == 0
+        assert attempt.backoff_before_attempt_ns == 0
+        assert attempt.gfe_latency_ns > 0 and attempt.gfe_latency_ns < attempt.duration_ns
+        assert attempt.application_blocking_time_ns == 0
+        assert attempt.grpc_throttling_time_ns == 0  # TODO: confirm
+
     @CrossSync.pytest
     async def test_read_modify_write(self, target, temp_rows, handler, cluster_config):
         from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule