[Cosmos] query_items with feed_range and continuation bug fix

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
 * Enhanced logging to ensure when a region is marked unavailable we have the proper context. See [PR 43602](https://github.com/Azure/azure-sdk-for-python/pull/43602)
@@ -40,7 +42,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).

sdk/cosmos/azure-cosmos/azure/cosmos/_base.py

@@ -945,3 +945,15 @@ 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: dict[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 "containerRID" in query_options:
+        pk_range_options["containerRID"] = query_options["containerRID"]
+    return pk_range_options

sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py

@@ -3259,9 +3259,9 @@ 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()
+            pk_range_options = base.format_pk_range_options(query_options=options)
             over_lapping_ranges = self._routing_map_provider.get_overlapping_ranges(resource_id, [feed_range_epk],
-                                                                                    options)
+                                                                                    pk_range_options)
             # It is possible to get more than one over lapping range. We need to get the query results for each one
             results: dict[str, Any] = {}
             # For each over lapping range we will take a sub range of the feed range EPK that overlaps with the over

sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py

@@ -3064,8 +3064,9 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                 feed_range_epk = partition_key_obj._get_epk_range_for_prefix_partition_key(partition_key_value)
 
         if feed_range_epk is not None:
+            pk_range_options = base.format_pk_range_options(query_options=options)
             over_lapping_ranges = await self._routing_map_provider.get_overlapping_ranges(id_, [feed_range_epk],
-                                                                                          options)
+                                                                                          pk_range_options)
             results: dict[str, Any] = {}
             # For each over lapping range we will take a sub range of the feed range EPK that overlaps with the over
             # lapping physical partition. The EPK sub range will be one of four: