fix - full partition key routing bug

sdk/cosmos/azure-cosmos/CHANGELOG.md

@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed complete-partition-key queries scanning documents instead of using partition-key routing, which caused excessive RU consumption and latency for aggregates such as `COUNT`. See [PR 48237](https://github.com/Azure/azure-sdk-for-python/pull/48237)
 
 #### Other Changes
 

sdk/cosmos/azure-cosmos/api.metadata.yml

@@ -1,3 +1,3 @@
 apiMdSha256: 1538a79b2c2da38fb83fd37ccea945314bf2f34218acda3d8f15e2f0268f3040
-parserVersion: 0.3.28
+parserVersion: 0.3.30
 pythonVersion: 3.13.14

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

@@ -93,7 +93,6 @@
 from ._inference_service import _InferenceService
 from .documents import ConnectionPolicy, DatabaseAccount
 from .partition_key import (
-    _build_partition_key_from_properties,
     _Undefined,
     _Empty,
     _PartitionKeyKind,
@@ -3368,9 +3367,21 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
             base.set_session_token_header(self, req_headers, path, request_params, options, partition_key_range_id)
 
         # Check if the overlapping ranges can be populated
+        #
+        # Complete partition keys are classified upstream in
+        # container.py::query_items and routed via the PartitionKey request
+        # header on __Post. Do NOT re-add an EPK conversion for complete keys
+        # here: it drops the PartitionKey header and defeats the server-side
+        # index-only aggregate (COUNT/etc.), causing a full partition scan.
+        # Only feed ranges and hierarchical-prefix keys belong on the EPK path
+        # below. Note the sync classifier lives in container.py while the async
+        # twin classifies inside __QueryFeed, so do not mirror the async branch
+        # into this file.
+        #
+        # The container_properties pop below is a no-op in production (nothing
+        # sets that kwarg anymore); it is retained solely as a test API shim.
         feed_range_epk = None
-        container_properties = kwargs.pop("container_properties", None)
-        is_full_pk_scope = False
+        kwargs.pop("container_properties", None)
         if "feed_range" in kwargs:
             feed_range = kwargs.pop("feed_range")
             feed_range_epk = FeedRangeInternalEpk.from_json(feed_range).get_normalized_range()
@@ -3379,27 +3390,6 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
             prefix_partition_key_value: _SequentialPartitionKeyType = kwargs.pop("prefix_partition_key_value")
             feed_range_epk = (
                 prefix_partition_key_obj._get_epk_range_for_prefix_partition_key(prefix_partition_key_value))
-        elif options.get("partitionKey") is not None and container_properties is not None:
-            partition_key_value = options["partitionKey"]
-            partition_key_obj = _build_partition_key_from_properties(container_properties)
-            if not partition_key_obj._is_prefix_partition_key(partition_key_value):
-                # Full-PK returns a single-value inclusive range; normalize to
-                # [min, max) before routing-map overlap resolution.
-                #
-                # NOTE: do NOT pop the PartitionKey header here. The pop is
-                # deferred to the `if pagination_state is not None:` block
-                # below, i.e. until we've confirmed the new feed-range
-                # routing path is actually taking over. If routing comes back
-                # with zero overlaps (stale cache, mid-split, etc.) we fall
-                # through to the regular __Post path, and that fallthrough
-                # must still carry the legacy PK header — otherwise the
-                # backend gets a request with no partition scoping and either
-                # raises BAD_REQUEST (cross-partition disabled) or silently
-                # runs an unscoped cross-partition query (wrong results).
-                feed_range_epk = partition_key_obj._get_epk_range_for_partition_key(
-                    partition_key_value
-                ).to_normalized_range()
-                is_full_pk_scope = True
 
         # If feed_range_epk exist, query with the range
         if feed_range_epk is not None:
@@ -3446,21 +3436,17 @@ def _is_input_scope_single_partition() -> bool:
                 return cached_is_single_partition
 
             if inbound_serialized_continuation and inbound_token_payload is None:
-                scope_is_single_partition = False
-                if not is_full_pk_scope:
-                    scope_is_single_partition = _is_input_scope_single_partition()
+                scope_is_single_partition = _is_input_scope_single_partition()
                 if _should_bridge_legacy_continuation(
                     inbound_serialized_continuation,
                     inbound_token_payload,
-                    is_full_pk_scope,
                     scope_is_single_partition,
                 ):
                     legacy_bridge_in_use = True
-                    # Hot path: legacy is the normal inbound shape for full-PK
-                    # and currently-single-partition feed-range queries (we
-                    # just emitted one). The bridge wires the legacy string
-                    # into the internal pagination queue; the outbound token
-                    # format on the next page is unchanged.
+                    # Hot path: legacy is the normal inbound shape for currently
+                    # single-partition feed-range queries. The bridge wires the
+                    # string into the internal pagination queue; the outbound
+                    # token format on the next page is unchanged.
                     _LOGGER.debug(
                         "Bridging inbound legacy continuation into internal pagination state; "
                         "outbound token format will remain unchanged (legacy single-string)."
@@ -3510,13 +3496,6 @@ def _is_input_scope_single_partition() -> bool:
                     )
 
             if pagination_state is not None:
-                if is_full_pk_scope:
-                    # Drop the legacy partition-key header now that the
-                    # feed-range routing path is taking over. The inner POSTs
-                    # in the loop set PartitionKeyRangeID / StartEpkString /
-                    # EndEpkString explicitly; sending both routing styles on
-                    # one request is undefined on the service side.
-                    req_headers.pop(http_constants.HttpHeaders.PartitionKey, None)
                 results: dict[str, Any] = {}
                 feedrange_response_headers: CaseInsensitiveDict = CaseInsensitiveDict()
                 consecutive_no_progress_pages = 0
@@ -3532,8 +3511,7 @@ def _checkpoint_and_reraise(error: Exception) -> NoReturn:
                             resource_id_str,
                             query,
                             feed_range_epk,
-                            is_full_pk_scope,
-                            (not is_full_pk_scope) and _is_input_scope_single_partition(),
+                            _is_input_scope_single_partition(),
                         )
                     except Exception as continuation_write_error:  # pylint: disable=broad-exception-caught
                         _LOGGER.warning(
@@ -3701,8 +3679,7 @@ def _checkpoint_and_reraise(error: Exception) -> NoReturn:
                     resource_id_str,
                     query,
                     feed_range_epk,
-                    is_full_pk_scope,
-                    (not is_full_pk_scope) and _is_input_scope_single_partition(),
+                    _is_input_scope_single_partition(),
                 )
                 # End feed_range pagination block.
                 self.last_response_headers = feedrange_response_headers

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

@@ -307,7 +307,6 @@ def _validate_token_identity(
 def _should_bridge_legacy_continuation(
     inbound_serialized_continuation: Optional[str],
     inbound_token_payload: Optional[dict],
-    is_full_pk_scope: bool,
     is_single_partition_scope: bool,
 ) -> bool:
     """Whether to bridge an inbound legacy continuation into pagination state.
@@ -316,16 +315,12 @@ def _should_bridge_legacy_continuation(
     structured ``v=1`` (legacy/opaque token), and the current request scope can
     be represented safely by a single legacy continuation slot:
 
-    * full-PK scope (structurally single-partition forever), or
-    * non-full-PK scope that currently maps to one physical partition.
+    * a feed-range or prefix scope that currently maps to one physical partition.
 
     :param inbound_serialized_continuation: Caller-supplied continuation string, if any.
     :type inbound_serialized_continuation: Optional[str]
     :param inbound_token_payload: Decoded structured payload, or ``None`` for legacy/absent token.
     :type inbound_token_payload: Optional[dict]
-    :param is_full_pk_scope: Whether request scope is a full partition-key query
-        (always emits legacy outbound regardless of partition count).
-    :type is_full_pk_scope: bool
     :param is_single_partition_scope: Whether the current input scope maps to one partition.
     :type is_single_partition_scope: bool
     :returns: ``True`` when the legacy continuation can safely be bridged.
@@ -334,7 +329,7 @@ def _should_bridge_legacy_continuation(
     return bool(
         inbound_serialized_continuation
         and inbound_token_payload is None
-        and (is_full_pk_scope or is_single_partition_scope)
+        and is_single_partition_scope
     )
 
 
@@ -505,10 +500,10 @@ def from_single_feedrange_with_continuation(
         """Build state for one feedrange where a backend continuation
         already exists.
 
-        Used for legacy-token compatibility on full-PK queries:
-        we keep the decoder strict, then bridge a legacy continuation
-        string into the queue head's ``bc`` slot for the single target
-        range.
+        Used for legacy-token compatibility when a feed-range or prefix query
+        currently maps to one physical partition. We keep the decoder strict,
+        then bridge a legacy continuation string into the queue head's ``bc``
+        slot for the single target range.
 
         :param feedrange: Single feedrange to seed.
         :type feedrange: ~azure.cosmos._routing.routing_range.Range
@@ -634,16 +629,13 @@ def _write_query_outbound_continuation(
     resource_id: str,
     query: Any,
     feed_range_epk: routing_range.Range,
-    is_full_pk_scope: bool,
     emit_legacy_for_single_partition: bool,
 ) -> None:
     """Write outbound continuation for feed-range pagination.
 
-    Full-PK queries always emit the legacy single-string continuation
-    so persisted bookmarks remain readable by older SDK versions.
-    Feed-range/prefix queries emit legacy continuation when the caller's
-    input scope currently maps to a single physical partition; otherwise
-    they emit the structured envelope.
+    Feed-range and prefix queries emit legacy continuation when the caller's
+    input scope currently maps to a single physical partition; otherwise they
+    emit the structured envelope.
 
     Defense in depth: even when the caller requests legacy emission, the
     writer verifies that the pagination queue can actually be represented
@@ -664,16 +656,13 @@ def _write_query_outbound_continuation(
     :type query: Any
     :param feed_range_epk: Original request feed range.
     :type feed_range_epk: ~azure.cosmos._routing.routing_range.Range
-    :param is_full_pk_scope: Whether request scope is a full partition-key query
-        (always emits legacy outbound regardless of partition count).
-    :type is_full_pk_scope: bool
-    :param emit_legacy_for_single_partition: Whether non-full-PK scope currently maps to a
+    :param emit_legacy_for_single_partition: Whether the input scope currently maps to a
         single physical partition and can safely emit legacy continuation.
     :type emit_legacy_for_single_partition: bool
     :returns: None. Mutates ``last_response_headers`` in place.
     :rtype: None
     """
-    if is_full_pk_scope or emit_legacy_for_single_partition:
+    if emit_legacy_for_single_partition:
         # A single legacy string can represent at most one queue entry's
         # backend continuation. If the queue grew past that (e.g. a
         # mid-page split exploded the head into children), emitting
@@ -690,12 +679,11 @@ def _write_query_outbound_continuation(
             return
         _LOGGER.warning(
             "Pagination queue has %d entries but caller requested legacy emission "
-            "(is_full_pk_scope=%s, emit_legacy_for_single_partition=%s). Falling "
+            "(emit_legacy_for_single_partition=%s). Falling "
             "through to structured envelope to preserve full pagination state; "
             "this indicates a caller-side single-partition classification that is "
             "out of sync with the actual queue shape.",
             len(pagination_state.queue),
-            is_full_pk_scope,
             emit_legacy_for_single_partition,
         )
     pagination_state.write_outbound_continuation(

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

@@ -3164,33 +3164,17 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]:
 
         # Check if the overlapping ranges can be populated
         feed_range_epk = None
-        is_full_pk_scope = False
         if "feed_range" in kwargs:
             feed_range = kwargs.pop("feed_range")
             feed_range_epk = FeedRangeInternalEpk.from_json(feed_range).get_normalized_range()
         elif options.get("partitionKey") is not None and container_property is not None:
             partition_key_value = options["partitionKey"]
             partition_key_obj = _build_partition_key_from_properties(container_property)
+            # Only HPK prefixes require EPK routing; complete keys keep the normal PartitionKey request path.
             if partition_key_obj._is_prefix_partition_key(partition_key_value):
                 req_headers.pop(http_constants.HttpHeaders.PartitionKey, None)
                 partition_key_value = cast(_SequentialPartitionKeyType, partition_key_value)
                 feed_range_epk = partition_key_obj._get_epk_range_for_prefix_partition_key(partition_key_value)
-            else:
-                # Full-PK returns a single-value inclusive range; normalize to
-                # [min, max) before routing-map overlap resolution.
-                #
-                # Do not pop the PartitionKey header here. The pop is deferred
-                # until we have confirmed the feed-range routing path is taking
-                # over (see the pop below, gated on `pagination_state is not None`
-                # and `is_full_pk_scope`). If routing returns zero overlaps we
-                # fall through to the regular __Post path, which still needs the
-                # legacy PK header — otherwise the backend receives an unscoped
-                # request and either raises BAD_REQUEST or silently runs an
-                # unscoped cross-partition query.
-                feed_range_epk = partition_key_obj._get_epk_range_for_partition_key(
-                    partition_key_value
-                ).to_normalized_range()
-                is_full_pk_scope = True
 
         if feed_range_epk is not None:
             if id_ is None:
@@ -3236,21 +3220,17 @@ async def _is_input_scope_single_partition() -> bool:
                 return cached_is_single_partition
 
             if inbound_serialized_continuation and inbound_token_payload is None:
-                scope_is_single_partition = False
-                if not is_full_pk_scope:
-                    scope_is_single_partition = await _is_input_scope_single_partition()
+                scope_is_single_partition = await _is_input_scope_single_partition()
                 if _should_bridge_legacy_continuation(
                     inbound_serialized_continuation,
                     inbound_token_payload,
-                    is_full_pk_scope,
                     scope_is_single_partition,
                 ):
                     legacy_bridge_in_use = True
-                    # Hot path: legacy is the normal inbound shape for full-PK
-                    # and currently-single-partition feed-range queries (we
-                    # just emitted one). The bridge wires the legacy string
-                    # into the internal pagination queue; the outbound token
-                    # format on the next page is unchanged.
+                    # Hot path: legacy is the normal inbound shape for currently
+                    # single-partition feed-range queries. The bridge wires the
+                    # string into the internal pagination queue; the outbound
+                    # token format on the next page is unchanged.
                     _LOGGER.debug(
                         "Bridging inbound legacy continuation into internal pagination state; "
                         "outbound token format will remain unchanged (legacy single-string)."
@@ -3296,13 +3276,6 @@ async def _is_input_scope_single_partition() -> bool:
                         )
 
             if pagination_state is not None:
-                if is_full_pk_scope:
-                    # Drop the legacy partition-key header now that the
-                    # feed-range routing path is taking over. The inner POSTs
-                    # in the loop set PartitionKeyRangeID / StartEpkString /
-                    # EndEpkString explicitly; sending both routing styles on
-                    # one request is undefined on the service side.
-                    req_headers.pop(http_constants.HttpHeaders.PartitionKey, None)
                 results: dict[str, Any] = {}
                 feedrange_response_headers: CaseInsensitiveDict = CaseInsensitiveDict()
                 consecutive_no_progress_pages = 0
@@ -3312,16 +3285,13 @@ async def _checkpoint_and_reraise(error: Exception) -> NoReturn:
                     # for any mid-page failure, then re-raise the original error.
                     self.last_response_headers = feedrange_response_headers
                     try:
-                        single_partition_scope_for_outbound = (
-                            (not is_full_pk_scope) and (await _is_input_scope_single_partition())
-                        )
+                        single_partition_scope_for_outbound = await _is_input_scope_single_partition()
                         _write_query_outbound_continuation(
                             feedrange_response_headers,
                             pagination_state,
                             resource_id_str,
                             query,
                             feed_range_epk,
-                            is_full_pk_scope,
                             single_partition_scope_for_outbound,
                         )
                     except Exception as continuation_write_error:  # pylint: disable=broad-exception-caught
@@ -3489,16 +3459,13 @@ async def _checkpoint_and_reraise(error: Exception) -> NoReturn:
                 # 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.
-                single_partition_scope_for_outbound = (
-                    (not is_full_pk_scope) and (await _is_input_scope_single_partition())
-                )
+                single_partition_scope_for_outbound = await _is_input_scope_single_partition()
                 _write_query_outbound_continuation(
                     feedrange_response_headers,
                     pagination_state,
                     resource_id_str,
                     query,
                     feed_range_epk,
-                    is_full_pk_scope,
                     single_partition_scope_for_outbound,
                 )
                 # End feed_range pagination block.