Skip to content

fix - full partition key routing bug #48237

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dibahlfi wants to merge 3 commits into
base: main
Choose a base branch
from
Open

fix - full partition key routing bug #48237

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7428aba
fix - full partition key routing bug
dibahlfi Jul 24, 2026
41bcb27
addressing agent comments
dibahlfi Jul 24, 2026
3ca40e5
added test coverage
dibahlfi Jul 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions sdk/cosmos/azure-cosmos/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Comment thread
tvaron3 marked this conversation as resolved.

#### Other Changes

Expand Down
2 changes: 1 addition & 1 deletion sdk/cosmos/azure-cosmos/api.metadata.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,3 @@
apiMdSha256: 1538a79b2c2da38fb83fd37ccea945314bf2f34218acda3d8f15e2f0268f3040
parserVersion: 0.3.28
parserVersion: 0.3.30
pythonVersion: 3.13.14
65 changes: 21 additions & 44 deletions sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Comment thread
dibahlfi marked this conversation as resolved.
if "feed_range" in kwargs:
feed_range = kwargs.pop("feed_range")
feed_range_epk = FeedRangeInternalEpk.from_json(feed_range).get_normalized_range()
Expand All @@ -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:
Expand Down Expand Up @@ -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)."
Expand Down Expand Up @@ -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
Expand All @@ -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(
Expand Down Expand Up @@ -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
Expand Down
36 changes: 12 additions & 24 deletions sdk/cosmos/azure-cosmos/azure/cosmos/_routing/feed_range_continuation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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.
Expand All @@ -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
)


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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(
Expand Down
49 changes: 8 additions & 41 deletions sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Comment thread
dibahlfi marked this conversation as resolved.
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.
Comment thread
dibahlfi marked this conversation as resolved.
Comment thread
dibahlfi marked this conversation as resolved.
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:
Expand Down Expand Up @@ -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)."
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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.
Expand Down
Loading
Loading