Visitar URL original
feat: Added a batch completed callback to the data client mutations b… · googleapis/google-cloud-python@e79c9a4 · GitHub
Skip to content

Commit e79c9a4

Browse files
gkevinzhengdaniel-sanche
gkevinzheng
authored and
daniel-sanche
committed
feat: Added a batch completed callback to the data client mutations batcher (#1308)
**Changes made:** - Refactored logic from `Table.mutate_rows` from producing a list of `Status` protos from a `MutationsExceptionGroup` - Added private keyword argument for a batch completion callback in the MutationsBatcher. - Added unit tests/system tests.
1 parent caea489 commit e79c9a4

10 files changed

Lines changed: 489 additions & 36 deletions

File tree

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,12 @@
2424
import warnings
2525
from functools import partial
2626
from typing import (
27+
<<<<<<< ours
2728
TYPE_CHECKING,
29+
=======
30+
Callable,
31+
cast,
32+
>>>>>>> theirs
2833
Any,
2934
AsyncIterable,
3035
Callable,
@@ -146,7 +151,13 @@
146151
)
147152

148153
if TYPE_CHECKING:
154+
<<<<<<< ours
149155
from google.cloud.bigtable.data._helpers import RowKeySamples, ShardedQuery
156+
=======
157+
from google.cloud.bigtable.data._helpers import RowKeySamples
158+
from google.cloud.bigtable.data._helpers import ShardedQuery
159+
from google.rpc import status_pb2
160+
>>>>>>> theirs
150161

151162
if CrossSync.is_async:
152163
from google.cloud.bigtable.data._async.mutations_batcher import (

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,12 +14,30 @@
1414
#
1515
from __future__ import annotations
1616

17+
<<<<<<< ours
18+
=======
19+
from typing import Callable, Optional, Sequence, TYPE_CHECKING, cast
20+
>>>>>>> theirs
1721
import atexit
1822
import concurrent.futures
1923
import time
2024
import warnings
2125
from collections import deque
26+
<<<<<<< ours
2227
from typing import TYPE_CHECKING, Sequence, cast
28+
=======
29+
import concurrent.futures
30+
31+
from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
32+
from google.cloud.bigtable.data.exceptions import FailedMutationEntryError
33+
from google.cloud.bigtable.data._helpers import _get_retryable_errors
34+
from google.cloud.bigtable.data._helpers import _get_timeouts
35+
from google.cloud.bigtable.data._helpers import (
36+
_get_statuses_from_mutations_exception_group,
37+
)
38+
39+
from google.cloud.bigtable.data._helpers import TABLE_DEFAULT
40+
>>>>>>> theirs
2341

2442
from google.cloud.bigtable.data._cross_sync import CrossSync
2543
from google.cloud.bigtable.data._helpers import (
@@ -37,6 +55,9 @@
3755
Mutation,
3856
)
3957

58+
from google.rpc import code_pb2
59+
from google.rpc import status_pb2
60+
4061
if TYPE_CHECKING:
4162
from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController
4263
from google.cloud.bigtable.data.mutations import RowMutationEntry
@@ -294,6 +315,7 @@ def __init__(
294315
self._newest_exceptions: deque[Exception] = deque(
295316
maxlen=self._exception_list_limit
296317
)
318+
self._user_batch_completed_callback = None
297319
# clean up on program exit
298320
atexit.register(self._on_exit)
299321

@@ -410,6 +432,7 @@ async def _execute_mutate_rows(
410432
list of FailedMutationEntryError objects for mutations that failed.
411433
FailedMutationEntryError objects will not contain index information
412434
"""
435+
statuses = [status_pb2.Status(code=code_pb2.Code.UNKNOWN)] * len(batch)
413436
try:
414437
operation = CrossSync._MutateRowsOperation(
415438
self._target.client._gapic_client,
@@ -422,13 +445,21 @@ async def _execute_mutate_rows(
422445
)
423446
await operation.start()
424447
except MutationsExceptionGroup as e:
448+
statuses = _get_statuses_from_mutations_exception_group(e, len(batch))
449+
425450
# strip index information from exceptions, since it is not useful in a batch context
426451
for subexc in e.exceptions:
427452
subexc.index = None
428453
return list(e.exceptions)
454+
else:
455+
statuses = [status_pb2.Status(code=code_pb2.Code.OK)] * len(batch)
429456
finally:
430457
# mark batch as complete in flow control
431458
await self._flow_control.remove_from_flow(batch)
459+
460+
# Call batch done callback with list of statuses.
461+
if self._user_batch_completed_callback:
462+
self._user_batch_completed_callback(statuses)
432463
return []
433464

434465
def _add_exceptions(self, excs: list[Exception]):

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

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,11 @@
1717

1818
from __future__ import annotations
1919

20+
<<<<<<< ours
21+
=======
22+
from typing import Callable, Sequence, List, Optional, Tuple, TYPE_CHECKING, Union
23+
import time
24+
>>>>>>> theirs
2025
import enum
2126
import time
2227
from collections import namedtuple
@@ -34,7 +39,14 @@
3439
from google.api_core.retry import RetryFailureReason, exponential_sleep_generator
3540

3641
from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
42+
<<<<<<< ours
3743
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
44+
=======
45+
from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
46+
from google.rpc import code_pb2
47+
from google.rpc import status_pb2
48+
49+
>>>>>>> theirs
3850

3951
if TYPE_CHECKING:
4052
import grpc
@@ -238,6 +250,66 @@ def _align_timeouts(operation: float, attempt: float | None) -> tuple[float, flo
238250
return operation, final_attempt
239251

240252

253+
def _get_statuses_from_mutations_exception_group(
254+
exc_group: MutationsExceptionGroup, batch_size: int
255+
) -> list[status_pb2.Status]:
256+
"""
257+
Helper function that populates a list of Status objects with exception information from
258+
the exception group.
259+
260+
Args:
261+
exc_group: The exception group from a mutate rows operation
262+
batch_size: How many RowMutationGroups were provided to the batch
263+
Returns:
264+
list[status_pb2.Status]: A list of Status proto objects
265+
"""
266+
# We exception handle as follows:
267+
#
268+
# 1. Each exception in the error group is a FailedMutationEntryError, and its
269+
# cause is either a singular exception or a RetryExceptionGroup consisting of
270+
# multiple exceptions.
271+
#
272+
# 2. In the case of a singular exception, if the error does not have a gRPC status
273+
# code, we return a status code of UNKNOWN.
274+
#
275+
# 3. In the case of a RetryExceptionGroup, we use terminal exception in the exception
276+
# group and process that.
277+
statuses = [status_pb2.Status(code=code_pb2.OK)] * batch_size
278+
for error in exc_group.exceptions:
279+
if isinstance(error.index, int) and 0 <= error.index < len(statuses):
280+
cause = error.__cause__
281+
if isinstance(cause, RetryExceptionGroup):
282+
statuses[error.index] = _get_status(cause.exceptions[-1])
283+
else:
284+
statuses[error.index] = _get_status(cause)
285+
return statuses
286+
287+
288+
def _get_status(exc: Optional[Exception]) -> status_pb2.Status:
289+
"""
290+
Helper function that returns a Status object corresponding to the given exception.
291+
292+
Args:
293+
exc: An exception to be converted into a Status.
294+
Returns:
295+
status_pb2.Status: A Status proto object.
296+
"""
297+
if (
298+
isinstance(exc, core_exceptions.GoogleAPICallError)
299+
and exc.grpc_status_code is not None
300+
):
301+
return status_pb2.Status( # type: ignore[unreachable]
302+
code=exc.grpc_status_code.value[0],
303+
message=exc.message,
304+
details=exc.details,
305+
)
306+
307+
return status_pb2.Status(
308+
code=code_pb2.Code.UNKNOWN,
309+
message=str(exc) if exc else "An unknown error has occurred",
310+
)
311+
312+
241313
def _validate_timeouts(
242314
operation_timeout: float, attempt_timeout: float | None, allow_none: bool = False
243315
):

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import time
2323
import warnings
2424
from collections import deque
25+
<<<<<<< ours
2526
from typing import TYPE_CHECKING, Sequence, cast
2627

2728
from google.cloud.bigtable.data._cross_sync import CrossSync
@@ -39,6 +40,22 @@
3940
_MUTATE_ROWS_REQUEST_MUTATION_LIMIT,
4041
Mutation,
4142
)
43+
=======
44+
import concurrent.futures
45+
from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
46+
from google.cloud.bigtable.data.exceptions import FailedMutationEntryError
47+
from google.cloud.bigtable.data._helpers import _get_retryable_errors
48+
from google.cloud.bigtable.data._helpers import _get_timeouts
49+
from google.cloud.bigtable.data._helpers import (
50+
_get_statuses_from_mutations_exception_group,
51+
)
52+
from google.cloud.bigtable.data._helpers import TABLE_DEFAULT
53+
from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT
54+
from google.cloud.bigtable.data.mutations import Mutation
55+
from google.cloud.bigtable.data._cross_sync import CrossSync
56+
from google.rpc import code_pb2
57+
from google.rpc import status_pb2
58+
>>>>>>> theirs
4259

4360
if TYPE_CHECKING:
4461
from google.cloud.bigtable.data._metrics import BigtableClientSideMetricsController
@@ -259,6 +276,7 @@ def __init__(
259276
self._newest_exceptions: deque[Exception] = deque(
260277
maxlen=self._exception_list_limit
261278
)
279+
self._user_batch_completed_callback = None
262280
atexit.register(self._on_exit)
263281

264282
def _timer_routine(self, interval: float | None) -> None:
@@ -355,6 +373,7 @@ def _execute_mutate_rows(
355373
list[FailedMutationEntryError]:
356374
list of FailedMutationEntryError objects for mutations that failed.
357375
FailedMutationEntryError objects will not contain index information"""
376+
statuses = [status_pb2.Status(code=code_pb2.Code.UNKNOWN)] * len(batch)
358377
try:
359378
operation = CrossSync._Sync_Impl._MutateRowsOperation(
360379
self._target.client._gapic_client,
@@ -367,11 +386,16 @@ def _execute_mutate_rows(
367386
)
368387
operation.start()
369388
except MutationsExceptionGroup as e:
389+
statuses = _get_statuses_from_mutations_exception_group(e, len(batch))
370390
for subexc in e.exceptions:
371391
subexc.index = None
372392
return list(e.exceptions)
393+
else:
394+
statuses = [status_pb2.Status(code=code_pb2.Code.OK)] * len(batch)
373395
finally:
374396
self._flow_control.remove_from_flow(batch)
397+
if self._user_batch_completed_callback:
398+
self._user_batch_completed_callback(statuses)
375399
return []
376400

377401
def _add_exceptions(self, excs: list[Exception]):

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

Lines changed: 26 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import warnings
1818
from typing import Set
1919

20+
<<<<<<< ours
2021
from google.api_core.exceptions import (
2122
Aborted,
2223
DeadlineExceeded,
@@ -25,6 +26,13 @@
2526
NotFound,
2627
ServiceUnavailable,
2728
)
29+
=======
30+
from google.api_core.exceptions import Aborted
31+
from google.api_core.exceptions import DeadlineExceeded
32+
from google.api_core.exceptions import NotFound
33+
from google.api_core.exceptions import ServiceUnavailable
34+
from google.api_core.exceptions import InternalServerError
35+
>>>>>>> theirs
2836
from google.api_core.gapic_v1.method import DEFAULT
2937
from google.api_core.retry import Retry, if_exception_type
3038
from google.cloud._helpers import _to_bytes # type: ignore
@@ -39,12 +47,23 @@
3947
)
4048
from google.cloud.bigtable.column_family import ColumnFamily, _gc_rule_from_pb
4149
from google.cloud.bigtable.data._helpers import TABLE_DEFAULT
50+
<<<<<<< ours
4251
from google.cloud.bigtable.data.exceptions import (
4352
MutationsExceptionGroup,
4453
RetryExceptionGroup,
4554
)
4655
from google.cloud.bigtable.data.mutations import RowMutationEntry
4756
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
57+
=======
58+
from google.cloud.bigtable.data._helpers import (
59+
_get_statuses_from_mutations_exception_group,
60+
)
61+
from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
62+
from google.cloud.bigtable.data.mutations import RowMutationEntry
63+
from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
64+
from google.cloud.bigtable.batcher import MutationsBatcher
65+
from google.cloud.bigtable.batcher import FLUSH_COUNT, MAX_MUTATION_SIZE
66+
>>>>>>> theirs
4867
from google.cloud.bigtable.encryption_info import EncryptionInfo
4968
from google.cloud.bigtable.policy import Policy
5069
from google.cloud.bigtable.row import (
@@ -796,9 +815,6 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
796815
% (row.row_key, row.table.name, self.name)
797816
)
798817
mutation_entries.append(RowMutationEntry(row.row_key, row._get_mutations()))
799-
return_statuses = [
800-
status_pb2.Status(code=code_pb2.OK) for _ in range(len(mutation_entries))
801-
] # By default, return status OKs for everything
802818

803819
try:
804820
self._table_impl.bulk_mutate_rows(
@@ -808,41 +824,15 @@ def mutate_rows(self, rows, retry=DEFAULT_RETRY, timeout=DEFAULT):
808824
retryable_errors=retryable_errors,
809825
)
810826
except MutationsExceptionGroup as mut_exc_group:
811-
# We exception handle as follows:
812-
#
813-
# 1. Each exception in the error group is a FailedMutationEntryError, and its
814-
# cause is either a singular exception or a RetryExceptionGroup consisting of
815-
# multiple exceptions.
816-
#
817-
# 2. In the case of a singular exception, if the error does not have a gRPC status
818-
# code, we return a status code of UNKNOWN.
819-
#
820-
# 3. In the case of a RetryExceptionGroup, we use terminal exception in the exception
821-
# group and process that.
822-
for error in mut_exc_group.exceptions:
823-
cause = error.__cause__
824-
if isinstance(cause, RetryExceptionGroup):
825-
return_statuses[error.index] = self._get_status(
826-
cause.exceptions[-1]
827-
)
828-
else:
829-
return_statuses[error.index] = self._get_status(cause)
830-
831-
return return_statuses
832-
833-
@staticmethod
834-
def _get_status(error):
835-
if isinstance(error, GoogleAPICallError) and error.grpc_status_code is not None:
836-
return status_pb2.Status(
837-
code=error.grpc_status_code.value[0],
838-
message=error.message,
839-
details=error.details,
827+
return_statuses = _get_statuses_from_mutations_exception_group(
828+
mut_exc_group, len(mutation_entries)
829+
)
830+
else:
831+
return_statuses = [status_pb2.Status(code=code_pb2.Code.OK)] * len(
832+
mutation_entries
840833
)
841834

842-
return status_pb2.Status(
843-
code=code_pb2.UNKNOWN,
844-
message=str(error),
845-
)
835+
return return_statuses
846836

847837
def sample_row_keys(self):
848838
"""Read a sample of row keys in the table.

0 commit comments

Comments
 (0)