Add retry for yield_rows

bigtable/google/cloud/bigtable/row.py

@@ -243,7 +243,8 @@ def _delete_cells(self, column_family_id, columns, time_range=None,
 def _retry_commit_exception(exc):
     if isinstance(exc, grpc.RpcError):
         exc = exceptions.from_grpc_error(exc)
-    return isinstance(exc, exceptions.ServiceUnavailable)
+    return isinstance(exc, (exceptions.ServiceUnavailable,
+                            exceptions.DeadlineExceeded))
 
 
 class DirectRow(_SetDeleteRow):

bigtable/google/cloud/bigtable/row_data.py

@@ -251,8 +251,6 @@ def __init__(self, response_iterator):
         self._response_iterator = response_iterator
         # Counter for responses pulled from iterator
         self._counter = 0
-        # Maybe cached from previous response
-        self._last_scanned_row_key = None
         # In-progress row, unset until first response, after commit/reset
         self._row = None
         # Last complete row, unset until first commit
@@ -261,6 +259,8 @@ def __init__(self, response_iterator):
         self._cell = None
         # Last complete cell, unset until first completion, after new row
         self._previous_cell = None
+        # May be cached from previous response
+        self.last_scanned_row_key = None
 
     @property
     def state(self):
@@ -270,7 +270,7 @@ def state(self):
         :returns:  name of state corresponding to current row / chunk
                    processing.
         """
-        if self._last_scanned_row_key is None:
+        if self.last_scanned_row_key is None:
             return self.START
         if self._row is None:
             assert self._cell is None
@@ -301,11 +301,11 @@ def read_rows(self):
 
             self._counter += 1
 
-            if self._last_scanned_row_key is None:  # first response
+            if self.last_scanned_row_key is None:  # first response
                 if response.last_scanned_row_key:
                     raise InvalidReadRowsResponse()
 
-            self._last_scanned_row_key = response.last_scanned_row_key
+            self.last_scanned_row_key = response.last_scanned_row_key
 
             row = self._row
             cell = self._cell
@@ -343,6 +343,7 @@ def read_rows(self):
 
                     yield self._row
 
+                    self.last_scanned_row_key = self._row.row_key
                     self._row, self._previous_row = None, self._row
                     self._previous_cell = None
                     row = cell = None

bigtable/google/cloud/bigtable/table.py

@@ -15,6 +15,11 @@
 """User-friendly container for Google Cloud Bigtable Table."""
 
 
+import functools
+import logging
+
+from grpc import StatusCode
+
 from google.api_core.exceptions import RetryError
 from google.api_core.retry import if_exception_type
 from google.api_core.retry import Retry
@@ -27,18 +32,19 @@
     table_pb2 as table_v2_pb2)
 from google.cloud.bigtable.column_family import _gc_rule_from_pb
 from google.cloud.bigtable.column_family import ColumnFamily
+from google.cloud.bigtable.row import _retry_commit_exception
 from google.cloud.bigtable.row import AppendRow
 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 YieldRowsData
-from grpc import StatusCode
 
 
 # Maximum number of mutations in bulk (MutateRowsRequest message):
 # (https://cloud.google.com/bigtable/docs/reference/data/rpc/
 #  google.bigtable.v2#google.bigtable.v2.MutateRowRequest)
 _MAX_BULK_MUTATIONS = 100000
+logging.getLogger().setLevel(logging.INFO)
 
 
 class _BigtableRetryableError(Exception):
@@ -317,7 +323,7 @@ def read_rows(self, start_key=None, end_key=None, limit=None,
         return PartialRowsData(response_iterator)
 
     def yield_rows(self, start_key=None, end_key=None, limit=None,
-                   filter_=None):
+                   filter_=None, retry=True):
         """Read rows from this table.
 
         :type start_key: bytes
@@ -340,17 +346,18 @@ def yield_rows(self, start_key=None, end_key=None, limit=None,
                         specified row(s). If unset, reads every column in
                         each row.
 
+        :type retry: bool
+        :param retry: (Optional) Apply retry on read rows.
+
         :rtype: :class:`.PartialRowData`
         :returns: A :class:`.PartialRowData` for each row returned
         """
-        request_pb = _create_row_request(
-            self.name, start_key=start_key, end_key=end_key, filter_=filter_,
-            limit=limit)
-        client = self._instance._client
-        response_iterator = client._data_stub.ReadRows(request_pb)
-        # We expect an iterator of `data_messages_v2_pb2.ReadRowsResponse`
-        generator = YieldRowsData(response_iterator)
-        for row in generator.read_rows():
+
+        retryable_read_rows = _RetryableReadRows(
+            self._instance._client, self.name, start_key,
+            end_key, filter_, limit, retry)
+
+        for row in retryable_read_rows():
             yield row
 
     def mutate_rows(self, rows, retry=DEFAULT_RETRY):
@@ -532,8 +539,59 @@ def _do_mutate_retryable_rows(self):
         return self.responses_statuses
 
 
+class _RetryableReadRows(object):
+    """A callable worker that can retry to read rows with transient errors.
+
+    This class is a callable that can retry reading rows that result in
+    transient errors.
+    """
+
+    def __init__(self, client, table_name, start_key,
+                 end_key, filter_, limit, retry):
+        self.generator = None
+        self.client = client
+        self.table_name = table_name
+        self.start_key = start_key
+        self.end_key = end_key
+        self.filter_ = filter_
+        self.limit = limit
+        self.retry = retry
+
+    def __call__(self):
+        if self.retry:
+            read_rows = functools.partial(self._do_read_retryable_rows)
+            retry_ = Retry(
+                predicate=_retry_commit_exception,
+                deadline=30)
+            return retry_(read_rows)()
+        else:
+            return self._do_read_retryable_rows()
+
+    def _do_read_retryable_rows(self):
+        if (self.generator and self.generator.last_scanned_row_key):
+            next_start_key = self.generator.last_scanned_row_key
+            logging.info('Start key is {} for retry read rows.'
+                         .format(next_start_key))
+        else:
+            next_start_key = self.start_key
+
+        request_pb = _create_row_request(
+            self.table_name,
+            start_key=next_start_key,
+            end_key=self.end_key,
+            filter_=self.filter_,
+            limit=self.limit,
+            start_inclusive=False)
+        client = self.client
+        response_iterator = client._data_stub.ReadRows(request_pb)
+        self.generator = YieldRowsData(response_iterator)
+        for row in self.generator.read_rows():
+            yield row
+
+
 def _create_row_request(table_name, row_key=None, start_key=None, end_key=None,
-                        filter_=None, limit=None, end_inclusive=False):
+                        filter_=None, limit=None,
+                        end_inclusive=False, start_inclusive=True):
     """Creates a request to read rows in a table.
 
     :type table_name: str
@@ -578,7 +636,10 @@ def _create_row_request(table_name, row_key=None, start_key=None, end_key=None,
     range_kwargs = {}
     if start_key is not None or end_key is not None:
         if start_key is not None:
-            range_kwargs['start_key_closed'] = _to_bytes(start_key)
+            start_key_key = 'start_key_open'
+            if start_inclusive:
+                start_key_key = 'start_key_closed'
+            range_kwargs[start_key_key] = _to_bytes(start_key)
         if end_key is not None:
             end_key_key = 'end_key_open'
             if end_inclusive:

bigtable/tests/unit/test_row_data.py

@@ -256,7 +256,7 @@ def test_state_start(self):
 
     def test_state_new_row_w_row(self):
         yrd = self._make_one([])
-        yrd._last_scanned_row_key = ''
+        yrd.last_scanned_row_key = ''
         yrd._row = object()
         self.assertEqual(yrd.state, yrd.NEW_ROW)
 
@@ -370,9 +370,9 @@ def test_valid_last_scanned_row_key_on_start(self):
             chunks=(), last_scanned_row_key='AFTER')
         iterator = _MockCancellableIterator(response)
         yrd = self._make_one(iterator)
-        yrd._last_scanned_row_key = 'BEFORE'
+        yrd.last_scanned_row_key = 'BEFORE'
         self._consume_all(yrd)
-        self.assertEqual(yrd._last_scanned_row_key, 'AFTER')
+        self.assertEqual(yrd.last_scanned_row_key, 'AFTER')
 
     def test_invalid_empty_chunk(self):
         from google.cloud.bigtable.row_data import InvalidChunk

bigtable/tests/unit/test_table.py

@@ -127,6 +127,8 @@ class TestTable(unittest.TestCase):
     TABLE_ID = 'table-id'
     TABLE_NAME = INSTANCE_NAME + '/tables/' + TABLE_ID
     ROW_KEY = b'row-key'
+    ROW_KEY_1 = b'row-key-1'
+    ROW_KEY_2 = b'row-key-2'
     FAMILY_NAME = u'family'
     QUALIFIER = b'qualifier'
     TIMESTAMP_MICROS = 100
@@ -541,18 +543,109 @@ def test_yield_rows(self):
         chunks = [chunk]
 
         response = _ReadRowsResponseV2(chunks)
-        response_iterator = _MockCancellableIterator(response)
+        response_iterator = _MockReadRowsIterator(response)
 
         # Patch the stub used by the API method.
         client._data_stub = _FakeStub(response_iterator)
 
+        generator = table.yield_rows(retry=False)
         rows = []
-        for row in table.yield_rows():
+        for row in generator:
             rows.append(row)
         result = rows[0]
 
         self.assertEqual(result.row_key, self.ROW_KEY)
 
+    def test_yield_retry_rows_with_response_exception(self):
+        # import grpc
+
+        client = _Client()
+        instance = _Instance(self.INSTANCE_NAME, client=client)
+        table = self._make_one(self.TABLE_ID, instance)
+
+        # class ErrorUnavailable(grpc.RpcError, grpc.Call):
+        #     """ErrorUnavailable exception"""
+        #
+        #     def code(self):
+        #         return grpc.StatusCode.UNAVAILABLE
+        #
+        #     def details(self):
+        #         return 'Endpoint read failed'
+        #
+        # class ErrorDeadlineExceeded(grpc.RpcError, grpc.Call):
+        #     """ErrorDeadlineExceeded exception"""
+        #
+        #     def code(self):
+        #         return grpc.StatusCode.DEADLINE_EXCEEDED
+        #
+        #     def details(self):
+        #         return 'Error while reading table'
+
+        # Create response_iterator
+        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
+        )
+
+        response = _ReadRowsResponseV2([chunk])
+        response_iterator= _MockReadRowsIterator(response)
+
+        # Patch the stub used by the API method.
+        client._data_stub = mock.MagicMock()
+        client._data_stub.ReadRows.side_effect = [response_iterator]
+
+        rows = []
+        for row in table.yield_rows():
+            rows.append(row)
+
+        result = rows[0]
+        self.assertEqual(result.row_key, self.ROW_KEY)
+
+    def test_yield_retry_rows(self):
+        client = _Client()
+        instance = _Instance(self.INSTANCE_NAME, client=client)
+        table = self._make_one(self.TABLE_ID, instance)
+
+        # 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_retryable_iterator = _MockRetryableIterator([response_1,
+                                                              response_2])
+
+        # Patch the stub used by the API method.
+        client._data_stub = mock.MagicMock()
+        client._data_stub.ReadRows.side_effect = [response_retryable_iterator]
+
+        rows = []
+        for row in table.yield_rows(start_key=self.ROW_KEY_1,
+                                    end_key=self.ROW_KEY_2):
+            rows.append(row)
+
+        result = rows[1]
+        self.assertEqual(result.row_key, self.ROW_KEY_2)
+
     def test_sample_row_keys(self):
         from tests.unit._testing import _FakeStub
 
