Repository: googleapis/google-cloud-python

Issue:
BigTable: Make read deadlines modifiable
The 0.31 bigtable client does a good job of retrying read requests, but the retry parameters are fixed (in `google/cloud/bigtable/row_data.py:408`). The defaults are fine for most uses, but not always, e.g. during long-running scans that really Must Work. (e.g., when doing a backup of an entire bigtable, which is where I ran into this)

It would be very helpful to give `table.read_rows` a retry argument analogous to `table.mutate_rows`.


Discussion:
@sumit-ql do you want to take a look at this?
@sduskis Yes i will work on this.

Resolving patch:
diff --git a/bigtable/google/cloud/bigtable/row_data.py b/bigtable/google/cloud/bigtable/row_data.py
--- a/bigtable/google/cloud/bigtable/row_data.py
+++ b/bigtable/google/cloud/bigtable/row_data.py
@@ -305,6 +305,20 @@ def _retry_read_rows_exception(exc):
                             exceptions.DeadlineExceeded))
 
 
+DEFAULT_RETRY_READ_ROWS = retry.Retry(
+    predicate=_retry_read_rows_exception,
+    initial=1.0,
+    maximum=15.0,
+    multiplier=2.0,
+    deadline=60.0,  # 60 seconds
+)
+"""The default retry strategy to be used on retry-able errors.
+
+Used by
+:meth:`~google.cloud.bigtable.row_data.PartialRowsData._read_next_response`.
+"""
+
+
 class PartialRowsData(object):
     """Convenience wrapper for consuming a ``ReadRows`` streaming response.
 
@@ -319,6 +333,14 @@ class PartialRowsData(object):
                     identified by self.last_scanned_row_key. The retry happens
                     inside of the Retry class, using a predicate for the
                     expected exceptions during iteration.
+
+    :type retry: :class:`~google.api_core.retry.Retry`
+    :param retry: (Optional) Retry delay and deadline arguments. To override,
+                  the default value :attr:`DEFAULT_RETRY_READ_ROWS` can be
+                  used and modified with the
+                  :meth:`~google.api_core.retry.Retry.with_delay` method
+                  or the
+                  :meth:`~google.api_core.retry.Retry.with_deadline` method.
     """
 
     NEW_ROW = 'New row'  # No cells yet complete for row
@@ -333,7 +355,8 @@ class PartialRowsData(object):
                    STATE_ROW_IN_PROGRESS: ROW_IN_PROGRESS,
                    STATE_CELL_IN_PROGRESS: CELL_IN_PROGRESS}
 
-    def __init__(self, read_method, request):
+    def __init__(self, read_method, request,
+                 retry=DEFAULT_RETRY_READ_ROWS):
         # Counter for rows returned to the user
         self._counter = 0
         # In-progress row, unset until first response, after commit/reset
@@ -349,6 +372,7 @@ def __init__(self, read_method, request):
         self.last_scanned_row_key = None
         self.read_method = read_method
         self.request = request
+        self.retry = retry
         self.response_iterator = read_method(request)
 
         self.rows = {}
@@ -405,10 +429,7 @@ def _read_next(self):
 
     def _read_next_response(self):
         """Helper for :meth:`__iter__`."""
-        retry_ = retry.Retry(
-            predicate=_retry_read_rows_exception,
-            deadline=60)
-        return retry_(self._read_next, on_error=self._on_error)()
+        return self.retry(self._read_next, on_error=self._on_error)()
 
     def __iter__(self):
         """Consume the ``ReadRowsResponse``s from the stream.
diff --git a/bigtable/google/cloud/bigtable/table.py b/bigtable/google/cloud/bigtable/table.py
--- a/bigtable/google/cloud/bigtable/table.py
+++ b/bigtable/google/cloud/bigtable/table.py
@@ -31,6 +31,7 @@
 from google.cloud.bigtable.row import ConditionalRow
 from google.cloud.bigtable.row import DirectRow
 from google.cloud.bigtable.row_data import PartialRowsData
+from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS
 from google.cloud.bigtable.row_set import RowSet
 from google.cloud.bigtable.row_set import RowRange
 from google.cloud.bigtable import enums
@@ -60,7 +61,7 @@ class _BigtableRetryableError(Exception):
     multiplier=2.0,
     deadline=120.0,  # 2 minutes
 )
-"""The default retry stategy to be used on retry-able errors.
+"""The default retry strategy to be used on retry-able errors.
 
 Used by :meth:`~google.cloud.bigtable.table.Table.mutate_rows`.
 """
