Skip to content

[Cosmos] Honor max_item_count for query_items(feed_range=...)#46469

Closed
tvaron3 wants to merge 2 commits into
Azure:mainfrom
tvaron3:fix/cosmos-feed-range-max-item-count
Closed

[Cosmos] Honor max_item_count for query_items(feed_range=...)#46469
tvaron3 wants to merge 2 commits into
Azure:mainfrom
tvaron3:fix/cosmos-feed-range-max-item-count

Conversation

@tvaron3

@tvaron3 tvaron3 commented Apr 22, 2026

Copy link
Copy Markdown
Member

Bug

When a user calls query_items(feed_range=..., max_item_count=N) and the supplied feed_range overlaps 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):

# feed_range was captured when the container had a single physical PK range;
# a split has since produced 2 physical PK ranges that overlap it.
for page in container.query_items(
    query="SELECT * FROM c",
    feed_range=feed_range,
    max_item_count=5,
).by_page():
    items = list(page)
    assert len(items) <= 5  # ❌ fails: 10 items returned

Root cause

In __QueryFeed (sync: azure/cosmos/_cosmos_client_connection.py, async: azure/cosmos/aio/_cosmos_client_connection_async.py), the feed_range_epk is not None branch resolves the user range to a list of overlapping physical PK ranges and issues one __Post per inner range:

for over_lapping_range in over_lapping_ranges:
    req_headers[PartitionKeyRangeID] = over_lapping_range["id"]
    partial_result, last_response_headers = self.__Post(...)   # honors max_item_count per inner range
    self.last_response_headers = last_response_headers          # overwritten every iteration
    results = base._merge_query_results(results, partial_result, query)  # accumulated without cap
return __GetBodiesFromQueryResult(results)

Each inner __Post honors x-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"] to options["maxItemCount"] (the user-requested page size) before returning. Applied symmetrically to the sync and async client connections.

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 cap > 0 and len(docs) > cap:
        results["Documents"] = docs[:cap]

The truthy guard (if max_item_count and ...) intentionally mirrors the contract used by _base.GetHeaders, which only emits the x-ms-max-item-count header when options['maxItemCount'] is truthy. This way maxItemCount=0/None/missing all 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'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. 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_range queries; 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 tests
  • tests/test_query_feed_range_max_item_count_async.py — 8 async tests

Each file constructs a bare CosmosClientConnection via object.__new__, patches base.GetHeaders / base.set_session_token_header[_async], mocks _routing_map_provider.get_overlapping_ranges to return overlapping PK ranges, and mocks __Post to return controlled partial results. Coverage:

Test What it pins
first_page_truncated_to_max_item_count K=3, N=5 → page is exactly 5, not 15 (the original bug). Also asserts the underlying merged dict is truncated, not just the projection.
truncation_to_one_across_three_ranges Tightest cap (K=3, N=1) — distinguishes truncation from "merge correctly".
no_truncation_when_under_cap merged ≤ cap → list returned unchanged.
boundary_exact_cap_no_slice merged == cap → list returned unchanged.
no_max_item_count_no_truncation maxItemCount unset → no-op.
max_item_count_zero_means_server_default_no_truncation maxItemCount=0 mirrors _base.GetHeaders truthy contract → no-op (does not return an empty page).
single_overlapping_range_unchanged K=1 → single POST, original behavior preserved.
missing_documents_key_does_not_crash partial result without Documents key → 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=0 test 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 to main baseline (no new warnings; remaining W0012 items are pre-existing plugin-version mismatches unrelated to this change).
  • python3 -c "import azure.cosmos; import azure.cosmos.aio" succeeds.
  • ✅ All 16 new tests pass (8 sync + 8 async).

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 unreleased 4.16.0b3 section, 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-continuation would 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 local last_response_headers dict, async via self.last_response_headers). The truncated page is therefore observed as terminal rather than producing wrong results on resume. Tests added: test_truncation_suppresses_continuation_header and test_no_truncation_preserves_continuation_header in both the sync and async test files. The full composite-continuation work is tracked separately as #46473.

@tvaron3 tvaron3 force-pushed the fix/cosmos-feed-range-max-item-count branch 2 times, most recently from 252b766 to 6db6860 Compare April 22, 2026 18:38
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>
@tvaron3 tvaron3 force-pushed the fix/cosmos-feed-range-max-item-count branch from 6db6860 to 9fa67f6 Compare April 22, 2026 19:23
@tvaron3 tvaron3 marked this pull request as ready for review April 22, 2026 20:41
@tvaron3 tvaron3 requested a review from a team as a code owner April 22, 2026 20:41
Copilot AI review requested due to automatic review settings April 22, 2026 20:41

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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-range feed_range branch (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.

Comment thread sdk/cosmos/azure-cosmos/azure/cosmos/_cosmos_client_connection.py
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>
Comment on lines +3205 to +3237
# 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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

waaaay too long of a comment - we probably only need the first paragraph and the last one at most

@Pilchie

Pilchie commented Apr 22, 2026

Copy link
Copy Markdown
Member

After the per-inner-range loop, truncate results["Documents"] to options["maxItemCount"] (the user-requested page size) before returning. Applied symmetrically to the sync and async client connections.

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.

@tvaron3 tvaron3 closed this Apr 22, 2026
tvaron3 added a commit that referenced this pull request Jun 1, 2026
* [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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

Status: Done

Development

Successfully merging this pull request may close these issues.

4 participants