Commit aff5f7a
[Cosmos] Fix /pkranges change-feed drain loop (#47245)
* [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>1 parent c8abdc1 commit aff5f7a
20 files changed
Lines changed: 2310 additions & 160 deletions
File tree
- sdk/cosmos/azure-cosmos
- azure/cosmos
- _routing
- aio
- aio
- tests
- routing
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1 | 1 | | |
2 | 2 | | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
| 7 | + | |
3 | 8 | | |
4 | 9 | | |
5 | 10 | | |
| |||
8 | 13 | | |
9 | 14 | | |
10 | 15 | | |
| 16 | + | |
11 | 17 | | |
12 | 18 | | |
13 | 19 | | |
14 | 20 | | |
15 | 21 | | |
16 | 22 | | |
17 | 23 | | |
18 | | - | |
19 | 24 | | |
20 | 25 | | |
21 | 26 | | |
| |||
Lines changed: 81 additions & 0 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
171 | 171 | | |
172 | 172 | | |
173 | 173 | | |
| 174 | + | |
174 | 175 | | |
175 | 176 | | |
176 | 177 | | |
| |||
295 | 296 | | |
296 | 297 | | |
297 | 298 | | |
| 299 | + | |
| 300 | + | |
| 301 | + | |
| 302 | + | |
| 303 | + | |
| 304 | + | |
| 305 | + | |
| 306 | + | |
| 307 | + | |
| 308 | + | |
| 309 | + | |
| 310 | + | |
| 311 | + | |
| 312 | + | |
| 313 | + | |
| 314 | + | |
| 315 | + | |
| 316 | + | |
| 317 | + | |
| 318 | + | |
| 319 | + | |
| 320 | + | |
| 321 | + | |
| 322 | + | |
| 323 | + | |
| 324 | + | |
| 325 | + | |
| 326 | + | |
| 327 | + | |
| 328 | + | |
| 329 | + | |
| 330 | + | |
| 331 | + | |
| 332 | + | |
| 333 | + | |
| 334 | + | |
| 335 | + | |
| 336 | + | |
| 337 | + | |
| 338 | + | |
| 339 | + | |
| 340 | + | |
| 341 | + | |
| 342 | + | |
| 343 | + | |
| 344 | + | |
| 345 | + | |
| 346 | + | |
| 347 | + | |
| 348 | + | |
| 349 | + | |
| 350 | + | |
| 351 | + | |
| 352 | + | |
| 353 | + | |
| 354 | + | |
| 355 | + | |
| 356 | + | |
| 357 | + | |
| 358 | + | |
| 359 | + | |
| 360 | + | |
| 361 | + | |
| 362 | + | |
| 363 | + | |
| 364 | + | |
| 365 | + | |
| 366 | + | |
| 367 | + | |
| 368 | + | |
| 369 | + | |
| 370 | + | |
| 371 | + | |
| 372 | + | |
| 373 | + | |
| 374 | + | |
| 375 | + | |
| 376 | + | |
| 377 | + | |
| 378 | + | |
298 | 379 | | |
299 | 380 | | |
300 | 381 | | |
| |||
Lines changed: 74 additions & 22 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
41 | 41 | | |
42 | 42 | | |
43 | 43 | | |
| 44 | + | |
| 45 | + | |
44 | 46 | | |
45 | 47 | | |
46 | 48 | | |
| |||
334 | 336 | | |
335 | 337 | | |
336 | 338 | | |
| 339 | + | |
337 | 340 | | |
338 | 341 | | |
339 | 342 | | |
| |||
377 | 380 | | |
378 | 381 | | |
379 | 382 | | |
380 | | - | |
381 | | - | |
382 | | - | |
383 | | - | |
384 | | - | |
| 383 | + | |
| 384 | + | |
| 385 | + | |
| 386 | + | |
| 387 | + | |
| 388 | + | |
| 389 | + | |
| 390 | + | |
| 391 | + | |
| 392 | + | |
| 393 | + | |
| 394 | + | |
| 395 | + | |
| 396 | + | |
| 397 | + | |
| 398 | + | |
| 399 | + | |
| 400 | + | |
| 401 | + | |
| 402 | + | |
| 403 | + | |
| 404 | + | |
| 405 | + | |
385 | 406 | | |
386 | | - | |
| 407 | + | |
387 | 408 | | |
| 409 | + | |
388 | 410 | | |
389 | | - | |
390 | | - | |
391 | | - | |
392 | | - | |
393 | | - | |
394 | | - | |
395 | | - | |
396 | | - | |
397 | | - | |
398 | | - | |
399 | | - | |
400 | | - | |
401 | | - | |
402 | | - | |
| 411 | + | |
| 412 | + | |
| 413 | + | |
| 414 | + | |
| 415 | + | |
| 416 | + | |
| 417 | + | |
| 418 | + | |
| 419 | + | |
| 420 | + | |
| 421 | + | |
| 422 | + | |
| 423 | + | |
| 424 | + | |
| 425 | + | |
| 426 | + | |
| 427 | + | |
| 428 | + | |
| 429 | + | |
| 430 | + | |
| 431 | + | |
| 432 | + | |
403 | 433 | | |
404 | | - | |
| 434 | + | |
| 435 | + | |
| 436 | + | |
| 437 | + | |
| 438 | + | |
| 439 | + | |
| 440 | + | |
| 441 | + | |
| 442 | + | |
| 443 | + | |
| 444 | + | |
| 445 | + | |
| 446 | + | |
| 447 | + | |
| 448 | + | |
| 449 | + | |
| 450 | + | |
| 451 | + | |
| 452 | + | |
| 453 | + | |
| 454 | + | |
| 455 | + | |
405 | 456 | | |
406 | 457 | | |
| 458 | + | |
407 | 459 | | |
408 | | - | |
| 460 | + | |
409 | 461 | | |
410 | 462 | | |
411 | 463 | | |
| |||
0 commit comments