addressing PR comments

Author: dibahlfi

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

@@ -211,8 +211,12 @@ def _merge_query_results(
             results_docs[0] = min(results_docs[0], partial_docs[0]) # type: ignore[index]
         elif aggregate_fn == "MAX":
             results_docs[0] = max(results_docs[0], partial_docs[0]) # type: ignore[index]
+        elif aggregate_fn == "AVG":
+            raise ValueError(
+                "VALUE AVG aggregate merge across partitions is not supported client-side."
+            )
         else:
-            # COUNT/SUM are additive; VALUE AVG is not fully supported client-side yet.
+            # COUNT/SUM are additive.
             results_docs[0] += partial_docs[0] # type: ignore[index]
         return results
 

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

@@ -3376,19 +3376,21 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
             # the feed_range helpers below read as ``resource_id_str`` instead
             # of the generic ``resource_id``.
             resource_id_str: str = resource_id
-            # (a) Look at the continuation the caller passed in.
-            #     - Empty or from a pre-fix SDK: start fresh.
-            #     - One of our v=1 envelopes: check the collection, query, and
-            #       feed_range still match before resuming from it.
-            #
-            # Shared state transitions (resume, split handling, page-item update,
-            # outbound token) live in _FeedRangePaginationState so sync/async
-            # stay behaviorally aligned.
+            # Decode the inbound continuation. Empty/legacy -> start fresh
+            # (``_decode_token`` returns ``None``); a valid v=1 envelope
+            # is checked against the current collection/query/feed_range
+            # before we resume from it. The shared
+            # ``_FeedRangePaginationState`` owns all state transitions
+            # (resume, split handling, page-item update, outbound token)
+            # so the sync and async loops below remain twin code paths
+            # — change one, change the other.
             items_left_in_page = _normalize_max_item_count(options.get("maxItemCount"))
-            inbound = _decode_token(options.get("continuation"))
-            if inbound is not None:
-                _validate_token_identity(inbound, resource_id_str, query, feed_range_epk)
-                pagination_state = _FeedRangePaginationState.from_inbound(inbound, items_left_in_page)
+            inbound_token_payload = _decode_token(options.get("continuation"))
+            if inbound_token_payload is not None:
+                _validate_token_identity(inbound_token_payload, resource_id_str, query, feed_range_epk)
+                pagination_state = _FeedRangePaginationState.from_inbound(
+                    inbound_token_payload, items_left_in_page
+                )
             else:
                 # First call (or legacy passthrough). Ask the routing map which
                 # partitions the input feed_range overlaps right now and turn
@@ -3432,30 +3434,10 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                         overlapping, head_feedrange
                     )
 
-                    # Handle the case where Cosmos split a partition between the
-                    # previous run and this one. Example: the saved
-                    # head_feedrange used to live inside one partition X, but X
-                    # has since been split into children X1 and X2. The routing
-                    # map now returns two partitions for the same feedrange. If
-                    # we sent one POST to X1 with X's full range as the EPK
-                    # filter, the backend would filter in-partition only and
-                    # silently drop every row living on X2 (resume after a
-                    # split would then come back short of ground truth).
-                    #
-                    # So when the lookup returns more than one partition, slice
-                    # the saved feedrange into one sub-feedrange per child
-                    # (intersection with the saved feedrange, ordered by EPK
-                    # min), make the first sub-feedrange the new current one,
-                    # put the rest in front of the remaining list, and clear the
-                    # saved backend continuation - it was issued by the old
-                    # parent partition and the children won't accept it. The next
-                    # loop iteration sees a single overlap and falls through to
-                    # the normal single-partition POST below.
-                    #
-                    # One edge case remains by design: if some rows were already
-                    # read from parent X before it split, those rows can show up
-                    # once more after resume when children X1/X2 restart from the
-                    # start of their slices.
+                    # If routing returns multiple overlaps, the head sub-range now spans a split
+                    # that occurred after the token was created. Re-slice and re-resolve until
+                    # each head maps to one partition. See
+                    # ``_FeedRangePaginationState.explode_on_multi_overlap`` for details.
                     while pagination_state.explode_on_multi_overlap(overlapping):
                         head_feedrange = pagination_state.head_range
                         if head_feedrange is None:
@@ -3467,31 +3449,22 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                             overlapping, head_feedrange
                         )
 
-                    backend_request_options = dict(options)
-                    if pagination_state.remaining_page_item_count is not None:
-                        backend_request_options["maxItemCount"] = pagination_state.remaining_page_item_count
-                    if pagination_state.head_bc is not None:
-                        backend_request_options["continuation"] = pagination_state.head_bc
-                    else:
-                        backend_request_options.pop("continuation", None)
-
                     # Populate request headers for this single backend POST.
                     # The shared helper handles partition routing (PKR id +
                     # optional EPK filter), page-size cap, and continuation
                     # set/clear so the same rules apply to sync and async.
-                    assert head_feedrange is not None  # narrowed by the loop guards above
                     _apply_feedrange_request_headers(
                         req_headers,
                         overlapping,
                         partition_scope,
                         head_feedrange,
                         pagination_state.remaining_page_item_count,
-                        backend_request_options.get("continuation"),
+                        pagination_state.head_bc,
                     )
                     # Use the session token for this specific partition so we don't
                     # send a compound token covering all partitions.
                     base.set_session_token_header(
-                        self, req_headers, path, request_params, backend_request_options, overlapping[0]["id"]
+                        self, req_headers, path, request_params, options, overlapping[0]["id"]
                     )
 
                     try:
