|
20 | 20 | from google.api_core.exceptions import NotFound |
21 | 21 | from google.api_core.exceptions import RetryError |
22 | 22 | from google.api_core.exceptions import ServiceUnavailable |
| 23 | +from google.api_core.gapic_v1.method import DEFAULT |
23 | 24 | from google.api_core.retry import if_exception_type |
24 | 25 | from google.api_core.retry import Retry |
25 | | -from google.api_core.gapic_v1.method import wrap_method |
26 | 26 | from google.cloud._helpers import _to_bytes |
27 | 27 | from google.cloud.bigtable.backup import Backup |
28 | 28 | from google.cloud.bigtable.column_family import _gc_rule_from_pb |
@@ -625,7 +625,7 @@ def yield_rows(self, **kwargs): |
625 | 625 | ) |
626 | 626 | return self.read_rows(**kwargs) |
627 | 627 |
|
628 | | - def mutate_rows(self, rows, retry=DEFAULT_RETRY): |
| 628 | + def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT): |
629 | 629 | """Mutates multiple rows in bulk. |
630 | 630 |
|
631 | 631 | For example: |
@@ -656,17 +656,23 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY): |
656 | 656 | the :meth:`~google.api_core.retry.Retry.with_delay` method or the |
657 | 657 | :meth:`~google.api_core.retry.Retry.with_deadline` method. |
658 | 658 |
|
| 659 | + :type timeout: float |
| 660 | + :param timeout: number of seconds bounding retries for the call |
| 661 | +
|
659 | 662 | :rtype: list |
660 | 663 | :returns: A list of response statuses (`google.rpc.status_pb2.Status`) |
661 | 664 | corresponding to success or failure of each row mutation |
662 | 665 | sent. These will be in the same order as the `rows`. |
663 | 666 | """ |
| 667 | + if timeout is DEFAULT: |
| 668 | + timeout = self.mutation_timeout |
| 669 | + |
664 | 670 | retryable_mutate_rows = _RetryableMutateRowsWorker( |
665 | 671 | self._instance._client, |
666 | 672 | self.name, |
667 | 673 | rows, |
668 | 674 | app_profile_id=self._app_profile_id, |
669 | | - timeout=self.mutation_timeout, |
| 675 | + timeout=timeout, |
670 | 676 | ) |
671 | 677 | return retryable_mutate_rows(retry=retry) |
672 | 678 |
|
@@ -1058,27 +1064,20 @@ def _do_mutate_retryable_rows(self): |
1058 | 1064 | # All mutations are either successful or non-retryable now. |
1059 | 1065 | return self.responses_statuses |
1060 | 1066 |
|
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) |
1064 | 1068 | 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) |
1078 | 1073 |
|
1079 | 1074 | 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 |
1082 | 1081 | ) |
1083 | 1082 | except (ServiceUnavailable, DeadlineExceeded, Aborted): |
1084 | 1083 | # If an exception, considered retryable by `RETRY_CODES`, is |
@@ -1260,38 +1259,38 @@ def _create_row_request( |
1260 | 1259 | return message |
1261 | 1260 |
|
1262 | 1261 |
|
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 |
1265 | 1264 |
|
1266 | 1265 | :type table_name: str |
1267 | 1266 | :param table_name: The name of the table to write to. |
1268 | 1267 |
|
1269 | 1268 | :type rows: list |
1270 | 1269 | :param rows: List or other iterable of :class:`.DirectRow` instances. |
1271 | 1270 |
|
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. |
1277 | 1273 | :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 |
1282 | 1277 | ) |
| 1278 | + entries = [] |
1283 | 1279 | mutations_count = 0 |
| 1280 | + entry_klass = data_messages_v2_pb2.MutateRowsRequest.Entry |
| 1281 | + |
1284 | 1282 | for row in rows: |
1285 | 1283 | _check_row_table_name(table_name, row) |
1286 | 1284 | _check_row_type(row) |
1287 | 1285 | 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)) |
1289 | 1287 | mutations_count += len(mutations) |
| 1288 | + |
1290 | 1289 | if mutations_count > _MAX_BULK_MUTATIONS: |
1291 | 1290 | raise TooManyMutationsError( |
1292 | 1291 | "Maximum number of mutations is %s" % (_MAX_BULK_MUTATIONS,) |
1293 | 1292 | ) |
1294 | | - return request_pb |
| 1293 | + return entries |
1295 | 1294 |
|
1296 | 1295 |
|
1297 | 1296 | def _check_row_table_name(table_name, row): |
|
0 commit comments