[Cosmos] query_items with feed_range and continuation bug fix (#43700)

simorenoh · web-flow · commit e4d4839c5586 · 2025-10-31T11:08:23.000-04:00
* changes, changelog

* Update CHANGELOG.md

* pylint, mypy, tests, samples, comments

* move logic

* Update _read_items_helper.py

* forgot to actually move the logic

* Update _base.py
diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md
@@ -3,11 +3,13 @@
 ### 4.14.1 (Unreleased)
 
 #### Features Added
+* Added merge support. See [PR 42924](https://github.com/Azure/azure-sdk-for-python/pull/42924).
 
 #### Breaking Changes
 
 #### Bugs Fixed
 * Fixed bug where customer provided excluded region was not always being honored during certain transient failures. See [PR 43602](https://github.com/Azure/azure-sdk-for-python/pull/43602)
+* Fixed bug where queries using `feed_range` and `continuation` options would not work as expected. See [PR 43700](https://github.com/Azure/azure-sdk-for-python/pull/43700).
 
 #### Other Changes
 * Further optimized health checks for sync and async clients. See [PR 43339](https://github.com/Azure/azure-sdk-for-python/pull/43339)
@@ -41,7 +43,6 @@ This version and all future versions will require Python 3.9+.
 #### Features Added
 * Added read_items API to provide an efficient method for retrieving multiple items in a single request. See [PR 42167](https://github.com/Azure/azure-sdk-for-python/pull/42167).
 * Added ability to replace a container's indexing policy if a vector embedding policy was present. See [PR 42810](https://github.com/Azure/azure-sdk-for-python/pull/42810).
-* Added merge support. See [PR 42924](https://github.com/Azure/azure-sdk-for-python/pull/42924).
 
 #### Bugs Fixed
 * Improved the resilience of Database Account Read metadata operation against short-lived network issues by increasing number of retries. See [PR 42525](https://github.com/Azure/azure-sdk-for-python/pull/42525).
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_base.py
@@ -945,3 +945,18 @@ def _build_properties_cache(properties: dict[str, Any], container_link: str) ->
         "_self": properties.get("_self", None), "_rid": properties.get("_rid", None),
         "partitionKey": properties.get("partitionKey", None), "container_link": container_link
     }
+
+def format_pk_range_options(query_options: Mapping[str, Any]) -> dict[str, Any]:
+    """Formats the partition key range options to be used internally from the query ones.
+
+    :param dict query_options: The query options being used.
+    :return: The relevant partition key range options.
+    :rtype: dict
+    """
+    pk_range_options: dict[str, Any] = {}
+    if query_options is not None:
+        if "containerRID" in query_options:
+            pk_range_options["containerRID"] = query_options["containerRID"]
+        if "excludedLocations" in query_options:
+            pk_range_options["excludedLocations"] = query_options["excludedLocations"]
+    return pk_range_options
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py
@@ -3258,7 +3258,6 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
 
         # If feed_range_epk exist, query with the range
         if feed_range_epk is not None:
-            last_response_headers = CaseInsensitiveDict()
             over_lapping_ranges = self._routing_map_provider.get_overlapping_ranges(resource_id, [feed_range_epk],
                                                                                     options)
             # It is possible to get more than one over lapping range. We need to get the query results for each one
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/hybrid_search_aggregator.py
@@ -3,7 +3,6 @@
 
 """Internal class for multi execution context aggregator implementation in the Azure Cosmos database service.
 """
-
 from azure.cosmos._execution_context.aio.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context.aio import document_producer
 from azure.cosmos._execution_context.hybrid_search_aggregator import _retrieve_component_scores, _rewrite_query_infos, \
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/multi_execution_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/multi_execution_aggregator.py
@@ -21,7 +21,6 @@
 
 """Internal class for multi execution context aggregator implementation in the Azure Cosmos database service.
 """
-
 from azure.cosmos._execution_context.aio.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context.aio import document_producer, _queue_async_helper
 from azure.cosmos._routing import routing_range
@@ -162,7 +161,6 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
         )
 
     async def _get_target_partition_key_range(self):
-
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()
         return await self._routing_provider.get_overlapping_ranges(
             self._resource_link,
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/non_streaming_order_by_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/aio/non_streaming_order_by_aggregator.py
@@ -3,7 +3,6 @@
 
 """Internal class for multi execution context aggregator implementation in the Azure Cosmos database service.
 """
-
 from azure.cosmos._execution_context.aio.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context.aio.multi_execution_aggregator import _MultiExecutionContextAggregator
 from azure.cosmos._execution_context.aio import document_producer
@@ -107,7 +106,6 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
         )
 
     async def _get_target_partition_key_range(self):
-
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()
         return await self._routing_provider.get_overlapping_ranges(
             self._resource_link,
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/multi_execution_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/multi_execution_aggregator.py
@@ -195,7 +195,6 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range
         )
 
     def _get_target_partition_key_range(self):
-
         query_ranges = self._partitioned_query_ex_info.get_query_ranges()
         return self._routing_provider.get_overlapping_ranges(
             self._resource_link,
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/non_streaming_order_by_aggregator.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_execution_context/non_streaming_order_by_aggregator.py
@@ -3,7 +3,6 @@
 
 """Internal class for multi execution context aggregator implementation in the Azure Cosmos database service.
 """
-
 from azure.cosmos._execution_context.base_execution_context import _QueryExecutionContextBase
 from azure.cosmos._execution_context.multi_execution_aggregator import _MultiExecutionContextAggregator
 from azure.cosmos._execution_context import document_producer
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/aio/routing_map_provider.py
@@ -61,7 +61,8 @@ async def get_overlapping_ranges(self, collection_link, partition_key_ranges, fe
         :rtype: list
         """
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
-        await self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
+        pk_range_options = _base.format_pk_range_options(feed_options)
+        await self.init_collection_routing_map_if_needed(collection_link, collection_id, pk_range_options, **kwargs)
 
         return self._collection_routing_map_by_item[collection_id].get_overlapping_ranges(partition_key_ranges)
 
@@ -95,7 +96,8 @@ async def get_range_by_partition_key_range_id(
             **kwargs: dict[str, Any]
     ) -> Optional[dict[str, Any]]:
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
-        await self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
+        pk_range_options = _base.format_pk_range_options(feed_options)
+        await self.init_collection_routing_map_if_needed(collection_link, collection_id, pk_range_options, **kwargs)
 
         return self._collection_routing_map_by_item[collection_id].get_range_by_partition_key_range_id(
             partition_key_range_id)
diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/routing_map_provider.py
@@ -84,7 +84,8 @@ def get_overlapping_ranges(self, collection_link, partition_key_ranges, feed_opt
         :rtype: list
         """
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
-        self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
+        pk_range_options = _base.format_pk_range_options(feed_options)
+        self.init_collection_routing_map_if_needed(collection_link, collection_id, pk_range_options, **kwargs)
 
         return self._collection_routing_map_by_item[collection_id].get_overlapping_ranges(partition_key_ranges)
 
@@ -96,7 +97,8 @@ def get_range_by_partition_key_range_id(
             **kwargs: dict[str, Any]
     ) -> Optional[dict[str, Any]]:
         collection_id = _base.GetResourceIdOrFullNameFromLink(collection_link)
-        self.init_collection_routing_map_if_needed(collection_link, collection_id, feed_options, **kwargs)
+        pk_range_options = _base.format_pk_range_options(feed_options)
+        self.init_collection_routing_map_if_needed(collection_link, collection_id, pk_range_options, **kwargs)
 
         return (self._collection_routing_map_by_item[collection_id]
                 .get_range_by_partition_key_range_id(partition_key_range_id))
diff --git a/sdk/cosmos/azure-cosmos/samples/document_management.py b/sdk/cosmos/azure-cosmos/samples/document_management.py
@@ -176,6 +176,44 @@ def query_items_cross_partition_with_pagination(container):
     print(f'  - Pages: {page_count}')
     print(f'  - Total items: {total_item_count}')
 
+def query_items_with_feed_ranges_and_pagination(container):
+    print('\n1.5c Querying with Feed Range and Pagination\n')
+    # We again use max_item_count to control page size
+    max_items_per_page = 3
+
+    # First we fetch the relevant feed ranges for the container - feed ranges represent logical partitions
+    # or ranges of partition key values, and can be used to query as well
+    feed_ranges = container.read_feed_ranges()
+
+    # For this example, we will just use the first feed range
+    feed_ranges_list = list(feed_ranges)
+    query_iterable = container.query_items(
+        query="SELECT * FROM c",
+        feed_range=feed_ranges_list[0],  # Query specific feed range
+        max_item_count=max_items_per_page
+    )
+
+    # Iterate through pages and count both pages and total items
+    total_item_count = 0
+    page_count = 0
+
+    item_pages = query_iterable.by_page()
+    for page in item_pages:
+        page_count += 1
+        items_in_page = list(page)
+        items_in_current_page = len(items_in_page)
+        total_item_count += items_in_current_page
+
+        print(f'Page {page_count}: Retrieved {items_in_current_page} items (max per page: {max_items_per_page})')
+
+        # Process items in this page
+        for item in items_in_page:
+            # Do something with each item
+            pass
+
+    print(f'\nTotal pages processed: {page_count}')
+    print(f'Total items retrieved: {total_item_count}')
+    print(f'Note: max_item_count limits items PER PAGE, not total results\n')
 
 def replace_item(container, doc_id):
     print('\n1.6 Replace an Item\n')
diff --git a/sdk/cosmos/azure-cosmos/samples/document_management_async.py b/sdk/cosmos/azure-cosmos/samples/document_management_async.py
@@ -192,6 +192,38 @@ async def query_items_cross_partition_with_pagination(container):
     print(f'  - Pages: {page_count}')
     print(f'  - Total items: {total_item_count}')
 
+async def query_items_with_feed_ranges_and_pagination(container):
+    print('\n1.5c Querying with Feed Range and Pagination\n')
+    # We again use max_item_count to control page size
+    max_items_per_page = 3
+
+    # First we fetch the relevant feed ranges for the container - feed ranges represent logical partitions
+    # or ranges of partition key values, and can be used to query as well
+    feed_ranges = container.read_feed_ranges()
+
+    # For this example, we will just use the first feed range
+    feed_ranges_list = [feed_range async for feed_range in feed_ranges]
+    query_iterable = container.query_items(
+        query="SELECT * FROM c",
+        feed_range=feed_ranges_list[0],  # Query specific feed range
+        max_item_count=max_items_per_page
+    )
+
+    # Iterate through pages and count both pages and total items
+    total_item_count = 0
+    page_count = 0
+
+    item_pages = query_iterable.by_page()
+    async for page in item_pages:
+        page_count += 1
+        items_in_page = [item async for item in page]
+        total_item_count += len(items_in_page)
+
+        print(f'Page {page_count}: {len(items_in_page)} items from across partitions')
+
+    print(f'\nCross-partition query completed:')
+    print(f'  - Pages: {page_count}')
+    print(f'  - Total items: {total_item_count}')
 
 async def replace_item(container, doc_id):
     print('\n1.6 Replace an Item\n')
diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range.py
@@ -52,7 +52,7 @@ def get_container(container_id: str):
     return db.get_container_client(container_id)
 
 @pytest.mark.cosmosQuery
-class TestQueryFeedRange():
+class TestQueryFeedRange:
     @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS)
     def test_query_with_feed_range_for_all_partitions(self, container_id):
         container = get_container(container_id)
@@ -123,5 +123,56 @@ def test_query_with_feed_range_for_a_full_range(self, container_id):
         add_all_pk_values_to_set(items, actual_pk_values)
         assert expected_pk_values.issubset(actual_pk_values)
 
+    def test_query_with_static_continuation(self):
+        container = get_container(SINGLE_PARTITION_CONTAINER_ID)
+        query = 'SELECT * from c'
+
+        # verify continuation token does not have any impact
+        for i in range(10):
+            query_by_page = container.query_items(
+                query=query,
+                feed_range={
+                    'Range': {'isMaxInclusive': False, 'isMinInclusive': True,
+                              'max': '1FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE',
+                              'min': '0FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF'}},
+                max_item_count=1,
+                continuation='-RID:~a0NPAOszCpOChB4AAAAAAA==#RT:1#TRC:2#ISV:2#IEO:65567#QCF:8').by_page()
+            for page in query_by_page:
+                items = list(page)
+                assert len(items) > 0
+
+    def test_query_with_continuation(self):
+        container = get_container(SINGLE_PARTITION_CONTAINER_ID)
+        query = 'SELECT * from c'
+
+        # go through all feed ranges using pagination
+        feed_ranges = container.read_feed_ranges()
+        for feed in feed_ranges:
+            query_kwargs = {
+                "query": query,
+                "feed_range": feed,
+                "priority": "Low",
+                "max_item_count": 1
+            }
+            query_results = container.query_items(**query_kwargs)
+            pager = query_results.by_page()
+            first_page = next(pager)
+            items = list(first_page)
+            assert len(items) > 0
+            continuation_token = pager.continuation_token
+            # use that continuation token to restart the query, and drain it from there
+            query_kwargs = {
+                "query": query,
+                "feed_range": feed,
+                "continuation": continuation_token,
+                "priority": "Low",
+                "max_item_count": 2
+            }
+            query_results = container.query_items(**query_kwargs)
+            pager = query_results.by_page(continuation_token=continuation_token)
+            for new_page in pager:
+                items = list(new_page)
+                assert len(items) > 0
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_async.py
@@ -58,7 +58,7 @@ async def get_container(container_id: str):
 @pytest.mark.usefixtures("setup_and_teardown_async")
 class TestQueryFeedRangeAsync:
     @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS)
-    async def test_query_with_feed_range_for_all_partitions(self, container_id):
+    async def test_query_with_feed_range_for_all_partitions_async(self, container_id):
         container = await get_container(container_id)
         query = 'SELECT * from c'
 
@@ -75,7 +75,7 @@ async def test_query_with_feed_range_for_all_partitions(self, container_id):
         assert expected_pk_values == actual_pk_values
 
     @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS)
-    async def test_query_with_feed_range_for_partition_key(self, container_id):
+    async def test_query_with_feed_range_for_partition_key_async(self, container_id):
         container = await get_container(container_id)
         query = 'SELECT * from c'
 
@@ -94,7 +94,7 @@ async def test_query_with_feed_range_for_partition_key(self, container_id):
             assert expected_pk_values == actual_pk_values
 
     @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS)
-    async def test_query_with_both_feed_range_and_partition_key(self, container_id):
+    async def test_query_with_both_feed_range_and_partition_key_async(self, container_id):
         container = await get_container(container_id)
 
         expected_error_message = "'feed_range' and 'partition_key' are exclusive parameters, please only set one of them."
@@ -112,7 +112,7 @@ async def test_query_with_both_feed_range_and_partition_key(self, container_id):
         assert str(e.value) == expected_error_message
 
     @pytest.mark.parametrize('container_id', TEST_CONTAINERS_IDS)
-    async def test_query_with_feed_range_for_a_full_range(self, container_id):
+    async def test_query_with_feed_range_for_a_full_range_async(self, container_id):
         container = await get_container(container_id)
         query = 'SELECT * from c'
 
@@ -134,5 +134,56 @@ async def test_query_with_feed_range_for_a_full_range(self, container_id):
         await add_all_pk_values_to_set_async(items, actual_pk_values)
         assert expected_pk_values.issubset(actual_pk_values)
 
+    async def test_query_with_static_continuation_async(self):
+        container = await get_container(SINGLE_PARTITION_CONTAINER_ID)
+        query = 'SELECT * from c'
+
+        # verify continuation token does not have any impact
+        for i in range(10):
+            query_by_page = container.query_items(
+                query=query,
+                feed_range={
+                    'Range': {'isMaxInclusive': False, 'isMinInclusive': True,
+                              'max': '1FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFE',
+                              'min': '0FFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF'}},
+                max_item_count=1,
+                continuation='-RID:~a0NPAOszCpOChB4AAAAAAA==#RT:1#TRC:2#ISV:2#IEO:65567#QCF:8').by_page()
+            async for page in query_by_page:
+                items = [item async for item in page]
+                assert len(items) > 0
+
+    async def test_query_with_continuation_async(self):
+        container = await get_container(SINGLE_PARTITION_CONTAINER_ID)
+        query = 'SELECT * from c'
+
+        # go through all feed ranges using pagination
+        feed_ranges = container.read_feed_ranges()
+        async for feed in feed_ranges:
+            query_kwargs = {
+                "query": query,
+                "feed_range": feed,
+                "priority": "Low",
+                "max_item_count": 1
+            }
+            query_results = container.query_items(**query_kwargs)
+            pager = query_results.by_page()
+            first_page = await pager.__anext__()
+            items = [item async for item in first_page]
+            assert len(items) > 0
+            continuation_token = pager.continuation_token
+            # use that continuation token to restart the query, and drain it from there
+            query_kwargs = {
+                "query": query,
+                "feed_range": feed,
+                "continuation": continuation_token,
+                "priority": "Low",
+                "max_item_count": 2
+            }
+            query_results = container.query_items(**query_kwargs)
+            pager = query_results.by_page(continuation_token=continuation_token)
+            async for new_page in pager:
+                items = [item async for item in new_page]
+                assert len(items) > 0
+
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -195,7 +195,6 @@ def _createTargetPartitionQueryExecutionContext(self, partition_key_target_range`
`195`	`195`	`)`
`196`	`196`
`197`	`197`	`def _get_target_partition_key_range(self):`
`198`		`-`
`199`	`198`	`query_ranges = self._partitioned_query_ex_info.get_query_ranges()`
`200`	`199`	`return self._routing_provider.get_overlapping_ranges(`
`201`	`200`	`self._resource_link,`