From 247667e32cd7789ab3018a5c86ad7f4fbeb72d0a Mon Sep 17 00:00:00 2001 From: Mattie Fu Date: Sat, 12 Sep 2026 02:55:50 +0000 Subject: [PATCH] feat(bigtable): stamp batcher mutate_rows RPCs with x-goog-api-client bigtable-batcher header Change-Id: I8d21a96552e9386d2441cc1b4f88d13592e1f5d8 --- .../google/cloud/bigtable/batcher.py | 7 +++++- .../bigtable/data/_async/_mutate_rows.py | 4 +++- .../cloud/bigtable/data/_async/client.py | 2 ++ .../bigtable/data/_async/mutations_batcher.py | 5 ++++ .../data/_sync_autogen/_mutate_rows.py | 4 +++- .../bigtable/data/_sync_autogen/client.py | 2 ++ .../data/_sync_autogen/mutations_batcher.py | 5 ++++ .../google/cloud/bigtable/table.py | 4 ++-- .../data/_async/test_mutations_batcher.py | 22 ++++++++++++++++++ .../_sync_autogen/test_mutations_batcher.py | 23 +++++++++++++++++++ .../tests/unit/v2_client/test_batcher.py | 23 ++++++++++++++++++- .../tests/unit/v2_client/test_table.py | 23 +++++++++++++++++++ 12 files changed, 118 insertions(+), 6 deletions(-) diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py index e69b46382eb0..3859b80560fc 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py @@ -22,6 +22,8 @@ from google.api_core.exceptions import from_grpc_status +from google.cloud.bigtable.gapic_version import __version__ as _bigtable_version + FLUSH_COUNT = 100 # after this many elements, send out the batch MAX_MUTATION_SIZE = 20 * 1024 * 1024 # 20MB # after this many bytes, send out the batch @@ -418,7 +420,10 @@ def _flush_rows(self, rows_to_flush): """ responses = [] if len(rows_to_flush) > 0: - response = self.table.mutate_rows(rows_to_flush) + response = self.table.mutate_rows( + rows_to_flush, + metadata=[("x-goog-api-client", f"bigtable-batcher/{_bigtable_version}")], + ) if self._user_batch_completed_callback: self._user_batch_completed_callback(response) 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 0007447a5505..c58067a3c63f 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 @@ -14,6 +14,7 @@ # from __future__ import annotations +import functools from typing import TYPE_CHECKING, Sequence from google.api_core import exceptions as core_exceptions @@ -85,6 +86,7 @@ def __init__( attempt_timeout: float | None, metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]] = (), + metadata: Sequence[tuple[str, str]] = (), ): # check that mutations are within limits total_mutations = sum(len(entry.mutations) for entry in mutation_entries) @@ -95,7 +97,7 @@ def __init__( f"all entries. Found {total_mutations}." ) self._target = target - self._gapic_fn = gapic_client.mutate_rows + self._gapic_fn = functools.partial(gapic_client.mutate_rows, metadata=metadata) # create predicate for determining which errors are retryable self.is_retryable = retries.if_exception_type( # RPC level errors 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 048171474648..5feff8b33707 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 @@ -1726,6 +1726,7 @@ async def bulk_mutate_rows( attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, retryable_errors: Sequence[type[Exception]] | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + metadata: Sequence[tuple[str, str]] = (), ): """ Applies mutations for multiple rows in a single batched request. @@ -1771,6 +1772,7 @@ async def bulk_mutate_rows( attempt_timeout, metric=self._create_operation(OperationType.BULK_MUTATE_ROWS), retryable_exceptions=retryable_excs, + metadata=metadata, ) 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 0019e34518fe..c0ff8d3ac1de 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 @@ -21,6 +21,8 @@ from collections import deque from typing import TYPE_CHECKING, Sequence, cast +from google.cloud.bigtable.gapic_version import __version__ as _bigtable_version + from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( TABLE_DEFAULT, @@ -419,6 +421,9 @@ async def _execute_mutate_rows( attempt_timeout=self._attempt_timeout, metric=metric, retryable_exceptions=self._retryable_errors, + metadata=[ + ("x-goog-api-client", f"bigtable-batcher/{_bigtable_version}") + ], ) await operation.start() except MutationsExceptionGroup as e: 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 8bb4e49e22eb..4b26600210bb 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 @@ -17,6 +17,7 @@ from __future__ import annotations +import functools from typing import TYPE_CHECKING, Sequence from google.api_core import exceptions as core_exceptions @@ -73,6 +74,7 @@ def __init__( attempt_timeout: float | None, metric: ActiveOperationMetric, retryable_exceptions: Sequence[type[Exception]] = (), + metadata: Sequence[tuple[str, str]] = (), ): total_mutations = sum((len(entry.mutations) for entry in mutation_entries)) if total_mutations > _MUTATE_ROWS_REQUEST_MUTATION_LIMIT: @@ -80,7 +82,7 @@ def __init__( f"mutate_rows requests can contain at most {_MUTATE_ROWS_REQUEST_MUTATION_LIMIT} mutations across all entries. Found {total_mutations}." ) self._target = target - self._gapic_fn = gapic_client.mutate_rows + self._gapic_fn = functools.partial(gapic_client.mutate_rows, metadata=metadata) self.is_retryable = retries.if_exception_type( *retryable_exceptions, bt_exceptions._MutateRowsIncomplete ) 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 24ece63767d4..16fc0ab5ec5d 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 @@ -1431,6 +1431,7 @@ def bulk_mutate_rows( attempt_timeout: float | None | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, retryable_errors: Sequence[type[Exception]] | TABLE_DEFAULT = TABLE_DEFAULT.MUTATE_ROWS, + metadata: Sequence[tuple[str, str]] = (), ): """Applies mutations for multiple rows in a single batched request. @@ -1473,6 +1474,7 @@ def bulk_mutate_rows( attempt_timeout, metric=self._create_operation(OperationType.BULK_MUTATE_ROWS), retryable_exceptions=retryable_excs, + metadata=metadata, ) 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 eb1f0055f5c9..4b572c09134b 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 @@ -24,6 +24,8 @@ from collections import deque from typing import TYPE_CHECKING, Sequence, cast +from google.cloud.bigtable.gapic_version import __version__ as _bigtable_version + from google.cloud.bigtable.data._cross_sync import CrossSync from google.cloud.bigtable.data._helpers import ( TABLE_DEFAULT, @@ -364,6 +366,9 @@ def _execute_mutate_rows( attempt_timeout=self._attempt_timeout, metric=metric, retryable_exceptions=self._retryable_errors, + metadata=[ + ("x-goog-api-client", f"bigtable-batcher/{_bigtable_version}") + ], ) operation.start() except MutationsExceptionGroup as e: diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 74f37e6ae1e8..f0888bb1dbb2 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -699,7 +699,7 @@ def yield_rows(self, **kwargs): ) return self.read_rows(**kwargs) - def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): + def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT, metadata=()): """Mutates multiple rows in bulk. For example: @@ -788,6 +788,7 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): operation_timeout=operation_timeout, attempt_timeout=attempt_timeout, retryable_errors=retryable_errors, + metadata=metadata, ) except MutationsExceptionGroup as mut_exc_group: # We exception handle as follows: @@ -1158,7 +1159,6 @@ def restore(self, new_table_id, cluster_id=None, backup_id=None, backup_name=Non } ) - class ClusterState(object): """Representation of a Cluster State. 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 2757e8a95cdb..5aa9bda0f776 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 @@ -956,6 +956,28 @@ async def test__execute_mutate_rows(self): assert kwargs["metric"] == expected_metric assert result == [] + @CrossSync.pytest + async def test__execute_mutate_rows_passes_batcher_metadata(self): + """_execute_mutate_rows constructs _MutateRowsOperation with the batcher version header.""" + from google.cloud.bigtable.gapic_version import __version__ as _bigtable_version + + with mock.patch.object(CrossSync, "_MutateRowsOperation") as mock_op_cls: + mock_op_cls.return_value = CrossSync.Mock() + mock_op_cls.return_value.start = CrossSync.Mock(return_value=None) + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 10 + table.default_mutate_rows_attempt_timeout = 8 + table.default_mutate_rows_retryable_errors = () + async with self._make_one(table) as instance: + await instance._execute_mutate_rows([self._make_mutation()], mock.Mock()) + _, kwargs = mock_op_cls.call_args + metadata = list(kwargs.get("metadata", [])) + assert any( + k == "x-goog-api-client" + and v == f"bigtable-batcher/{_bigtable_version}" + for k, v in metadata + ) + @CrossSync.pytest async def test__execute_mutate_rows_returns_errors(self): """Errors from operation should be retruned as list""" 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 ad50533fe57c..3807eb6d8414 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 @@ -836,6 +836,29 @@ def test__execute_mutate_rows(self): assert kwargs["metric"] == expected_metric assert result == [] + def test__execute_mutate_rows_passes_batcher_metadata(self): + """_execute_mutate_rows constructs _MutateRowsOperation with the batcher version header.""" + from google.cloud.bigtable.gapic_version import __version__ as _bigtable_version + + with mock.patch.object( + CrossSync._Sync_Impl, "_MutateRowsOperation" + ) as mock_op_cls: + mock_op_cls.return_value = CrossSync.Mock() + mock_op_cls.return_value.start = CrossSync.Mock(return_value=None) + table = mock.Mock() + table.default_mutate_rows_operation_timeout = 10 + table.default_mutate_rows_attempt_timeout = 8 + table.default_mutate_rows_retryable_errors = () + with self._make_one(table) as instance: + instance._execute_mutate_rows([self._make_mutation()], mock.Mock()) + _, kwargs = mock_op_cls.call_args + metadata = list(kwargs.get("metadata", [])) + assert any( + k == "x-goog-api-client" + and v == f"bigtable-batcher/{_bigtable_version}" + for k, v in metadata + ) + def test__execute_mutate_rows_returns_errors(self): """Errors from operation should be retruned as list""" from google.cloud.bigtable.data.exceptions import ( diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py index 758ca226a12d..b76342c3d999 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_batcher.py @@ -331,6 +331,25 @@ def test_flush_async_batch_count(mocked_executor_submit): assert mocked_executor_submit.call_count == 3 +def test_flush_rows_passes_batcher_header(): + """_flush_rows calls table.mutate_rows with the batcher version header.""" + from google.cloud.bigtable.gapic_version import __version__ as _bigtable_version + + table = _Table(TABLE_NAME) + with MutationsBatcher(table=table) as batcher: + row = DirectRow(row_key=b"row_key") + row.set_cell("cf1", b"c1", 1) + batcher.mutate(row) + + assert table.mutation_calls == 1 + metadata = table.last_mutate_rows_kwargs.get("metadata", []) + assert any( + k == "x-goog-api-client" and v == f"bigtable-batcher/{_bigtable_version}" + for k, v in metadata + ) + + + class _Instance(object): def __init__(self, client=None): self._client = client @@ -341,10 +360,12 @@ def __init__(self, name, client=None): self.name = name self._instance = _Instance(client) self.mutation_calls = 0 + self.last_mutate_rows_kwargs = {} - def mutate_rows(self, rows): + def mutate_rows(self, rows, **kwargs): from google.rpc.status_pb2 import Status self.mutation_calls += 1 + self.last_mutate_rows_kwargs = kwargs return [Status(code=0) for _ in rows] diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py index a163f0a2a341..b21ca4c6a68d 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py @@ -838,6 +838,7 @@ def _table_mutate_rows_helper( operation_timeout=expected_operation_timeout, attempt_timeout=expected_attempt_timeout, retryable_errors=expected_retryable_errors, + metadata=(), ) # Check that mutation entries are in order @@ -1619,6 +1620,28 @@ def test_table_restore_table_w_backup_name(): _table_restore_helper(backup_name=BACKUP_NAME) +def test_table_mutate_rows_no_batcher_header(): + """Direct table.mutate_rows calls do not add the batcher header.""" + from google.cloud.bigtable.row import DirectRow + + credentials = _make_credentials() + client = _make_client(project="project-id", credentials=credentials, admin=True) + instance = client.instance(instance_id=INSTANCE_ID) + table = _make_table(TABLE_ID, instance) + + row = DirectRow(row_key=b"row_key", table=table) + row.set_cell("cf", b"col", b"value") + + with mock.patch.object(table._table_impl, "bulk_mutate_rows"): + table.mutate_rows([row], retry=None) + _, kwargs = table._table_impl.bulk_mutate_rows.call_args + metadata = list(kwargs.get("metadata", [])) + assert not any( + k == "x-goog-api-client" and "bigtable-batcher" in v + for k, v in metadata + ) + + def test__create_row_request_table_name_only(): from google.cloud.bigtable.table import _create_row_request