From 9fa67f6d38c9d0b06e269c423acdc05e5f5df948 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 22 Apr 2026 11:19:38 -0700 Subject: [PATCH 1/2] [Cosmos] Honor max_item_count for query_items(feed_range=...) When a user-supplied feed_range overlaps K physical partition key ranges (for example, after a server-side split), __QueryFeed issues one POST per overlapping range and merges the partial results. Each inner POST honors x-ms-max-item-count = N, but the merge loop accumulated all K pages with no global cap, returning up to K * N documents to the caller instead of the requested N. Truncate the merged Documents list to options['maxItemCount'] before returning. Apply the fix to both the sync and async client connections. Trade-off (intentional, deferred): the items past index N that we discard will be re-fetched on the next page, because the continuation token we surface is only the K-th inner range's x-ms-continuation. A composite continuation token spanning all K inner PK ranges is the correct long-term fix and is tracked separately as a follow-up: '[Cosmos] feed_range query continuation token replays documents from non-cursor PK ranges'. Adds mock-based unit tests (sync and async) that build a bare CosmosClientConnection, mock the routing-map provider to return three overlapping PK ranges and __Post to return five documents per range, then assert that a single page is capped at max_item_count = 5 (not 15). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 + .../azure/cosmos/_cosmos_client_connection.py | 39 +++ .../aio/_cosmos_client_connection_async.py | 38 +++ .../test_query_feed_range_max_item_count.py | 238 ++++++++++++++++++ ...t_query_feed_range_max_item_count_async.py | 188 ++++++++++++++ 5 files changed, 505 insertions(+) create mode 100644 sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py create mode 100644 sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index fd8e170900ed..30711bdb2276 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -7,6 +7,8 @@ #### Breaking Changes #### Bugs Fixed +* Fixed bug where `query_items(feed_range=..., max_item_count=N)` could return up to `K * N` documents per page when the supplied feed range overlapped `K` physical partition key ranges (for example, after a server-side split). The page returned to the caller is now truncated to the requested `max_item_count`. + * Known limitation (deferred): when a `feed_range` overlaps multiple PK ranges, only the last inner range's `x-ms-continuation` is surfaced as the page continuation token. Round-tripping that token sends it to every inner range on the next page, which is undefined server-side and can produce duplicates, missing documents, or non-terminating iteration on subsequent pages. A composite continuation token across overlapping inner ranges is tracked separately as a follow-up; until then, only the *first* page of a multi-range `feed_range` query is reliable. * Fixed bug where `CosmosClient` construction with AAD credentials would crash at startup if the semantic reranking inference endpoint environment variable was not set, even when semantic reranking was not being used. The inference service is now lazily initialized on first use. See [PR 46243](https://github.com/Azure/azure-sdk-for-python/pull/46243) #### 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 37c8bf219306..5a901d277e48 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3408,6 +3408,45 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: response_hook(last_response_headers, partial_result) # if the prefix partition query has results lets return it if results: + # Honor the user-requested page size (maxItemCount) across the K overlapping + # inner physical PK ranges. Each inner __Post above honors max_item_count per + # range, so the merged result can hold up to K * max_item_count documents. + # Truncate to the user-requested cap so a single page never exceeds it. + # + # NOTE: Use a truthy check (mirrors the contract used by _base.GetHeaders, + # which only emits the x-ms-max-item-count header when options['maxItemCount'] + # is truthy). max_item_count=0/None/missing all mean "use the server default + # page size" and must be a no-op here, otherwise we would silently return + # empty pages while the server-side default page actually returned data. + # + # Known limitation (intentionally deferred — tracked separately as a + # follow-up: "[Cosmos] feed_range query continuation token replays documents + # from non-cursor PK ranges"): + # The continuation token surfaced to the caller (last_response_headers + # below is the K-th iteration's headers) is only the K-th inner range's + # x-ms-continuation. When the caller round-trips that token on the next + # page, __QueryFeed re-resolves the K overlapping ranges and sends the + # same caller-supplied continuation to every inner POST. Range K-1's + # continuation against ranges 0..K-2 is undefined server-side (may error, + # may restart from the beginning, may return undefined slices), so callers + # paginating a feed_range that overlaps multiple PK ranges may observe + # duplicates, missing documents, or non-terminating iteration. The correct + # fix is a composite continuation token spanning all K inner PK ranges. + # Until that lands, this branch only delivers a correct *first* page for + # multi-range feed_range queries. + max_item_count = options.get("maxItemCount") + docs = results.get("Documents") + if max_item_count and isinstance(docs, list): + try: + cap = int(max_item_count) + except (TypeError, ValueError): + cap = 0 + if 0 < cap < len(docs): + results["Documents"] = docs[:cap] + # Keep the internal _count field consistent with the truncated + # Documents list so any downstream consumer that introspects + # the merged dict observes a coherent shape. + results["_count"] = cap if last_response_headers.get(http_constants.HttpHeaders.IndexUtilization) is not None: index_metrics_raw = last_response_headers[http_constants.HttpHeaders.IndexUtilization] last_response_headers[http_constants.HttpHeaders.IndexUtilization] = ( 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 3be6fecdc0f9..53aaa0dd78f5 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 @@ -3202,6 +3202,44 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: response_hook(self.last_response_headers, partial_result) # if the prefix partition query has results lets return it if results: + # Honor the user-requested page size (maxItemCount) across the K overlapping + # inner physical PK ranges. Each inner __Post above honors max_item_count per + # range, so the merged result can hold up to K * max_item_count documents. + # Truncate to the user-requested cap so a single page never exceeds it. + # + # NOTE: Use a truthy check (mirrors the contract used by _base.GetHeaders, + # which only emits the x-ms-max-item-count header when options['maxItemCount'] + # is truthy). max_item_count=0/None/missing all mean "use the server default + # page size" and must be a no-op here, otherwise we would silently return + # empty pages while the server-side default page actually returned data. + # + # Known limitation (intentionally deferred — tracked separately as a + # follow-up: "[Cosmos] feed_range query continuation token replays documents + # from non-cursor PK ranges"): + # The continuation token surfaced to the caller is only the K-th inner + # range's x-ms-continuation. When the caller round-trips that token on the + # next page, _CosmosClientConnection__QueryFeed re-resolves the K overlapping + # ranges and sends the same caller-supplied continuation to every inner POST. + # Range K-1's continuation against ranges 0..K-2 is undefined server-side + # (may error, may restart from the beginning, may return undefined slices), + # so callers paginating a feed_range that overlaps multiple PK ranges may + # observe duplicates, missing documents, or non-terminating iteration. The + # correct fix is a composite continuation token spanning all K inner PK + # ranges. Until that lands, this branch only delivers a correct *first* page + # for multi-range feed_range queries. + max_item_count = options.get("maxItemCount") + docs = results.get("Documents") + if max_item_count and isinstance(docs, list): + try: + cap = int(max_item_count) + except (TypeError, ValueError): + cap = 0 + if 0 < cap < len(docs): + results["Documents"] = docs[:cap] + # Keep the internal _count field consistent with the truncated + # Documents list so any downstream consumer that introspects + # the merged dict observes a coherent shape. + results["_count"] = cap if self.last_response_headers.get(http_constants.HttpHeaders.IndexUtilization) is not None: index_metrics_raw = self.last_response_headers[http_constants.HttpHeaders.IndexUtilization] self.last_response_headers[http_constants.HttpHeaders.IndexUtilization] = ( diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py new file mode 100644 index 000000000000..2480b41f9e8f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py @@ -0,0 +1,238 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Sync unit test for the ``feed_range`` query page-size honoring fix. + +When a user-supplied ``feed_range`` overlaps multiple physical PK ranges (for +example, after a server-side split), ``__QueryFeed`` issues one POST per +overlapping range and merges the partial results. The user-requested +``max_item_count`` was previously honored *per inner range*, so a single page +could return up to ``K * max_item_count`` documents (where ``K`` is the number +of overlapping physical ranges). + +This test pins the post-merge truncation that caps the page at the +user-requested ``max_item_count``. + +Note: these tests reach into the name-mangled +``_CosmosClientConnection__QueryFeed`` / ``_CosmosClientConnection__Post`` +members. If ``__QueryFeed`` is renamed or moved off +``CosmosClientConnection``, move these tests with it. +""" + +import unittest +from unittest.mock import MagicMock, patch + +from azure.cosmos._cosmos_client_connection import CosmosClientConnection +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk +from azure.cosmos._routing.routing_range import Range + + +def _build_client_connection(overlapping_ranges=None): + """Build a bare ``CosmosClientConnection`` instance with only the attributes + referenced by ``__QueryFeed``'s feed_range branch. + + We deliberately bypass ``__init__`` so the test does not require an + emulator or any network setup. + """ + client = object.__new__(CosmosClientConnection) + client.default_headers = {} + client._query_compatibility_mode = CosmosClientConnection._QueryCompatibilityMode.Default + client.availability_strategy = None + client.availability_strategy_executor = None + client.availability_strategy_max_concurrency = None + client.last_response_headers = {} + if overlapping_ranges is None: + overlapping_ranges = [ + {"id": "0", "minInclusive": "", "maxExclusive": "55"}, + {"id": "1", "minInclusive": "55", "maxExclusive": "AA"}, + {"id": "2", "minInclusive": "AA", "maxExclusive": "FF"}, + ] + client._routing_map_provider = MagicMock() + client._routing_map_provider.get_overlapping_ranges.return_value = overlapping_ranges + client._UpdateSessionIfRequired = MagicMock() + return client + + +def _make_feed_range_dict(): + """Return a feed_range JSON-serializable dict that spans the full hash space.""" + return FeedRangeInternalEpk( + Range(range_min="", range_max="FF", isMinInclusive=True, isMaxInclusive=False) + ).to_dict() + + +def _docs(n, prefix="d"): + return {"Documents": [{"id": f"{prefix}-{i}"} for i in range(n)]} + + +def _capture_result_fn(): + """A ``result_fn`` that records the dict it is called with so tests can assert + that the *underlying merged dict* (not just the projection) was truncated.""" + captured = {} + + def fn(result): + captured["result"] = result + return result["Documents"] + return captured, fn + + +@patch("azure.cosmos._cosmos_client_connection.base.set_session_token_header", + lambda *args, **kwargs: None) +@patch("azure.cosmos._cosmos_client_connection.base.GetHeaders", + side_effect=lambda *args, **kwargs: {}) +class TestQueryFeedRangeMaxItemCount(unittest.TestCase): + + def _query(self, client, options, post_side_effect): + post_mock = MagicMock(side_effect=post_side_effect) + client._CosmosClientConnection__Post = post_mock + captured, result_fn = _capture_result_fn() + docs, _headers = client._CosmosClientConnection__QueryFeed( + path="/dbs/db1/colls/coll1/docs", + resource_type="docs", + resource_id="coll1", + result_fn=result_fn, + create_fn=None, + query={"query": "SELECT * FROM c"}, + options=options, + feed_range=_make_feed_range_dict(), + ) + return docs, post_mock, captured + + def test_first_page_truncated_to_max_item_count(self, _mock_get_headers): + """A single page must not exceed ``max_item_count`` even when multiple + physical PK ranges overlap the requested feed_range.""" + client = _build_client_connection() + page_size = 5 + docs, post_mock, captured = self._query( + client, + options={"maxItemCount": page_size}, + post_side_effect=lambda *a, **kw: (_docs(page_size), {}), + ) + # All three inner ranges queried (intentional — see the follow-up note + # about composite continuation tokens). + self.assertEqual(post_mock.call_count, 3) + # Both the projection and the merged dict are capped. + self.assertEqual(len(docs), page_size) + self.assertEqual(len(captured["result"]["Documents"]), page_size) + + def test_truncation_to_one_across_three_ranges(self, _mock_get_headers): + """Tightest cap: K=3, N=1 — proves we truncate, not "merge correctly".""" + client = _build_client_connection() + docs, _post_mock, captured = self._query( + client, + options={"maxItemCount": 1}, + post_side_effect=lambda *a, **kw: (_docs(5), {}), + ) + self.assertEqual(len(docs), 1) + self.assertEqual(len(captured["result"]["Documents"]), 1) + + def test_no_truncation_when_under_cap(self, _mock_get_headers): + """If the merged result is already <= max_item_count, nothing is dropped.""" + client = _build_client_connection() + docs, _post_mock, _captured = self._query( + client, + options={"maxItemCount": 10}, + post_side_effect=lambda *a, **kw: (_docs(1), {}), + ) + self.assertEqual(len(docs), 3) + + def test_boundary_exact_cap_no_slice(self, _mock_get_headers): + """When merged length == cap, the list is returned unchanged.""" + client = _build_client_connection() + # 3 ranges * 1 doc = 3 merged; cap = 3. + docs, _post_mock, _captured = self._query( + client, + options={"maxItemCount": 3}, + post_side_effect=lambda *a, **kw: (_docs(1), {}), + ) + self.assertEqual(len(docs), 3) + + def test_no_max_item_count_no_truncation(self, _mock_get_headers): + """When no maxItemCount is supplied, the merged page is returned in full.""" + client = _build_client_connection() + docs, _post_mock, _captured = self._query( + client, + options={}, + post_side_effect=lambda *a, **kw: (_docs(4), {}), + ) + # 3 ranges * 4 docs each = 12, no truncation since maxItemCount is unset. + self.assertEqual(len(docs), 12) + + def test_max_item_count_zero_means_server_default_no_truncation(self, _mock_get_headers): + """maxItemCount=0 mirrors _base.GetHeaders' truthy contract: it means + "use the server default page size", not "return zero items". The + truncation block must be a no-op so we don't silently empty a page + whose docs were actually fetched at server cost.""" + client = _build_client_connection() + docs, _post_mock, _captured = self._query( + client, + options={"maxItemCount": 0}, + post_side_effect=lambda *a, **kw: (_docs(7), {}), + ) + # 3 ranges * 7 docs each = 21, no truncation since cap is non-positive. + self.assertEqual(len(docs), 21) + + def test_single_overlapping_range_unchanged(self, _mock_get_headers): + """Single-range feed_range case: the truncation must not regress the + existing behavior (one POST, return the partial result as-is).""" + client = _build_client_connection(overlapping_ranges=[ + {"id": "0", "minInclusive": "", "maxExclusive": "FF"}, + ]) + docs, post_mock, _captured = self._query( + client, + options={"maxItemCount": 5}, + post_side_effect=lambda *a, **kw: (_docs(5), {}), + ) + self.assertEqual(post_mock.call_count, 1) + self.assertEqual(len(docs), 5) + + def test_missing_documents_key_does_not_crash(self, _mock_get_headers): + """A partial result missing the Documents key entirely must not raise + from the truncation block; the ``isinstance(docs, list)`` guard + rejects ``None`` and the block is a no-op.""" + client = _build_client_connection(overlapping_ranges=[ + {"id": "0", "minInclusive": "", "maxExclusive": "FF"}, + ]) + post_mock = MagicMock(side_effect=lambda *a, **kw: ({"some_other_field": 42}, {})) + client._CosmosClientConnection__Post = post_mock + captured = {} + + def lenient_result_fn(result): + captured["result"] = result + # Mimic real-world result_fns that defensively project; the point + # of this test is that the truncation block itself does not raise + # when Documents is missing. + return result.get("Documents") or [] + + # Should not raise. + docs, _headers = client._CosmosClientConnection__QueryFeed( + path="/dbs/db1/colls/coll1/docs", + resource_type="docs", + resource_id="coll1", + result_fn=lenient_result_fn, + create_fn=None, + query={"query": "SELECT * FROM c"}, + options={"maxItemCount": 5}, + feed_range=_make_feed_range_dict(), + ) + self.assertEqual(docs, []) + self.assertNotIn("Documents", captured["result"]) + + + def test_truncation_keeps_count_field_consistent(self, _mock_get_headers): + """After truncation, ``results['_count']`` (set by _merge_query_results) + must be updated to match the truncated Documents length so any + downstream introspection sees a coherent shape.""" + client = _build_client_connection() + docs, _post_mock, captured = self._query( + client, + options={"maxItemCount": 5}, + post_side_effect=lambda *a, **kw: (_docs(5), {}), + ) + self.assertEqual(len(docs), 5) + self.assertEqual(captured["result"].get("_count"), 5, + "_count must be updated alongside Documents") + + +if __name__ == "__main__": + unittest.main() + diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py new file mode 100644 index 000000000000..a42d62de31dc --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py @@ -0,0 +1,188 @@ +# The MIT License (MIT) +# Copyright (c) Microsoft Corporation. All rights reserved. + +"""Async unit test for the ``feed_range`` query page-size honoring fix. + +Mirror of ``test_query_feed_range_max_item_count.py`` for the async +``CosmosClientConnection`` in ``azure.cosmos.aio``. + +Note: these tests reach into the name-mangled +``_CosmosClientConnection__QueryFeed`` / ``_CosmosClientConnection__Post`` +members. If ``__QueryFeed`` is renamed or moved off the async +``CosmosClientConnection``, move these tests with it. +""" + +import unittest +from unittest.mock import MagicMock, patch, AsyncMock + +import pytest + +from azure.cosmos.aio._cosmos_client_connection_async import CosmosClientConnection +from azure.cosmos._change_feed.feed_range_internal import FeedRangeInternalEpk +from azure.cosmos._routing.routing_range import Range + + +def _build_async_client_connection(overlapping_ranges=None): + """Build a bare async ``CosmosClientConnection`` instance with only the + attributes referenced by ``__QueryFeed``'s feed_range branch.""" + client = object.__new__(CosmosClientConnection) + client.default_headers = {} + client._query_compatibility_mode = CosmosClientConnection._QueryCompatibilityMode.Default + client.availability_strategy = None + client.availability_strategy_executor = None + client.availability_strategy_max_concurrency = None + client.last_response_headers = {} + if overlapping_ranges is None: + overlapping_ranges = [ + {"id": "0", "minInclusive": "", "maxExclusive": "55"}, + {"id": "1", "minInclusive": "55", "maxExclusive": "AA"}, + {"id": "2", "minInclusive": "AA", "maxExclusive": "FF"}, + ] + client._routing_map_provider = MagicMock() + client._routing_map_provider.get_overlapping_ranges = AsyncMock(return_value=overlapping_ranges) + client._UpdateSessionIfRequired = MagicMock() + return client + + +def _make_feed_range_dict(): + return FeedRangeInternalEpk( + Range(range_min="", range_max="FF", isMinInclusive=True, isMaxInclusive=False) + ).to_dict() + + +def _docs(n, prefix="d"): + return {"Documents": [{"id": f"{prefix}-{i}"} for i in range(n)]} + + +def _capture_result_fn(): + captured = {} + + def fn(result): + captured["result"] = result + return result["Documents"] + return captured, fn + + +@pytest.mark.asyncio +@patch("azure.cosmos.aio._cosmos_client_connection_async.base.set_session_token_header_async", + new=AsyncMock(return_value=None)) +@patch("azure.cosmos.aio._cosmos_client_connection_async.base.GetHeaders", + side_effect=lambda *args, **kwargs: {}) +class TestQueryFeedRangeMaxItemCountAsync: + + async def _query(self, client, options, post_side_effect): + post_mock = AsyncMock(side_effect=post_side_effect) + client._CosmosClientConnection__Post = post_mock + captured, result_fn = _capture_result_fn() + docs = await client._CosmosClientConnection__QueryFeed( + path="/dbs/db1/colls/coll1/docs", + resource_type="docs", + id_="coll1", + result_fn=result_fn, + create_fn=None, + query={"query": "SELECT * FROM c"}, + options=options, + feed_range=_make_feed_range_dict(), + ) + return docs, post_mock, captured + + async def test_first_page_truncated_to_max_item_count(self, _mock_get_headers): + client = _build_async_client_connection() + page_size = 5 + docs, post_mock, captured = await self._query( + client, + options={"maxItemCount": page_size}, + post_side_effect=lambda *a, **kw: (_docs(page_size), {}), + ) + assert post_mock.call_count == 3 + assert len(docs) == page_size + assert len(captured["result"]["Documents"]) == page_size + + async def test_truncation_to_one_across_three_ranges(self, _mock_get_headers): + client = _build_async_client_connection() + docs, _post_mock, captured = await self._query( + client, + options={"maxItemCount": 1}, + post_side_effect=lambda *a, **kw: (_docs(5), {}), + ) + assert len(docs) == 1 + assert len(captured["result"]["Documents"]) == 1 + + async def test_no_truncation_when_under_cap(self, _mock_get_headers): + client = _build_async_client_connection() + docs, _post_mock, _captured = await self._query( + client, + options={"maxItemCount": 10}, + post_side_effect=lambda *a, **kw: (_docs(1), {}), + ) + assert len(docs) == 3 + + async def test_boundary_exact_cap_no_slice(self, _mock_get_headers): + client = _build_async_client_connection() + docs, _post_mock, _captured = await self._query( + client, + options={"maxItemCount": 3}, + post_side_effect=lambda *a, **kw: (_docs(1), {}), + ) + assert len(docs) == 3 + + async def test_no_max_item_count_no_truncation(self, _mock_get_headers): + client = _build_async_client_connection() + docs, _post_mock, _captured = await self._query( + client, + options={}, + post_side_effect=lambda *a, **kw: (_docs(4), {}), + ) + assert len(docs) == 12 + + async def test_max_item_count_zero_means_server_default_no_truncation(self, _mock_get_headers): + """maxItemCount=0 means "use the server default page size", not + "return zero items". See the corresponding sync test for the + full rationale.""" + client = _build_async_client_connection() + docs, _post_mock, _captured = await self._query( + client, + options={"maxItemCount": 0}, + post_side_effect=lambda *a, **kw: (_docs(7), {}), + ) + assert len(docs) == 21 + + async def test_single_overlapping_range_unchanged(self, _mock_get_headers): + client = _build_async_client_connection(overlapping_ranges=[ + {"id": "0", "minInclusive": "", "maxExclusive": "FF"}, + ]) + docs, post_mock, _captured = await self._query( + client, + options={"maxItemCount": 5}, + post_side_effect=lambda *a, **kw: (_docs(5), {}), + ) + assert post_mock.call_count == 1 + assert len(docs) == 5 + + async def test_missing_documents_key_does_not_crash(self, _mock_get_headers): + """A partial result missing the Documents key entirely must not raise + from the truncation block.""" + client = _build_async_client_connection(overlapping_ranges=[ + {"id": "0", "minInclusive": "", "maxExclusive": "FF"}, + ]) + post_mock = AsyncMock(side_effect=lambda *a, **kw: ({"some_other_field": 42}, {})) + client._CosmosClientConnection__Post = post_mock + captured = {} + + def lenient_result_fn(result): + captured["result"] = result + return result.get("Documents") or [] + + docs = await client._CosmosClientConnection__QueryFeed( + path="/dbs/db1/colls/coll1/docs", + resource_type="docs", + id_="coll1", + result_fn=lenient_result_fn, + create_fn=None, + query={"query": "SELECT * FROM c"}, + options={"maxItemCount": 5}, + feed_range=_make_feed_range_dict(), + ) + assert docs == [] + assert "Documents" not in captured["result"] + From 2b36cc912924630671f37a3a08cee8ab6d6d856b Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 22 Apr 2026 13:55:27 -0700 Subject: [PATCH 2/2] Suppress continuation token when feed_range page is truncated Address Copilot review on PR #46469: truncating the merged page while surfacing the last inner range's x-ms-continuation can cause silent data loss on resume (the token has advanced past truncated documents from earlier ranges). Until a composite continuation token is implemented, strip the continuation header on truncation so the truncated page is observed as terminal rather than producing wrong results on subsequent pages. - _cosmos_client_connection.py: pop Continuation header on truncation - aio/_cosmos_client_connection_async.py: mirror on self.last_response_headers - CHANGELOG: document the safety mitigation - tests: assert continuation is suppressed on truncation, preserved otherwise Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure-cosmos/CHANGELOG.md | 2 +- .../azure/cosmos/_cosmos_client_connection.py | 16 +++++++++ .../aio/_cosmos_client_connection_async.py | 16 +++++++++ .../test_query_feed_range_max_item_count.py | 35 +++++++++++++++++++ ...t_query_feed_range_max_item_count_async.py | 31 ++++++++++++++++ 5 files changed, 99 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index 30711bdb2276..70896754862f 100644 --- a/sdk/cosmos/azure-cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure-cosmos/CHANGELOG.md @@ -8,7 +8,7 @@ #### Bugs Fixed * Fixed bug where `query_items(feed_range=..., max_item_count=N)` could return up to `K * N` documents per page when the supplied feed range overlapped `K` physical partition key ranges (for example, after a server-side split). The page returned to the caller is now truncated to the requested `max_item_count`. - * Known limitation (deferred): when a `feed_range` overlaps multiple PK ranges, only the last inner range's `x-ms-continuation` is surfaced as the page continuation token. Round-tripping that token sends it to every inner range on the next page, which is undefined server-side and can produce duplicates, missing documents, or non-terminating iteration on subsequent pages. A composite continuation token across overlapping inner ranges is tracked separately as a follow-up; until then, only the *first* page of a multi-range `feed_range` query is reliable. + * Known limitation (deferred): when a `feed_range` overlaps multiple PK ranges, only the last inner range's `x-ms-continuation` is surfaced as the page continuation token. Round-tripping that token sends it to every inner range on the next page, which is undefined server-side and can produce duplicates, missing documents, or non-terminating iteration on subsequent pages. As a safety mitigation, when the merged page is actually truncated the continuation header is suppressed so the truncated page is observed as terminal rather than producing wrong results on resume. A composite continuation token across overlapping inner ranges is tracked separately as a follow-up. * Fixed bug where `CosmosClient` construction with AAD credentials would crash at startup if the semantic reranking inference endpoint environment variable was not set, even when semantic reranking was not being used. The inference service is now lazily initialized on first use. See [PR 46243](https://github.com/Azure/azure-sdk-for-python/pull/46243) #### 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 5a901d277e48..25ec836b6454 100644 --- a/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py +++ b/sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py @@ -3434,6 +3434,14 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: # fix is a composite continuation token spanning all K inner PK ranges. # Until that lands, this branch only delivers a correct *first* page for # multi-range feed_range queries. + # + # Mitigation for the deferred limitation: when truncation actually + # discards documents, the surfaced continuation token would describe + # the wrong cursor (it is the last inner range's token, but documents + # from earlier ranges and from the truncated tail were dropped). To + # avoid silent data loss, we strip the continuation header below when + # truncation occurs, so the truncated page is observed as terminal + # rather than producing wrong results on resume. max_item_count = options.get("maxItemCount") docs = results.get("Documents") if max_item_count and isinstance(docs, list): @@ -3447,6 +3455,14 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: # Documents list so any downstream consumer that introspects # the merged dict observes a coherent shape. results["_count"] = cap + # The merged page was assembled from multiple inner PK ranges and + # then truncated. The continuation token from the last inner range + # only describes progress for that range, not for the truncated + # tail or for ranges past the cursor, so resuming with it would + # silently skip documents. Until a composite continuation token is + # implemented, suppress the misleading token so callers see the + # truncated page as terminal rather than experiencing data loss. + last_response_headers.pop(http_constants.HttpHeaders.Continuation, None) if last_response_headers.get(http_constants.HttpHeaders.IndexUtilization) is not None: index_metrics_raw = last_response_headers[http_constants.HttpHeaders.IndexUtilization] last_response_headers[http_constants.HttpHeaders.IndexUtilization] = ( 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 53aaa0dd78f5..ce5169aafc2e 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 @@ -3227,6 +3227,14 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: # correct fix is a composite continuation token spanning all K inner PK # ranges. Until that lands, this branch only delivers a correct *first* page # for multi-range feed_range queries. + # + # Mitigation for the deferred limitation: when truncation actually + # discards documents, the surfaced continuation token would describe + # the wrong cursor (it is the last inner range's token, but documents + # from earlier ranges and from the truncated tail were dropped). To + # avoid silent data loss, we strip the continuation header below when + # truncation occurs, so the truncated page is observed as terminal + # rather than producing wrong results on resume. max_item_count = options.get("maxItemCount") docs = results.get("Documents") if max_item_count and isinstance(docs, list): @@ -3240,6 +3248,14 @@ def __GetBodiesFromQueryResult(result: dict[str, Any]) -> list[dict[str, Any]]: # Documents list so any downstream consumer that introspects # the merged dict observes a coherent shape. results["_count"] = cap + # The merged page was assembled from multiple inner PK ranges and + # then truncated. The continuation token from the last inner range + # only describes progress for that range, not for the truncated + # tail or for ranges past the cursor, so resuming with it would + # silently skip documents. Until a composite continuation token is + # implemented, suppress the misleading token so callers see the + # truncated page as terminal rather than experiencing data loss. + self.last_response_headers.pop(http_constants.HttpHeaders.Continuation, None) if self.last_response_headers.get(http_constants.HttpHeaders.IndexUtilization) is not None: index_metrics_raw = self.last_response_headers[http_constants.HttpHeaders.IndexUtilization] self.last_response_headers[http_constants.HttpHeaders.IndexUtilization] = ( diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py index 2480b41f9e8f..a4e11664129f 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py @@ -232,6 +232,41 @@ def test_truncation_keeps_count_field_consistent(self, _mock_get_headers): self.assertEqual(captured["result"].get("_count"), 5, "_count must be updated alongside Documents") + def test_truncation_suppresses_continuation_header(self, _mock_get_headers): + """When the merged page is truncated, the surfaced continuation token + only describes the last inner PK range and would silently skip documents + on resume. It must be suppressed on the page that returns to the caller.""" + client = _build_client_connection() + # Each inner POST returns 5 docs and a continuation token. + # K=3 -> merged 15 docs, capped at 5 -> truncation occurs. + docs, _post_mock, _captured = self._query( + client, + options={"maxItemCount": 5}, + post_side_effect=lambda *a, **kw: ( + _docs(5), {"x-ms-continuation": "inner-token"} + ), + ) + self.assertEqual(len(docs), 5) + self.assertNotIn("x-ms-continuation", client.last_response_headers, + "continuation header must be stripped on truncation") + + def test_no_truncation_preserves_continuation_header(self, _mock_get_headers): + """When the merged page fits within the cap, no truncation happens, so + the inner continuation must be left intact (today's pre-fix behavior).""" + client = _build_client_connection() + docs, _post_mock, _captured = self._query( + client, + # 3 ranges * 2 docs = 6, cap=10 -> no truncation + options={"maxItemCount": 10}, + post_side_effect=lambda *a, **kw: ( + _docs(2), {"x-ms-continuation": "inner-token"} + ), + ) + self.assertEqual(len(docs), 6) + self.assertEqual(client.last_response_headers.get("x-ms-continuation"), + "inner-token", + "continuation must be preserved when no truncation") + if __name__ == "__main__": unittest.main() diff --git a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py index a42d62de31dc..e8802a5dec78 100644 --- a/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py @@ -186,3 +186,34 @@ def lenient_result_fn(result): assert docs == [] assert "Documents" not in captured["result"] + async def test_truncation_suppresses_continuation_header(self, _mock_get_headers): + """When the merged page is truncated, the surfaced continuation token + only describes the last inner PK range and would silently skip + documents on resume. It must be stripped from last_response_headers.""" + client = _build_async_client_connection() + docs, _post_mock, _captured = await self._query( + client, + options={"maxItemCount": 5}, + post_side_effect=lambda *a, **kw: ( + _docs(5), {"x-ms-continuation": "inner-token"} + ), + ) + assert len(docs) == 5 + assert "x-ms-continuation" not in client.last_response_headers, \ + "continuation header must be stripped on truncation" + + async def test_no_truncation_preserves_continuation_header(self, _mock_get_headers): + """When the merged page fits within the cap, no truncation happens, + so the inner continuation must be left intact.""" + client = _build_async_client_connection() + docs, _post_mock, _captured = await self._query( + client, + # 3 ranges * 2 docs = 6, cap=10 -> no truncation + options={"maxItemCount": 10}, + post_side_effect=lambda *a, **kw: ( + _docs(2), {"x-ms-continuation": "inner-token"} + ), + ) + assert len(docs) == 6 + assert client.last_response_headers.get("x-ms-continuation") == "inner-token" +