From ebe0815a72d72e8eb3af477b6878d23486fd4f07 Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Mon, 24 Aug 2020 14:40:52 -0700 Subject: [PATCH 1/6] fixed bug in querying by page using continuation token --- .../base_execution_context.py | 2 +- sdk/cosmos/azure-cosmos/test/test_query.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py index c125e870b3a4..758553dd1ae0 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py @@ -44,7 +44,7 @@ def __init__(self, client, options): 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: + if "continuation" in options: self._continuation = options["continuation"] self._has_started = False self._has_finished = False diff --git a/sdk/cosmos/azure-cosmos/test/test_query.py b/sdk/cosmos/azure-cosmos/test/test_query.py index 59204ef863ea..cec9b2bdf0d0 100644 --- a/sdk/cosmos/azure-cosmos/test/test_query.py +++ b/sdk/cosmos/azure-cosmos/test/test_query.py @@ -522,6 +522,30 @@ 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 _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 From 205475d3bdf8da36df702477263b544aa06afeb9 Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Mon, 24 Aug 2020 15:07:09 -0700 Subject: [PATCH 2/6] updated changelog --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index e2940d239b72..0b0f25909a76 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -1,3 +1,7 @@ +**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. From 0c964948ccfaab501fdd396b6515f4a90c73a216 Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Mon, 24 Aug 2020 16:39:43 -0700 Subject: [PATCH 3/6] modified changelog --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 0b0f25909a76..90015159104c 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -1,3 +1,5 @@ +## 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. From 98d1776af32684350234d054b68714d0913462ea Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Wed, 26 Aug 2020 07:25:27 -0700 Subject: [PATCH 4/6] added support from single partition query continuation tokens --- .../base_execution_context.py | 18 +++++++++++++----- sdk/cosmos/azure-cosmos/test/test_query.py | 12 ++++++++++++ 2 files changed, 25 insertions(+), 5 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py index 758553dd1ae0..8d5bd2c53fdf 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py @@ -26,6 +26,7 @@ from collections import deque from .. import _retry_utility from .. import http_constants +import copy # pylint: disable=protected-access @@ -43,13 +44,19 @@ 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: - 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 AttributeError("continuation tokens are not supported for cross-partition queries.") + else: + return self._options["continuation"] + return None + def _has_more_pages(self): return not self._has_started or self._continuation @@ -112,8 +119,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: diff --git a/sdk/cosmos/azure-cosmos/test/test_query.py b/sdk/cosmos/azure-cosmos/test/test_query.py index cec9b2bdf0d0..ca28e70a81fc 100644 --- a/sdk/cosmos/azure-cosmos/test/test_query.py +++ b/sdk/cosmos/azure-cosmos/test/test_query.py @@ -546,6 +546,18 @@ def test_paging_with_continuation_token(self): 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(AttributeError): + 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 From a1414701e4a82d3260f1cc0b7a3da0ed8c4db278 Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Wed, 26 Aug 2020 11:43:34 -0700 Subject: [PATCH 5/6] fixed linting error --- .../cosmos/_execution_context/base_execution_context.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py index 8d5bd2c53fdf..93f1e514c666 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py @@ -24,9 +24,9 @@ """ from collections import deque +import copy from .. import _retry_utility from .. import http_constants -import copy # pylint: disable=protected-access @@ -53,8 +53,7 @@ def _get_initial_continuation(self): if "continuation" in self._options: if "enableCrossPartitionQuery" in self._options: raise AttributeError("continuation tokens are not supported for cross-partition queries.") - else: - return self._options["continuation"] + return self._options["continuation"] return None def _has_more_pages(self): From a0973da8e0dbfa31908f6f0874049c9664f5d5ff Mon Sep 17 00:00:00 2001 From: Srinath Narayanan Date: Wed, 26 Aug 2020 16:53:43 -0700 Subject: [PATCH 6/6] addressed PR comments --- .../azure/cosmos/_execution_context/base_execution_context.py | 2 +- sdk/cosmos/azure-cosmos/test/test_query.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py index 93f1e514c666..3897e5a8a6b1 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/base_execution_context.py @@ -52,7 +52,7 @@ def __init__(self, client, options): def _get_initial_continuation(self): if "continuation" in self._options: if "enableCrossPartitionQuery" in self._options: - raise AttributeError("continuation tokens are not supported for cross-partition queries.") + raise ValueError("continuation tokens are not supported for cross-partition queries.") return self._options["continuation"] return None diff --git a/sdk/cosmos/azure-cosmos/test/test_query.py b/sdk/cosmos/azure-cosmos/test/test_query.py index ca28e70a81fc..44d59975c46c 100644 --- a/sdk/cosmos/azure-cosmos/test/test_query.py +++ b/sdk/cosmos/azure-cosmos/test/test_query.py @@ -555,7 +555,7 @@ def test_cross_partition_query_with_continuation_token_fails(self): max_item_count=1, ) - with self.assertRaises(AttributeError): + 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):