Implementing consume_*() methods on Bigtable PartialRowsData.

These methods read from a stream of ReadRowsResponse's,
parse them and then store them on the PartialRowsData
instance.

Author: dhermes

gcloud/bigtable/row_data.py

@@ -147,6 +147,19 @@ def clear(self):
         self._chunks_encountered = False
         self._cells.clear()
 
+    def update_from_read_rows(self, read_rows_response_pb):
+        """Updates the current row from a ``ReadRows`` response.
+
+        :type read_rows_response_pb:
+            :class:`._generated.bigtable_service_messages_pb2.ReadRowsResponse`
+        :param read_rows_response_pb: A response streamed back as part of a
+                                      ``ReadRows`` request.
+
+        :raises: :class:`NotImplementedError <exceptions.NotImplementedError>`
+                 always
+        """
+        raise NotImplementedError
+
 
 class PartialRowsData(object):
     """Convenience wrapper for consuming a ``ReadRows`` streaming response.
@@ -180,3 +193,45 @@ def rows(self):
         # NOTE: To avoid duplicating large objects, this is just the
         #       mutable private data.
         return self._rows
+
+    def cancel(self):
+        """Cancels the iterator, closing the stream."""
+        self._response_iterator.cancel()
+
+    def consume_next(self):
+        """Consumes the next ``ReadRowsResponse`` from the stream.
+
+        Parses the response and stores it as a :class:`PartialRowData`
+        in a dictionary owned by this object.
+
+        :raises: :class:`StopIteration <exceptions.StopIteration>` if the
+                 response iterator has no more responses to stream.
+        """
+        read_rows_response = self._response_iterator.next()
+        row_key = read_rows_response.row_key
+        partial_row = self._rows.get(row_key)
+        if partial_row is None:
+            partial_row = self._rows[row_key] = PartialRowData(row_key)
+        # NOTE: This is not atomic in the case of failures.
+        partial_row.update_from_read_rows(read_rows_response)
+
+    def consume_all(self, max_loops=None):
+        """Consume the streamed responses until there are no more.
+
+        This simply calls :meth:`consume_next` until there are no
+        more to consume.
+
+        :type max_loops: int
+        :param max_loops: (Optional) Maximum number of times to try to consume
+                          an additional ``ReadRowsResponse``. You can use this
+                          to avoid long wait times.
+        """
+        curr_loop = 0
+        if max_loops is None:
+            max_loops = float('inf')
+        while curr_loop < max_loops:
+            curr_loop += 1
+            try:
+                self.consume_next()
+            except StopIteration:
+                break

gcloud/bigtable/test_row_data.py

@@ -207,13 +207,34 @@ def test_clear(self):
         self.assertFalse(partial_row_data._chunks_encountered)
         self.assertEqual(partial_row_data.cells, {})
 
+    def test_update_from_read_rows(self):
+        partial_row_data = self._makeOne(None)
+        with self.assertRaises(NotImplementedError):
+            partial_row_data.update_from_read_rows(None)
+
 
 class TestPartialRowsData(unittest2.TestCase):
 
     def _getTargetClass(self):
         from gcloud.bigtable.row_data import PartialRowsData
         return PartialRowsData
 
+    def _getDoNothingClass(self):
+        klass = self._getTargetClass()
+
+        class FakePartialRowsData(klass):
+
+            def __init__(self, *args, **kwargs):
+                super(FakePartialRowsData, self).__init__(*args, **kwargs)
+                self._consumed = []
+
+            def consume_next(self):
+                value = self._response_iterator.next()
+                self._consumed.append(value)
+                return value
+
+        return FakePartialRowsData
+
     def _makeOne(self, *args, **kwargs):
         return self._getTargetClass()(*args, **kwargs)
 
@@ -253,3 +274,101 @@ def test_rows_getter(self):
         partial_rows_data = self._makeOne(None)
         partial_rows_data._rows = value = object()
         self.assertTrue(partial_rows_data.rows is value)
+
+    def test_cancel(self):
+        response_iterator = _MockCancellableIterator()
+        partial_rows_data = self._makeOne(response_iterator)
+        self.assertEqual(response_iterator.cancel_calls, 0)
+        partial_rows_data.cancel()
+        self.assertEqual(response_iterator.cancel_calls, 1)
+
+    def test_consume_next(self):
+        from gcloud._testing import _Monkey
+        from gcloud.bigtable._generated import (
+            bigtable_service_messages_pb2 as messages_pb2)
+        from gcloud.bigtable import row_data as MUT
+
+        row_key = b'row-key'
+        value_pb = messages_pb2.ReadRowsResponse(row_key=row_key)
+        response_iterator = _MockCancellableIterator(value_pb)
+        partial_rows_data = self._makeOne(response_iterator)
+        self.assertEqual(partial_rows_data.rows, {})
+        with _Monkey(MUT, PartialRowData=_MockPartialRowData):
+            partial_rows_data.consume_next()
+        self.assertEqual(len(partial_rows_data.rows), 1)
+        partial_row = partial_rows_data.rows[row_key]
+        self.assertTrue(isinstance(partial_row, _MockPartialRowData))
+        self.assertEqual(partial_row.row_key, row_key)
+        self.assertEqual(partial_row.read_rows_updates, [value_pb])
+
+    def test_consume_next_row_exists(self):
+        from gcloud._testing import _Monkey
+        from gcloud.bigtable._generated import (
+            bigtable_service_messages_pb2 as messages_pb2)
+        from gcloud.bigtable import row_data as MUT
+
+        row_key = b'row-key'
+        chunk = messages_pb2.ReadRowsResponse.Chunk(commit_row=True)
+        value_pb = messages_pb2.ReadRowsResponse(row_key=row_key,
+                                                 chunks=[chunk])
+        response_iterator = _MockCancellableIterator(value_pb)
+        partial_rows_data = self._makeOne(response_iterator)
+        existing_values = _MockPartialRowData(row_key)
+        partial_rows_data._rows[row_key] = existing_values
+        self.assertFalse(existing_values.committed)
+        with _Monkey(MUT, PartialRowData=_MockPartialRowData):
+            partial_rows_data.consume_next()
+        self.assertEqual(existing_values.read_rows_updates, [value_pb])
+
+    def test_consume_next_empty_iter(self):
+        response_iterator = _MockCancellableIterator()
+        partial_rows_data = self._makeOne(response_iterator)
+        with self.assertRaises(StopIteration):
+            partial_rows_data.consume_next()
+
+    def test_consume_all(self):
+        klass = self._getDoNothingClass()
+
+        value1, value2, value3 = object(), object(), object()
+        response_iterator = _MockCancellableIterator(value1, value2, value3)
+        partial_rows_data = klass(response_iterator)
+        self.assertEqual(partial_rows_data._consumed, [])
+        partial_rows_data.consume_all()
+        self.assertEqual(partial_rows_data._consumed, [value1, value2, value3])
+
+    def test_consume_all_with_max_loops(self):
+        klass = self._getDoNothingClass()
+
+        value1, value2, value3 = object(), object(), object()
+        response_iterator = _MockCancellableIterator(value1, value2, value3)
+        partial_rows_data = klass(response_iterator)
+        self.assertEqual(partial_rows_data._consumed, [])
+        partial_rows_data.consume_all(max_loops=1)
+        self.assertEqual(partial_rows_data._consumed, [value1])
+        # Make sure the iterator still has the remaining values.
+        self.assertEqual(list(response_iterator.iter_values), [value2, value3])
+
+
+class _MockCancellableIterator(object):
+
+    cancel_calls = 0
+
+    def __init__(self, *values):
+        self.iter_values = iter(values)
+
+    def cancel(self):
+        self.cancel_calls += 1
+
+    def next(self):
+        return next(self.iter_values)
+
+
+class _MockPartialRowData(object):
+
+    def __init__(self, row_key):
+        self.row_key = row_key
+        self.committed = False
+        self.read_rows_updates = []
+
+    def update_from_read_rows(self, read_rows_response):
+        self.read_rows_updates.append(read_rows_response)