@@ -3500,14 +3473,14 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                         )
                     except Exception:  # pylint: disable=broad-exception-caught
                         # Preserve resume progress if a later POST fails mid-page.
+                        self.last_response_headers = feedrange_response_headers
                         try:
                             pagination_state.write_outbound_continuation(
                                 feedrange_response_headers,
                                 resource_id_str,
                                 query,
                                 feed_range_epk,
                             )
-                            self.last_response_headers = feedrange_response_headers
                         except Exception as continuation_write_error:  # pylint: disable=broad-exception-caught
                             _LOGGER.warning(
                                 "Failed to write continuation while handling query POST failure: %s",
@@ -3571,8 +3544,9 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
                             )
                         )
 
-                # (c) Build the outbound token. Clear the continuation header if
-                # there is no work left at all.
+                # Pagination loop is done — write the final outbound
+                # continuation (or clear the header if the queue is fully
+                # drained) so the caller's ``by_page`` loop terminates.
                 pagination_state.write_outbound_continuation(
                     feedrange_response_headers,
                     resource_id_str,

sdk/cosmos/azure-cosmos/azure/cosmos/_query_aggregate_utils.py

@@ -39,9 +39,10 @@ def _extract_query_text(query: Optional[Union[str, dict[str, Any]]]) -> Optional
 def _get_select_value_aggregate_function(query: Optional[Union[str, dict[str, Any]]]) -> Optional[str]:
     """Identify the aggregate function for ``SELECT VALUE`` aggregate queries.
 
-    This is a lightweight text heuristic over the full query string (not a SQL
-    parser). Aggregate function tokens inside subqueries can therefore produce
-    false positives for the outer query.
+    This is a lightweight text heuristic (not a SQL parser). It extracts only
+    the OUTER ``SELECT VALUE`` projection and then matches aggregate function
+    names in that projection so nested subqueries do not drive outer
+    classification.
 
     :param query: Query text or query spec dictionary.
     :type query: Optional[Union[str, dict[str, Any]]]
@@ -53,18 +54,52 @@ def _get_select_value_aggregate_function(query: Optional[Union[str, dict[str, An
         return None
 
     normalized = " ".join(query_text.upper().split())
-    if "SELECT VALUE" not in normalized:
+    projection = _extract_outer_select_value_projection(normalized)
+    if projection is None:
         return None
 
-    # NOTE: This checks the full normalized query text, so aggregate function
-    # names inside subqueries can still be matched as false positives.
     # Match whole function names only (avoid MYCOUNT) and allow COUNT (1).
     for aggregate_fn in ("COUNT", "SUM", "MIN", "MAX", "AVG"):
-        if re.search(rf"(?<![A-Z0-9_]){aggregate_fn}\s*\(", normalized):
+        if re.search(rf"(?<![A-Z0-9_]){aggregate_fn}\s*\(", projection):
             return aggregate_fn
     return None
 
 
+def _extract_outer_select_value_projection(normalized_query: str) -> Optional[str]:
+    """Return the outer ``SELECT VALUE`` projection text up to the outer ``FROM``.
+
+    Uses a lightweight parenthesis-depth scan so nested subqueries do not
+    influence outer aggregate detection.
+    """
+    select_value = "SELECT VALUE"
+    start_idx = normalized_query.find(select_value)
+    if start_idx < 0:
+        return None
+
+    projection_start = start_idx + len(select_value)
+    if projection_start < len(normalized_query) and normalized_query[projection_start] == " ":
+        projection_start += 1
+
+    depth = 0
+    index = projection_start
+    while index <= len(normalized_query) - 4:
+        ch = normalized_query[index]
+        if ch == "(":
+            depth += 1
+        elif ch == ")" and depth > 0:
+            depth -= 1
+
+        if depth == 0 and normalized_query[index:index + 4] == "FROM":
+            prev_char = normalized_query[index - 1] if index > 0 else " "
+            next_char = normalized_query[index + 4] if index + 4 < len(normalized_query) else " "
+            if not (prev_char.isalnum() or prev_char == "_") and not (next_char.isalnum() or next_char == "_"):
+                projection = normalized_query[projection_start:index].strip()
+                return projection or None
+        index += 1
+
+    return None
+
+
 def _classify_aggregate_partial(
     docs: Any,
     query: Optional[Union[str, dict[str, Any]]]

sdk/cosmos/azure-cosmos/azure/cosmos/_routing/feed_range_continuation.py

@@ -434,6 +434,10 @@ class _FeedRangePaginationState:
     structure scale uniformly to non-sequential merges / parallel
     fan-out: every entry is structurally equal, and any subset of
     entries may carry a non-null backend continuation simultaneously.
+
+    Not thread-safe. One instance is created per ``query_items`` call
+    and is mutated only by that call's pagination loop (sync or async)
+    — never shared across threads or concurrent tasks.
     """
 
     def __init__(
@@ -790,4 +794,3 @@ def _apply_feedrange_request_headers(
         req_headers[http_constants.HttpHeaders.Continuation] = inbound_continuation
     else:
         req_headers.pop(http_constants.HttpHeaders.Continuation, None)
-