fix(bigtable): preserve resource labels on batcher metrics

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py

@@ -43,7 +43,6 @@
 )
 
 if TYPE_CHECKING:
-    from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController
     from google.cloud.bigtable.data.mutations import RowMutationEntry
 
     if CrossSync.is_async:
@@ -194,12 +193,12 @@ async def add_to_flow(self, mutations: RowMutationEntry | list[RowMutationEntry]
     async def add_to_flow_with_metrics(
         self,
         mutations: RowMutationEntry | list[RowMutationEntry],
-        metrics_controller: BigtableClientSideMetricsController,
+        target: TargetType,
     ):
         inner_generator = self.add_to_flow(mutations)
         while True:
             # start a new metric
-            metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS)
+            metric = target._create_operation(OperationType.BULK_MUTATE_ROWS)
             flow_start_time = time.monotonic_ns()
             try:
                 value = await inner_generator.__anext__()
@@ -394,7 +393,7 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
         ] = []
         found_exceptions: list[FailedMutationEntryError] = []
         async for batch, metric in self._flow_control.add_to_flow_with_metrics(
-            new_entries, self._target.client._metrics
+            new_entries, self._target
         ):
             try:
                 batch_task = CrossSync.create_task(

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py

@@ -46,7 +46,6 @@
 )
 
 if TYPE_CHECKING:
-    from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController
     from google.cloud.bigtable.data._sync_autogen.client import (
         _DataApiTarget as TargetType,
     )
@@ -164,13 +163,11 @@ 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,
+        self, mutations: RowMutationEntry | list[RowMutationEntry], target: TargetType
     ):
         inner_generator = self.add_to_flow(mutations)
         while True:
-            metric = metrics_controller.create_operation(OperationType.BULK_MUTATE_ROWS)
+            metric = target._create_operation(OperationType.BULK_MUTATE_ROWS)
             flow_start_time = time.monotonic_ns()
             try:
                 value = inner_generator.__next__()
@@ -342,7 +339,7 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]):
         ] = []
         found_exceptions: list[FailedMutationEntryError] = []
         for batch, metric in self._flow_control.add_to_flow_with_metrics(
-            new_entries, self._target.client._metrics
+            new_entries, self._target
         ):
             try:
                 batch_task = CrossSync._Sync_Impl.create_task(

packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py

@@ -297,6 +297,39 @@ async def test_add_to_flow_oversize(self):
         ]
         assert len(count_results) == 1
 
+    @CrossSync.pytest
+    async def test_add_to_flow_with_metrics_throttling(self):
+        from google.cloud.bigtable.data._metrics import ActiveOperationMetric
+
+        instance = self._make_one()
+        target = mock.Mock()
+        target._create_operation.side_effect = ActiveOperationMetric
+        mutations = [self._make_mutation(), self._make_mutation()]
+        with mock.patch("time.monotonic_ns", return_value=100) as clock:
+
+            async def flow(entries):
+                clock.return_value = 125
+                yield entries[:1]
+                clock.return_value = 240
+                yield entries[1:]
+
+            with mock.patch.object(instance, "add_to_flow", side_effect=flow):
+                results = []
+                async for batch, metric in instance.add_to_flow_with_metrics(
+                    mutations, target
+                ):
+                    results.append((batch, metric))
+                    clock.return_value = 200
+
+        assert len(results) == 2
+        assert results[0][0] == mutations[:1]
+        assert results[1][0] == mutations[1:]
+        assert results[0][1] is not results[1][1]
+        assert results[0][1].start_time_ns == 100
+        assert results[0][1].flow_throttling_time_ns == 25
+        assert results[1][1].start_time_ns == 200
+        assert results[1][1].flow_throttling_time_ns == 40
+
 
 @CrossSync.convert_class(sync_name="TestMutationsBatcher")
 class TestMutationsBatcherAsync:
@@ -322,6 +355,7 @@ def _make_one(self, table=None, **kwargs):
                 ServiceUnavailable,
             )
             table.client._metrics = BigtableClientSideMetricsController([])
+            table._create_operation.side_effect = table.client._metrics.create_operation
 
         return self._get_target_class()(table, **kwargs)
 
@@ -817,6 +851,45 @@ async def gen(x):
                     instance._oldest_exceptions.clear()
                     instance._newest_exceptions.clear()
 
