diff --git a/sdk/cosmos/azure-cosmos/CHANGELOG.md b/sdk/cosmos/azure-cosmos/CHANGELOG.md index fd8e170900ed..70896754862f 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. 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 37c8bf219306..25ec836b6454 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,61 @@ 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. + # + # 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): + 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 + # 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 3be6fecdc0f9..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 @@ -3202,6 +3202,60 @@ 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. + # + # 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): + 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 + # 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 new file mode 100644 index 000000000000..a4e11664129f --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count.py @@ -0,0 +1,273 @@ +# 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") + + 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 new file mode 100644 index 000000000000..e8802a5dec78 --- /dev/null +++ b/sdk/cosmos/azure-cosmos/tests/test_query_feed_range_max_item_count_async.py @@ -0,0 +1,219 @@ +# 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"] + + 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" +