added start of system tests for metrics

Author: daniel-sanche

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

@@ -187,6 +187,7 @@ async def _run_attempt(self):
                 ),
                 timeout=next(self.timeout_generator),
                 retry=None,
+                metadata=[self._operation_metric.interceptor_metadata],
             )
             async for result_list in result_generator:
                 for result in result_list.entries:

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

@@ -156,6 +156,7 @@ def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
             self.request,
             timeout=next(self.attempt_timeout_gen),
             retry=None,
+            metadata=[self._operation_metric.interceptor_metadata],
         )
         chunked_stream = self.chunk_stream(gapic_stream)
         return self.merge_rows(chunked_stream)

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]
 
@@ -1488,6 +1489,7 @@ async def mutate_row(
                 ),
                 timeout=attempt_timeout,
                 retry=None,
+                metadata=[operation_metric.interceptor_metadata],
             )
             return await CrossSync.retry_target(
                 target,
@@ -1608,7 +1610,7 @@ async def check_and_mutate_row(
             false_case_mutations = [false_case_mutations]
         false_case_list = [m._to_pb() for m in false_case_mutations or []]
 
-        with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE):
+        with self._metrics.create_operation(OperationType.CHECK_AND_MUTATE) as op:
             result = await self.client._gapic_client.check_and_mutate_row(
                 request=CheckAndMutateRowRequest(
                     true_mutations=true_case_list,
@@ -1624,6 +1626,7 @@ async def check_and_mutate_row(
                 ),
                 timeout=operation_timeout,
                 retry=None,
+                metadata=[op.interceptor_metadata],
             )
             return result.predicate_matched
 
@@ -1666,7 +1669,7 @@ async def read_modify_write_row(
         if not rules:
             raise ValueError("rules must contain at least one item")
 
-        with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE):
+        with self._metrics.create_operation(OperationType.READ_MODIFY_WRITE) as op:
             result = await self.client._gapic_client.read_modify_write_row(
                 request=ReadModifyWriteRowRequest(
                     rules=[rule._to_pb() for rule in rules],
@@ -1678,6 +1681,7 @@ async def read_modify_write_row(
                 ),
                 timeout=operation_timeout,
                 retry=None,
+                metadata=[op.interceptor_metadata],
             )
             # construct Row from result
             return Row._from_pb(result.row)

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/__init__.py

@@ -13,7 +13,67 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+import asyncio
+import pytest
+import uuid
+from google.cloud.bigtable.data._cross_sync import CrossSync
 
 TEST_FAMILY = "test-family"
 TEST_FAMILY_2 = "test-family-2"
 TEST_AGGREGATE_FAMILY = "test-aggregate-family"
+
+class SystemTestRunner:
+    """
+    configures a system test class with configuration for clusters/tables/etc
+
+    used by standard system tests, and metrics tests
+    """
+
+    @CrossSync.drop
+    @pytest.fixture(scope="session")
+    def event_loop(self):
+        loop = asyncio.get_event_loop()
+        yield loop
+        loop.stop()
+        loop.close()
+
+    @pytest.fixture(scope="session")
+    def init_table_id(self):
+        """
+        The table_id to use when creating a new test table
+        """
+        return f"test-table-{uuid.uuid4().hex}"
+
+    @pytest.fixture(scope="session")
+    def cluster_config(self, project_id):
+        """
+        Configuration for the clusters to use when creating a new instance
+        """
+        from google.cloud.bigtable_admin_v2 import types
+
+        cluster = {
+            "test-cluster": types.Cluster(
+                location=f"projects/{project_id}/locations/us-central1-b",
+                serve_nodes=1,
+            )
+        }
+        return cluster
+
+    @pytest.fixture(scope="session")
+    def column_family_config(self):
+        """
+        specify column families to create when creating a new test table
+        """
+        from google.cloud.bigtable_admin_v2 import types
+
+        int_aggregate_type = types.Type.Aggregate(
+            input_type=types.Type(int64_type={"encoding": {"big_endian_bytes": {}}}),
+            sum={},
+        )
+        return {
+            TEST_FAMILY: types.ColumnFamily(),
+            TEST_FAMILY_2: types.ColumnFamily(),
+            TEST_AGGREGATE_FAMILY: types.ColumnFamily(
+                value_type=types.Type(aggregate_type=int_aggregate_type)
+            ),
+        }