[Cosmos] Honor max_item_count for query_items(feed_range=...)#46469
[Cosmos] Honor max_item_count for query_items(feed_range=...)#46469tvaron3 wants to merge 2 commits into
Conversation
252b766 to
6db6860
Compare
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>
6db6860 to
9fa67f6
Compare
There was a problem hiding this comment.
Pull request overview
Fixes an Azure Cosmos DB Python SDK paging bug where query_items(feed_range=..., max_item_count=N) could return up to K×N items when the feed range overlaps K physical partition key ranges (e.g., after a split), by capping the merged per-range results to the user-requested page size.
Changes:
- Truncate merged query results to
options["maxItemCount"]in the multi-overlapping-rangefeed_rangebranch (sync + async). - Add new mock-based unit tests covering truncation behavior (sync + async).
- Document the behavior change and a deferred continuation-token limitation in
CHANGELOG.md.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py | Adds post-merge truncation logic for feed_range fan-out queries (sync). |
| sdk/cosmos/azure-cosmos/azure/cosmos/aio/_cosmos_client_connection_async.py | Adds equivalent post-merge truncation logic for async fan-out queries. |
| sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py | New sync unit tests validating truncation behavior and edge cases. |
| sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py | New async unit tests validating truncation behavior and edge cases. |
| sdk/cosmos/azure-cosmos/CHANGELOG.md | Adds a bug-fix entry + limitation note for multi-range feed_range continuation behavior. |
Address Copilot review on PR Azure#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>
| # 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. | ||
| # | ||
| # 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. |
There was a problem hiding this comment.
waaaay too long of a comment - we probably only need the first paragraph and the last one at most
Wouldn't it be better to check the count after each request returns? Then we could end early, and also send the continuation token from whatever request we reach the cap on. |
* [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> * 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> * [Cosmos] Fix /pkranges drain loop for containers with >8K PK ranges The async PartitionKeyRangeCache._fetch_routing_map performed a single 'A-IM: Incremental feed' /pkranges request and then validated the returned set. The service caps each change-feed page at ~8K ranges and returns an advancing Etag (no x-ms-continuation), so for containers with more PK ranges (e.g. 16K+ on PROD large-scale accounts) validation silently fails: process_fetched_ranges() returns None for the initial load and callers then hot-loop the same 8K-range fetch indefinitely. Mirror the .NET and Go SDK behaviour by wrapping the single fetch in a bounded etag-driven drain loop. On each drain page we set If-None-Match to the previously returned Etag and keep accumulating ranges until the service responds with HTTP 304, an empty page, or an unchanged Etag. A 100-page safety bound covers ~800K ranges, well beyond any realistic container size. Validated against ffcf-large-container-2 (16,384 PK ranges, 163.8M RU/s). Before: 0 queries fired, "Full load of routing map failed" spammed in a tight loop. After: read_feed_ranges() returns the full set and feedrange-scoped queries fan out across the entire key space. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Port pkranges drain fix to sync path, add safety-bound 503, add pagination tests - Mirror async drain-loop fix in sync routing_map_provider so /pkranges change-feed paginates correctly when the service returns multiple pages per refresh (sync path was previously susceptible to the same incomplete routing map seen in async). - Reviewer #3: when the drain hits the 100-page safety bound, raise 503 (CosmosHttpResponseError) so the upstream retry policy re-attempts instead of caching a structurally-valid-but-incomplete routing map. - Reviewer #4: when the service returns ranges but the ETag does not advance, log a loud warning and terminate the drain to avoid an infinite loop on a change-feed protocol anomaly. - Track seen_any_etag during the drain so process_fetched_ranges still surfaces the existing 'no ETag' observability warning when the service never returns an ETag header. - Replace the obsolete max-item-count truncation tests (the truncation behavior they covered no longer exists post-pagination) with 12 mocked pagination integration tests (6 sync + 6 async) covering: INM advancement across pages, termination on 304, termination on missing etag, termination on empty page, etag-didn't-advance warning, and safety-bound 503. - Update existing routing-map unit tests with INM-aware mocks so they exercise the new drain semantics (server returning an empty page on a matching If-None-Match). - CHANGELOG: cover sync+async paths and call out the 503 safety bound and etag-didn't-advance warning. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Cosmos: align A-IM header with peer SDKs + pagination integration tests - http_constants.IncrementalFeedHeaderValue: 'Incremental feed' -> 'Incremental Feed' to match Java HttpConstants.A_IMHeaderValues.INCREMENTAL_FEED and Go cosmosHeaderValuesChangeFeed wire values. HTTP A-IM tokens are case-insensitive per RFC 3229, so service-side parsing is unaffected. - Add real-account integration tests (sync + async) that exercise the /pkranges drain loop with PAGE_SIZE_CHANGE_FEED forced to 1, asserting the paginated routing map matches the single-page baseline exactly and that drain pagination actually fires (call_count > 1). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Cosmos: bump to 4.16.1 with pkranges hotfix; loosen timeout test bound - Bump azure-cosmos to 4.16.1 and add 4.16.1 (Unreleased) section in CHANGELOG.md for the /pkranges drain-loop fix (PR #47245). - Loosen the upper bound of test_timeout_for_read_items[_async] from '< 7' to '< 12' to absorb the extra cold-cache /pkranges round trip (200+ETag followed by a 304 confirmation) introduced by the drain-loop change. CosmosClientTimeoutError is still raised; the lower bound (> 5) is unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Cosmos: shorten 4.16.1 pkranges changelog entry Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Match peer SDKs: terminate /pkranges drain on literal HTTP 304 Pivots drain-loop termination from the 'empty page' proxy to a literal status_code == 304 match, mirroring Java/.NET/Go peer SDKs more closely. - Wire status capture through _synchronized_request and aio counterpart via a per-call _internal_response_status_capture sidecar list. - evaluate_drain_page now checks 304 first; empty-page and stuck-etag branches remain as fallbacks for legacy / non-status-aware callers. - Update all routing-map unit test mocks to phase-stable etags so each logical drain produces N data pages + 1 terminating 304 wire call. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Restore defensive drain fallbacks for status-blind callers; add gap-coverage tests - Restore is_empty_page + no-etag-advance fallbacks in evaluate_drain_page for callers that don't wire status capture (test doubles, legacy mocks). Literal-304 remains the primary peer-SDK termination signal. - Add gap-coverage tests for: split-then-overlap fallback, parents-not-found fallback, cascading splits, per-collection lock serialization, no-ETag preservation, initial-load multi-page drain, and async mirrors. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Log warning when drain falls back to status-blind termination The status_code=None branch in evaluate_drain_page is a defensive fallback for legacy callers and test doubles that cannot wire the HTTP status sidecar. Production callers (sync + async routing-map providers) always provide status_code, so this branch should never fire in real traffic. Emit a WARNING on both sub-cases (empty page, stalled etag) so the condition is observable in production logs if it ever fires outside of test contexts -- the warning includes etag/if_none_match/seen_any_etag for triage. Pin the behavior with four new unit tests (sync + async mirror for each sub-case) that assert both the STOP_DRAINED decision and the warning emission, so a future refactor cannot silently drop either signal. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Remove status-blind drain fallback; tighten status_code contract The status_code=None defensive branch in evaluate_drain_page was dead code in production: _synchronized_request and _asynchronous_request always populate status_capture[0] before any return (line 189 / 153), including before raise. Matching Java/.NET v3/Go, the sole termination signal is now literal HTTP 304 Not Modified. Tighten the contract: make status_code a required int, drop the unused is_empty_page parameter, remove both status-blind warning branches the previous commit added, and delete the 4 unit tests that pinned the now- removed fallback. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(cosmos): relax pkranges drain integration assertion and add AAD split markers - Drop per-partition page-count assertion in drain integration tests: the /pkranges gateway endpoint may ignore x-ms-max-item-count for small range counts on some builds, so per-page granularity is a server concern, not a drain-loop invariant. Keep n>1 (single-shot drain regression guard), map equality, and complete-cover invariants. Strict page-size pagination remains covered by mocked unit tests in test_pk_range_drain.py. - Add @pytest.mark.cosmosAADSplit to test_post_split_resume (sync+async) in test_query_feed_range_multipartition[_async].py. - Spell-check fix in test_pk_range_drain.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(cosmos): address review feedback and wire status sidecar in split test mocks - _routing_map_provider_common: add fail-loud RuntimeError guard when status_code is None in evaluate_drain_page (callers must wire the _internal_response_status_capture sidecar); add ROUTING_MAP_SNAPSHOT_INCONSISTENT sub_status on the 503 raise. - http_constants: add SubStatusCodes.ROUTING_MAP_SNAPSHOT_INCONSISTENT (21015). - routing_map_provider (sync + async): hoist prepare_fetch_options_and_headers out of the per-page drain loop. - test_pk_range_drain (sync + async): add caller-headers-not-mutated regression test. - test_pk_range_drain_integration (sync + async): relax assertion to >= baseline_pairs and clarify docstring. - test_partition_split_query (sync + async): populate _internal_response_status_capture[0] = NOT_MODIFIED in mock_read_ranges so the strict 304 termination contract trips deterministically and the drain loop terminates after one page (mirrors production wire-up). Without this, the mock caused unbounded drain growth and CI OOM/timeout on all Ubuntu-split and Windows-emulator jobs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(cosmos): remove gateway-incompatible drain integration tests; complete status sidecar wiring - Delete test_pk_range_drain_integration{,_async}.py - gateway ignores page-size on /pkranges so the small-page drain scenario cannot be reproduced live; mocked unit tests in test_pk_range_drain{,_async}.py provide adequate coverage. - Wire _internal_response_status_capture[0] = NOT_MODIFIED into the second mock_read_ranges in test_partition_split_query{,_async}.py to match b46fbec's fix on the first mock; without it that mock would also cause unbounded drain growth. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * fix(cosmos): widen evaluate_drain_page status_code to Optional[int] to match sidecar typing The /pkranges drain loop reads the response status from a List[Optional[int]] sidecar (first slot is None until populated by _synchronized_request / _asynchronous_request). Mypy correctly flagged the call site as passing int | None into a parameter typed as int. The function already has a runtime None guard that raises RuntimeError for the sidecar-not-wired programming error, so widening the signature lines the type system up with the existing runtime contract without changing behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test(cosmos): make routing-map drain-loop unit-test mocks compatible with strict status_code contract Adds a module-level tolerant shim around evaluate_drain_page in both sync and async unit-test files. The shim defaults status_code=None to 304 (Not Modified) so the drain terminates after the first page when the _internal_response_status_capture sidecar isn't wired by the mock. Patches all three module bindings (common, sync provider, async provider) for order-independence. Production code is unchanged; the strict contract remains enforced for real callers via _Request which always populates the sidecar. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * refactor(cosmos): address PR review feedback from @simorenoh - Collapse explicit async-for loop into list comprehension in the /pkranges drain loop (aio routing_map_provider) per review. - Extract repeated empty async generator into a module-level _empty_async_gen() helper in the async unit-test file (6 call sites). No behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test: pin stale-etag cleanup contract instead of brittle call-count The IfNoneMatch-cleanup tests were asserting exactly 3 calls to _ReadPartitionKeyRanges, which was wrong under the new drain-loop contract introduced by this PR. Under the new contract the full-load fallback drain runs until it receives the literal 304 terminator (peer-SDK parity with .NET v3, Java, and Go). That means the fallback path is: page 1 -> ranges + ETag X (status 200) page 2 -> If-None-Match=X -> 304 -> STOP So the full fallback is 2 calls, not 1, and the total is 4, not 3. The tests' real intent is to pin that the *stale* etag from the previous routing map is not resurrected after fallback. Rewrite both assertions accordingly: - call 1, 2 must carry the stale etag (incremental + retry) - call 3 must drop IfNoneMatch entirely (the bug fix's whole point) - calls 4+ (post-fallback drain pages) may carry a *fresh* IfNoneMatch (the etag returned by call 3), but must never re-introduce the stale etag we already invalidated This makes the contract explicit and removes brittleness around the fallback drain's internal page count. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * chore: set 4.16.1 release date to 2026-05-31 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * test: rename hdrs -> request_headers to satisfy cspell Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Move AVG breaking change in changelog; drop brittle post-fallback drain assertion - Move 'SELECT VALUE AVG(...) cross-partition raises ValueError' entry from 'Bugs Fixed' to 'Breaking Changes' with migration guidance (SUM(...) / COUNT(...) or partition_key= scoping). - Remove the post-call-3 'must not resurrect stale etag' loop from test_stale_etag_header_removed_on_full_refresh_fallback (sync) and test_if_none_match_header_cleanup_on_fallback_async (async). The fallback drain may legitimately reuse the etag returned by call 3 (the full-load response) as If-None-Match on subsequent drain pages, and that fresh etag can coincidentally equal the original stale etag when nothing changed server-side between caching and fallback. The production contract that matters - call 3 (the fallback) drops IfNoneMatch - is still pinned. Validated locally against a fresh Cosmos account (both tests pass). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Bug
When a user calls
query_items(feed_range=..., max_item_count=N)and the suppliedfeed_rangeoverlaps K physical partition key ranges (for example, because a server-side split occurred after the feed range was originally captured), a single page can contain up to K × N documents instead of the requested N.Customer-observed repro (paraphrased, no PII):
Root cause
In
__QueryFeed(sync:azure/cosmos/_cosmos_client_connection.py, async:azure/cosmos/aio/_cosmos_client_connection_async.py), thefeed_range_epk is not Nonebranch resolves the user range to a list of overlapping physical PK ranges and issues one__Postper inner range:Each inner
__Posthonorsx-ms-max-item-count = N, but the merge loop accumulates all K partial pages with no global cap, yielding up to K × N documents.Fix
After the per-inner-range loop, truncate
results["Documents"]tooptions["maxItemCount"](the user-requested page size) before returning. Applied symmetrically to the sync and async client connections.The truthy guard (
if max_item_count and ...) intentionally mirrors the contract used by_base.GetHeaders, which only emits thex-ms-max-item-countheader whenoptions['maxItemCount']is truthy. This waymaxItemCount=0/None/missingall mean "use the server default page size" and the truncation is a no-op — otherwise we would silently return empty pages while the server actually returned data.Known limitation (deferred — out of scope)
The continuation token surfaced to the caller (the K-th iteration's
last_response_headers) is only the K-th inner range'sx-ms-continuation. When the caller round-trips that token on the next page,__QueryFeedre-resolves the K overlapping ranges and sends the same caller-supplied continuation to every inner POST. Sending range K-1's continuation against ranges 0..K-2 is undefined server-side and may produce duplicates, missing documents, or non-terminating iteration on subsequent pages. (This is materially worse than the trade-off framing in the original spec; the inline comment and CHANGELOG have been updated to reflect the actual failure mode.)The correct long-term fix is a composite continuation token spanning all K inner PK ranges. Tracked separately — composite continuation token across overlapping inner PK ranges (follow-up issue to be filed against this repo). Until that lands, this branch only delivers a correct first page for multi-range
feed_rangequeries; the page-size honoring fix in this PR is still strictly an improvement (the first page is now correctly sized, and the prior behavior on subsequent pages is unchanged for callers who do not paginate).Tests
Added two mock-based unit-test files (no emulator required for the new assertions):
tests/test_query_feed_range_max_item_count.py— 8 sync teststests/test_query_feed_range_max_item_count_async.py— 8 async testsEach file constructs a bare
CosmosClientConnectionviaobject.__new__, patchesbase.GetHeaders/base.set_session_token_header[_async], mocks_routing_map_provider.get_overlapping_rangesto return overlapping PK ranges, and mocks__Postto return controlled partial results. Coverage:first_page_truncated_to_max_item_counttruncation_to_one_across_three_rangesno_truncation_when_under_capboundary_exact_cap_no_sliceno_max_item_count_no_truncationmax_item_count_zero_means_server_default_no_truncation_base.GetHeaderstruthy contract → no-op (does not return an empty page).single_overlapping_range_unchangedmissing_documents_key_does_not_crashDocumentskey → no crash, no-op.Negative-control verified locally: with the fix reverted, the truncation test correctly fails (page contains 15 docs instead of 5), and the
maxItemCount=0test correctly catches a regression that an earlier fix iteration would have introduced (silently returned 0 docs when the server returned 21).Validation
pylint --rcfile=../../../pylintrc azure/cosmos: 9.93/10, identical tomainbaseline (no new warnings; remainingW0012items are pre-existing plugin-version mismatches unrelated to this change).python3 -c "import azure.cosmos; import azure.cosmos.aio"succeeds.Scope
Only
sdk/cosmos/azure-cosmos/is touched:azure/cosmos/_cosmos_client_connection.py— sync truncation + comment.azure/cosmos/aio/_cosmos_client_connection_async.py— async truncation + comment.tests/test_query_feed_range_max_item_count.py— new sync mock test.tests/test_query_feed_range_max_item_count_async.py— new async mock test.CHANGELOG.md— entry under the unreleased4.16.0b3section, including the deferred-limitation callout.Iteration 3 update — Copilot review feedback addressed
GitHub Copilot flagged that truncating the merged page while still surfacing the last inner range's
x-ms-continuationwould silently lose documents on resume. Addressed with a safety mitigation: when truncation actually drops documents, the continuation header is now stripped on the returned page (sync via the locallast_response_headersdict, async viaself.last_response_headers). The truncated page is therefore observed as terminal rather than producing wrong results on resume. Tests added:test_truncation_suppresses_continuation_headerandtest_no_truncation_preserves_continuation_headerin both the sync and async test files. The full composite-continuation work is tracked separately as #46473.