From 7428abaf165c23cc0609bc219ac84e2c66811e7b Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Thu, 23 Jul 2026 23:45:45 -0500 Subject: [PATCH 1/3] fix - full partition key routing bug --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 1 + .../azure/cosmos/_cosmos_client_connection.py | 52 ++----- .../_routing/feed_range_continuation.py | 36 ++--- .../aio/_cosmos_client_connection_async.py | 49 ++----- .../azure-cosmos/azure/cosmos/container.py | 2 - .../test_feed_range_continuation_token.py | 79 ++--------- .../tests/test_partition_split_retry_unit.py | 131 +++++++++++++++--- .../test_partition_split_retry_unit_async.py | 130 ++++++++++++++--- sdk/cosmos/azure-cosmos/tests/test_query.py | 41 ++++++ .../azure-cosmos/tests/test_query_async.py | 44 ++++++ .../test_query_feed_range_multipartition.py | 31 +++++ ...t_query_feed_range_multipartition_async.py | 34 ++++- 12 files changed, 406 insertions(+), 224 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 9091743794c9..ff319c9010f3 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/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`. #### Other Changes diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index a5d01a7a12c9..acced14e1c4b 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/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, @@ -3369,8 +3368,7 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: # Check if the overlapping ranges can be populated 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 +3377,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 +3423,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 +3483,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 +3498,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 +3666,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 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/feed_range_continuation.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/feed_range_continuation.py index 6a89fb211a24..f3ede33cca40 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_routing/feed_range_continuation.py +++ b/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( diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py b/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py index 1f9f2ca87369..92e9c107d390 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py +++ b/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. diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py index a18a9ca2be5b..cc5f5ad7210c 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/container.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/container.py @@ -1034,8 +1034,6 @@ def query_items( # pylint:disable=docstring-missing-param,too-many-statements # Get container property and init client container caches container_properties = self._get_properties_with_options(feed_options) - kwargs["container_properties"] = container_properties - # Update 'feed_options' from 'kwargs' if utils.valid_key_value_exist(kwargs, "enable_cross_partition_query"): feed_options["enableCrossPartitionQuery"] = kwargs.pop("enable_cross_partition_query") diff --git a/sdk/cosmos/azure-cosmos/tests/test_feed_range_continuation_token.py b/sdk/cosmos/azure-cosmos/tests/test_feed_range_continuation_token.py index 3b3c2d588ca7..c655bfa6714c 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_feed_range_continuation_token.py +++ b/sdk/cosmos/azure-cosmos/tests/test_feed_range_continuation_token.py @@ -1699,7 +1699,7 @@ def test_drained_state_clears_continuation_header(self): class TestWriteQueryOutboundContinuation: """Outbound continuation format selection should match request scope policy.""" - def test_full_pk_scope_always_emits_legacy(self): + def test_single_partition_scope_emits_legacy(self): state = _FeedRangePaginationState.from_single_feedrange_with_continuation( _HEAD_FEEDRANGE, _BACKEND_CONT, @@ -1713,33 +1713,12 @@ def test_full_pk_scope_always_emits_legacy(self): _RID, _QUERY, _FEED_RANGE, - is_full_pk_scope=True, - emit_legacy_for_single_partition=False, - ) - - assert headers[http_constants.HttpHeaders.Continuation] == _BACKEND_CONT - - def test_non_full_pk_single_partition_scope_emits_legacy(self): - state = _FeedRangePaginationState.from_single_feedrange_with_continuation( - _HEAD_FEEDRANGE, - _BACKEND_CONT, - page_size_hint=5, - ) - headers: dict = {} - - _write_query_outbound_continuation( - headers, - state, - _RID, - _QUERY, - _FEED_RANGE, - is_full_pk_scope=False, emit_legacy_for_single_partition=True, ) assert headers[http_constants.HttpHeaders.Continuation] == _BACKEND_CONT - def test_non_full_pk_multi_partition_scope_emits_structured(self): + def test_multi_partition_scope_emits_structured(self): state = _FeedRangePaginationState( [(_HEAD_FEEDRANGE, _BACKEND_CONT), (_REMAINING_FEEDRANGE, None)], page_size_hint=5, @@ -1752,7 +1731,6 @@ def test_non_full_pk_multi_partition_scope_emits_structured(self): _RID, _QUERY, _FEED_RANGE, - is_full_pk_scope=False, emit_legacy_for_single_partition=False, ) @@ -1760,41 +1738,6 @@ def test_non_full_pk_multi_partition_scope_emits_structured(self): assert decoded is not None assert decoded[_FIELD_VERSION] == _TOKEN_VERSION - def test_full_pk_scope_with_multi_entry_queue_falls_through_to_structured(self, caplog): - """Defense in depth: if the queue somehow has >1 entries while the - caller claims ``is_full_pk_scope=True`` (a structural impossibility - in normal operation, but a possible caller-side bug surface), the - writer must NOT silently discard tail entries via legacy emission. - It falls through to the structured envelope and logs a warning. - """ - state = _FeedRangePaginationState( - [(_HEAD_FEEDRANGE, _BACKEND_CONT), (_REMAINING_FEEDRANGE, None)], - page_size_hint=5, - ) - headers: dict = {} - - with caplog.at_level("WARNING", logger="azure.cosmos._routing.feed_range_continuation"): - _write_query_outbound_continuation( - headers, - state, - _RID, - _QUERY, - _FEED_RANGE, - is_full_pk_scope=True, - emit_legacy_for_single_partition=False, - ) - - # Defense: queue has 2 entries, so the writer falls through to structured. - decoded = _decode_token(headers[http_constants.HttpHeaders.Continuation]) - assert decoded is not None - assert decoded[_FIELD_VERSION] == _TOKEN_VERSION - # And it surfaces the caller-side inconsistency as a WARNING. - assert any( - "Pagination queue has 2 entries" in record.getMessage() - and record.levelname == "WARNING" - for record in caplog.records - ) - def test_emit_legacy_with_multi_entry_queue_falls_through_to_structured(self, caplog): """Defense in depth: a stale single-partition cache after a mid-page split could set ``emit_legacy_for_single_partition=True`` on a @@ -1814,7 +1757,6 @@ def test_emit_legacy_with_multi_entry_queue_falls_through_to_structured(self, ca _RID, _QUERY, _FEED_RANGE, - is_full_pk_scope=False, emit_legacy_for_single_partition=True, ) @@ -1832,31 +1774,26 @@ class TestLegacyBridgeDecision: """Legacy inbound continuation is bridged only when scope is safely single-partition.""" @pytest.mark.parametrize( - "inbound_serialized_continuation,inbound_token_payload,is_full_pk_scope," - "is_single_partition_scope,expected", + "inbound_serialized_continuation,inbound_token_payload,is_single_partition_scope,expected", [ - (None, None, False, True, False), - ("", None, False, True, False), - ("legacy", {"v": 1}, False, True, False), - ("legacy", None, True, False, True), - ("legacy", None, False, True, True), - ("legacy", None, False, False, False), + (None, None, True, False), + ("", None, True, False), + ("legacy", {"v": 1}, True, False), + ("legacy", None, True, True), + ("legacy", None, False, False), ], ) def test_should_bridge_legacy_continuation_policy( self, inbound_serialized_continuation, inbound_token_payload, - is_full_pk_scope, is_single_partition_scope, expected, ): assert _should_bridge_legacy_continuation( inbound_serialized_continuation, inbound_token_payload, - is_full_pk_scope, is_single_partition_scope, ) is expected - diff --git a/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit.py b/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit.py index d4ab355ed253..dd6e36b79908 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit.py +++ b/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit.py @@ -13,7 +13,7 @@ import pytest from azure.core.exceptions import ServiceRequestError -from azure.cosmos import exceptions +from azure.cosmos import exceptions, PartitionKey from azure.cosmos import _retry_utility from azure.cosmos._cosmos_client_connection import CosmosClientConnection from azure.cosmos._execution_context.base_execution_context import _DefaultQueryExecutionContext @@ -1231,41 +1231,128 @@ def test_queryfeed_populates_capture_dict_from_options(self): assert result == [{"id": "1"}] assert headers is canned_headers - def test_queryfeed_full_pk_no_overlap_fallback_preserves_partition_key_header(self): - """Full-PK no-overlap fallback must retain legacy PartitionKey header on __Post.""" + def test_queryfeed_full_partition_key_count_uses_partition_key_routing(self): + """Complete single-path and hierarchical keys must avoid EPK feed-range routing.""" client = self._create_minimal_connection() client._query_compatibility_mode = client._QueryCompatibilityMode.Default client._routing_map_provider = MagicMock() - client._routing_map_provider.get_overlapping_ranges.return_value = [] + client._routing_map_provider.get_overlapping_ranges.side_effect = AssertionError( + "complete partition keys must not resolve EPK overlaps" + ) + + for partition_key, partition_key_header, container_properties in ( + ( + "PowerCut", + '["PowerCut"]', + {"partitionKey": {"paths": ["/documentType"], "kind": "Hash", "version": 2}}, + ), + ( + ["tenant", "user"], + '["tenant","user"]', + { + "partitionKey": { + "paths": ["/tenant", "/user"], + "kind": "MultiHash", + "version": 2, + } + }, + ), + ): + with self.subTest(partition_key=partition_key): + seen_headers = [] - seen_partition_key_headers = [] + def post_side_effect(_path, _request_params, _query, req_headers, **_kwargs): + seen_headers.append(dict(req_headers)) + return {"Documents": [30735986]}, {} - def post_side_effect(_path, _request_params, _query, req_headers, **_kwargs): - seen_partition_key_headers.append(req_headers.get(HttpHeaders.PartitionKey)) - return {"Documents": [{"id": "doc-1"}]}, {} + with patch( + "azure.cosmos._cosmos_client_connection.base.GetHeaders", + return_value={HttpHeaders.PartitionKey: partition_key_header}, + ): + with patch( + "azure.cosmos._cosmos_client_connection.base.set_session_token_header", + return_value=None, + ): + with patch.object( + client, + "_CosmosClientConnection__Post", + side_effect=post_side_effect, + ): + docs, headers = client.QueryFeed( + path="/dbs/db/colls/c1/docs", + collection_id="rid-c1", + query="SELECT VALUE COUNT(1) FROM c", + options={ + "partitionKey": partition_key, + "enableCrossPartitionQuery": False, + "populateQueryMetrics": True, + }, + container_properties=container_properties, + ) + + assert docs == [30735986] + assert HttpHeaders.Continuation not in headers + assert len(seen_headers) == 1 + assert seen_headers[0][HttpHeaders.PartitionKey] == partition_key_header + assert HttpHeaders.PartitionKeyRangeID not in seen_headers[0] + assert HttpHeaders.ReadFeedKeyType not in seen_headers[0] + assert HttpHeaders.StartEpkString not in seen_headers[0] + assert HttpHeaders.EndEpkString not in seen_headers[0] + + client._routing_map_provider.get_overlapping_ranges.assert_not_called() + + def test_queryfeed_hpk_prefix_count_keeps_split_safe_epk_routing(self): + """A hierarchical-key prefix still spans logical partitions and uses EPK pagination.""" + client = self._create_minimal_connection() + client._query_compatibility_mode = client._QueryCompatibilityMode.Default + client._routing_map_provider = MagicMock() + + def overlap_side_effect(_rid, ranges, _options): + requested = ranges[0] + return [{ + "id": "0", + "minInclusive": requested.min, + "maxExclusive": requested.max, + }] + + client._routing_map_provider.get_overlapping_ranges.side_effect = overlap_side_effect + seen_headers = [] - container_properties = {"partitionKey": {"paths": ["/pk"], "kind": "Hash", "version": 2}} - options = {"partitionKey": ["mypk"]} + def post_side_effect(_path, _request_params, _query, req_headers, **_kwargs): + seen_headers.append(dict(req_headers)) + return {"Documents": [42]}, {} with patch( "azure.cosmos._cosmos_client_connection.base.GetHeaders", - return_value={HttpHeaders.PartitionKey: '["mypk"]'}, + return_value={}, ): - with patch("azure.cosmos._cosmos_client_connection.base.set_session_token_header", return_value=None): - with patch.object(client, "_CosmosClientConnection__Post", side_effect=post_side_effect): - docs, _headers = client.QueryFeed( + with patch( + "azure.cosmos._cosmos_client_connection.base.set_session_token_header", + return_value=None, + ): + with patch.object( + client, + "_CosmosClientConnection__Post", + side_effect=post_side_effect, + ): + docs, headers = client.QueryFeed( path="/dbs/db/colls/c1/docs", collection_id="rid-c1", - query="SELECT * FROM c", - options=options, - container_properties=container_properties, + query="SELECT VALUE COUNT(1) FROM c", + options={}, + prefix_partition_key_object=PartitionKey( + path=["/tenant", "/user"], kind="MultiHash" + ), + prefix_partition_key_value=["tenant"], ) - assert docs == [{"id": "doc-1"}] - assert seen_partition_key_headers == ['["mypk"]'], ( - "When full-PK routing finds no overlaps and falls back to __Post, " - "the legacy PartitionKey header must be preserved." - ) + assert docs == [42] + assert HttpHeaders.Continuation not in headers + assert len(seen_headers) == 1 + assert HttpHeaders.PartitionKey not in seen_headers[0] + assert seen_headers[0][HttpHeaders.PartitionKeyRangeID] == "0" + assert seen_headers[0][HttpHeaders.ReadFeedKeyType] == "EffectivePartitionKeyRange" + assert client._routing_map_provider.get_overlapping_ranges.call_count >= 2 def test_queryfeed_feed_range_legacy_inbound_single_partition_honors_and_emits_legacy(self): """Legacy inbound continuation is honored when feed_range currently maps to one partition.""" diff --git a/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit_async.py b/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit_async.py index 417e141dafab..a75279a0f744 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_partition_split_retry_unit_async.py @@ -1175,28 +1175,118 @@ async def _noop_set_session(*args, **kwargs): assert mock_post.await_count == 1 assert result == ([], {}) - async def test_queryfeed_full_pk_no_overlap_fallback_preserves_partition_key_header_async(self): - """Async full-PK no-overlap fallback must retain legacy PartitionKey header on __Post.""" + async def test_queryfeed_full_partition_key_count_uses_partition_key_routing_async(self): + """Complete single-path and hierarchical keys must avoid EPK feed-range routing.""" client = self._create_minimal_connection() client._query_compatibility_mode = client._QueryCompatibilityMode.Default client._routing_map_provider = MagicMock() - client._routing_map_provider.get_overlapping_ranges = AsyncMock(return_value=[]) + client._routing_map_provider.get_overlapping_ranges = AsyncMock( + side_effect=AssertionError("complete partition keys must not resolve EPK overlaps") + ) + + async def _noop_set_session(*args, **kwargs): + return None + + for partition_key, partition_key_header, container_properties in ( + ( + "PowerCut", + '["PowerCut"]', + {"partitionKey": {"paths": ["/documentType"], "kind": "Hash", "version": 2}}, + ), + ( + ["tenant", "user"], + '["tenant","user"]', + { + "partitionKey": { + "paths": ["/tenant", "/user"], + "kind": "MultiHash", + "version": 2, + } + }, + ), + ): + with self.subTest(partition_key=partition_key): + seen_headers = [] - seen_partition_key_headers = [] + async def post_side_effect(_path, _request_params, _query, req_headers, **_kwargs): + seen_headers.append(dict(req_headers)) + return {"Documents": [30735986]}, {} + + async def get_container_properties(_options): + return container_properties + + with patch( + "azure.cosmos.aio._cosmos_client_connection_async.base.GetHeaders", + return_value={HttpHeaders.PartitionKey: partition_key_header}, + ): + with patch( + "azure.cosmos.aio._cosmos_client_connection_async.base.set_session_token_header_async", + side_effect=_noop_set_session, + ): + with patch.object( + client, + "_CosmosClientConnection__Post", + side_effect=post_side_effect, + ): + docs, headers = await client.QueryFeed( + path="/dbs/db/colls/c1/docs", + collection_id="rid-c1", + query="SELECT VALUE COUNT(1) FROM c", + options={ + "partitionKey": partition_key, + "enableCrossPartitionQuery": False, + "populateQueryMetrics": True, + }, + containerProperties=get_container_properties, + ) + + assert docs == [30735986] + assert HttpHeaders.Continuation not in headers + assert len(seen_headers) == 1 + assert seen_headers[0][HttpHeaders.PartitionKey] == partition_key_header + assert HttpHeaders.PartitionKeyRangeID not in seen_headers[0] + assert HttpHeaders.ReadFeedKeyType not in seen_headers[0] + assert HttpHeaders.StartEpkString not in seen_headers[0] + assert HttpHeaders.EndEpkString not in seen_headers[0] + + client._routing_map_provider.get_overlapping_ranges.assert_not_awaited() + + async def test_queryfeed_hpk_prefix_count_keeps_split_safe_epk_routing_async(self): + """A hierarchical-key prefix still spans logical partitions and uses EPK pagination.""" + client = self._create_minimal_connection() + client._query_compatibility_mode = client._QueryCompatibilityMode.Default + client._routing_map_provider = MagicMock() + + async def overlap_side_effect(_rid, ranges, _options): + requested = ranges[0] + return [{ + "id": "0", + "minInclusive": requested.min, + "maxExclusive": requested.max, + }] + + client._routing_map_provider.get_overlapping_ranges = AsyncMock(side_effect=overlap_side_effect) + seen_headers = [] async def post_side_effect(_path, _request_params, _query, req_headers, **_kwargs): - seen_partition_key_headers.append(req_headers.get(HttpHeaders.PartitionKey)) - return {"Documents": [{"id": "doc-1"}]}, {} + seen_headers.append(dict(req_headers)) + return {"Documents": [42]}, {} + + async def get_container_properties(_options): + return { + "partitionKey": { + "paths": ["/tenant", "/user"], + "kind": "MultiHash", + "version": 2, + } + } async def _noop_set_session(*args, **kwargs): return None - container_properties = {"partitionKey": {"paths": ["/pk"], "kind": "Hash", "version": 2}} - options = {"partitionKey": ["mypk"]} - with patch( "azure.cosmos.aio._cosmos_client_connection_async.base.GetHeaders", - return_value={HttpHeaders.PartitionKey: '["mypk"]'}, + return_value={HttpHeaders.PartitionKey: '["tenant"]'}, ): with patch( "azure.cosmos.aio._cosmos_client_connection_async.base.set_session_token_header_async", @@ -1207,19 +1297,21 @@ async def _noop_set_session(*args, **kwargs): "_CosmosClientConnection__Post", side_effect=post_side_effect, ): - docs, _headers = await client.QueryFeed( + docs, headers = await client.QueryFeed( path="/dbs/db/colls/c1/docs", collection_id="rid-c1", - query="SELECT * FROM c", - options=options, - container_property=container_properties, + query="SELECT VALUE COUNT(1) FROM c", + options={"partitionKey": ["tenant"]}, + containerProperties=get_container_properties, ) - assert docs == [{"id": "doc-1"}] - assert seen_partition_key_headers == ['["mypk"]'], ( - "When async full-PK routing finds no overlaps and falls back to __Post, " - "the legacy PartitionKey header must be preserved." - ) + assert docs == [42] + assert HttpHeaders.Continuation not in headers + assert len(seen_headers) == 1 + assert HttpHeaders.PartitionKey not in seen_headers[0] + assert seen_headers[0][HttpHeaders.PartitionKeyRangeID] == "0" + assert seen_headers[0][HttpHeaders.ReadFeedKeyType] == "EffectivePartitionKeyRange" + assert client._routing_map_provider.get_overlapping_ranges.await_count >= 2 async def test_queryfeed_feed_range_legacy_inbound_single_partition_honors_and_emits_legacy_async(self): """Async: legacy inbound continuation is honored when feed_range maps to one partition.""" diff --git a/sdk/cosmos/azure-cosmos/tests/test_query.py b/sdk/cosmos/azure-cosmos/tests/test_query.py index 1c46c470af4c..a008fac5abd2 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query.py @@ -101,6 +101,47 @@ def test_populate_query_metrics(self): self.assertTrue(all(['=' in x for x in metrics])) self._delete_container_for_test(created_collection.id) + def test_partition_scoped_count_is_index_only(self): + container_id = "query_count_metrics_test_" + str(uuid.uuid4()) + created_collection = self._create_container_for_test(container_id, PartitionKey(path="/documentType")) + try: + for i in range(25): + created_collection.create_item({ + "id": f"power-cut-{i}", + "documentType": "PowerCut", + "value": i, + }) + created_collection.create_item({ + "id": "other-document-type", + "documentType": "Other", + }) + + query_iterable = created_collection.query_items( + query="SELECT VALUE COUNT(1) FROM c", + partition_key="PowerCut", + enable_cross_partition_query=False, + populate_query_metrics=True, + ) + pager = query_iterable.by_page() + pages = [list(page) for page in pager] + + self.assertEqual([[25]], pages) + self.assertIsNone(pager.continuation_token) + + response_headers = query_iterable.get_response_headers() + metrics_header = response_headers[http_constants.HttpHeaders.QueryMetrics] + metrics = dict( + entry.split("=", 1) + for entry in metrics_header.split(";") + if entry + ) + self.assertEqual(0, int(metrics["retrievedDocumentCount"])) + self.assertEqual(0, int(metrics["retrievedDocumentSize"])) + self.assertEqual(1.0, float(metrics["indexUtilizationRatio"])) + self.assertLess(float(response_headers[http_constants.HttpHeaders.RequestCharge]), 20.0) + finally: + self._delete_container_for_test(container_id) + def test_populate_index_metrics(self): created_collection = self._create_container_for_test("query_index_test", PartitionKey(path="/pk")) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_async.py index 969b5e0fc724..1c308dcdd205 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_async.py @@ -140,6 +140,50 @@ async def test_populate_query_metrics_async(self): self._delete_container_for_test(container_id) + async def test_partition_scoped_count_is_index_only_async(self): + container_id = "query_count_metrics_test_" + str(uuid.uuid4()) + created_collection = self._create_container_for_test( + container_id, PartitionKey(path="/documentType") + ) + try: + for i in range(25): + await created_collection.create_item({ + "id": f"power-cut-{i}", + "documentType": "PowerCut", + "value": i, + }) + await created_collection.create_item({ + "id": "other-document-type", + "documentType": "Other", + }) + + query_iterable = created_collection.query_items( + query="SELECT VALUE COUNT(1) FROM c", + partition_key="PowerCut", + populate_query_metrics=True, + ) + pager = query_iterable.by_page() + pages = [] + async for page in pager: + pages.append([item async for item in page]) + + assert pages == [[25]] + assert pager.continuation_token is None + + response_headers = query_iterable.get_response_headers() + metrics_header = response_headers[http_constants.HttpHeaders.QueryMetrics] + metrics = dict( + entry.split("=", 1) + for entry in metrics_header.split(";") + if entry + ) + assert int(metrics["retrievedDocumentCount"]) == 0 + assert int(metrics["retrievedDocumentSize"]) == 0 + assert float(metrics["indexUtilizationRatio"]) == 1.0 + assert float(response_headers[http_constants.HttpHeaders.RequestCharge]) < 20.0 + finally: + self._delete_container_for_test(container_id) + async def test_populate_index_metrics_async(self): container_id = "index_metrics_test" + str(uuid.uuid4()) created_collection = self._create_container_for_test(container_id, PartitionKey(path="/pk")) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py index 13f485362a1b..3feb28a85e9f 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py @@ -308,6 +308,37 @@ def test_full_partition_key_query_pagination_resume(self): ] fetched_ids = [item['id'] for item in first_page] + resumed_remaining_ids assert baseline_ids == fetched_ids + + # A single-partition feed-range query emits a legacy opaque + # continuation token; that token must still resume on the + # partition-key request path and return the same remaining items. + feed_range = created_container.feed_range_from_partition_key(full_key) + epk_pager = created_container.query_items( + query=query, + feed_range=feed_range, + max_item_count=7, + ).by_page() + epk_first_page = list(next(epk_pager)) + assert epk_first_page + legacy_token = epk_pager.continuation_token + assert legacy_token + assert _decode_token(legacy_token) is None + + expected_epk_remaining_ids = [ + item['id'] + for page in epk_pager + for item in page + ] + resumed_via_partition_key_ids = [ + item['id'] + for page in created_container.query_items( + query=query, + partition_key=full_key, + max_item_count=7, + ).by_page(legacy_token) + for item in page + ] + assert expected_epk_remaining_ids == resumed_via_partition_key_ids finally: db.delete_container(created_container.id) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py index de4279fa12eb..c1ab68cab09a 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py @@ -873,6 +873,39 @@ async def test_full_partition_key_query_pagination_resume_async(self): assert baseline_ids == fetched_ids, ( "Page 1 plus the resumed pages must equal the baseline " "documents for this partition key in the same order") + + # A single-partition feed-range query emits a legacy opaque + # continuation token; that token must still resume on the + # partition-key request path and return the same remaining items. + feed_range = await created_container.feed_range_from_partition_key(full_key) + epk_pager = created_container.query_items( + query=query, + feed_range=feed_range, + max_item_count=7, + ).by_page() + epk_first_page_iter = await epk_pager.__anext__() + epk_first_page = [item async for item in epk_first_page_iter] + assert epk_first_page + legacy_token = epk_pager.continuation_token + assert legacy_token + assert _decode_token(legacy_token) is None + + expected_epk_remaining_ids: List[str] = [] + async for page in epk_pager: + expected_epk_remaining_ids.extend( + [item['id'] async for item in page]) + + resumed_via_partition_key_ids: List[str] = [] + partition_key_pager = created_container.query_items( + query=query, + partition_key=full_key, + max_item_count=7, + ).by_page(legacy_token) + async for page in partition_key_pager: + resumed_via_partition_key_ids.extend( + [item['id'] async for item in page]) + + assert expected_epk_remaining_ids == resumed_via_partition_key_ids finally: try: await db.delete_container(created_container.id) @@ -1080,4 +1113,3 @@ async def _stalled_post(*_args, **_kwargs): if __name__ == "__main__": unittest.main() - From 41bcb27a00c1318b1dca4bc74fa9dad3310378ef Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Fri, 24 Jul 2026 11:00:28 -0500 Subject: [PATCH 2/3] addressing agent comments --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- sdk/cosmos/azure-cosmos/api.metadata.yml | 2 +- .../azure/cosmos/_cosmos_client_connection.py | 13 +++++++++++++ 3 files changed, 15 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index ff319c9010f3..64214d37b36c 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,7 +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`. +* 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 diff --git a/sdk/cosmos/azure-cosmos/api.metadata.yml b/sdk/cosmos/azure-cosmos/api.metadata.yml index 51c24d28144b..fd25fd0f7de8 100644 --- a/sdk/cosmos/azure-cosmos/api.metadata.yml +++ b/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 diff --git a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py index acced14e1c4b..3fcfdff03c49 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3367,6 +3367,19 @@ 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 kwargs.pop("container_properties", None) if "feed_range" in kwargs: From 3ca40e5f40136a4bb1b75947157b016371fa7cb4 Mon Sep 17 00:00:00 2001 From: dibahlfi <106994927+dibahlfi@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:40:17 -0500 Subject: [PATCH 3/3] added test coverage --- .../test_query_feed_range_multipartition.py | 28 ++++++++++++++++-- ...t_query_feed_range_multipartition_async.py | 29 +++++++++++++++++-- 2 files changed, 51 insertions(+), 6 deletions(-) diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py index 3feb28a85e9f..97515905a237 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition.py @@ -247,12 +247,12 @@ def test_single_partition_feed_range(self): # ------------------------------------------------------------------ # # Partition-key caller shapes (full key and prefix key) # ------------------------------------------------------------------ # - def test_full_partition_key_query_pagination_resume(self): - """Full hierarchical partition-key query resumes correctly by continuation. + def test_full_partition_key_query_pagination_resume_and_count_is_index_only(self): + """Full hierarchical partition-key query resumes and remains index-only. This uses a dedicated MultiHash container and a full key value so the query stays scoped to one logical partition while still exercising - pagination + resume on the partition_key path. + pagination, resume, and aggregate query metrics on the partition_key path. """ db = _client().get_database_client(DATABASE_ID) container_id = "FeedRangeMultiPartitionFullPK-" + str(uuid.uuid4()) @@ -309,6 +309,28 @@ def test_full_partition_key_query_pagination_resume(self): fetched_ids = [item['id'] for item in first_page] + resumed_remaining_ids assert baseline_ids == fetched_ids + count_iterable = created_container.query_items( + query="SELECT VALUE COUNT(1) FROM c", + partition_key=full_key, + enable_cross_partition_query=False, + populate_query_metrics=True, + ) + count_pager = count_iterable.by_page() + count_pages = [list(page) for page in count_pager] + assert count_pages == [[25]] + assert count_pager.continuation_token is None + + count_headers = count_iterable.get_response_headers() + count_metrics = dict( + entry.split("=", 1) + for entry in count_headers[http_constants.HttpHeaders.QueryMetrics].split(";") + if entry + ) + assert int(count_metrics["retrievedDocumentCount"]) == 0 + assert int(count_metrics["retrievedDocumentSize"]) == 0 + assert float(count_metrics["indexUtilizationRatio"]) == 1.0 + assert float(count_headers[http_constants.HttpHeaders.RequestCharge]) < 20.0 + # A single-partition feed-range query emits a legacy opaque # continuation token; that token must still resume on the # partition-key request path and return the same remaining items. diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py index c1ab68cab09a..7fc2e644085d 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_multipartition_async.py @@ -798,11 +798,12 @@ async def _drain(p): finally: await client.close() - async def test_full_partition_key_query_pagination_resume_async(self): + async def test_full_partition_key_query_pagination_resume_and_count_is_index_only_async(self): """Query with a full partition key on a hierarchical container, drain page 1, resume from the returned continuation token, and confirm the resumed pages match the remaining pages of a fresh iterator and - cover the same documents as a baseline scan in the same order. + cover the same documents as a baseline scan in the same order. Also + confirm a partition-scoped aggregate remains index-only. """ client = _client() try: @@ -874,6 +875,29 @@ async def test_full_partition_key_query_pagination_resume_async(self): "Page 1 plus the resumed pages must equal the baseline " "documents for this partition key in the same order") + count_iterable = created_container.query_items( + query="SELECT VALUE COUNT(1) FROM c", + partition_key=full_key, + populate_query_metrics=True, + ) + count_pager = count_iterable.by_page() + count_pages = [] + async for page in count_pager: + count_pages.append([item async for item in page]) + assert count_pages == [[25]] + assert count_pager.continuation_token is None + + count_headers = count_iterable.get_response_headers() + count_metrics = dict( + entry.split("=", 1) + for entry in count_headers[http_constants.HttpHeaders.QueryMetrics].split(";") + if entry + ) + assert int(count_metrics["retrievedDocumentCount"]) == 0 + assert int(count_metrics["retrievedDocumentSize"]) == 0 + assert float(count_metrics["indexUtilizationRatio"]) == 1.0 + assert float(count_headers[http_constants.HttpHeaders.RequestCharge]) < 20.0 + # A single-partition feed-range query emits a legacy opaque # continuation token; that token must still resume on the # partition-key request path and return the same remaining items. @@ -1112,4 +1136,3 @@ async def _stalled_post(*_args, **_kwargs): if __name__ == "__main__": unittest.main() -