fix: fixing 410 SDK refresh for feed ranges (#44425)

* fix: fixing 410 for feed ranges

* fix: adding sync support for handling 410 bug

* fix: fixing Pylint errors

* fix: addressing comments

* fix: adding tests

* fix: addressing comments

* fix: fixing pylint warnings

* fix: updating version date

* fix: addressing MyPy errors

* fix: combining tests to avoid timeout pipeline errors

* fix: fixing pipeline timeout issue

* fix: fixing infinite loop bug

Author: dibahlfi

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -1,14 +1,9 @@
 ## Release History
 
-### 4.15.0b2 (Unreleased)
-
-#### Features Added
-
-#### Breaking Changes
+### 4.15.0b2 (2025-12-16)
 
 #### Bugs Fixed
-
-#### Other Changes
+* Fixed bug where sdk was not properly retrying requests in some edge cases after partition splits.See [PR 44425](https://github.com/Azure/azure-sdk-for-python/pull/44425)
 
 ### 4.15.0b1 (2025-11-26)
 

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

@@ -129,6 +129,115 @@ def build_options(kwargs: dict[str, Any]) -> dict[str, Any]:
         options['accessCondition'] = {'type': 'IfNoneMatch', 'condition': if_none_match}
     return options
 
+def _merge_query_results(
+        results: dict[str, Any],
+        partial_result: dict[str, Any],
+        query: Optional[Union[str, dict[str, Any]]]
+) -> dict[str, Any]:
+    """Merges partial query results from different partitions.
+
+    This method is required for queries that are manually fanned out to multiple
+    partitions or ranges within the SDK, such as prefix partition key queries.
+    For non-aggregated queries, results from each partition are simply concatenated.
+    However, for aggregate queries (COUNT, SUM, MIN, MAX, AVG), each partition
+    returns a partial aggregate. This method merges these partial results to compute
+    the final, correct aggregate value.
+
+    TODO:This client-side aggregation is a temporary workaround. Ideally, this logic
+    should be integrated into the core pipeline as aggregate queries are handled by DefaultExecutionContext,
+    not MultiAggregatorExecutionContext, which is not split proof until the logic is moved to the core pipeline.
+    This method handles the aggregation of results when a query spans multiple
+    partitions. It specifically handles:
+    1. Standard queries: Appends documents from partial_result to results.
+    2. Aggregate queries that return a JSON object (e.g., `SELECT COUNT(1) FROM c`, `SELECT MIN(c.field) FROM c`).
+    3. VALUE queries with aggregation that return a scalar value (e.g., `SELECT VALUE COUNT(1) FROM c`).
+
+    :param dict[str, Any] results: The accumulated result's dictionary.
+    :param dict[str, Any] partial_result: The new partial result dictionary to merge.
+    :param query: The query being executed.
+    :type query: str or dict[str, Any]
+    :return: The merged result's dictionary.
+    :rtype: dict[str, Any]
+    """
+    if not results:
+        return partial_result
+
+    partial_docs = partial_result.get("Documents")
+    if not partial_docs:
+        return results
+
+    results_docs = results.get("Documents")
+
+    # Check if both results are aggregate queries
+    is_partial_agg = (
+            isinstance(partial_docs, list)
+            and len(partial_docs) == 1
+            and isinstance(partial_docs[0], dict)
+            and partial_docs[0].get("_aggregate") is not None
+    )
+    is_results_agg = (
+            results_docs
+            and isinstance(results_docs, list)
+            and len(results_docs) == 1
+            and isinstance(results_docs[0], dict)
+            and results_docs[0].get("_aggregate") is not None
+    )
+
+    if is_partial_agg and is_results_agg:
+        agg_results = results_docs[0]["_aggregate"] # type: ignore[index]
+        agg_partial = partial_docs[0]["_aggregate"]
+        for key in agg_partial:
+            if key not in agg_results:
+                agg_results[key] = agg_partial[key]
+            elif isinstance(agg_partial.get(key), dict) and "count" in agg_partial[key]:  # AVG
+                if isinstance(agg_results.get(key), dict):
+                    agg_results[key]["sum"] += agg_partial[key]["sum"]
+                    agg_results[key]["count"] += agg_partial[key]["count"]
+            elif key.lower().startswith("min"):
+                agg_results[key] = min(agg_results[key], agg_partial[key])
+            elif key.lower().startswith("max"):
+                agg_results[key] = max(agg_results[key], agg_partial[key])
+            else:  # COUNT, SUM
+                agg_results[key] += agg_partial[key]
+        return results
+
+    # Check if both are VALUE aggregate queries
+    is_partial_value_agg = (
+            isinstance(partial_docs, list)
+            and len(partial_docs) == 1
+            and isinstance(partial_docs[0], (int, float))
+    )
+    is_results_value_agg = (
+            results_docs
+            and isinstance(results_docs, list)
+            and len(results_docs) == 1
+            and isinstance(results_docs[0], (int, float))
+    )
+
+    if is_partial_value_agg and is_results_value_agg:
+        query_text = query.get("query") if isinstance(query, dict) else query
+        if query_text:
+            query_upper = query_text.upper()
+            # For MIN/MAX, we find the min/max of the partial results.
+            # For COUNT/SUM, we sum the partial results.
+            # Without robust query parsing, we can't distinguish them reliably.
+            # Defaulting to sum for COUNT/SUM. MIN/MAX VALUE queries are not fully supported client-side.
+            if " SELECT VALUE MIN" in query_upper:
+                results_docs[0] = min(results_docs[0], partial_docs[0]) # type: ignore[index]
+            elif " SELECT VALUE MAX" in query_upper:
+                results_docs[0] = max(results_docs[0], partial_docs[0]) # type: ignore[index]
+            else:  # For COUNT/SUM, we sum the partial results
+                results_docs[0] += partial_docs[0] # type: ignore[index]
+        return results
+
+    # Standard query, append documents
+    if results_docs is None:
+        results["Documents"] = partial_docs
+    elif isinstance(results_docs, list) and isinstance(partial_docs, list):
+        results_docs.extend(partial_docs)
+    results["_count"] = len(results["Documents"])
+    return results
+
 
 def GetHeaders(  # pylint: disable=too-many-statements,too-many-branches
         cosmos_client_connection: Union["CosmosClientConnection", "AsyncClientConnection"],

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

@@ -3353,11 +3353,16 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                 )
                 self.last_response_headers = last_response_headers
                 self._UpdateSessionIfRequired(req_headers, partial_result, last_response_headers)
-                if results:
-                    # add up all the query results from all over lapping ranges
-                    results["Documents"].extend(partial_result["Documents"])
-                else:
-                    results = partial_result
+                # Introducing a temporary complex function into a critical path to handle aggregated queries
+                # during splits, as a precaution falling back to the original logic if anything goes wrong
+                try:
+                    results = base._merge_query_results(results, partial_result, query)
+                except Exception: # pylint: disable=broad-exception-caught
+                    # If the new merge logic fails, fall back to the original logic.
+                    if results:
+                        results["Documents"].extend(partial_result["Documents"])
+                    else:
+                        results = partial_result
                 if response_hook:
                     response_hook(last_response_headers, partial_result)
             # if the prefix partition query has results lets return it

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

@@ -27,7 +27,8 @@
 import copy
 
 from ...aio import _retry_utility_async
-from ... import http_constants
+from ... import http_constants, exceptions
+
 
 # pylint: disable=protected-access
 
@@ -136,12 +137,30 @@ async def _fetch_items_helper_with_retries(self, fetch_function):
         # ExecuteAsync passes retry context parameters (timeout, operation start time, logger, etc.)
         # The callback need to accept these parameters even if unused
         # Removing **kwargs results in a TypeError when ExecuteAsync tries to pass these parameters
-        async def callback(**kwargs):  # pylint: disable=unused-argument
-            return await self._fetch_items_helper_no_retries(fetch_function)
-
-        return await _retry_utility_async.ExecuteAsync(
-            self._client, self._client._global_endpoint_manager, callback, **self._options
-        )
+        async def execute_fetch():
+            async def callback(**kwargs):  # pylint: disable=unused-argument
+                return await self._fetch_items_helper_no_retries(fetch_function)
+
+            return await _retry_utility_async.ExecuteAsync(
+                self._client, self._client._global_endpoint_manager, callback, **self._options
+            )
+
+        max_retries = 3
+        attempt = 0
+        while attempt <= max_retries:
+            try:
+                return await execute_fetch()
+            except exceptions.CosmosHttpResponseError as e:
+                if exceptions._partition_range_is_gone(e):
+                    attempt += 1
+                    if attempt > max_retries:
+                        raise  # Exhausted retries, propagate error
+
+                    # Refresh routing map to get new partition key ranges
+                    self._client.refresh_routing_map_provider()
+                    # Retry immediately (no backoff needed for partition splits)
+                    continue
+                raise  # Not a partition split error, propagate immediately
 
 
 class _DefaultQueryExecutionContext(_QueryExecutionContextBase):

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

@@ -25,7 +25,8 @@
 
 from collections import deque
 import copy
-from .. import _retry_utility, http_constants
+from .. import _retry_utility, http_constants, exceptions
+
 
 # pylint: disable=protected-access
 
@@ -131,14 +132,34 @@ def _fetch_items_helper_no_retries(self, fetch_function):
     def _fetch_items_helper_with_retries(self, fetch_function):
         # TODO: Properly propagate kwargs from retry utility to fetch function
         # the callback keep the **kwargs parameter to maintain compatibility with the retry utility's execution pattern.
-        # ExecuteAsync passes retry context parameters (timeout, operation start time, logger, etc.)
+        # Execute passes retry context parameters (timeout, operation start time, logger, etc.)
         # The callback need to accept these parameters even if unused
-        # Removing **kwargs results in a TypeError when ExecuteAsync tries to pass these parameters
-        def callback(**kwargs): # pylint: disable=unused-argument
-            return self._fetch_items_helper_no_retries(fetch_function)
-
-        return _retry_utility.Execute(self._client, self._client._global_endpoint_manager, callback, **self._options)
-
+        # Removing **kwargs results in a TypeError when Execute tries to pass these parameters
+        def execute_fetch():
+            def callback(**kwargs):  # pylint: disable=unused-argument
+                return self._fetch_items_helper_no_retries(fetch_function)
+
+            return _retry_utility.Execute(
+                self._client, self._client._global_endpoint_manager, callback, **self._options
+            )
+
+        max_retries = 3
+        attempt = 0
+
+        while attempt <= max_retries:
+            try:
+                return execute_fetch()
+            except exceptions.CosmosHttpResponseError as e:
+                if exceptions._partition_range_is_gone(e):
+                    attempt += 1
+                    if attempt > max_retries:
+                        raise  # Exhausted retries, propagate error
+
+                    # Refresh routing map to get new partition key ranges
+                    self._client.refresh_routing_map_provider()
+                    # Retry immediately (no backoff needed for partition splits)
+                    continue
+                raise  # Not a partition split error, propagate immediately
     next = __next__  # Python 2 compatibility.
 
 

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

@@ -3149,11 +3149,18 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                 )
                 self.last_response_headers = last_response_headers
                 self._UpdateSessionIfRequired(req_headers, partial_result, last_response_headers)
-                if results:
-                    # add up all the query results from all over lapping ranges
-                    results["Documents"].extend(partial_result["Documents"])
-                else:
-                    results = partial_result
+
+                # Introducing a temporary complex function into a critical path to handle aggregated queries,
+                # during splits as a precaution falling back to the original logic if anything goes wrong
+                try:
+                    results = base._merge_query_results(results, partial_result, query)
+                except Exception: # pylint: disable=broad-exception-caught
+                    # If the new merge logic fails, fall back to the original logic.
+                    if results:
+                        results["Documents"].extend(partial_result["Documents"])
+                    else:
+                        results = partial_result
+
                 if response_hook:
                     response_hook(self.last_response_headers, partial_result)
             # if the prefix partition query has results lets return it