@@ -298,7 +299,8 @@ def read_row(self, row_key, filter_=None):
         return row
 
     def read_rows(self, start_key=None, end_key=None, limit=None,
-                  filter_=None, end_inclusive=False, row_set=None):
+                  filter_=None, end_inclusive=False, row_set=None,
+                  retry=DEFAULT_RETRY_READ_ROWS):
         """Read rows from this table.
 
         :type start_key: bytes
@@ -329,6 +331,14 @@ def read_rows(self, start_key=None, end_key=None, limit=None,
         :param filter_: (Optional) The row set containing multiple row keys and
                         row_ranges.
 
+        :type retry: :class:`~google.api_core.retry.Retry`
+        :param retry:
+            (Optional) Retry delay and deadline arguments. To override, the
+            default value :attr:`DEFAULT_RETRY_READ_ROWS` can be used and
+            modified with the :meth:`~google.api_core.retry.Retry.with_delay`
+            method or the :meth:`~google.api_core.retry.Retry.with_deadline`
+            method.
+
         :rtype: :class:`.PartialRowsData`
         :returns: A :class:`.PartialRowsData` a generator for consuming
                   the streamed results.
@@ -340,7 +350,7 @@ def read_rows(self, start_key=None, end_key=None, limit=None,
         data_client = self._instance._client.table_data_client
         return PartialRowsData(
             data_client.transport.read_rows,
-            request_pb)
+            request_pb, retry)
 
     def yield_rows(self, **kwargs):
         """Read rows from this table.


Test patch:
diff --git a/bigtable/tests/unit/test_row_data.py b/bigtable/tests/unit/test_row_data.py
--- a/bigtable/tests/unit/test_row_data.py
+++ b/bigtable/tests/unit/test_row_data.py
@@ -14,9 +14,9 @@
 
 
 import unittest
-
 import mock
 
+from google.api_core.exceptions import DeadlineExceeded
 from ._testing import _make_credentials
 from google.cloud.bigtable.row_set import RowRange
 from google.cloud.bigtable_v2.proto import (
@@ -343,6 +343,7 @@ def _make_one(self, *args, **kwargs):
         return self._get_target_class()(*args, **kwargs)
 
     def test_constructor(self):
+        from google.cloud.bigtable.row_data import DEFAULT_RETRY_READ_ROWS
         client = _Client()
         client._data_stub = mock.MagicMock()
         request = object()
@@ -350,6 +351,19 @@ def test_constructor(self):
                                            request)
         self.assertIs(partial_rows_data.request, request)
         self.assertEqual(partial_rows_data.rows, {})
+        self.assertEqual(partial_rows_data.retry,
+                         DEFAULT_RETRY_READ_ROWS)
+
+    def test_constructor_with_retry(self):
+        client = _Client()
+        client._data_stub = mock.MagicMock()
+        request = retry = object()
+        partial_rows_data = self._make_one(client._data_stub.ReadRows,
+                                           request, retry)
+        self.assertIs(partial_rows_data.request, request)
+        self.assertEqual(partial_rows_data.rows, {})
+        self.assertEqual(partial_rows_data.retry,
+                         retry)
 
     def test___eq__(self):
         client = _Client()
@@ -642,6 +656,40 @@ def test_yield_rows_data(self):
 
         self.assertEqual(result, self.ROW_KEY)
 