+    @CrossSync.pytest
+    @pytest.mark.parametrize("app_profile_id", [None, "test-profile"])
+    @pytest.mark.parametrize("authorized_view", [False, True])
+    async def test_flush_metric_resource_labels(self, app_profile_id, authorized_view):
+        from google.cloud.bigtable.data._metrics import (
+            BigtableClientSideMetricsController,
+            OperationType,
+        )
+
+        client = mock.Mock(project="test-project")
+        client._metrics = BigtableClientSideMetricsController([])
+        with mock.patch.object(CrossSync, "create_task"):
+            if authorized_view:
+                target = CrossSync.AuthorizedView(
+                    client, "test-instance", "test-table", "test-view", app_profile_id
+                )
+            else:
+                target = CrossSync.Table(
+                    client, "test-instance", "test-table", app_profile_id
+                )
+
+        with mock.patch.object(CrossSync, "_MutateRowsOperation") as operation:
+            operation.return_value.start = CrossSync.Mock()
+            async with self._make_one(
+                target, flush_interval=None, flow_control_max_mutation_count=1
+            ) as instance:
+                await instance.append(self._make_mutation())
+                await instance.append(self._make_mutation())
+
+        assert operation.call_count == 2
+        metrics = [call.kwargs["metric"] for call in operation.call_args_list]
+        assert metrics[0] is not metrics[1]
+        for metric in metrics:
+            assert metric.op_type == OperationType.BULK_MUTATE_ROWS
+            assert metric.project_id == "test-project"
+            assert metric.instance_id == "test-instance"
+            assert metric.table_id == "test-table"
+            assert metric.app_profile_id == app_profile_id
+
     @CrossSync.pytest
     async def test_flush_clears_job_list(self):
         """

packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py

@@ -250,6 +250,37 @@ def test_add_to_flow_oversize(self):
         count_results = [out for out in instance.add_to_flow(large_count_mutation)]
         assert len(count_results) == 1
 
+    def test_add_to_flow_with_metrics_throttling(self):
+        from google.cloud.bigtable.data._metrics import ActiveOperationMetric
+
+        instance = self._make_one()
+        target = mock.Mock()
+        target._create_operation.side_effect = ActiveOperationMetric
+        mutations = [self._make_mutation(), self._make_mutation()]
+        with mock.patch("time.monotonic_ns", return_value=100) as clock:
+
+            def flow(entries):
+                clock.return_value = 125
+                yield entries[:1]
+                clock.return_value = 240
+                yield entries[1:]
+
+            with mock.patch.object(instance, "add_to_flow", side_effect=flow):
+                results = []
+                for batch, metric in instance.add_to_flow_with_metrics(
+                    mutations, target
+                ):
+                    results.append((batch, metric))
+                    clock.return_value = 200
+        assert len(results) == 2
+        assert results[0][0] == mutations[:1]
+        assert results[1][0] == mutations[1:]
+        assert results[0][1] is not results[1][1]
+        assert results[0][1].start_time_ns == 100
+        assert results[0][1].flow_throttling_time_ns == 25
+        assert results[1][1].start_time_ns == 200
+        assert results[1][1].flow_throttling_time_ns == 40
+
 
 class TestMutationsBatcher:
     def _get_target_class(self):
@@ -273,6 +304,7 @@ def _make_one(self, table=None, **kwargs):
                 ServiceUnavailable,
             )
             table.client._metrics = BigtableClientSideMetricsController([])
+            table._create_operation.side_effect = table.client._metrics.create_operation
         return self._get_target_class()(table, **kwargs)
 
     @staticmethod
@@ -715,6 +747,44 @@ def gen(x):
                     instance._oldest_exceptions.clear()
                     instance._newest_exceptions.clear()
 
+    @pytest.mark.parametrize("app_profile_id", [None, "test-profile"])
+    @pytest.mark.parametrize("authorized_view", [False, True])
+    def test_flush_metric_resource_labels(self, app_profile_id, authorized_view):
+        from google.cloud.bigtable.data._metrics import (
+            BigtableClientSideMetricsController,
+            OperationType,
+        )
+
+        client = mock.Mock(project="test-project")
+        client._metrics = BigtableClientSideMetricsController([])
+        with mock.patch.object(CrossSync._Sync_Impl, "create_task"):
+            if authorized_view:
+                target = CrossSync._Sync_Impl.AuthorizedView(
+                    client, "test-instance", "test-table", "test-view", app_profile_id
+                )
+            else:
+                target = CrossSync._Sync_Impl.Table(
+                    client, "test-instance", "test-table", app_profile_id
+                )
+        with mock.patch.object(
+            CrossSync._Sync_Impl, "_MutateRowsOperation"
+        ) as operation:
+            operation.return_value.start = CrossSync._Sync_Impl.Mock()
+            with self._make_one(
+                target, flush_interval=None, flow_control_max_mutation_count=1
+            ) as instance:
+                instance.append(self._make_mutation())
+                instance.append(self._make_mutation())
+        assert operation.call_count == 2
+        metrics = [call.kwargs["metric"] for call in operation.call_args_list]
+        assert metrics[0] is not metrics[1]
+        for metric in metrics:
+            assert metric.op_type == OperationType.BULK_MUTATE_ROWS
+            assert metric.project_id == "test-project"
+            assert metric.instance_id == "test-instance"
+            assert metric.table_id == "test-table"
+            assert metric.app_profile_id == app_profile_id
+
     def test_flush_clears_job_list(self):
         """a job should be added to _flush_jobs when _schedule_flush is called,
         and removed when it completes"""