Do not pass 'offset' once the query iterator has a cursor (googleapis#5503)

mcdonc · tseaver · commit fde9b59d8f5a · 2018-06-15T16:48:49.000-04:00
Closes googleapis#4675.
diff --git a/datastore/google/cloud/datastore/query.py b/datastore/google/cloud/datastore/query.py
@@ -452,10 +452,12 @@ def _build_protobuf(self):
         if self.max_results is not None:
             pb.limit.value = self.max_results - self.num_results
 
-        if self._offset is not None:
-            # NOTE: The offset goes down relative to the location
-            #       because we are updating the cursor each time.
-            pb.offset = self._offset - self._skipped_results
+        if start_cursor is None and self._offset is not None:
+            # NOTE: We don't need to add an offset to the request protobuf
+            #       if we are using an existing cursor, because the offset
+            #       is only relative to the start of the result set, not
+            #       relative to each page (this method is called per-page)
+            pb.offset = self._offset
 
         return pb
 
diff --git a/datastore/tests/system/test_system.py b/datastore/tests/system/test_system.py
@@ -397,6 +397,18 @@ def test_query_paginate_simple_timestamp_keys(self):
 
         self.assertTrue(page_count > 1)
 
+    def test_query_offset_timestamp_keys(self):
+        # See issue #4675
+        max_all = 10000
+        offset = 1
+        max_offset = max_all - offset
+        query = self.CLIENT.query(kind='timestamp_key')
+        all_w_limit = list(query.fetch(limit=max_all))
+        self.assertEqual(len(all_w_limit), max_all)
+
+        offset_w_limit = list(query.fetch(offset=offset, limit=max_offset))
+        self.assertEqual(offset_w_limit, all_w_limit[offset:])
+
     def test_query_paginate_with_offset(self):
         page_query = self._base_query()
         page_query.order = 'appearances'
diff --git a/datastore/tests/unit/test_query.py b/datastore/tests/unit/test_query.py
@@ -404,21 +404,27 @@ def test__build_protobuf_empty(self):
         expected_pb = query_pb2.Query()
         self.assertEqual(pb, expected_pb)
 
-    def test__build_protobuf_all_values(self):
+    def test__build_protobuf_all_values_except_offset(self):
+        # this test and the following (all_values_except_start_and_end_cursor)
+        # test mutually exclusive states; the offset is ignored
+        # if a start_cursor is supplied
         from google.cloud.datastore_v1.proto import query_pb2
         from google.cloud.datastore.query import Query
 
         client = _Client(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._make_one(
-            query, client, limit=limit, offset=offset,
-            start_cursor=start_cursor, end_cursor=end_cursor)
+            query,
+            client,
+            limit=limit,
+            start_cursor=start_cursor,
+            end_cursor=end_cursor
+        )
         self.assertEqual(iterator.max_results, limit)
         iterator.num_results = 4
         iterator._skipped_results = 1
@@ -427,6 +433,32 @@ def test__build_protobuf_all_values(self):
         expected_pb = query_pb2.Query(
             start_cursor=start_bytes,
             end_cursor=end_bytes,
+        )
+        expected_pb.limit.value = limit - iterator.num_results
+        self.assertEqual(pb, expected_pb)
+
+    def test__build_protobuf_all_values_except_start_and_end_cursor(self):
+        # this test and the previous (all_values_except_start_offset)
+        # test mutually exclusive states; the offset is ignored
+        # if a start_cursor is supplied
+        from google.cloud.datastore_v1.proto import query_pb2
+        from google.cloud.datastore.query import Query
+
+        client = _Client(None)
+        query = Query(client)
+        limit = 15
+        offset = 9
+        iterator = self._make_one(
+            query,
+            client,
+            limit=limit,
+            offset=offset,
+        )
+        self.assertEqual(iterator.max_results, limit)
+        iterator.num_results = 4
+
+        pb = iterator._build_protobuf()
+        expected_pb = query_pb2.Query(
             offset=offset - iterator._skipped_results,
         )
         expected_pb.limit.value = limit - iterator.num_results