@@ -1218,10 +1311,7 @@ def __init__(self, name, client=None):
         self._client = client
 
 
-class _MockCancellableIterator(object):
-
-    cancel_calls = 0
-
+class _MockReadRowsIterator(object):
     def __init__(self, *values):
         self.iter_values = iter(values)
 
@@ -1232,8 +1322,41 @@ def __next__(self):  # pragma: NO COVER Py3k
         return self.next()
 
 
+class _MockRetryableIterator(object):
+
+    def __init__(self, *values):
+        self.iter_values = values[0]
+        self.calls = 0
+
+    def next(self):
+        # import grpc
+        #
+        # class ErrorUnavailable(grpc.RpcError, grpc.Call):
+        #     """ErrorUnavailable exception"""
+        #
+        #     def code(self):
+        #         return grpc.StatusCode.UNAVAILABLE
+        #
+        #     def details(self):
+        #         return 'Endpoint read failed'
+
+        self.calls += 1
+        if self.calls == 1:
+            first = self.iter_values[0]
+            return first
+        # elif self.calls == 2:
+        #     raise ErrorUnavailable()
+        elif self.calls == 2:
+            return self.iter_values[1]
+        else:
+            raise StopIteration()
+
+    def __next__(self):  # pragma: NO COVER Py3k
+        return self.next()
+
+
 class _ReadRowsResponseV2(object):
 
     def __init__(self, chunks, last_scanned_row_key=''):
         self.chunks = chunks
-        self.last_scanned_row_key = last_scanned_row_key
+        self.last_scanned_row_key = last_scanned_row_key