Bigtable: Add retry parameter to 'Table.read_rows()'.

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.

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.