+    def test_yield_retry_rows_data(self):
+        from google.api_core import retry
+        client = _Client()
+
+        retry_read_rows = retry.Retry(
+            predicate=_read_rows_retry_exception,
+        )
+
+        chunk = _ReadRowsResponseCellChunkPB(
+            row_key=self.ROW_KEY,
+            family_name=self.FAMILY_NAME,
+            qualifier=self.QUALIFIER,
+            timestamp_micros=self.TIMESTAMP_MICROS,
+            value=self.VALUE,
+            commit_row=True,
+        )
+        chunks = [chunk]
+
+        response = _ReadRowsResponseV2(chunks)
+        failure_iterator = _MockFailureIterator_1()
+        iterator = _MockCancellableIterator(response)
+        client._data_stub = mock.MagicMock()
+        client._data_stub.ReadRows.side_effect = [failure_iterator,
+                                                  iterator]
+
+        request = object()
+
+        yrd = self._make_one(client._data_stub.ReadRows, request,
+                             retry_read_rows)
+
+        result = self._consume_all(yrd)[0]
+
+        self.assertEqual(result, self.ROW_KEY)
+
     def _consume_all(self, yrd):
         return [row.row_key for row in yrd]
 
@@ -1141,6 +1189,14 @@ def next(self):
     __next__ = next
 
 
+class _MockFailureIterator_1(object):
+
+    def next(self):
+        raise DeadlineExceeded("Failed to read from server")
+
+    __next__ = next
+
+
 class _PartialCellData(object):
 
     row_key = b''
@@ -1221,3 +1277,7 @@ def _ReadRowsRequestPB(*args, **kw):
         bigtable_pb2 as messages_v2_pb2)
 
     return messages_v2_pb2.ReadRowsRequest(*args, **kw)
+
+
+def _read_rows_retry_exception(exc):
+    return isinstance(exc, DeadlineExceeded)
diff --git a/bigtable/tests/unit/test_table.py b/bigtable/tests/unit/test_table.py
--- a/bigtable/tests/unit/test_table.py
+++ b/bigtable/tests/unit/test_table.py
@@ -15,9 +15,9 @@
 
 import unittest
 
-import grpc
 import mock
 from ._testing import _make_credentials
+from google.api_core.exceptions import DeadlineExceeded
 
 
 class Test___mutate_rows_request(unittest.TestCase):
@@ -629,7 +629,7 @@ def test_read_rows(self):
                                app_profile_id=app_profile_id)
 
         # Create request_pb
-        request = object()  # Returned by our mock.
+        request = retry = object()  # Returned by our mock.
         mock_created = []
 
         def mock_create_row_request(table_name, **kwargs):
@@ -639,7 +639,7 @@ def mock_create_row_request(table_name, **kwargs):
         # Create expected_result.
         expected_result = PartialRowsData(
             client._table_data_client.transport.read_rows,
-            request)
+            request, retry)
 
         # Perform the method and check the result.
         start_key = b'start-key'
@@ -649,9 +649,10 @@ def mock_create_row_request(table_name, **kwargs):
         with _Monkey(MUT, _create_row_request=mock_create_row_request):
             result = table.read_rows(
                 start_key=start_key, end_key=end_key, filter_=filter_obj,
-                limit=limit)
+                limit=limit, retry=retry)
 
         self.assertEqual(result.rows, expected_result.rows)
+        self.assertEqual(result.retry, expected_result.retry)
         created_kwargs = {
             'start_key': start_key,
             'end_key': end_key,
@@ -663,6 +664,67 @@ def mock_create_row_request(table_name, **kwargs):
         }
         self.assertEqual(mock_created, [(table.name, created_kwargs)])
 
