Skip to content

Commit 0b488a3

Browse files
daniel-sanchegkevinzheng
daniel-sanche
and
gkevinzheng
authored
feat(bigtable): Reroute Mutations Batcher to use data client (#18200)
Migrating over @gkevinzheng PR from bigtable monorepo googleapis/python-bigtable#1309 ### Original description: > **Changes Made:** > > - Replaced mutations batcher implementation with one based off of the data client. > - Reworked unit tests. > - Added additional system tests. > ### Additional Changes: - 785f138: fix references in propertoes to point to data client - ac62f4c: removed deprecation language around flush_interval, since it is supported in the data client Note to reviewers: This PR has already been reviewed and merged to a staging branch, with the intention of doing a single merge to main. We are now planning to slowly rollout these changes back to the main branch. Minimal re-review should be necessary --------- Co-authored-by: Kevin Zheng <[email protected]>
1 parent 30b4f44 commit 0b488a3

11 files changed

Lines changed: 493 additions & 574 deletions

File tree

‎packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py‎

Lines changed: 59 additions & 297 deletions
Large diffs are not rendered by default.

‎packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/_mutate_rows.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -139,7 +139,7 @@ async def start(self):
139139
self._handle_entry_error(idx, exc)
140140
finally:
141141
# raise exception detailing incomplete mutations
142-
all_errors: list[Exception] = []
142+
all_errors: list[bt_exceptions.FailedMutationEntryError] = []
143143
for idx, exc_list in self.errors.items():
144144
if len(exc_list) == 0:
145145
raise core_exceptions.ClientError(

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

Lines changed: 32 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -296,10 +296,11 @@ def __init__(
296296
self._exceptions_since_last_raise: int = 0
297297
# keep track of the first and last _exception_list_limit exceptions
298298
self._exception_list_limit: int = 10
299-
self._oldest_exceptions: list[Exception] = []
300-
self._newest_exceptions: deque[Exception] = deque(
299+
self._oldest_exceptions: list[FailedMutationEntryError] = []
300+
self._newest_exceptions: deque[FailedMutationEntryError] = deque(
301301
maxlen=self._exception_list_limit
302302
)
303+
# only used by the shim right now.
303304
self._user_batch_completed_callback: (
304305
Callable[[list[status_pb2.Status]], Any] | None
305306
) = None
@@ -386,7 +387,11 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
386387
new_entries list of RowMutationEntry objects to flush
387388
"""
388389
# flush new entries
389-
in_process_requests: list[CrossSync.Future[list[FailedMutationEntryError]]] = []
390+
in_process_requests: list[
391+
tuple[
392+
CrossSync.Future[list[FailedMutationEntryError]], list[RowMutationEntry]
393+
]
394+
] = []
390395
async for batch, metric in self._flow_control.add_to_flow_with_metrics(
391396
new_entries, self._target.client._metrics
392397
):
@@ -396,7 +401,7 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
396401
metric,
397402
sync_executor=self._sync_rpc_executor,
398403
)
399-
in_process_requests.append(batch_task)
404+
in_process_requests.append((batch_task, batch))
400405
# wait for all inflight requests to complete
401406
found_exceptions = await self._wait_for_batch_results(*in_process_requests)
402407
# update exception data to reflect any new errors
@@ -466,7 +471,7 @@ async def _execute_mutate_rows(
466471
)
467472
return []
468473

469-
def _add_exceptions(self, excs: list[Exception]):
474+
def _add_exceptions(self, excs: list[FailedMutationEntryError]):
470475
"""
471476
Add new list of exceptions to internal store. To avoid unbounded memory,
472477
the batcher will store the first and last _exception_list_limit exceptions,
@@ -566,26 +571,28 @@ def _on_exit(self):
566571
@staticmethod
567572
@CrossSync.convert
568573
async def _wait_for_batch_results(
569-
*tasks: CrossSync.Future[list[FailedMutationEntryError]]
570-
| CrossSync.Future[None],
571-
) -> list[Exception]:
574+
*tasks: tuple[
575+
CrossSync.Future[list[FailedMutationEntryError]] | CrossSync.Future[None],
576+
list[RowMutationEntry],
577+
],
578+
) -> list[FailedMutationEntryError]:
572579
"""
573580
Takes in a list of futures representing _execute_mutate_rows tasks,
574581
waits for them to complete, and returns a list of errors encountered.
575582
576583
Args:
577-
*tasks: futures representing _execute_mutate_rows or _flush_internal tasks
584+
*tasks: Tuples of futures representing _execute_mutate_rows or
585+
_flush_internal tasks, and their associated batches
578586
Returns:
579-
list[Exception]:
580-
list of Exceptions encountered by any of the tasks. Errors are expected
581-
to be FailedMutationEntryError, representing a failed mutation operation.
582-
If a task fails with a different exception, it will be included in the
583-
output list. Successful tasks will not be represented in the output list.
587+
list[FailedMutationEntryError]:
588+
list of FailedMutationEntryError encountered by any of the tasks,
589+
representing a failed mutation operation.
590+
Successful tasks will not be represented in the output list.
584591
"""
585592
if not tasks:
586593
return []
587-
exceptions: list[Exception] = []
588-
for task in tasks:
594+
exceptions: list[FailedMutationEntryError] = []
595+
for task, batch in tasks:
589596
if CrossSync.is_async:
590597
# futures don't need to be awaited in sync mode
591598
await task
@@ -598,5 +605,13 @@ async def _wait_for_batch_results(
598605
exc.index = None
599606
exceptions.extend(exc_list)
600607
except Exception as e:
601-
exceptions.append(e)
608+
exceptions.extend(
609+
[
610+
FailedMutationEntryError(
611+
failed_idx=None, failed_mutation_entry=entry, cause=e
612+
)
613+
for entry in batch
614+
]
615+
)
616+
602617
return exceptions

‎packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/_mutate_rows.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,7 @@ def start(self):
112112
for idx in incomplete_indices:
113113
self._handle_entry_error(idx, exc)
114114
finally:
115-
all_errors: list[Exception] = []
115+
all_errors: list[bt_exceptions.FailedMutationEntryError] = []
116116
for idx, exc_list in self.errors.items():
117117
if len(exc_list) == 0:
118118
raise core_exceptions.ClientError(

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

Lines changed: 30 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -261,8 +261,8 @@ def __init__(
261261
self._entries_processed_since_last_raise: int = 0
262262
self._exceptions_since_last_raise: int = 0
263263
self._exception_list_limit: int = 10
264-
self._oldest_exceptions: list[Exception] = []
265-
self._newest_exceptions: deque[Exception] = deque(
264+
self._oldest_exceptions: list[FailedMutationEntryError] = []
265+
self._newest_exceptions: deque[FailedMutationEntryError] = deque(
266266
maxlen=self._exception_list_limit
267267
)
268268
self._user_batch_completed_callback: (
@@ -335,7 +335,10 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]):
335335
Args:
336336
new_entries list of RowMutationEntry objects to flush"""
337337
in_process_requests: list[
338-
CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
338+
tuple[
339+
CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]],
340+
list[RowMutationEntry],
341+
]
339342
] = []
340343
for batch, metric in self._flow_control.add_to_flow_with_metrics(
341344
new_entries, self._target.client._metrics
@@ -346,7 +349,7 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]):
346349
metric,
347350
sync_executor=self._sync_rpc_executor,
348351
)
349-
in_process_requests.append(batch_task)
352+
in_process_requests.append((batch_task, batch))
350353
found_exceptions = self._wait_for_batch_results(*in_process_requests)
351354
self._entries_processed_since_last_raise += len(new_entries)
352355
self._add_exceptions(found_exceptions)
@@ -400,7 +403,7 @@ def _execute_mutate_rows(
400403
)
401404
return []
402405

403-
def _add_exceptions(self, excs: list[Exception]):
406+
def _add_exceptions(self, excs: list[FailedMutationEntryError]):
404407
"""Add new list of exceptions to internal store. To avoid unbounded memory,
405408
the batcher will store the first and last _exception_list_limit exceptions,
406409
and discard any in between.
@@ -479,30 +482,40 @@ def _on_exit(self):
479482

480483
@staticmethod
481484
def _wait_for_batch_results(
482-
*tasks: CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
483-
| CrossSync._Sync_Impl.Future[None],
484-
) -> list[Exception]:
485+
*tasks: tuple[
486+
CrossSync._Sync_Impl.Future[list[FailedMutationEntryError]]
487+
| CrossSync._Sync_Impl.Future[None],
488+
list[RowMutationEntry],
489+
],
490+
) -> list[FailedMutationEntryError]:
485491
"""Takes in a list of futures representing _execute_mutate_rows tasks,
486492
waits for them to complete, and returns a list of errors encountered.
487493
488494
Args:
489-
*tasks: futures representing _execute_mutate_rows or _flush_internal tasks
495+
*tasks: Tuples of futures representing _execute_mutate_rows or
496+
_flush_internal tasks, and their associated batches
490497
Returns:
491-
list[Exception]:
492-
list of Exceptions encountered by any of the tasks. Errors are expected
493-
to be FailedMutationEntryError, representing a failed mutation operation.
494-
If a task fails with a different exception, it will be included in the
495-
output list. Successful tasks will not be represented in the output list."""
498+
list[FailedMutationEntryError]:
499+
list of FailedMutationEntryError encountered by any of the tasks,
500+
representing a failed mutation operation.
501+
Successful tasks will not be represented in the output list."""
496502
if not tasks:
497503
return []
498-
exceptions: list[Exception] = []
499-
for task in tasks:
504+
exceptions: list[FailedMutationEntryError] = []
505+
for task, batch in tasks:
500506
try:
501507
exc_list = task.result()
502508
if exc_list:
503509
for exc in exc_list:
504510
exc.index = None
505511
exceptions.extend(exc_list)
506512
except Exception as e:
507-
exceptions.append(e)
513+
exceptions.extend(
514+
[
515+
FailedMutationEntryError(
516+
failed_idx=None, failed_mutation_entry=entry, cause=e
517+
)
518+
for entry in batch
519+
]
520+
)
508521
return exceptions

‎packages/google-cloud-bigtable/google/cloud/bigtable/data/exceptions.py‎

Lines changed: 14 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -141,19 +141,18 @@ def __repr__(self):
141141
return f"{self.__class__.__name__}({message!r}, {self.exceptions!r})"
142142

143143

144-
# TODO: When working on mutations batcher, rework exception handling to guarantee that
145-
# MutationsExceptionGroup only stores FailedMutationEntryErrors.
146144
class MutationsExceptionGroup(_BigtableExceptionGroup):
147145
"""
148146
Represents one or more exceptions that occur during a bulk mutation operation
149147
150-
Exceptions will typically be of type FailedMutationEntryError, but other exceptions may
151-
be included if they are raised during the mutation operation
148+
Exceptions will be of type FailedMutationEntryError.
152149
"""
153150

154151
@staticmethod
155152
def _format_message(
156-
excs: list[Exception], total_entries: int, exc_count: int | None = None
153+
excs: list[FailedMutationEntryError],
154+
total_entries: int,
155+
exc_count: int | None = None,
157156
) -> str:
158157
"""
159158
Format a message for the exception group
@@ -171,7 +170,10 @@ def _format_message(
171170
return f"{exc_count} failed {entry_str} from {total_entries} attempted."
172171

173172
def __init__(
174-
self, excs: list[Exception], total_entries: int, message: str | None = None
173+
self,
174+
excs: list[FailedMutationEntryError],
175+
total_entries: int,
176+
message: str | None = None,
175177
):
176178
"""
177179
Args:
@@ -189,7 +191,10 @@ def __init__(
189191
self.total_entries_attempted = total_entries
190192

191193
def __new__(
192-
cls, excs: list[Exception], total_entries: int, message: str | None = None
194+
cls,
195+
excs: list[FailedMutationEntryError],
196+
total_entries: int,
197+
message: str | None = None,
193198
):
194199
"""
195200
Args:
@@ -209,8 +214,8 @@ def __new__(
209214
@classmethod
210215
def from_truncated_lists(
211216
cls,
212-
first_list: list[Exception],
213-
last_list: list[Exception],
217+
first_list: list[FailedMutationEntryError],
218+
last_list: list[FailedMutationEntryError],
214219
total_excs: int,
215220
entry_count: int,
216221
) -> MutationsExceptionGroup:

0 commit comments

Comments
 (0)