Skip to content
This repository was archived by the owner on Aug 30, 2026. It is now read-only.

Commit 6d597a1

Browse files
authored
feat: add 'timeout' arg to 'Table.mutate_rows' (#157)
Also, call data client's 'mutate_rows' directly -- do *not* scribble on its internal API wrappers. See: #7 (comment) Closes #7
1 parent 44932cb commit 6d597a1

2 files changed

Lines changed: 172 additions & 91 deletions

File tree

google/cloud/bigtable/table.py

Lines changed: 33 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,9 @@
2020
from google.api_core.exceptions import NotFound
2121
from google.api_core.exceptions import RetryError
2222
from google.api_core.exceptions import ServiceUnavailable
23+
from google.api_core.gapic_v1.method import DEFAULT
2324
from google.api_core.retry import if_exception_type
2425
from google.api_core.retry import Retry
25-
from google.api_core.gapic_v1.method import wrap_method
2626
from google.cloud._helpers import _to_bytes
2727
from google.cloud.bigtable.backup import Backup
2828
from google.cloud.bigtable.column_family import _gc_rule_from_pb
@@ -625,7 +625,7 @@ def yield_rows(self, **kwargs):
625625
)
626626
return self.read_rows(**kwargs)
627627

628-
def mutate_rows(self, rows, retry=DEFAULT_RETRY):
628+
def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
629629
"""Mutates multiple rows in bulk.
630630
631631
For example:
@@ -656,17 +656,23 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY):
656656
the :meth:`~google.api_core.retry.Retry.with_delay` method or the
657657
:meth:`~google.api_core.retry.Retry.with_deadline` method.
658658
659+
:type timeout: float
660+
:param timeout: number of seconds bounding retries for the call
661+
659662
:rtype: list
660663
:returns: A list of response statuses (`google.rpc.status_pb2.Status`)
661664
corresponding to success or failure of each row mutation
662665
sent. These will be in the same order as the `rows`.
663666
"""
667+
if timeout is DEFAULT:
668+
timeout = self.mutation_timeout
669+
664670
retryable_mutate_rows = _RetryableMutateRowsWorker(
665671
self._instance._client,
666672
self.name,
667673
rows,
668674
app_profile_id=self._app_profile_id,
669-
timeout=self.mutation_timeout,
675+
timeout=timeout,
670676
)
671677
return retryable_mutate_rows(retry=retry)
672678

@@ -1058,27 +1064,20 @@ def _do_mutate_retryable_rows(self):
10581064
# All mutations are either successful or non-retryable now.
10591065
return self.responses_statuses
10601066

1061-
mutate_rows_request = _mutate_rows_request(
1062-
self.table_name, retryable_rows, app_profile_id=self.app_profile_id
1063-
)
1067+
entries = _compile_mutation_entries(self.table_name, retryable_rows)
10641068
data_client = self.client.table_data_client
1065-
inner_api_calls = data_client._inner_api_calls
1066-
if "mutate_rows" not in inner_api_calls:
1067-
default_retry = (data_client._method_configs["MutateRows"].retry,)
1068-
if self.timeout is None:
1069-
default_timeout = data_client._method_configs["MutateRows"].timeout
1070-
else:
1071-
default_timeout = timeout.ExponentialTimeout(deadline=self.timeout)
1072-
data_client._inner_api_calls["mutate_rows"] = wrap_method(
1073-
data_client.transport.mutate_rows,
1074-
default_retry=default_retry,
1075-
default_timeout=default_timeout,
1076-
client_info=data_client._client_info,
1077-
)
1069+
1070+
kwargs = {}
1071+
if self.timeout is not None:
1072+
kwargs["timeout"] = timeout.ExponentialTimeout(deadline=self.timeout)
10781073

10791074
try:
1080-
responses = data_client._inner_api_calls["mutate_rows"](
1081-
mutate_rows_request, retry=None
1075+
responses = data_client.mutate_rows(
1076+
self.table_name,
1077+
entries,
1078+
app_profile_id=self.app_profile_id,
1079+
retry=None,
1080+
**kwargs
10821081
)
10831082
except (ServiceUnavailable, DeadlineExceeded, Aborted):
10841083
# If an exception, considered retryable by `RETRY_CODES`, is
@@ -1260,38 +1259,38 @@ def _create_row_request(
12601259
return message
12611260

12621261

1263-
def _mutate_rows_request(table_name, rows, app_profile_id=None):
1264-
"""Creates a request to mutate rows in a table.
1262+
def _compile_mutation_entries(table_name, rows):
1263+
"""Create list of mutation entries
12651264
12661265
:type table_name: str
12671266
:param table_name: The name of the table to write to.
12681267
12691268
:type rows: list
12701269
:param rows: List or other iterable of :class:`.DirectRow` instances.
12711270
1272-
:type: app_profile_id: str
1273-
:param app_profile_id: (Optional) The unique name of the AppProfile.
1274-
1275-
:rtype: :class:`data_messages_v2_pb2.MutateRowsRequest`
1276-
:returns: The ``MutateRowsRequest`` protobuf corresponding to the inputs.
1271+
:rtype: List[:class:`data_messages_v2_pb2.MutateRowsRequest.Entry`]
1272+
:returns: entries corresponding to the inputs.
12771273
:raises: :exc:`~.table.TooManyMutationsError` if the number of mutations is
1278-
greater than 100,000
1279-
"""
1280-
request_pb = data_messages_v2_pb2.MutateRowsRequest(
1281-
table_name=table_name, app_profile_id=app_profile_id
1274+
greater than the max ({})
1275+
""".format(
1276+
_MAX_BULK_MUTATIONS
12821277
)
1278+
entries = []
12831279
mutations_count = 0
1280+
entry_klass = data_messages_v2_pb2.MutateRowsRequest.Entry
1281+
12841282
for row in rows:
12851283
_check_row_table_name(table_name, row)
12861284
_check_row_type(row)
12871285
mutations = row._get_mutations()
1288-
request_pb.entries.add(row_key=row.row_key, mutations=mutations)
1286+
entries.append(entry_klass(row_key=row.row_key, mutations=mutations))
12891287
mutations_count += len(mutations)
1288+
12901289
if mutations_count > _MAX_BULK_MUTATIONS:
12911290
raise TooManyMutationsError(
12921291
"Maximum number of mutations is %s" % (_MAX_BULK_MUTATIONS,)
12931292
)
1294-
return request_pb
1293+
return entries
12951294

12961295

12971296
def _check_row_table_name(table_name, row):

0 commit comments

Comments
 (0)