fix: retry if failure occurs on initial call in MutateRows (#123)

* fix: Retry if failure occurs on initial call in MutateRows

* fix: use exception clases to DRY RETRY_CODES

Co-authored-by: Tres Seaver <tseaver@palladion.com>

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