Migrate datastore iterator to use base iterator class

datastore/google/cloud/datastore/client.py

@@ -456,19 +456,30 @@ def query(self, **kwargs):
  >>> query = client.query(kind='MyKind')
  >>> query.add_filter('property', '=', 'val')
 
- Using the query iterator's
- :meth:`~google.cloud.datastore.query.Iterator.next_page` method:
+ Using the query iterator
 
  .. code-block:: python
 
  >>> query_iter = query.fetch()
- >>> entities, more_results, cursor = query_iter.next_page()
- >>> entities
- [<list of Entity unmarshalled from protobuf>]
- >>> more_results
- <boolean of more results>
- >>> cursor
- <string containing cursor where fetch stopped>
+ >>> for entity in query_iter:
+ ... do_something(entity)
+
+ or manually page through results
+
+ .. code-block:: python
+
+ >>> query_iter = query.fetch(start_cursor='2mdd223i944')
+ >>> pages = query_iter.pages
+ >>>
+ >>> first_page = next(pages)
+ >>> first_page_entities = list(first_page)
+ >>> query_iter.next_page_token
+ 'abc-some-cursor'
+ >>>
+ >>> second_page = next(pages)
+ >>> second_page_entities = list(second_page)
+ >>> query_iter.next_page_token is None
+ True
 
  Under the hood this is doing:
 

datastore/google/cloud/datastore/query.py

@@ -17,11 +17,23 @@
 import base64
 
 from google.cloud._helpers import _ensure_tuple_or_list
+from google.cloud.iterator import Iterator as BaseIterator
+from google.cloud.iterator import Page
+
 from google.cloud.datastore._generated import query_pb2 as _query_pb2
 from google.cloud.datastore import helpers
 from google.cloud.datastore.key import Key
 
 
+_NOT_FINISHED = _query_pb2.QueryResultBatch.NOT_FINISHED
+
+_FINISHED = (
+ _query_pb2.QueryResultBatch.NO_MORE_RESULTS,
+ _query_pb2.QueryResultBatch.MORE_RESULTS_AFTER_LIMIT,
+ _query_pb2.QueryResultBatch.MORE_RESULTS_AFTER_CURSOR,
+)
+
+
 class Query(object):
  """A Query against the Cloud Datastore.
 
@@ -355,18 +367,19 @@ def fetch(self, limit=None, offset=0, start_cursor=None, end_cursor=None,
  client = self._client
 
  return Iterator(
- self, client, limit, offset, start_cursor, end_cursor)
+ self, client, limit=limit, offset=offset,
+ start_cursor=start_cursor, end_cursor=end_cursor)
 
 
-class Iterator(object):
+class Iterator(BaseIterator):
  """Represent the state of a given execution of a Query.
 
- :type query: :class:`google.cloud.datastore.query.Query`
+ :type query: :class:`~google.cloud.datastore.query.Query`
  :param query: Query object holding permanent configuration (i.e.
  things that don't change on with each page in
  a results set).
 
- :type client: :class:`google.cloud.datastore.client.Client`
+ :type client: :class:`~google.cloud.datastore.client.Client`
  :param client: The client used to make a request.
 
  :type limit: int
@@ -384,97 +397,109 @@ class Iterator(object):
  query results.
  """
 
- _NOT_FINISHED = _query_pb2.QueryResultBatch.NOT_FINISHED
-
- _FINISHED = (
- _query_pb2.QueryResultBatch.NO_MORE_RESULTS,
- _query_pb2.QueryResultBatch.MORE_RESULTS_AFTER_LIMIT,
- _query_pb2.QueryResultBatch.MORE_RESULTS_AFTER_CURSOR,
- )
+ next_page_token = None
 
  def __init__(self, query, client, limit=None, offset=None,
  start_cursor=None, end_cursor=None):
+ super(Iterator, self).__init__(
+ client=client, item_to_value=_item_to_entity,
+ page_token=start_cursor, max_results=limit)
  self._query = query
- self._client = client
- self._limit = limit
  self._offset = offset
- self._start_cursor = start_cursor
  self._end_cursor = end_cursor
- self._page = self._more_results = None
- self._skipped_results = None
+ # The attributes below will change over the life of the iterator.
+ self._more_results = True
+ self._skipped_results = 0
 
- def next_page(self):
- """Fetch a single "page" of query results.
+ def _build_protobuf(self):
+ """Build a query protobuf.
 
- Low-level API for fine control: the more convenient API is
- to iterate on the current Iterator.
+ Relies on the current state of the iterator.
 
- :rtype: tuple, (entities, more_results, cursor)
- :returns: The next page of results.
+ :rtype:
+ :class:`google.cloud.datastore._generated.query_pb2.Query`
+ :returns: The query protobuf object for the current
+ state of the iterator.
  """
  pb = _pb_from_query(self._query)
 
