fixed bug in querying by page using continuation token (#13298)

* fixed bug in querying by page using continuation token

* updated changelog

* modified changelog

* added support from single partition query continuation tokens

* fixed linting error

* addressed PR comments

Author: Srinath Narayanan

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -1,3 +1,9 @@
+## 4.1.1 (unreleased)
+
+**Bug fixes**
+- Fixed bug where continuation token is not honored when query_iterable is used to get results by page. Issue #13265.
+
+
 ## 4.1.0 (2020-08-10)
 
 - Added deprecation warning for "lazy" indexing mode. The backend no longer allows creating containers with this mode and will set them to consistent instead.

sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py

@@ -24,6 +24,7 @@
 """
 
 from collections import deque
+import copy
 from .. import _retry_utility
 from .. import http_constants
 
@@ -43,13 +44,18 @@ def __init__(self, client, options):
         self._client = client
         self._options = options
         self._is_change_feed = "changeFeed" in options and options["changeFeed"] is True
-        self._continuation = None
-        if "continuation" in options and self._is_change_feed:
-            self._continuation = options["continuation"]
+        self._continuation = self._get_initial_continuation()
         self._has_started = False
         self._has_finished = False
         self._buffer = deque()
 
+    def _get_initial_continuation(self):
+        if "continuation" in self._options:
+            if "enableCrossPartitionQuery" in self._options:
+                raise ValueError("continuation tokens are not supported for cross-partition queries.")
+            return self._options["continuation"]
+        return None
+
     def _has_more_pages(self):
         return not self._has_started or self._continuation
 
@@ -112,8 +118,9 @@ def _fetch_items_helper_no_retries(self, fetch_function):
         while self._continuation or not self._has_started:
             if not self._has_started:
                 self._has_started = True
-            self._options["continuation"] = self._continuation
-            (fetched_items, response_headers) = fetch_function(self._options)
+            new_options = copy.deepcopy(self._options)
+            new_options["continuation"] = self._continuation
+            (fetched_items, response_headers) = fetch_function(new_options)
             continuation_key = http_constants.HttpHeaders.Continuation
             # Use Etag as continuation token for change feed queries.
             if self._is_change_feed:

sdk/cosmos/azure-cosmos/test/test_query.py

@@ -522,6 +522,42 @@ def test_distinct_on_different_types_and_field_orders(self):
         _QueryExecutionContextBase.__next__ = self.OriginalExecuteFunction
         _QueryExecutionContextBase.next = self.OriginalExecuteFunction
 
+    def test_paging_with_continuation_token(self):
+        created_collection = self.config.create_multi_partition_collection_with_custom_pk_if_not_exist(self.client)
+
+        document_definition = {'pk': 'pk', 'id': '1'}
+        created_collection.create_item(body=document_definition)
+        document_definition = {'pk': 'pk', 'id': '2'}
+        created_collection.create_item(body=document_definition)
+
+        query = 'SELECT * from c'
+        query_iterable = created_collection.query_items(
+            query=query,
+            partition_key='pk',
+            max_item_count=1
+        )
+        pager = query_iterable.by_page()
+        pager.next()
+        token = pager.continuation_token
+        second_page = list(pager.next())[0]
+
+        pager = query_iterable.by_page(token)
+        second_page_fetched_with_continuation_token = list(pager.next())[0]
+
+        self.assertEqual(second_page['id'], second_page_fetched_with_continuation_token['id'])
+
+    def test_cross_partition_query_with_continuation_token_fails(self):
+        created_collection = self.config.create_multi_partition_collection_with_custom_pk_if_not_exist(self.client)
+        query = 'SELECT * from c'
+        query_iterable = created_collection.query_items(
+            query=query,
+            enable_cross_partition_query=True,
+            max_item_count=1,
+        )
+
+        with self.assertRaises(ValueError):
+            pager = query_iterable.by_page("fake_continuation_token")
+
     def _validate_distinct_on_different_types_and_field_orders(self, collection, query, expected_results, get_mock_result):
         self.count = 0
         self.get_mock_result = get_mock_result