fix: Correctly handle limit and offset when batching query results. (#237)

Fixes #236. It was found that when a result set spanned multiple
batches, we weren't updating `limit` and `offset` on subsequent queries.
Now we do.

Author: Chris Rossi

packages/google-cloud-ndb/google/cloud/ndb/_datastore_query.py

@@ -311,8 +311,11 @@ def _next_batch(self):
 
         if more_results:
             # Fix up query for next batch
+            limit = self._query.limit
+            if limit is not None:
+                limit -= len(self._batch)
             self._query = self._query.copy(
-                start_cursor=Cursor(batch.end_cursor)
+                start_cursor=Cursor(batch.end_cursor), offset=None, limit=limit
             )
 
     def next(self):

packages/google-cloud-ndb/tests/system/test_query.py

@@ -96,6 +96,33 @@ def make_entities():
     assert [entity.foo for entity in results][:5] == [0, 1, 2, 3, 4]
 
 
+@pytest.mark.usefixtures("client_context")
+def test_high_limit(dispose_of):
+    """Regression test for Issue #236
+
+    https://github.com/googleapis/python-ndb/issues/236
+    """
+    n_entities = 500
+
+    class SomeKind(ndb.Model):
+        foo = ndb.IntegerProperty()
+
+    @ndb.toplevel
+    def make_entities():
+        entities = [SomeKind(foo=i) for i in range(n_entities)]
+        keys = yield [entity.put_async() for entity in entities]
+        raise ndb.Return(keys)
+
+    for key in make_entities():
+        dispose_of(key._key)
+
+    query = SomeKind.query()
+    eventually(query.fetch, _length_equals(n_entities))
+    results = query.fetch(limit=400)
+
+    assert len(results) == 400
+
+
 @pytest.mark.usefixtures("client_context")
 def test_fetch_and_immediately_cancel(dispose_of):
     # Make a lot of entities so the query call won't complete before we get to

packages/google-cloud-ndb/tests/unit/test__datastore_query.py

@@ -356,6 +356,43 @@ def test__next_batch_has_more(_datastore_run_query):
         assert iterator._has_next_batch
         assert iterator._query.start_cursor.cursor == b"abc"
 
+    @staticmethod
+    @pytest.mark.usefixtures("in_context")
+    @mock.patch("google.cloud.ndb._datastore_query._datastore_run_query")
+    def test__next_batch_has_more_w_offset_and_limit(_datastore_run_query):
+        """Regression test for Issue #236
+
+        https://github.com/googleapis/python-ndb/issues/236
+        """
+        entity_results = [
+            mock.Mock(entity="entity1", cursor=b"a"),
+            mock.Mock(entity="entity2", cursor=b"b"),
+            mock.Mock(entity="entity3", cursor=b"c"),
+        ]
+        _datastore_run_query.return_value = utils.future_result(
+            mock.Mock(
+                batch=mock.Mock(
+                    entity_result_type=query_pb2.EntityResult.FULL,
+                    entity_results=entity_results,
+                    end_cursor=b"abc",
+                    more_results=query_pb2.QueryResultBatch.NOT_FINISHED,
+                )
+            )
+        )
+
+        query = query_module.QueryOptions(offset=5, limit=5)
+        iterator = _datastore_query._QueryIteratorImpl(query)
+        assert iterator._next_batch().result() is None
+        assert iterator._index == 0
+        assert len(iterator._batch) == 3
+        assert iterator._batch[0].result_pb.entity == "entity1"
+        assert iterator._batch[0].result_type == query_pb2.EntityResult.FULL
+        assert iterator._batch[0].order_by is None
+        assert iterator._has_next_batch
+        assert iterator._query.start_cursor.cursor == b"abc"
+        assert iterator._query.offset is None
+        assert iterator._query.limit == 2
+
     @staticmethod
     def test_next_done():
         iterator = _datastore_query._QueryIteratorImpl("foo")