Reducing limit or the life of a datastore query iterator.

Fixes #1763.
googleapis · May 12, 2016 · 2a72031 · 2a72031
1 parent 6064270
commit 2a72031
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 2 deletions.
diff --git a/gcloud/datastore/query.py b/gcloud/datastore/query.py
@@ -463,6 +463,8 @@ def __iter__(self):
                 yield entity
             if not self._more_results:
                 break
+            if self._limit is not None:
+                self._limit -= len(self._page)
             self.next_page()
 
 

diff --git a/gcloud/datastore/test_query.py b/gcloud/datastore/test_query.py
@@ -523,9 +523,7 @@ def test___iter___w_more(self):
                 [{'kind': self._KIND, 'id': self._ID}])
             self.assertEqual(entities[1]['foo'], u'Foo')
         qpb1 = _pb_from_query(query)
-        qpb1.offset = 0
         qpb2 = _pb_from_query(query)
-        qpb2.offset = 0
         qpb2.start_cursor = self._END
         EXPECTED1 = {
             'project': self._PROJECT,
@@ -543,6 +541,45 @@ def test___iter___w_more(self):
         self.assertEqual(connection._called_with[0], EXPECTED1)
         self.assertEqual(connection._called_with[1], EXPECTED2)
 
+    def test___iter___w_limit(self):
+        from gcloud.datastore.query import _pb_from_query
+
+        connection = _Connection()
+        client = self._makeClient(connection)
+        query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE)
+        self._addQueryResults(connection, more=True)
+        self._addQueryResults(connection)
+        iterator = self._makeOne(query, client, limit=2)
+        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)
+        qpb1.limit.value = 2
+        qpb2 = _pb_from_query(query)
+        qpb2.start_cursor = self._END
+        qpb2.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,
+        }
+        self.assertEqual(len(connection._called_with), 2)
+        self.assertEqual(connection._called_with[0], EXPECTED1)
+        self.assertEqual(connection._called_with[1], EXPECTED2)
+
 
 class Test__pb_from_query(unittest2.TestCase):