- start_cursor = self._start_cursor
+ start_cursor = self.next_page_token
  if start_cursor is not None:
  pb.start_cursor = base64.urlsafe_b64decode(start_cursor)
 
  end_cursor = self._end_cursor
  if end_cursor is not None:
  pb.end_cursor = base64.urlsafe_b64decode(end_cursor)
 
- if self._limit is not None:
- pb.limit.value = self._limit
+ if self.max_results is not None:
+ pb.limit.value = self.max_results - self.num_results
 
  if self._offset is not None:
- pb.offset = self._offset
+ # NOTE: The offset goes down relative to the location
+ # because we are updating the cursor each time.
+ pb.offset = self._offset - self._skipped_results
 
-transaction = self._client.current_transaction
+return pb
 
- query_results = self._client.connection.run_query(
- query_pb=pb,
- project=self._query.project,
- namespace=self._query.namespace,
- transaction_id=transaction and transaction.id,
- )
- (entity_pbs, cursor_as_bytes,
- more_results_enum, self._skipped_results) = query_results
+ def _process_query_results(self, entity_pbs, cursor_as_bytes,
+ more_results_enum, skipped_results):
+ """Process the response from a datastore query.
+
+ :type entity_pbs: iterable
+ :param entity_pbs: The entities returned in the current page.
+
+ :type cursor_as_bytes: bytes
+ :param cursor_as_bytes: The end cursor of the query.
+
+ :type more_results_enum:
+ :class:`._generated.query_pb2.QueryResultBatch.MoreResultsType`
+ :param more_results_enum: Enum indicating if there are more results.
+
+ :type skipped_results: int
+ :param skipped_results: The number of skipped results.
 
- if cursor_as_bytes == b'':
- self._start_cursor = None
+ :rtype: iterable
+ :returns: The next page of entity results.
+ :raises ValueError: If ``more_results`` is an unexpected value.
+ """
+ self._skipped_results = skipped_results
+
+ if cursor_as_bytes == b'': # Empty-value for bytes.
+ self.next_page_token = None
  else:
- self._start_cursor = base64.urlsafe_b64encode(cursor_as_bytes)
+ self.next_page_token = base64.urlsafe_b64encode(cursor_as_bytes)
  self._end_cursor = None
 
- if more_results_enum == self._NOT_FINISHED:
+ if more_results_enum == _NOT_FINISHED:
  self._more_results = True
- elif more_results_enum in self._FINISHED:
+ elif more_results_enum in _FINISHED:
  self._more_results = False
  else:
  raise ValueError('Unexpected value returned for `more_results`.')
 
- self._page = [
- helpers.entity_from_protobuf(entity)
- for entity in entity_pbs]
- return self._page, self._more_results, self._start_cursor
+ return entity_pbs
 
- def __iter__(self):
- """Generator yielding all results matching our query.
+ def _next_page(self):
+ """Get the next page in the iterator.
 
- :rtype: sequence of :class:`google.cloud.datastore.entity.Entity`
+ :rtype: :class:`~google.cloud.iterator.Page`
+ :returns: The next page in the iterator (or :data:`None` if
+ there are no pages left).
  """
- while True:
- self.next_page()
- for entity in self._page:
- yield entity
- if not self._more_results:
- break
- num_results = len(self._page)
- if self._limit is not None:
- self._limit -= num_results
- if self._offset is not None and self._skipped_results is not None:
- # NOTE: The offset goes down relative to the location
- # because we are updating the cursor each time.
- self._offset -= self._skipped_results
+ if not self._more_results:
+ return None
+
+ pb = self._build_protobuf()
+ transaction = self.client.current_transaction
+
+ query_results = self.client.connection.run_query(
+ query_pb=pb,
+ project=self._query.project,
+ namespace=self._query.namespace,
+ transaction_id=transaction and transaction.id,
+ )
+ entity_pbs = self._process_query_results(*query_results)
+ return Page(self, entity_pbs, self._item_to_value)
 
 
 def _pb_from_query(query):
@@ -540,3 +565,21 @@ def _pb_from_query(query):
  pb.distinct_on.add().name = distinct_on_name
 
  return pb
+
+
+# pylint: disable=unused-argument
+def _item_to_entity(iterator, entity_pb):
+ """Convert a raw protobuf entity to the native object.
+
+ :type iterator: :class:`~google.cloud.iterator.Iterator`
+ :param iterator: The iterator that is currently in use.
+
+ :type entity_pb:
+ :class:`google.cloud.datastore._generated.entity_pb2.Entity`
+ :param entity_pb: An entity protobuf to convert to a native entity.
+
+ :rtype: :class:`~google.cloud.datastore.entity.Entity`
+ :returns: The next entity in the page.
+ """
+ return helpers.entity_from_protobuf(entity_pb)
+# pylint: enable=unused-argument