+    def test_read_retry_rows(self):
+        from google.cloud.bigtable_v2.gapic import bigtable_client
+        from google.cloud.bigtable_admin_v2.gapic import (
+            bigtable_table_admin_client)
+        from google.api_core import retry
+
+        data_api = bigtable_client.BigtableClient(mock.Mock())
+        table_api = bigtable_table_admin_client.BigtableTableAdminClient(
+            mock.Mock())
+        credentials = _make_credentials()
+        client = self._make_client(project='project-id',
+                                   credentials=credentials, admin=True)
+        client._table_data_client = data_api
+        client._table_admin_client = table_api
+        instance = client.instance(instance_id=self.INSTANCE_ID)
+        table = self._make_one(self.TABLE_ID, instance)
+
+        retry_read_rows = retry.Retry(
+            predicate=_read_rows_retry_exception,
+        )
+
+        # Create response_iterator
+        chunk_1 = _ReadRowsResponseCellChunkPB(
+            row_key=self.ROW_KEY_1,
+            family_name=self.FAMILY_NAME,
+            qualifier=self.QUALIFIER,
+            timestamp_micros=self.TIMESTAMP_MICROS,
+            value=self.VALUE,
+            commit_row=True
+        )
+
+        chunk_2 = _ReadRowsResponseCellChunkPB(
+            row_key=self.ROW_KEY_2,
+            family_name=self.FAMILY_NAME,
+            qualifier=self.QUALIFIER,
+            timestamp_micros=self.TIMESTAMP_MICROS,
+            value=self.VALUE,
+            commit_row=True
+        )
+
+        response_1 = _ReadRowsResponseV2([chunk_1])
+        response_2 = _ReadRowsResponseV2([chunk_2])
+        response_failure_iterator_1 = _MockFailureIterator_1()
+        response_failure_iterator_2 = _MockFailureIterator_2([response_1])
+        response_iterator = _MockReadRowsIterator(response_2)
+
+        # Patch the stub used by the API method.
+        client._table_data_client.transport.read_rows = mock.Mock(
+            side_effect=[
+                response_failure_iterator_1, response_failure_iterator_2,
+                response_iterator])
+
+        rows = []
+        for row in table.read_rows(start_key=self.ROW_KEY_1,
+                                   end_key=self.ROW_KEY_2,
+                                   retry=retry_read_rows):
+            rows.append(row)
+
+        result = rows[1]
+        self.assertEqual(result.row_key, self.ROW_KEY_2)
+
     def test_yield_retry_rows(self):
         from google.cloud.bigtable_v2.gapic import bigtable_client
         from google.cloud.bigtable_admin_v2.gapic import (
@@ -1719,16 +1781,7 @@ def next(self):
 class _MockFailureIterator_1(object):
 
     def next(self):
-        class DeadlineExceeded(grpc.RpcError, grpc.Call):
-            """ErrorDeadlineExceeded exception"""
-
-            def code(self):
-                return grpc.StatusCode.DEADLINE_EXCEEDED
-
-            def details(self):
-                return "Failed to read from server"
-
-        raise DeadlineExceeded()
+        raise DeadlineExceeded("Failed to read from server")
 
     __next__ = next
 
@@ -1740,20 +1793,11 @@ def __init__(self, *values):
         self.calls = 0
 
     def next(self):
-        class DeadlineExceeded(grpc.RpcError, grpc.Call):
-            """ErrorDeadlineExceeded exception"""
-
-            def code(self):
-                return grpc.StatusCode.DEADLINE_EXCEEDED
-
-            def details(self):
-                return "Failed to read from server"
-
         self.calls += 1
         if self.calls == 1:
             return self.iter_values[0]
         else:
-            raise DeadlineExceeded()
+            raise DeadlineExceeded("Failed to read from server")
 
     __next__ = next
 
@@ -1786,3 +1830,7 @@ def _ClusterStatePB(replication_state):
     return table_v2_pb2.Table.ClusterState(
         replication_state=replication_state
     )
+
+
+def _read_rows_retry_exception(exc):
+    return isinstance(exc, DeadlineExceeded)