From 5dc927c09f6a123ab68315d9dd0c4b8dcb533712 Mon Sep 17 00:00:00 2001 From: Ali Abouelazm Date: Fri, 16 Oct 2026 00:16:07 -0700 Subject: [PATCH] fix(bigtable): preserve resource labels on batcher metrics Create batcher operation metrics through the target helper so project, instance, table, and app-profile labels reach the metrics exporter. Keep metric creation before flow-control waiting to preserve latency timing. Add sync and async regression tests for tables, authorized views, app profiles, multiple batches, and flow-control timing. Regenerate sync code and tests with CrossSync. Fixes #18546 --- .../bigtable/data/_async/mutations_batcher.py | 7 +- .../data/_sync_autogen/mutations_batcher.py | 9 +-- .../data/_async/test_mutations_batcher.py | 73 +++++++++++++++++++ .../_sync_autogen/test_mutations_batcher.py | 70 ++++++++++++++++++ 4 files changed, 149 insertions(+), 10 deletions(-) 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 b49a09555b79..0527925b908a 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 @@ -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( 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 f451f0249e85..a7cc84b8c9e4 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 @@ -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( 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 a6b7194cfbbb..0649481156c7 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 @@ -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): """ 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 339323e30514..733d91c16e83 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 @@ -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"""