From 49d950300bdc7a3225bd6c564e53c86e4d877cd4 Mon Sep 17 00:00:00 2001
From: Danny Hermes <daniel.j.hermes@gmail.com>
Date: Thu, 12 May 2016 16:50:20 -0700
Subject: [PATCH] Updating offset during life of a datastore query iterator.

---
 gcloud/datastore/query.py      | 13 ++++++++-----
 gcloud/datastore/test_query.py |  6 ++++--
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/gcloud/datastore/query.py b/gcloud/datastore/query.py
index 38248638c6989..6261da3742e4a 100644
--- a/gcloud/datastore/query.py
+++ b/gcloud/datastore/query.py
@@ -363,8 +363,7 @@ class Iterator(object):
     :param limit: (Optional) Limit the number of results returned.
 
     :type offset: integer
-    :param offset: (Optional) Defaults to 0. Offset used to begin
-                   a query.
+    :param offset: (Optional) Offset used to begin a query.
 
     :type start_cursor: bytes
     :param start_cursor: (Optional) Cursor to begin paging through
@@ -382,7 +381,7 @@ class Iterator(object):
         _query_pb2.QueryResultBatch.MORE_RESULTS_AFTER_LIMIT,
     )
 
-    def __init__(self, query, client, limit=None, offset=0,
+    def __init__(self, query, client, limit=None, offset=None,
                  start_cursor=None, end_cursor=None):
         self._query = query
         self._client = client
@@ -413,7 +412,8 @@ def next_page(self):
         if self._limit is not None:
             pb.limit.value = self._limit
 
-        pb.offset = self._offset
+        if self._offset is not None:
+            pb.offset = self._offset
 
         transaction = self._client.current_transaction
 
@@ -463,8 +463,11 @@ def __iter__(self):
                 yield entity
             if not self._more_results:
                 break
+            num_results = len(self._page)
             if self._limit is not None:
-                self._limit -= len(self._page)
+                self._limit -= num_results
+            if self._offset is not None:
+                self._offset += num_results
             self.next_page()
 
 
diff --git a/gcloud/datastore/test_query.py b/gcloud/datastore/test_query.py
index de75871a1b502..02ebdea7eea1e 100644
--- a/gcloud/datastore/test_query.py
+++ b/gcloud/datastore/test_query.py
@@ -374,7 +374,7 @@ def test_ctor_defaults(self):
         iterator = self._makeOne(query, connection)
         self.assertTrue(iterator._query is query)
         self.assertEqual(iterator._limit, None)
-        self.assertEqual(iterator._offset, 0)
+        self.assertEqual(iterator._offset, None)
 
     def test_ctor_explicit(self):
         client = self._makeClient()
@@ -549,7 +549,7 @@ def test___iter___w_limit(self):
         query = _Query(client, self._KIND, self._PROJECT, self._NAMESPACE)
         self._addQueryResults(connection, more=True)
         self._addQueryResults(connection)
-        iterator = self._makeOne(query, client, limit=2)
+        iterator = self._makeOne(query, client, limit=2, offset=13)
         entities = list(iterator)
 
         self.assertFalse(iterator._more_results)
@@ -561,9 +561,11 @@ def test___iter___w_limit(self):
             self.assertEqual(entities[1]['foo'], u'Foo')
         qpb1 = _pb_from_query(query)
         qpb1.limit.value = 2
+        qpb1.offset = 13
         qpb2 = _pb_from_query(query)
         qpb2.start_cursor = self._END
         qpb2.limit.value = 1
+        qpb2.offset = 14
         EXPECTED1 = {
             'project': self._PROJECT,
             'query_pb': qpb1,