refactor retry logic on mutate row

bigtable/google/cloud/bigtable/row.py

@@ -15,10 +15,14 @@
 """User-friendly container for Google Cloud Bigtable Row."""
 
 
+import functools
 import struct
 
+import grpc
 import six
 
+from google.api_core import exceptions
+from google.api_core import retry
 from google.cloud._helpers import _datetime_from_microseconds
 from google.cloud._helpers import _microseconds_from_datetime
 from google.cloud._helpers import _to_bytes
@@ -236,6 +240,12 @@ def _delete_cells(self, column_family_id, columns, time_range=None,
             mutations_list.extend(to_append)
 
 
+def _retry_commit_exception(exc):
+    if isinstance(exc, grpc.RpcError):
+        exc = exceptions.from_grpc_error(exc)
+    return isinstance(exc, exceptions.ServiceUnavailable)
+
+
 class DirectRow(_SetDeleteRow):
     """Google Cloud Bigtable Row for sending "direct" mutations.
 
@@ -412,9 +422,15 @@ def commit(self):
             row_key=self._row_key,
             mutations=mutations_list,
         )
-        # We expect a `google.protobuf.empty_pb2.Empty`
-        client = self._table._instance._client
-        client._data_stub.MutateRow(request_pb)
+
+        commit = functools.partial(
+            self._table._instance._client._data_stub.MutateRow,
+            request_pb)
+        retry_ = retry.Retry(
+            predicate=_retry_commit_exception,
+            deadline=30)
+        retry_(commit)()
+
         self.clear()
 
     def clear(self):

bigtable/tests/unit/test_row.py

@@ -367,6 +367,26 @@ def test_commit(self):
         )])
         self.assertEqual(row._pb_mutations, [])
 
+    def test_retry_commit_exception(self):
+        import grpc
+        import mock
+
+        from google.cloud.bigtable.row import _retry_commit_exception
+
+        class ErrorUnavailable(grpc.RpcError, grpc.Call):
+            """ErrorUnavailable exception"""
+
+        message = 'Endpoint read failed'
+        error = mock.create_autospec(ErrorUnavailable, instance=True)
+        error.code.return_value = grpc.StatusCode.UNAVAILABLE
+        error.details.return_value = message
+
+        result = _retry_commit_exception(error)
+        self.assertEqual(result, True)
+
+        result = _retry_commit_exception(ValueError)
+        self.assertNotEqual(result, True)
+
     def test_commit_too_many_mutations(self):
         from google.cloud._testing import _Monkey
         from google.cloud.bigtable import row as MUT