Bigtable: Add 'PartialRowsData.cancel'.

bigtable/google/cloud/bigtable/row_data.py

@@ -376,6 +376,9 @@ def __init__(self, read_method, request, retry=DEFAULT_RETRY_READ_ROWS):
         self.rows = {}
         self._state = self.STATE_NEW_ROW
 
+        # Flag to stop iteration, for any reason not related to self.retry()
+        self.stop = False
+
     @property
     def state(self):
         """State machine state.
@@ -439,18 +442,19 @@ def __iter__(self):
         while True:
             try:
                 response = self._read_next_response()
+                for chunk in response.chunks:
+                    if self.stop:
+                        raise StopIteration
+                    self._process_chunk(chunk)
+                    if chunk.commit_row:
+                        self.last_scanned_row_key = self._previous_row.row_key
+                        self._counter += 1
+                        yield self._previous_row
             except StopIteration:
                 if self.state != self.NEW_ROW:
                     raise ValueError("The row remains partial / is not committed.")
                 break
 
-            for chunk in response.chunks:
-                self._process_chunk(chunk)
-                if chunk.commit_row:
-                    self.last_scanned_row_key = self._previous_row.row_key
-                    self._counter += 1
-                    yield self._previous_row
-
             resp_last_key = response.last_scanned_row_key
             if resp_last_key and resp_last_key > self.last_scanned_row_key:
                 self.last_scanned_row_key = resp_last_key