diff --git a/datastore/google/cloud/datastore/client.py b/datastore/google/cloud/datastore/client.py index 4b4a53f70d83..fc9cbf2ea321 100644 --- a/datastore/google/cloud/datastore/client.py +++ b/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 - [] - >>> more_results - - >>> cursor - + >>> 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: diff --git a/datastore/google/cloud/datastore/query.py b/datastore/google/cloud/datastore/query.py index 8a7807496b18..a9488db725af 100644 --- a/datastore/google/cloud/datastore/query.py +++ b/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,37 +397,33 @@ 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) @@ -422,59 +431,75 @@ def next_page(self): 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 diff --git a/datastore/unit_tests/test_query.py b/datastore/unit_tests/test_query.py index 0fa333a83e7c..7c2d2238410d 100644 --- a/datastore/unit_tests/test_query.py +++ b/datastore/unit_tests/test_query.py @@ -309,34 +309,35 @@ def test_distinct_on_multiple_calls(self): self.assertEqual(query.distinct_on, _DISTINCT_ON2) def test_fetch_defaults_w_client_attr(self): + from google.cloud.datastore.query import Iterator + connection = _Connection() client = self._makeClient(connection) query = self._makeOne(client) iterator = query.fetch() + + self.assertIsInstance(iterator, Iterator) self.assertIs(iterator._query, query) - self.assertIs(iterator._client, client) - self.assertIsNone(iterator._limit) + self.assertIs(iterator.client, client) + self.assertIsNone(iterator.max_results) self.assertEqual(iterator._offset, 0) def test_fetch_w_explicit_client(self): + from google.cloud.datastore.query import Iterator + connection = _Connection() client = self._makeClient(connection) other_client = self._makeClient(connection) query = self._makeOne(client) iterator = query.fetch(limit=7, offset=8, client=other_client) + self.assertIsInstance(iterator, Iterator) self.assertIs(iterator._query, query) - self.assertIs(iterator._client, other_client) - self.assertEqual(iterator._limit, 7) + self.assertIs(iterator.client, other_client) + self.assertEqual(iterator.max_results, 7) self.assertEqual(iterator._offset, 8) class TestIterator(unittest.TestCase): - _PROJECT = 'PROJECT' - _NAMESPACE = 'NAMESPACE' - _KIND = 'KIND' - _ID = 123 - _START = b'\x00' - _END = b'\xFF' def _getTargetClass(self): from google.cloud.datastore.query import Iterator @@ -345,270 +346,195 @@ def _getTargetClass(self): def _makeOne(self, *args, **kw): return self._getTargetClass()(*args, **kw) - def _addQueryResults(self, connection, cursor=_END, more=False, - skipped_results=None, no_entity=False): - from google.cloud.datastore._generated import entity_pb2 - from google.cloud.datastore._generated import query_pb2 - from google.cloud.datastore.helpers import _new_value_pb - - if more: - more_enum = query_pb2.QueryResultBatch.NOT_FINISHED - else: - more_enum = query_pb2.QueryResultBatch.MORE_RESULTS_AFTER_LIMIT - _ID = 123 - if no_entity: - entities = [] - else: - entity_pb = entity_pb2.Entity() - entity_pb.key.partition_id.project_id = self._PROJECT - path_element = entity_pb.key.path.add() - path_element.kind = self._KIND - path_element.id = _ID - value_pb = _new_value_pb(entity_pb, 'foo') - value_pb.string_value = u'Foo' - entities = [entity_pb] - - connection._results.append( - (entities, cursor, more_enum, skipped_results)) - - def _makeClient(self, connection=None): - if connection is None: - connection = _Connection() - return _Client(self._PROJECT, connection) - - def test_ctor_defaults(self): - connection = _Connection() + def test_constructor_defaults(self): query = object() - iterator = self._makeOne(query, connection) + client = object() + iterator = self._makeOne(query, client) + + self.assertFalse(iterator._started) + self.assertIs(iterator.client, client) + self.assertIsNotNone(iterator._item_to_value) + self.assertIsNone(iterator.max_results) + self.assertEqual(iterator.page_number, 0) + self.assertIsNone(iterator.next_page_token,) + self.assertEqual(iterator.num_results, 0) self.assertIs(iterator._query, query) - self.assertIsNone(iterator._limit) self.assertIsNone(iterator._offset) - self.assertIsNone(iterator._skipped_results) + self.assertIsNone(iterator._end_cursor) + self.assertTrue(iterator._more_results) - def test_ctor_explicit(self): - client = self._makeClient() - query = _Query(client) - iterator = self._makeOne(query, client, 13, 29) + def test_constructor_explicit(self): + query = object() + client = object() + limit = 43 + offset = 9 + start_cursor = b'8290\xff' + end_cursor = b'so20rc\ta' + iterator = self._makeOne( + query, client, limit=limit, offset=offset, + start_cursor=start_cursor, end_cursor=end_cursor) + + self.assertFalse(iterator._started) + self.assertIs(iterator.client, client) + self.assertIsNotNone(iterator._item_to_value) + self.assertEqual(iterator.max_results, limit) + self.assertEqual(iterator.page_number, 0) + self.assertEqual(iterator.next_page_token, start_cursor) + self.assertEqual(iterator.num_results, 0) self.assertIs(iterator._query, query) - self.assertEqual(iterator._limit, 13) - self.assertEqual(iterator._offset, 29) + self.assertEqual(iterator._offset, offset) + self.assertEqual(iterator._end_cursor, end_cursor) + self.assertTrue(iterator._more_results) - def test_next_page_no_cursors_no_more(self): - from google.cloud.datastore.query import _pb_from_query - connection = _Connection() - client = self._makeClient(connection) - query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE) - self._addQueryResults(connection, cursor=b'') + def test__build_protobuf_empty(self): + from google.cloud.datastore._generated import query_pb2 + from google.cloud.datastore.query import Query + + client = _Client(None, None) + query = Query(client) iterator = self._makeOne(query, client) - entities, more_results, cursor = iterator.next_page() - self.assertIsNone(iterator._skipped_results) - self.assertIsNone(cursor) - self.assertFalse(more_results) - self.assertFalse(iterator._more_results) - self.assertEqual(len(entities), 1) - self.assertEqual(entities[0].key.path, - [{'kind': self._KIND, 'id': self._ID}]) - self.assertEqual(entities[0]['foo'], u'Foo') - qpb = _pb_from_query(query) - qpb.offset = 0 - EXPECTED = { - 'project': self._PROJECT, - 'query_pb': qpb, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - self.assertEqual(connection._called_with, [EXPECTED]) + pb = iterator._build_protobuf() + expected_pb = query_pb2.Query() + self.assertEqual(pb, expected_pb) - def test_next_page_no_cursors_no_more_w_offset_and_limit(self): - from google.cloud.datastore.query import _pb_from_query - connection = _Connection() - client = self._makeClient(connection) - query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE) - skipped_results = object() - self._addQueryResults(connection, cursor=b'', - skipped_results=skipped_results) - iterator = self._makeOne(query, client, 13, 29) - entities, more_results, cursor = iterator.next_page() - - self.assertIsNone(cursor) - self.assertFalse(more_results) - self.assertFalse(iterator._more_results) - self.assertEqual(iterator._skipped_results, skipped_results) - self.assertEqual(len(entities), 1) - self.assertEqual(entities[0].key.path, - [{'kind': self._KIND, 'id': self._ID}]) - self.assertEqual(entities[0]['foo'], u'Foo') - qpb = _pb_from_query(query) - qpb.limit.value = 13 - qpb.offset = 29 - EXPECTED = { - 'project': self._PROJECT, - 'query_pb': qpb, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - self.assertEqual(connection._called_with, [EXPECTED]) + def test__build_protobuf_all_values(self): + from google.cloud.datastore._generated import query_pb2 + from google.cloud.datastore.query import Query - def test_next_page_w_cursors_w_more(self): - from base64 import urlsafe_b64decode - from base64 import urlsafe_b64encode - from google.cloud.datastore.query import _pb_from_query - connection = _Connection() - client = self._makeClient(connection) - query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE) - self._addQueryResults(connection, cursor=self._END, more=True) - iterator = self._makeOne(query, client) - iterator._start_cursor = self._START - iterator._end_cursor = self._END - entities, more_results, cursor = iterator.next_page() + client = _Client(None, None) + query = Query(client) + limit = 15 + offset = 9 + start_bytes = b'i\xb7\x1d' + start_cursor = 'abcd' + end_bytes = b'\xc3\x1c\xb3' + end_cursor = 'wxyz' + iterator = self._makeOne( + query, client, limit=limit, offset=offset, + start_cursor=start_cursor, end_cursor=end_cursor) + self.assertEqual(iterator.max_results, limit) + iterator.num_results = 4 + iterator._skipped_results = 1 + + pb = iterator._build_protobuf() + expected_pb = query_pb2.Query( + start_cursor=start_bytes, + end_cursor=end_bytes, + offset=offset - iterator._skipped_results, + ) + expected_pb.limit.value = limit - iterator.num_results + self.assertEqual(pb, expected_pb) + + def test__process_query_results(self): + from google.cloud.datastore._generated import query_pb2 + + iterator = self._makeOne(None, None, + end_cursor='abcd') + self.assertIsNotNone(iterator._end_cursor) - self.assertEqual(cursor, urlsafe_b64encode(self._END)) - self.assertTrue(more_results) + entity_pbs = object() + cursor_as_bytes = b'\x9ai\xe7' + cursor = b'mmnn' + skipped_results = 4 + more_results_enum = query_pb2.QueryResultBatch.NOT_FINISHED + result = iterator._process_query_results( + entity_pbs, cursor_as_bytes, + more_results_enum, skipped_results) + self.assertIs(result, entity_pbs) + + self.assertEqual(iterator._skipped_results, skipped_results) + self.assertEqual(iterator.next_page_token, cursor) self.assertTrue(iterator._more_results) - self.assertIsNone(iterator._skipped_results) - self.assertIsNone(iterator._end_cursor) - self.assertEqual(urlsafe_b64decode(iterator._start_cursor), self._END) - self.assertEqual(len(entities), 1) - self.assertEqual(entities[0].key.path, - [{'kind': self._KIND, 'id': self._ID}]) - self.assertEqual(entities[0]['foo'], u'Foo') - qpb = _pb_from_query(query) - qpb.offset = 0 - qpb.start_cursor = urlsafe_b64decode(self._START) - qpb.end_cursor = urlsafe_b64decode(self._END) - EXPECTED = { - 'project': self._PROJECT, - 'query_pb': qpb, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - self.assertEqual(connection._called_with, [EXPECTED]) - def test_next_page_w_cursors_w_bogus_more(self): - connection = _Connection() - client = self._makeClient(connection) - query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE) - self._addQueryResults(connection, cursor=self._END, more=True) - epb, cursor, _, _ = connection._results.pop() - connection._results.append((epb, cursor, 5, None)) # invalid enum - iterator = self._makeOne(query, client) - self.assertRaises(ValueError, iterator.next_page) + def test__process_query_results_done(self): + from google.cloud.datastore._generated import query_pb2 - def test___iter___no_more(self): - from google.cloud.datastore.query import _pb_from_query - connection = _Connection() - client = self._makeClient(connection) - query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE) - self._addQueryResults(connection) - iterator = self._makeOne(query, client) - entities = list(iterator) + iterator = self._makeOne(None, None, + end_cursor='abcd') + self.assertIsNotNone(iterator._end_cursor) + + entity_pbs = object() + cursor_as_bytes = b'' + skipped_results = 44 + more_results_enum = query_pb2.QueryResultBatch.NO_MORE_RESULTS + result = iterator._process_query_results( + entity_pbs, cursor_as_bytes, + more_results_enum, skipped_results) + self.assertIs(result, entity_pbs) + self.assertEqual(iterator._skipped_results, skipped_results) + self.assertIsNone(iterator.next_page_token) self.assertFalse(iterator._more_results) - self.assertEqual(len(entities), 1) - self.assertEqual(entities[0].key.path, - [{'kind': self._KIND, 'id': self._ID}]) - self.assertEqual(entities[0]['foo'], u'Foo') - qpb = _pb_from_query(query) - qpb.offset = 0 - EXPECTED = { - 'project': self._PROJECT, - 'query_pb': qpb, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - self.assertEqual(connection._called_with, [EXPECTED]) - def test___iter___w_more(self): - from google.cloud.datastore.query import _pb_from_query + def test__process_query_results_bad_enum(self): + iterator = self._makeOne(None, None) + more_results_enum = 999 + with self.assertRaises(ValueError): + iterator._process_query_results( + None, b'', more_results_enum, None) + + def test__next_page(self): + from google.cloud.iterator import Page + from google.cloud.datastore._generated import query_pb2 + from google.cloud.datastore.query import Query + connection = _Connection() - client = self._makeClient(connection) - query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE) - self._addQueryResults(connection, cursor=self._END, more=True) - self._addQueryResults(connection) + more_enum = query_pb2.QueryResultBatch.NOT_FINISHED + result = ([], b'', more_enum, 0) + connection._results = [result] + project = 'prujekt' + client = _Client(project, connection) + query = Query(client) iterator = self._makeOne(query, client) - entities = list(iterator) - self.assertFalse(iterator._more_results) - self.assertEqual(len(entities), 2) - for entity in entities: - self.assertEqual( - entity.key.path, - [{'kind': self._KIND, 'id': self._ID}]) - self.assertEqual(entities[1]['foo'], u'Foo') - qpb1 = _pb_from_query(query) - qpb2 = _pb_from_query(query) - qpb2.start_cursor = self._END - EXPECTED1 = { - 'project': self._PROJECT, - 'query_pb': qpb1, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - EXPECTED2 = { - 'project': self._PROJECT, - 'query_pb': qpb2, - 'namespace': self._NAMESPACE, + page = iterator._next_page() + self.assertIsInstance(page, Page) + self.assertIs(page._parent, iterator) + + self.assertEqual(connection._called_with, [{ + 'query_pb': query_pb2.Query(), + 'project': project, + 'namespace': None, 'transaction_id': None, - } - self.assertEqual(len(connection._called_with), 2) - self.assertEqual(connection._called_with[0], EXPECTED1) - self.assertEqual(connection._called_with[1], EXPECTED2) + }]) - def test___iter___w_limit(self): - from google.cloud.datastore.query import _pb_from_query + def test__next_page_no_more(self): + from google.cloud.datastore.query import Query connection = _Connection() - client = self._makeClient(connection) - query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE) - skip1 = 4 - skip2 = 9 - self._addQueryResults(connection, more=True, skipped_results=skip1, - no_entity=True) - self._addQueryResults(connection, more=True, skipped_results=skip2) - self._addQueryResults(connection) - offset = skip1 + skip2 - iterator = self._makeOne(query, client, limit=2, offset=offset) - entities = list(iterator) + client = _Client(None, connection) + query = Query(client) + iterator = self._makeOne(query, client) + iterator._more_results = False - self.assertFalse(iterator._more_results) - self.assertEqual(len(entities), 2) - for entity in entities: - self.assertEqual( - entity.key.path, - [{'kind': self._KIND, 'id': self._ID}]) - qpb1 = _pb_from_query(query) - qpb1.limit.value = 2 - qpb1.offset = offset - qpb2 = _pb_from_query(query) - qpb2.start_cursor = self._END - qpb2.limit.value = 2 - qpb2.offset = offset - skip1 - qpb3 = _pb_from_query(query) - qpb3.start_cursor = self._END - qpb3.limit.value = 1 - EXPECTED1 = { - 'project': self._PROJECT, - 'query_pb': qpb1, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - EXPECTED2 = { - 'project': self._PROJECT, - 'query_pb': qpb2, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - EXPECTED3 = { - 'project': self._PROJECT, - 'query_pb': qpb3, - 'namespace': self._NAMESPACE, - 'transaction_id': None, - } - self.assertEqual(len(connection._called_with), 3) - self.assertEqual(connection._called_with[0], EXPECTED1) - self.assertEqual(connection._called_with[1], EXPECTED2) - self.assertEqual(connection._called_with[2], EXPECTED3) + page = iterator._next_page() + self.assertIsNone(page) + self.assertEqual(connection._called_with, []) + + +class Test__item_to_entity(unittest.TestCase): + + def _callFUT(self, iterator, entity_pb): + from google.cloud.datastore.query import _item_to_entity + return _item_to_entity(iterator, entity_pb) + + def test_it(self): + from google.cloud._testing import _Monkey + from google.cloud.datastore import helpers + + result = object() + entities = [] + + def mocked(entity_pb): + entities.append(entity_pb) + return result + + entity_pb = object() + with _Monkey(helpers, entity_from_protobuf=mocked): + self.assertIs(result, self._callFUT(None, entity_pb)) + + self.assertEqual(entities, [entity_pb]) class Test__pb_from_query(unittest.TestCase): diff --git a/system_tests/clear_datastore.py b/system_tests/clear_datastore.py index 19ce60725e76..9dc24a49dc28 100644 --- a/system_tests/clear_datastore.py +++ b/system_tests/clear_datastore.py @@ -45,9 +45,8 @@ def fetch_keys(kind, client, fetch_max=FETCH_MAX, query=None, cursor=None): query.keys_only() iterator = query.fetch(limit=fetch_max, start_cursor=cursor) - - entities, _, cursor = iterator.next_page() - return query, entities, cursor + page = six.next(iterator.pages) + return query, list(page), iterator.next_page_token def get_ancestors(entities): diff --git a/system_tests/datastore.py b/system_tests/datastore.py index 18346cf10219..4508b280e649 100644 --- a/system_tests/datastore.py +++ b/system_tests/datastore.py @@ -17,6 +17,7 @@ import unittest import httplib2 +import six from google.cloud._helpers import UTC from google.cloud import datastore @@ -260,14 +261,16 @@ def test_limit_queries(self): # Fetch characters. iterator = query.fetch(limit=limit) - character_entities, _, cursor = iterator.next_page() + page = six.next(iterator.pages) + character_entities = list(page) + cursor = iterator.next_page_token self.assertEqual(len(character_entities), limit) # Check cursor after fetch. self.assertIsNotNone(cursor) - # Fetch remaining of characters. - new_character_entities = list(iterator) + # Fetch remaining characters. + new_character_entities = list(query.fetch(start_cursor=cursor)) characters_remaining = len(self.CHARACTERS) - limit self.assertEqual(len(new_character_entities), characters_remaining) @@ -362,7 +365,9 @@ def test_query_paginate_with_offset(self): iterator = page_query.fetch(limit=limit, offset=offset) # Fetch characters. - entities, _, cursor = iterator.next_page() + page = six.next(iterator.pages) + entities = list(page) + cursor = iterator.next_page_token self.assertEqual(len(entities), limit) self.assertEqual(entities[0]['name'], 'Robb') self.assertEqual(entities[1]['name'], 'Bran') @@ -385,7 +390,9 @@ def test_query_paginate_with_start_cursor(self): iterator = page_query.fetch(limit=limit, offset=offset) # Fetch characters. - entities, _, cursor = iterator.next_page() + page = six.next(iterator.pages) + entities = list(page) + cursor = iterator.next_page_token self.assertEqual(len(entities), limit) # Use cursor to create a fresh query.