diff --git a/google/cloud/bigtable/table.py b/google/cloud/bigtable/table.py index 199269013..950a8c3fe 100644 --- a/google/cloud/bigtable/table.py +++ b/google/cloud/bigtable/table.py @@ -14,11 +14,12 @@ """User-friendly container for Google Cloud Bigtable Table.""" -from grpc import StatusCode - from google.api_core import timeout -from google.api_core.exceptions import RetryError +from google.api_core.exceptions import Aborted +from google.api_core.exceptions import DeadlineExceeded from google.api_core.exceptions import NotFound +from google.api_core.exceptions import RetryError +from google.api_core.exceptions import ServiceUnavailable from google.api_core.retry import if_exception_type from google.api_core.retry import Retry from google.api_core.gapic_v1.method import wrap_method @@ -986,15 +987,12 @@ class _RetryableMutateRowsWorker(object): are retryable, any subsequent call on this callable will be a no-op. """ - # pylint: disable=unsubscriptable-object RETRY_CODES = ( - StatusCode.DEADLINE_EXCEEDED.value[0], - StatusCode.ABORTED.value[0], - StatusCode.UNAVAILABLE.value[0], + Aborted.grpc_status_code.value[0], + DeadlineExceeded.grpc_status_code.value[0], + ServiceUnavailable.grpc_status_code.value[0], ) - # pylint: enable=unsubscriptable-object - def __init__(self, client, table_name, rows, app_profile_id=None, timeout=None): self.client = client self.table_name = table_name @@ -1078,9 +1076,15 @@ def _do_mutate_retryable_rows(self): client_info=data_client._client_info, ) - responses = data_client._inner_api_calls["mutate_rows"]( - mutate_rows_request, retry=None - ) + try: + responses = data_client._inner_api_calls["mutate_rows"]( + mutate_rows_request, retry=None + ) + except (ServiceUnavailable, DeadlineExceeded, Aborted): + # If an exception, considered retryable by `RETRY_CODES`, is + # returned from the initial call, consider + # it to be retryable. Wrap as a Bigtable Retryable Error. + raise _BigtableRetryableError num_responses = 0 num_retryable_responses = 0