From 955c77239ddbfd91d876045b73ed657c4511a9c4 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Thu, 10 Sep 2026 17:26:50 -0400 Subject: [PATCH 1/5] Gate unbounded buffered Cosmos queries Require a finite global TOP/LIMIT for unordered DISTINCT and non-streaming ORDER BY, with a layered allow_unbounded_queries opt-out and shared 400/20126 admission status. Preserve bounded execution and continuation restrictions. Cover admission, option precedence, routing, encoding, and buffered execution; update SDK and driver APIs, changelogs, and query specifications. Native pagination remains out of scope. Fixes https://github.com/Azure/azure-sdk-for-rust/issues/5122 Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 3 + sdk/cosmos/azure_data_cosmos/api/API.md | 7 +- .../azure_data_cosmos/api/API.metadata.yml | 2 +- .../azure_data_cosmos/src/options/feed.rs | 9 + .../tests/binary_roundtrip_fuzzer.rs | 12 +- .../tests/emulator_tests/cosmos_hpk.rs | 12 +- .../tests/emulator_tests/cosmos_query.rs | 149 +++++++- .../emulator_tests/cosmos_vector_query.rs | 60 +++ .../query_comparison.rs | 86 +++++ .../cosmos_query_distinct_split.rs | 5 +- .../azure_data_cosmos_driver/CHANGELOG.md | 5 +- .../azure_data_cosmos_driver/api/API.md | 15 +- .../api/API.metadata.yml | 2 +- .../src/driver/cosmos_driver.rs | 8 + .../dataflow/non_streaming_ordered_merge.rs | 181 +++++++-- .../src/driver/dataflow/planner.rs | 348 ++++++++++++++++-- .../src/error/cosmos_status.rs | 43 ++- .../src/options/operation_options.rs | 51 +++ .../in_memory_emulator_tests/distinct.rs | 260 +++++++++++-- .../docs/specs/0001-configuration-options.md | 34 +- .../specs/0006-error-codes-and-retries.md | 16 + sdk/cosmos/docs/specs/0013-query-engine.md | 35 ++ 22 files changed, 1215 insertions(+), 128 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 193c38c759..728581bcb8 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -4,12 +4,15 @@ ### Features Added +- Added layered `OperationOptions::allow_unbounded_queries` and `QueryOptions::with_allow_unbounded_queries` to explicitly admit unbounded client-buffered queries. - Extended Cosmos binary JSON encoding to cross-partition `DISTINCT` query pages. ([#5070](https://github.com/Azure/azure-sdk-for-rust/pull/5070)) - Added `QueryPlanMode::{LocalPreferred, GatewayOnly}` to `OperationOptions`, allowing applications to force Gateway query planning globally or for an individual query as a livesite mitigation. ([#5181](https://github.com/Azure/azure-sdk-for-rust/pull/5181)) - Added finite cross-partition `ORDER BY VectorDistance(...)` queries with `TOP` or `OFFSET`/`LIMIT`. Results are fully buffered before the first page and cannot be resumed from continuation tokens. Hybrid/full-text vector ranking remains unsupported. ([#5130](https://github.com/Azure/azure-sdk-for-rust/pull/5130)) ### Breaking Changes +- Unordered cross-partition DISTINCT now requires a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`; non-streaming ORDER BY shares this admission policy, without a fixed numeric ceiling. + ### Bugs Fixed - Name-based container clients now automatically recover when a container is deleted and recreated. ([#5219](https://github.com/Azure/azure-sdk-for-rust/pull/5219)) diff --git a/sdk/cosmos/azure_data_cosmos/api/API.md b/sdk/cosmos/azure_data_cosmos/api/API.md index 1556efc15e..d8adb61a25 100644 --- a/sdk/cosmos/azure_data_cosmos/api/API.md +++ b/sdk/cosmos/azure_data_cosmos/api/API.md @@ -1208,6 +1208,7 @@ pub mod models { impl CosmosStatus { const AUTHENTICATION_TOKEN_ACQUISITION_FAILED: CosmosStatus = _; const CLIENT_BAD_REQUEST: CosmosStatus = _; + const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: CosmosStatus = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: CosmosStatus = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: CosmosStatus = _; @@ -1241,7 +1242,7 @@ pub mod models { const CLIENT_MIXED_NAME_RID_ADDRESSING: CosmosStatus = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: CosmosStatus = _; const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; - const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; + const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: CosmosStatus = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: CosmosStatus = _; @@ -2456,6 +2457,7 @@ pub mod options { #[derive(Clone, Debug, Default)] #[non_exhaustive] pub struct OperationOptions { + pub allow_unbounded_queries: Option, pub query_plan_mode: Option, pub patch_strategy: Option, pub read_consistency_strategy: Option, @@ -2489,6 +2491,7 @@ pub mod options { #[must_use] fn build(self) -> OperationOptions; fn new() -> Self; + fn with_allow_unbounded_queries(self, value: bool) -> Self; fn with_availability_strategy(self, value: AvailabilityStrategy) -> Self; fn with_binary_encoding(self, value: BinaryEncodingOptions) -> Self; fn with_content_response_on_write(self, value: ContentResponseOnWrite) -> Self; @@ -2513,6 +2516,7 @@ pub mod options { #[doc(inline)] #[automatically_derived] impl<'a> OperationOptionsView<'a> { + fn allow_unbounded_queries(&self) -> Option<&bool>; fn availability_strategy(&self) -> Option<&AvailabilityStrategy>; fn binary_encoding(&self) -> Option<&BinaryEncodingOptions>; fn content_response_on_write(&self) -> Option<&ContentResponseOnWrite>; @@ -2623,6 +2627,7 @@ pub mod options { pub populate_query_metrics: Option, } impl QueryOptions { + fn with_allow_unbounded_queries(self, allow: bool) -> Self; fn with_continuation_token(self, continuation_token: ContinuationToken) -> Self; fn with_feed_options(self, feed: FeedOptions) -> Self; fn with_max_item_count(self, max_item_count: MaxItemCountHint) -> Self; diff --git a/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml b/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml index 73115fbadc..015f89ab48 100644 --- a/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml +++ b/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml @@ -1,4 +1,4 @@ -apiMdSha256: b4b18c17377dd7ee0ebddfa2a12c6319da074a34bb3b8d23dd6a6d971381a166 +apiMdSha256: 3e5a5c6223f651ed143d5e8a1ac53b82feb30e92e4ed33208be78950887baddb packageVersion: 0.39.0 parserVersion: 2.2.1 rustVersion: 1.97.0-nightly diff --git a/sdk/cosmos/azure_data_cosmos/src/options/feed.rs b/sdk/cosmos/azure_data_cosmos/src/options/feed.rs index 7bca2495fd..55e3ecfd0e 100644 --- a/sdk/cosmos/azure_data_cosmos/src/options/feed.rs +++ b/sdk/cosmos/azure_data_cosmos/src/options/feed.rs @@ -145,6 +145,15 @@ pub struct QueryOptions { } impl QueryOptions { + /// Allows client-buffered queries without a global finite TOP or LIMIT. + /// + /// Overrides the client default, including when false. Enabling this can + /// consume unbounded memory; service and continuation restrictions still apply. + pub fn with_allow_unbounded_queries(mut self, allow: bool) -> Self { + self.operation.allow_unbounded_queries = Some(allow); + self + } + /// Sets the session token for this request. pub fn with_session_token(mut self, session_token: impl Into) -> Self { self.session_token = Some(session_token.into()); diff --git a/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs b/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs index 12c7b860f2..5a62357b57 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/binary_roundtrip_fuzzer.rs @@ -41,7 +41,7 @@ use azure_data_cosmos::clients::ContainerClient; use azure_data_cosmos::models::ContainerProperties; use azure_data_cosmos::options::{ BinaryEncodingOptions, ConnectionPoolOptions, ContentResponseOnWrite, ItemWriteOptions, - OperationOptions, Region, ServerCertificateValidation, + OperationOptions, QueryOptions, Region, ServerCertificateValidation, }; use azure_data_cosmos::{ AccountEndpoint, AccountReference, CosmosClient, CosmosRuntime, FeedScope, Query, @@ -2012,13 +2012,18 @@ async fn query_values( sql: &str, run_id: &str, context: &str, + allow_unbounded_queries: bool, ) -> Result, Box> { let mut attempt = 0; loop { attempt += 1; let query = Query::from(sql).with_parameter("@run", run_id)?; let result = match container - .query_items(query, FeedScope::full_container(), None) + .query_items( + query, + FeedScope::full_container(), + Some(QueryOptions::default().with_allow_unbounded_queries(allow_unbounded_queries)), + ) .await { Ok(iterator) => Box::pin(iterator.try_collect()).await, @@ -2410,7 +2415,7 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { .await?; let context = format!("iter={iter} config={label} query={phase} seed={}", cfg.seed); let actual = canonical_query_results( - query_values(&container, sql, &run_id, &context).await?, + query_values(&container, sql, &run_id, &context, phase == "distinct").await?, ordered, ); if let Some(expected) = &expected { @@ -2440,6 +2445,7 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { "SELECT VALUE {\"int\": 7} FROM c WHERE c.fuzzRun = @run", &run_id, &context, + false, ) .await?; assert!( diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs index d9b5cd4a1f..0be82ee27e 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs @@ -21,6 +21,7 @@ use azure_data_cosmos::feed::FeedScope; use azure_data_cosmos::models::{ ContainerProperties, PartitionKeyKind, PatchInstructions, PatchOperation, }; +use azure_data_cosmos::options::QueryOptions; use azure_data_cosmos::{PartitionKey, Query, SubStatusCode, TransactionalBatch}; use framework::{TestClient, TestOptions, TestRunContext}; use futures::{StreamExt, TryStreamExt}; @@ -834,15 +835,12 @@ pub async fn hpk_query_cross_partition_advanced_not_servable() -> Result<(), Box // Servable: DISTINCT has a client-side stage, and it must // deduplicate correctly across the container's physical partitions. - let mut countries = collect_query::( - &container, - "SELECT DISTINCT VALUE c.country FROM c", - FeedScope::full_container(), + let mut countries = container.query_items::( + "SELECT DISTINCT VALUE c.country FROM c", FeedScope::full_container(), + Some(QueryOptions::default().with_allow_unbounded_queries(true)), ) .await? - .into_iter() - .map(|v| v.as_str().unwrap_or_default().to_owned()) - .collect::>(); + .try_collect::>().await?; countries.sort(); assert_eq!( countries, diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs index 047912f996..f51b05506b 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs @@ -11,8 +11,11 @@ use azure_data_cosmos::{ clients::{ContainerClient, DatabaseClient}, feed::FeedScope, models::ThroughputProperties, - options::{MaxItemCountHint, QueryOptions}, - Query, + options::{ + BinaryEncodingOptions, MaxItemCountHint, OperationOptions, QueryOptions, QueryPlanMode, + Region, + }, + AccountReference, CosmosClient, CosmosStatus, Query, RoutingStrategy, }; use framework::{test_data, MockItem, TestClient, TestOptions}; use futures::StreamExt; @@ -79,6 +82,120 @@ fn unordered_query_results_preserve_multiplicity() { assert_query_results(vec![1, 1], vec![1, 2], QueryResultOrder::Unordered); } +#[tokio::test] +#[cfg_attr(not(test_category = "emulator"), ignore = "requires live account")] +async fn live_distinct_admission_and_option_precedence() -> Result<(), Box> { + if framework::targets_emulator() { + eprintln!("live DISTINCT admission coverage requires a live account"); + return Ok(()); + } + TestClient::run_with_unique_db( + async |run_context, db_client| { + println!("Live DISTINCT test database: {}", run_context.db_name()); + test_data::create_container_with_items( + db_client, + test_data::generate_mock_items(4, 3), + None, + ) + .await?; + let connection = framework::resolve_connection_string() + .expect("the live harness already resolved a connection string"); + let account = AccountReference::with_authentication_key( + connection.account_endpoint().parse()?, + connection.account_key().clone(), + ); + let expected: Vec = (0..4).map(|i| format!("partition{i}")).collect(); + let unbounded = "SELECT DISTINCT VALUE c.partitionKey FROM c"; + for mode in [QueryPlanMode::LocalPreferred, QueryPlanMode::GatewayOnly] { + for client_allow in [None, Some(false), Some(true)] { + let mut defaults = OperationOptions::default(); + defaults.query_plan_mode = Some(mode); + defaults.allow_unbounded_queries = client_allow; + let client = CosmosClient::builder() + .with_default_operation_options(defaults) + .build(account.clone(), RoutingStrategy::ProximityTo(Region::EAST_US)) + .await?; + let container = client + .database_client(run_context.db_name()) + .container_client("TestContainer", None) + .await?; + for binary in [false, true] { + let mut operation = OperationOptions::default(); + operation.binary_encoding = + Some(BinaryEncodingOptions::new().with_enabled(binary)); + let options = QueryOptions::default() + .with_operation_options(operation) + .with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(1).unwrap(), + )); + for request_allow in [None, Some(false), Some(true)] { + let mut options = options.clone(); + if let Some(allow) = request_allow { + options = options.with_allow_unbounded_queries(allow); + } + let result = container + .query_items::( + unbounded, + FeedScope::full_container(), + Some(options), + ) + .await; + if !request_allow.or(client_allow).unwrap_or(false) { + let error = match result { + Err(error) => error, + Ok(_) => panic!("unbounded DISTINCT must fail at admission"), + }; + assert_eq!( + error.status(), + CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW, + "{mode:?}, client={client_allow:?}, request={request_allow:?}, binary={binary}" + ); + } else { + let mut pages = result?.into_pages(); + assert_eq!( + pages.to_continuation_token().unwrap_err().status(), + CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED + ); + let mut actual = Vec::new(); + while let Some(page) = pages.next().await { + actual.extend(page?.into_items()); + } + actual.sort(); + assert_eq!(actual, expected); + } + } + for query in [ + "SELECT DISTINCT TOP @bound VALUE c.partitionKey FROM c", + "SELECT DISTINCT VALUE c.partitionKey FROM c OFFSET 1 LIMIT @bound", + ] { + let mut pages = container + .query_items::( + Query::from(query).with_parameter("@bound", 2)?, + FeedScope::full_container(), + Some(options.clone().with_allow_unbounded_queries(false)), + ) + .await? + .into_pages(); + let mut actual = Vec::new(); + while let Some(page) = pages.next().await { + actual.extend(page?.into_items()); + } + assert_eq!(actual.len(), 2); + actual.sort(); + actual.dedup(); + assert_eq!(actual.len(), 2); + assert!(actual.iter().all(|value| expected.contains(value))); + } + } + } + } + Ok(()) + }, + Some(TestOptions::default()), + ) + .await +} + async fn execute_query_test( db_client: &DatabaseClient, items: Vec, @@ -407,9 +524,11 @@ pub async fn cross_partition_query_with_unordered_distinct() -> Result<(), Box Result<(), Box "select distinct value c.partitionKey from c", FeedScope::full_container(), Some( - QueryOptions::default().with_max_item_count(MaxItemCountHint::Limit( - std::num::NonZeroU32::new(1).unwrap(), - )), + QueryOptions::default() + .with_allow_unbounded_queries(true) + .with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(1).unwrap(), + )), ), ) .await? @@ -1116,7 +1237,11 @@ pub async fn distinct_projection_shapes() -> Result<(), Box> { scope: FeedScope, ) -> Result> { let mut pages = container - .query_items::(query, scope, None) + .query_items::( + query, + scope, + Some(QueryOptions::default().with_allow_unbounded_queries(true)), + ) .await? .into_pages(); let mut n = 0; @@ -1228,7 +1353,11 @@ pub async fn distinct_combined_with_unsupported_stages_is_rejected() -> Result<( for query in unsupported { let outcome = container - .query_items::(query, FeedScope::full_container(), None) + .query_items::( + query, + FeedScope::full_container(), + Some(QueryOptions::default().with_allow_unbounded_queries(true)), + ) .await; let error = match outcome { Err(error) => error, diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs index f06fe69f2e..a19195f658 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs @@ -666,6 +666,66 @@ pub async fn cross_partition_vector_search() -> Result<(), Box> { .await } +#[tokio::test] +#[cfg_attr( + not(test_category = "emulator"), + ignore = "requires live vector-enabled account" +)] +pub async fn unbounded_vector_query_admission_and_execution() -> Result<(), Box> { + if framework::targets_emulator() { + eprintln!("live vector admission coverage unavailable on local emulators"); + return Ok(()); + } + TestClient::run_with_unique_db( + async |run_context, db_client| { + let container = seed_vector_container( + run_context, db_client, Some(CROSS_PARTITION_THROUGHPUT), + ).await?; + assert_seeded_across_physical_partitions(&container, vector_documents().len()).await?; + let query = Query::from( + "SELECT c.id, VectorDistance(c.embedding, @queryVector, true) AS score \ + FROM c WHERE c.active = true \ + ORDER BY VectorDistance(c.embedding, @queryVector, true)", + ).with_parameter("@queryVector", QUERY_VECTOR.as_slice())?; + let denied = container.query_items::( + query.clone(), FeedScope::full_container(), None, + ).await; + let error = match denied { + Err(error) => error, + Ok(_) => panic!("missing global bound must be rejected"), + }; + // A service rejection is not evidence of client admission or executable opt-out. + assert_eq!(error.status(), CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW, + "the service must supply a no-TOP vector plan to validate client admission: {error}"); + let mut pages = container.query_items::( + query.clone(), FeedScope::full_container(), + Some(QueryOptions::default().with_allow_unbounded_queries(true).with_max_item_count( + MaxItemCountHint::Limit(NonZeroU32::new(2).unwrap()), + )), + ).await?.into_pages(); + assert_eq!(pages.to_continuation_token().unwrap_err().status(), + CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED); + let mut ids = Vec::new(); + while let Some(page) = pages.next().await { + ids.extend(page?.into_items().into_iter().map(|item| item.id)); + } + assert_eq!(ids, ["origin", "other-partition-origin", "near", + "other-partition-near", "far", "farthest"]); + let mut pages = container.query_items::( + query, FeedScope::partition(SEARCH_PARTITION), + Some(QueryOptions::default().with_allow_unbounded_queries(false)), + ).await?.into_pages(); + let mut ids = Vec::new(); + while let Some(page) = pages.next().await { + ids.extend(page?.into_items().into_iter().map(|item| item.id)); + } + assert_eq!(ids, ["origin", "near", "far", "farthest"]); + Ok(()) + }, + Some(TestOptions::default()), + ).await +} + #[test] fn precomputed_vector_fixture_has_expected_shape() { precomputed_vector_fixture(); diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/query_comparison.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/query_comparison.rs index c76944d0b4..69a169db59 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/query_comparison.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/query_comparison.rs @@ -310,6 +310,92 @@ enum FixtureKind { Hpk, } +#[tokio::test] +async fn buffered_query_policy_sdk_overrides_and_hierarchical_routing() -> Result<(), Box> +{ + let harness = QueryComparisonHarness::setup_in_memory_only().await?; + let handles = provision_fixture_with_topology( + &harness, + "buffered-policy", + FixtureKind::Hpk, + Some(ContainerConfig::new().with_partition_count(1).build()?), + &[], + ) + .await?; + let query = "SELECT DISTINCT VALUE c.value FROM c"; + let ranges = handles.emulator_container.read_feed_ranges(None).await?; + assert_eq!(ranges.len(), 1); + for client_allow in [false, true] { + let mut defaults = OperationOptions::default(); + defaults.allow_unbounded_queries = Some(client_allow); + let client = CosmosClientBuilder::new() + .with_runtime( + CosmosRuntimeBuilder::from(harness.emulator_http.runtime_builder()) + .build() + .await?, + ) + .with_default_operation_options(defaults) + .build( + AccountReference::with_authentication_key( + EMULATOR_GATEWAY_URL.parse::()?, + Secret::new("dGVzdGtleQ=="), + ), + RoutingStrategy::ProximityTo(Region::EAST_US), + ) + .await?; + let container = client + .database_client("buffered-policy") + .container_client("hpk", None) + .await?; + for (scope, buffered, expected_count) in [ + (FeedScope::full_container(), true, handles.documents.len()), + ( + FeedScope::range(ranges[0].clone()), + true, + handles.documents.len(), + ), + (FeedScope::partition("tenant-a"), true, 7), + ( + FeedScope::partition(("tenant-a", "user-1", "session-1")), + false, + 1, + ), + ] { + for request in [None, Some(false), Some(true)] { + let mut options = QueryOptions::default() + .with_max_item_count(MaxItemCountHint::Limit(NonZeroU32::new(1).unwrap())); + if let Some(allow) = request { + options = options.with_allow_unbounded_queries(allow); + } + let result = container + .query_items::(query, scope.clone(), Some(options)) + .await; + if buffered && !request.unwrap_or(client_allow) { + let error = match result { + Err(error) => error, + Ok(_) => { + panic!("unbounded buffered query must fail before creating a pager") + } + }; + assert_eq!(error.status(), azure_data_cosmos::models::CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW); + } else { + let mut pages = result?.into_pages(); + if buffered { + assert_eq!(pages.to_continuation_token().unwrap_err().status(), + azure_data_cosmos::models::CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED); + } + let mut values = Vec::new(); + while let Some(page) = pages.next().await { + values.extend(page?.into_items()); + } + assert_eq!(values.len(), expected_count); + } + } + } + } + Ok(()) +} + impl FixtureKind { fn container_name(self) -> &'static str { match self { diff --git a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs index febb5e1969..f6691a0315 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/split_tests/cosmos_query_distinct_split.rs @@ -110,7 +110,10 @@ async fn start_query( .query_items::( query, FeedScope::full_container(), - Some(query_options(binary, page_size)), + Some( + query_options(binary, page_size) + .with_allow_unbounded_queries(query == UNORDERED_QUERY), + ), ) .await? .into_pages(); diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 6cde51bfc5..a267d33ce9 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -4,12 +4,15 @@ ### Features Added +- Added layered `OperationOptions::allow_unbounded_queries` and incremental unbounded non-streaming ORDER BY execution, preserving finite top-k execution and continuation restrictions. - Extended Cosmos binary JSON query-page handling to cross-partition `DISTINCT`, including composition with streaming `ORDER BY` and `OFFSET`/`LIMIT`/`TOP`. ([#5070](https://github.com/Azure/azure-sdk-for-rust/pull/5070)) - Added local Rust query planning for supported cross-partition queries, avoiding Gateway query-plan requests while retaining native and Gateway fallbacks for advanced query shapes. Added layered `OperationOptions::query_plan_mode`, `QueryPlanMode::{LocalPreferred, GatewayOnly}`, and the authoritative `AZURE_COSMOS_QUERY_PLAN_MODE_OVERRIDE=gateway` break-glass setting to force Gateway planning globally or per operation. ([#5181](https://github.com/Azure/azure-sdk-for-rust/pull/5181)) -- Added a fully buffered cross-partition merge for finite non-streaming `ORDER BY` plans, including `VectorDistance(...)`. Unbounded, resumed, DISTINCT, and hybrid non-streaming plans are rejected with typed statuses. ([#5130](https://github.com/Azure/azure-sdk-for-rust/pull/5130)) +- Added a fully buffered cross-partition merge for finite non-streaming `ORDER BY` plans, including `VectorDistance(...)`. Resumed, DISTINCT, and hybrid non-streaming plans are rejected with typed statuses. ([#5130](https://github.com/Azure/azure-sdk-for-rust/pull/5130)) ### Breaking Changes +- Unordered cross-partition DISTINCT and non-streaming ORDER BY require a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`, sharing 400/20126 (`CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`); the existing non-streaming status constant remains an alias. + - `error::cosmos_status` is no longer a public module; `CosmosStatus` and `SubStatusCode` remain available as re-exports from `error`. The internal-only `query` module (gated behind the `__internal_testing` feature) is now `#[doc(hidden)]` so it no longer appears as an empty public module in generated API surfaces. ([#5205](https://github.com/Azure/azure-sdk-for-rust/pull/5205)) ### Bugs Fixed diff --git a/sdk/cosmos/azure_data_cosmos_driver/api/API.md b/sdk/cosmos/azure_data_cosmos_driver/api/API.md index 3b5b1e895e..bb3c1526a2 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/api/API.md +++ b/sdk/cosmos/azure_data_cosmos_driver/api/API.md @@ -636,6 +636,7 @@ pub mod error { impl CosmosStatus { const AUTHENTICATION_TOKEN_ACQUISITION_FAILED: CosmosStatus = _; const CLIENT_BAD_REQUEST: CosmosStatus = _; + const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: CosmosStatus = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: CosmosStatus = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: CosmosStatus = _; @@ -669,7 +670,7 @@ pub mod error { const CLIENT_MIXED_NAME_RID_ADDRESSING: CosmosStatus = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: CosmosStatus = _; const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; - const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; + const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: CosmosStatus = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: CosmosStatus = _; @@ -792,6 +793,7 @@ pub mod error { const CANNOT_ACQUIRE_PKRANGE_LOCK: SubStatusCode = _; const CHANNEL_CLOSED: SubStatusCode = _; const CHECKPOINT_QUEUE_DEPTH_BACKPRESSURE: SubStatusCode = _; + const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: SubStatusCode = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: SubStatusCode = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: SubStatusCode = _; @@ -839,7 +841,7 @@ pub mod error { const CLIENT_MIXED_NAME_RID_ADDRESSING: SubStatusCode = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: SubStatusCode = _; const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = _; - const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = _; + const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: SubStatusCode = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: SubStatusCode = _; @@ -1796,6 +1798,7 @@ pub mod models { impl CosmosStatus { const AUTHENTICATION_TOKEN_ACQUISITION_FAILED: CosmosStatus = _; const CLIENT_BAD_REQUEST: CosmosStatus = _; + const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: CosmosStatus = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: CosmosStatus = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: CosmosStatus = _; @@ -1829,7 +1832,7 @@ pub mod models { const CLIENT_MIXED_NAME_RID_ADDRESSING: CosmosStatus = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: CosmosStatus = _; const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; - const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; + const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: CosmosStatus = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: CosmosStatus = _; @@ -2353,6 +2356,7 @@ pub mod models { const CANNOT_ACQUIRE_PKRANGE_LOCK: SubStatusCode = _; const CHANNEL_CLOSED: SubStatusCode = _; const CHECKPOINT_QUEUE_DEPTH_BACKPRESSURE: SubStatusCode = _; + const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: SubStatusCode = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: SubStatusCode = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: SubStatusCode = _; @@ -2400,7 +2404,7 @@ pub mod models { const CLIENT_MIXED_NAME_RID_ADDRESSING: SubStatusCode = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: SubStatusCode = _; const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = _; - const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = _; + const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: SubStatusCode = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: SubStatusCode = _; @@ -3261,6 +3265,7 @@ pub mod options { #[derive(Clone, Debug, Default)] #[non_exhaustive] pub struct OperationOptions { + pub allow_unbounded_queries: Option, pub query_plan_mode: Option, pub patch_strategy: Option, pub read_consistency_strategy: Option, @@ -3291,6 +3296,7 @@ pub mod options { #[must_use] fn build(self) -> OperationOptions; fn new() -> Self; + fn with_allow_unbounded_queries(self, value: bool) -> Self; fn with_availability_strategy(self, value: AvailabilityStrategy) -> Self; fn with_binary_encoding(self, value: BinaryEncodingOptions) -> Self; fn with_content_response_on_write(self, value: ContentResponseOnWrite) -> Self; @@ -3313,6 +3319,7 @@ pub mod options { } #[automatically_derived] impl<'a> OperationOptionsView<'a> { + fn allow_unbounded_queries(&self) -> Option<&bool>; fn availability_strategy(&self) -> Option<&AvailabilityStrategy>; fn binary_encoding(&self) -> Option<&BinaryEncodingOptions>; fn content_response_on_write(&self) -> Option<&ContentResponseOnWrite>; diff --git a/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml b/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml index 635ca58868..51e45e2887 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml +++ b/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml @@ -1,4 +1,4 @@ -apiMdSha256: 91c979e114416e320babcb9f8ac6d660448c660c134e62583d735a140f98b757 +apiMdSha256: 328e7c147d92350f50e91f76b4152e77af17801c5e415673352eb56cf90c310d packageVersion: 0.8.0 parserVersion: 2.2.1 rustVersion: 1.97.0-nightly diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 90e017eccf..61bf37303c 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -4290,6 +4290,14 @@ impl CosmosDriver { ResolvedQueryPlan::Plan(plan) => *plan, }; + planner::validate_buffered_query( + &query_plan, + self.operation_options_view(options) + .allow_unbounded_queries() + .copied() + .unwrap_or(false), + )?; + // Build the fan-out pipeline using the query plan. let container_ref = container.clone(); let mut topology = CachedTopologyProvider::new( diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs index 8b37b5a9e2..d1aea00571 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -//! Buffered merge for finite, non-streaming ORDER BY queries. +//! Buffered merge for admitted non-streaming ORDER BY queries. use super::{ binary_heap, @@ -25,13 +25,13 @@ struct RetainedRow { ordinal: u64, } -/// Drains rewritten partition queries before emitting the finite globally ordered window. +/// Drains rewritten partition queries before emitting the globally ordered window. pub(crate) struct NonStreamingOrderedMerge { child: Box, directions: Arc<[SortOrder]>, - retention_limit: usize, + retention_limit: Option, skip: usize, - take: usize, + take: Option, page_size: usize, emit_binary: bool, retained: Vec, @@ -47,9 +47,9 @@ impl NonStreamingOrderedMerge { pub(crate) fn new( child: Box, directions: Vec, - retention_limit: usize, + retention_limit: Option, skip: usize, - take: usize, + take: Option, max_item_count: Option, emit_binary: bool, ) -> Self { @@ -66,7 +66,7 @@ impl NonStreamingOrderedMerge { take, page_size, emit_binary, - retained: Vec::with_capacity(retention_limit), + retained: Vec::new(), next_ordinal: 0, results: VecDeque::new(), aggregator: Some(PageAggregator::new(emit_binary)), @@ -94,12 +94,25 @@ impl NonStreamingOrderedMerge { .build() })?; - if self.retention_limit == 0 { + if self.retention_limit == Some(0) { return Ok(()); } let candidate = RetainedRow { row, ordinal }; - if self.retained.len() < self.retention_limit { + if self + .retention_limit + .is_none_or(|limit| self.retained.len() < limit) + { + self.retained.try_reserve(1).map_err(|_| { + CosmosError::builder() + .with_status(CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE) + .with_message("non-streaming ORDER BY candidate storage could not be allocated") + .build() + })?; + if self.retention_limit.is_none() { + self.retained.push(candidate); + return Ok(()); + } let directions = &self.directions; binary_heap::push_by(&mut self.retained, candidate, |left, right| { compare_key_tuples(left.row.keys.as_ref(), right.row.keys.as_ref(), directions) @@ -120,17 +133,19 @@ impl NonStreamingOrderedMerge { fn finish_buffering(&mut self) { let mut retained = mem::take(&mut self.retained); let directions = &self.directions; - retained.sort_by(|left, right| { + retained.sort_unstable_by(|left, right| { compare_key_tuples(left.row.keys.as_ref(), right.row.keys.as_ref(), directions) .then_with(|| left.ordinal.cmp(&right.ordinal)) }); - self.results = retained + let results = retained .into_iter() .skip(self.skip) - .take(self.take) - .map(|retained| retained.row.payload) - .collect(); + .map(|retained| retained.row.payload); + self.results = match self.take { + Some(take) => results.take(take).collect(), + None => results.collect(), + }; self.session_token = self .aggregator .as_ref() @@ -281,9 +296,9 @@ mod tests { NonStreamingOrderedMerge::new( Box::new(MockLeaf::with_pages(pages.into_iter().map(Ok).collect())), vec![SortOrder::Ascending], - retention_limit, + Some(retention_limit), skip, - take, + Some(take), page_size .map(|value| MaxItemCountHint::Limit(std::num::NonZeroU32::new(value).unwrap())), false, @@ -411,9 +426,9 @@ mod tests { true, ))])), vec![SortOrder::Ascending], - 1, + Some(1), 0, - 1, + Some(1), None, true, ); @@ -453,9 +468,9 @@ mod tests { is_terminal: true, })])), vec![SortOrder::Ascending], - 1, + Some(1), 0, - 1, + Some(1), None, true, ); @@ -480,6 +495,128 @@ mod tests { ); } + #[tokio::test] + async fn unbounded_merge_orders_all_rows_with_offset_and_encoding_parity() { + for emit_binary in [false, true] { + for skip in [0, 2, 20] { + let rows = [ + ("a", 2, "a"), + ("b", 1, "b"), + ("c", 1, "c"), + ("d", 2, "z"), + ("e", 1, "c"), + ]; + let mut pages = Vec::new(); + for chunk in rows.chunks(2) { + let body = serde_json::to_vec(&json!({ + "Documents": chunk.iter().map(|(id, key, secondary)| json!({ + "_rid": id, + "orderByItems": [{"item": key}, {"item": secondary}], + "payload": {"id": id} + })).collect::>() + })) + .unwrap(); + pages.push(Ok(PageResult::Page { + response: response_with_charge(&body, 1.0), + is_terminal: false, + })); + pages.push(Ok(page(&[], 0.5, false))); + } + pages.push(Ok(PageResult::Drained)); + let mut node = NonStreamingOrderedMerge::new( + Box::new(MockLeaf::with_pages(pages)), + vec![SortOrder::Ascending, SortOrder::Descending], + None, + skip, + None, + Some(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(2).unwrap(), + )), + emit_binary, + ); + assert_eq!(node.retained.capacity(), 0); + assert_eq!( + node.snapshot_state().unwrap_err().status(), + CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED + ); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = context(&mut executor, &mut topology); + let mut actual = Vec::new(); + let mut charge = 0.0; + let mut terminal = false; + while let PageResult::Page { + response, + is_terminal, + } = node.next_page(&mut context).await.unwrap() + { + assert!(!terminal); + terminal = is_terminal; + charge += response.headers().request_charge.unwrap().value(); + let ResponseBody::Items(items) = response.body() else { + panic!("expected items"); + }; + assert!(items.len() <= 2); + for item in items { + assert_eq!(crate::binary_json::is_binary(&item), emit_binary); + let value: serde_json::Value = if emit_binary { + crate::binary_json::from_slice(&item).unwrap() + } else { + serde_json::from_slice(&item).unwrap() + }; + actual.push(value["id"].as_str().unwrap().to_owned()); + } + } + assert_eq!( + actual, + ["c", "e", "b", "d", "a"] + .into_iter() + .skip(skip) + .map(str::to_owned) + .collect::>() + ); + assert!(terminal); + assert_eq!(charge, 4.5); + assert!(matches!( + node.next_page(&mut context).await.unwrap(), + PageResult::Drained + )); + } + } + } + + #[tokio::test] + async fn buffering_preserves_upstream_failure_without_partial_results() { + for limit in [None, Some(2)] { + let error = CosmosError::builder() + .with_status(CosmosStatus::CLIENT_UNSUPPORTED_QUERY_FEATURE) + .with_message("upstream failure") + .build(); + let mut node = NonStreamingOrderedMerge::new( + Box::new(MockLeaf::with_pages(vec![ + Ok(page(&[("a", 1.0, "a")], 1.0, false)), + Err(error), + ])), + vec![SortOrder::Ascending], + limit, + 0, + limit, + None, + false, + ); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = context(&mut executor, &mut topology); + let error = node.next_page(&mut context).await.unwrap_err(); + assert_eq!( + error.status(), + CosmosStatus::CLIENT_UNSUPPORTED_QUERY_FEATURE + ); + assert!(error.to_string().contains("upstream failure")); + assert!(node.results.is_empty()); + } + } + #[test] fn comparison_uses_all_order_by_items() { let left = RetainedRow { @@ -507,9 +644,9 @@ mod tests { let node = NonStreamingOrderedMerge::new( Box::new(MockLeaf::with_pages(Vec::new())), vec![SortOrder::Ascending, SortOrder::Ascending], - 1, + Some(1), 0, - 1, + Some(1), None, false, ); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs index 4177ad91d0..c2c429262e 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs @@ -357,7 +357,33 @@ pub(crate) fn is_non_streaming_order_by(info: &QueryInfo) -> bool { info.has_non_streaming_order_by } -/// Builds a bounded, fully buffered merge for a finite non-streaming ORDER BY query. +/// Validates admission using global, normalized metadata, never partition rewrites. +pub(crate) fn validate_buffered_query( + query_plan: &QueryPlan, + allow_unbounded_queries: bool, +) -> crate::error::Result<()> { + let Some(info) = query_plan.query_info.as_ref() else { + return Ok(()); + }; + if allow_unbounded_queries || combine_take(info).is_some() { + return Ok(()); + } + let shape = if is_non_streaming_order_by(info) { + "non-streaming ORDER BY (including buffered vector search)" + } else if info.distinct_type == DistinctType::Unordered { + "unordered DISTINCT" + } else { + return Ok(()); + }; + Err(crate::error::CosmosError::builder() + .with_status(crate::error::CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW) + .with_message(format!( + "cross-partition {shape} requires a finite global TOP or LIMIT; set allow_unbounded_queries to true to allow unbounded client buffering" + )) + .build()) +} + +/// Builds a fully buffered merge for an admitted non-streaming ORDER BY query. pub(crate) async fn build_non_streaming_ordered_merge( query_plan: &QueryPlan, topology_provider: &mut dyn TopologyProvider, @@ -430,40 +456,50 @@ pub(crate) async fn build_non_streaming_ordered_merge( } let skip = info.offset.unwrap_or(0); - let take = combine_take(info).ok_or_else(|| { - crate::error::CosmosError::builder() - .with_status( - crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW, - ) - .with_message( - "cross-partition non-streaming ORDER BY requires a finite TOP or OFFSET/LIMIT window", - ) - .build() - })?; - let retention_limit = skip.checked_add(take).ok_or_else(|| { - crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE) - .with_message("non-streaming ORDER BY OFFSET plus take overflows the supported window") - .build() - })?; - let retention_limit = usize::try_from(retention_limit).map_err(|_| { - crate::error::CosmosError::builder() + let take = combine_take(info); + let retention_limit = take + .map(|take| { + skip.checked_add(take).ok_or_else(|| { + crate::error::CosmosError::builder() + .with_status( + crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE, + ) + .with_message( + "non-streaming ORDER BY OFFSET plus take overflows the supported window", + ) + .build() + }) + }) + .transpose()?; + let retention_limit = + retention_limit + .map(|limit| { + usize::try_from(limit).map_err(|_| { + crate::error::CosmosError::builder() .with_status(crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE) .with_message("non-streaming ORDER BY candidate window does not fit in memory") .build() - })?; + }) + }) + .transpose()?; let skip = usize::try_from(skip).map_err(|_| { crate::error::CosmosError::builder() .with_status(crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE) .with_message("non-streaming ORDER BY OFFSET does not fit in memory") .build() })?; - let take = usize::try_from(take).map_err(|_| { - crate::error::CosmosError::builder() - .with_status(crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE) - .with_message("non-streaming ORDER BY take does not fit in memory") - .build() - })?; + let take = take + .map(|take| { + usize::try_from(take).map_err(|_| { + crate::error::CosmosError::builder() + .with_status( + crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE, + ) + .with_message("non-streaming ORDER BY take does not fit in memory") + .build() + }) + }) + .transpose()?; let effective_operation = rewritten_operation(operation, query_plan)?; let request_nodes = plan_fresh(query_plan, topology_provider, &effective_operation).await?; @@ -1817,6 +1853,7 @@ fn render_feed_range_for_error(range: &FeedRange) -> String { mod tests { use std::borrow::Cow; + use super::super::mocks::{response_with_continuation, MockRequestExecutor}; use super::*; use crate::{ driver::dataflow::{ @@ -1825,8 +1862,9 @@ mod tests { }, models::{ effective_partition_key::EffectivePartitionKey, AccountReference, ContainerProperties, - ContainerReference, DatabaseReference, ItemReference, OperationType, PartitionKey, - PartitionKeyDefinition, ResourceType, SystemProperties, + ContainerReference, CosmosResponse, DatabaseReference, ItemReference, MaxItemCountHint, + OperationType, PartitionKey, PartitionKeyDefinition, RequestCharge, ResourceType, + ResponseBody, SystemProperties, }, }; @@ -3834,19 +3872,255 @@ mod tests { assert_eq!(drain.into_children().len(), 2); } - #[tokio::test] - async fn build_non_streaming_ordered_merge_requires_finite_window() { - let operation = Arc::new(non_streaming_order_by_operation()); + #[test] + fn non_streaming_ordered_merge_requires_finite_window_or_opt_out() { let mut plan = non_streaming_order_by_plan(); plan.query_info.as_mut().unwrap().top = None; - let mut topology = MockTopologyProvider::new(Vec::new()); - - let err = build_non_streaming_ordered_merge(&plan, &mut topology, &operation, None) - .await - .unwrap_err(); + let err = validate_buffered_query(&plan, false).unwrap_err(); assert_eq!( err.status(), - crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW + crate::error::CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW + ); + validate_buffered_query(&plan, true).unwrap(); + } + + #[test] + fn buffered_admission_uses_only_global_normalized_bounds() { + for non_streaming in [false, true] { + for distinct_type in [ + DistinctType::None, + DistinctType::Ordered, + DistinctType::Unordered, + ] { + for order_by in [Vec::new(), vec![SortOrder::Ascending]] { + for (top, limit) in [ + (None, None), + (Some(0), None), + (None, Some(0)), + (Some(1), None), + (None, Some(1)), + (Some(50_003), None), + (Some(10), Some(3)), + (Some(u64::MAX), None), + ] { + for allow in [false, true] { + let plan = QueryPlan { + query_info: Some(QueryInfo { + top, + limit, + offset: Some(50_000), + has_non_streaming_order_by: non_streaming, + distinct_type, + order_by: order_by.clone(), + rewritten_query: Some( + "SELECT TOP 1 VALUE 'private SQL' FROM c".into(), + ), + ..Default::default() + }), + ..Default::default() + }; + let denied = !allow + && top.is_none() + && limit.is_none() + && (non_streaming || distinct_type == DistinctType::Unordered); + let result = validate_buffered_query(&plan, allow); + assert_eq!(result.is_err(), denied); + if let Err(error) = result { + let message = error.to_string(); + assert!(message.contains("TOP or LIMIT")); + assert!(message.contains("allow_unbounded_queries")); + assert!(!message.contains("private SQL")); + assert!(message.contains(if non_streaming { + "non-streaming ORDER BY" + } else { + "unordered DISTINCT" + })); + assert_eq!( + error.status(), + crate::error::CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW + ); + } + } + } + } + } + } + } + + #[test] + fn global_take_preserves_zero_and_selects_minimum() { + for (top, limit, expected) in [ + (None, None, None), + (Some(0), None, Some(0)), + (None, Some(0), Some(0)), + (Some(7), Some(3), Some(3)), + (Some(3), Some(7), Some(3)), + ] { + assert_eq!( + combine_take(&QueryInfo { + top, + limit, + ..Default::default() + }), + expected + ); + } + } + + #[test] + fn native_and_gateway_metadata_share_buffered_admission() { + for native in [false, true] { + for top in [None, Some(0), Some(10)] { + let wire = serde_json::json!({ + "queryInfo": { + "hasNonStreamingOrderBy": if native { serde_json::json!(1) } else { serde_json::json!(true) }, + "orderBy": ["Ascending"], + "top": top, + "rewrittenQuery": "SELECT TOP 1 c._rid, [{\"item\": 0}] AS orderByItems, c AS payload FROM c" + }, + "queryRanges": [{ + "min": "", "max": "FF", + "isMinInclusive": if native { serde_json::json!(1) } else { serde_json::json!(true) }, + "isMaxInclusive": if native { serde_json::json!(0) } else { serde_json::json!(false) } + }] + }); + let plan: QueryPlan = if native { + serde_json::from_value(wire).unwrap() + } else { + let raw: super::super::query_plan::RawQueryPlan = + serde_json::from_value(wire).unwrap(); + raw.resolve(&test_partition_key_definition()).unwrap() + }; + assert_eq!(validate_buffered_query(&plan, false).is_ok(), top.is_some()); + validate_buffered_query(&plan, true).unwrap(); + } + } + } + + #[tokio::test] + async fn admitted_unbounded_non_streaming_plan_preserves_partition_rewrite() { + for input_binary in [false, true] { + for output_binary in [false, true] { + for (offset, limit, expected) in [ + (0, None, vec!["a", "b", "c", "d", "e", "f"]), + (1, None, vec!["b", "c", "d", "e", "f"]), + (1, Some(3), vec!["b", "c", "d"]), + (20, None, vec![]), + ] { + let mut plan = non_streaming_order_by_plan(); + let info = plan.query_info.as_mut().unwrap(); + info.top = None; + info.offset = Some(offset); + info.limit = limit; + let rewrite = info.rewritten_query.clone().unwrap(); + let parameters = serde_json::json!([{"name":"@floor","value":-1}]); + let operation = Arc::new( + CosmosOperation::query_items(test_container(), Some(FeedRange::full())) + .with_body(serde_json::to_vec(&serde_json::json!({ + "query":"SELECT c.id FROM c WHERE c.rank > @floor ORDER BY c.rank, c.tie DESC", + "parameters":parameters + })).unwrap()) + .with_max_item_count(MaxItemCountHint::Limit(2.try_into().unwrap())) + .with_supported_serialization_formats(if output_binary { "CosmosBinary" } else { "JsonText" }), + ); + validate_buffered_query(&plan, true).unwrap(); + let mut topology = MockTopologyProvider::new(vec![Ok(vec![ + rr("", "80", "a"), + rr("80", "FF", "b"), + ])]); + let mut pipeline = + build_non_streaming_ordered_merge(&plan, &mut topology, &operation, None) + .await + .unwrap(); + assert_eq!(pipeline.fan_out_width(), 2); + let responses = [ + (vec![("e", 3, 0), ("b", 1, 9)], Some("a-next"), "a:1#2"), + (vec![("d", 2, 0)], None, "a:1#4"), + (vec![("f", 4, 0), ("c", 1, 1)], Some("b-next"), "b:1#3"), + (vec![("a", 0, 0)], None, "b:1#5"), + ].into_iter().map(|(rows, continuation, session)| { + let envelope = serde_json::json!({"Documents": rows.into_iter().map(|(id, rank, tie)| { + serde_json::json!({"_rid":id,"orderByItems":[{"item":rank},{"item":tie}],"payload":id}) + }).collect::>()}); + let body = if input_binary { + crate::binary_json::to_vec(&envelope).unwrap() + } else { + serde_json::to_vec(&envelope).unwrap() + }; + let response = response_with_continuation(&body, continuation); + let mut headers = response.headers().clone(); + headers.request_charge = Some(RequestCharge::new(1.25)); + headers.session_token = Some(crate::models::SessionToken(session.into())); + Ok(CosmosResponse::new(body, headers, response.status(), response.diagnostics())) + }).collect(); + let mut executor = MockRequestExecutor::new(responses); + let mut context = PipelineContext::new(&mut executor, Some(&mut topology)); + let mut actual = Vec::new(); + let mut charge = 0.0; + while let Some(page) = pipeline.next_page(&mut context).await.unwrap() { + charge += page.headers().request_charge.unwrap().value(); + let session = page.headers().session_token.as_ref().unwrap().to_string(); + assert!(session.contains("a:1#4"), "{session}"); + assert!(session.contains("b:1#5"), "{session}"); + let ResponseBody::Items(items) = page.body() else { + panic!("expected items") + }; + assert!(items.len() <= 2); + for item in items { + assert_eq!(crate::binary_json::is_binary(item), output_binary); + let value: String = if output_binary { + crate::binary_json::from_slice(item).unwrap() + } else { + serde_json::from_slice(item).unwrap() + }; + actual.push(value); + } + } + assert_eq!(actual, expected); + assert_eq!(charge, 5.0); + assert!(pipeline.next_page(&mut context).await.unwrap().is_none()); + assert_eq!(pipeline.snapshot_state().unwrap_err().status(), + crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED); + assert_eq!( + executor.continuation_calls, + vec![None, Some("a-next".into()), None, Some("b-next".into())] + ); + assert_eq!(executor.target_calls.len(), 4); + assert_eq!(executor.target_calls[0], executor.target_calls[1]); + assert_eq!(executor.target_calls[2], executor.target_calls[3]); + assert_ne!(executor.target_calls[0], executor.target_calls[2]); + assert!(executor.responses.is_empty()); + for body in &executor.query_bodies { + let body: serde_json::Value = + serde_json::from_slice(body.as_ref().unwrap()).unwrap(); + assert_eq!( + body, + serde_json::json!({"query":rewrite,"parameters":parameters}) + ); + assert!(!body["query"].as_str().unwrap().contains("TOP")); + } + } + } + } + } + + #[tokio::test] + async fn non_streaming_window_overflow_is_reported_before_topology() { + let mut plan = non_streaming_order_by_plan(); + let info = plan.query_info.as_mut().unwrap(); + info.top = Some(u64::MAX); + info.offset = Some(1); + let error = build_non_streaming_ordered_merge( + &plan, + &mut NoopTopologyProvider, + &Arc::new(non_streaming_order_by_operation()), + None, + ) + .await + .unwrap_err(); + assert_eq!( + error.status(), + crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE ); } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs index 9d4e090af9..8f1edc50da 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs @@ -503,7 +503,7 @@ impl SubStatusCode { 20123 => Some("ClientDistinctValueTooDeeplyNested"), 20124 => Some("ClientDistinctContinuationUnsupported"), 20125 => Some("ClientNonStreamingOrderByContinuationUnsupported"), - 20126 => Some("ClientNonStreamingOrderByRequiresFiniteWindow"), + 20126 => Some("ClientBufferedQueryRequiresFiniteWindow"), 20127 => Some("ClientNonStreamingOrderByWindowTooLarge"), 20150 => Some("ClientDuplicateFaultInjectionRuleId"), 20151 => Some("ClientThroughputControlGroupRegistrationFailed"), @@ -1406,10 +1406,13 @@ impl SubStatusCode { pub const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = SubStatusCode(20125); - /// A non-streaming `ORDER BY` query did not contain a finite `TOP` or - /// `OFFSET`/`LIMIT` window (20126). + /// A buffered query requires a finite global TOP/LIMIT or explicit opt-out + /// (20126), including non-streaming ORDER BY and unordered DISTINCT. + pub const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = SubStatusCode(20126); + + /// Compatibility alias for [`Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`]. pub const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = - SubStatusCode(20126); + Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; /// A non-streaming `ORDER BY` query's candidate window cannot be represented /// by the current process (20127). @@ -2394,12 +2397,17 @@ impl CosmosStatus { sub_status: Some(SubStatusCode::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED), }; - /// 400 / 20126 — non-streaming `ORDER BY` requires a finite result window. - pub const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = CosmosStatus { + /// 400 / 20126 — a buffered query requires a finite global TOP/LIMIT or + /// explicit opt-out, including non-streaming ORDER BY and unordered DISTINCT. + pub const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = CosmosStatus { status_code: StatusCode::BadRequest, - sub_status: Some(SubStatusCode::CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW), + sub_status: Some(SubStatusCode::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW), }; + /// Compatibility alias for [`Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`]. + pub const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = + Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; + /// 400 / 20127 — the non-streaming `ORDER BY` candidate window cannot be /// represented by the current process. pub const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = CosmosStatus { @@ -2767,6 +2775,27 @@ mod tests { ); } + #[test] + fn buffered_query_status_preserves_existing_code() { + let status = CosmosStatus::new(StatusCode::BadRequest).with_sub_status(20126); + assert_eq!( + status, + CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW + ); + assert_eq!( + status, + CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW + ); + assert_eq!( + status.sub_status(), + Some(SubStatusCode::CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW) + ); + assert_eq!( + status.name(), + Some("ClientBufferedQueryRequiresFiniteWindow") + ); + } + #[test] fn with_sub_status_unambiguous() { let status = CosmosStatus::new(StatusCode::TooManyRequests).with_sub_status(3200); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs index cc13f6a602..a29499306e 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/operation_options.rs @@ -40,6 +40,14 @@ use crate::{ #[options(layers(runtime, account, operation))] #[non_exhaustive] pub struct OperationOptions { + /// Allows client-buffered queries without a global finite TOP or LIMIT. + /// + /// `None` inherits; the final default is false. Explicit false overrides a + /// client opt-out. Enabling this can consume unbounded memory and does not + /// relax service restrictions or enable buffered-query continuation tokens. + #[option(env = "AZURE_COSMOS_ALLOW_UNBOUNDED_QUERIES")] + pub allow_unbounded_queries: Option, + /// Query-plan provider selection for query operations. /// /// `None` inherits from a lower layer (default: @@ -397,6 +405,49 @@ mod tests { assert!(view.max_session_retry_count().is_none()); } + #[test] + fn unbounded_admission_resolves_each_layer_and_explicit_false() { + use crate::driver::dataflow::{ + planner::validate_buffered_query, + query_plan::{DistinctType, QueryInfo, QueryPlan}, + }; + use std::sync::Arc; + + let plan = QueryPlan { + query_info: Some(QueryInfo { + distinct_type: DistinctType::Unordered, + ..Default::default() + }), + ..Default::default() + }; + for env in [None, Some(false), Some(true)] { + for runtime in [None, Some(false), Some(true)] { + for account in [None, Some(false), Some(true)] { + for operation in [None, Some(false), Some(true)] { + let layer = |value| { + Arc::new(OperationOptions { + allow_unbounded_queries: value, + ..Default::default() + }) + }; + let request = layer(operation); + let view = OperationOptionsView::new( + Some(layer(env)), + Some(layer(runtime)), + Some(layer(account)), + Some(&request), + ); + let allowed = view.allow_unbounded_queries().copied().unwrap_or(false); + assert_eq!( + validate_buffered_query(&plan, allowed).is_ok(), + operation.or(account).or(runtime).or(env).unwrap_or(false) + ); + } + } + } + } + } + /// Rule 2 + Rule 3 (RCS resolution): /// An explicit per-request `Default` overrides a client-level non-`Default`, /// resulting in no RCS being emitted on the wire. diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs index 23f842647a..27c211a25b 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/distinct.rs @@ -28,6 +28,7 @@ use azure_data_cosmos_driver::models::{ }; use azure_data_cosmos_driver::options::{ BinaryEncodingOptions, DriverOptions, OperationOptions, OperationOptionsBuilder, PlanOptions, + QueryPlanMode, }; const GATEWAY_URL: &str = "https://eastus.emulator.local"; @@ -122,6 +123,14 @@ async fn setup() -> (Arc, Arc) { async fn setup_with_observer( observer: Option>, +) -> (Arc, Arc) { + setup_with_policy(observer, OperationOptions::default(), 2).await +} + +async fn setup_with_policy( + observer: Option>, + client_options: OperationOptions, + partition_count: u32, ) -> (Arc, Arc) { let config = VirtualAccountConfig::new(vec![VirtualRegion::new( "East US", @@ -139,7 +148,7 @@ async fn setup_with_observer( let store = emulator.store(); store.create_database("testdb"); let container_config = ContainerConfig::new() - .with_partition_count(2) + .with_partition_count(partition_count) .build() .unwrap(); store.create_container_with_config( @@ -159,7 +168,11 @@ async fn setup_with_observer( "ZW11bGF0b3Ita2V5", ); let driver = runtime - .create_driver(DriverOptions::builder(account).build()) + .create_driver( + DriverOptions::builder(account) + .with_operation_options(client_options) + .build(), + ) .await .expect("driver initializes against the emulator"); (emulator, driver) @@ -176,6 +189,166 @@ async fn setup_with_query_recorder() -> ( (emulator, driver, recorder) } +#[tokio::test] +async fn buffered_admission_precedes_items_and_resolves_client_query_overrides() { + use azure_data_cosmos_driver::error::CosmosStatus; + + for partition_count in [1, 2] { + for mode in [QueryPlanMode::LocalPreferred, QueryPlanMode::GatewayOnly] { + for client in [None, Some(false), Some(true)] { + let recorder = Arc::new(QueryRequestRecorder::default()); + let mut defaults = OperationOptions::default(); + defaults.allow_unbounded_queries = client; + defaults.query_plan_mode = Some(mode); + let (_, driver) = + setup_with_policy(Some(recorder.clone()), defaults, partition_count).await; + let container = driver + .resolve_container("testdb", "testcoll", OperationOptions::default()) + .await + .unwrap(); + let documents: Vec<_> = (0..18) + .map(|n| { + serde_json::json!({ + "id": format!("id-{n}"), "pk": format!("pk-{n}"), "value": n % 6, + }) + }) + .collect(); + seed(&driver, &container, &documents).await; + + for request in [None, Some(false), Some(true)] { + for (sql, parameters, bound, expected) in [ + ("SELECT DISTINCT VALUE c.value FROM c", vec![], false, 6), + ( + "SELECT DISTINCT TOP @take VALUE c.value FROM c", + vec![serde_json::json!({"name":"@take","value":3})], + true, + 3, + ), + ( + "SELECT DISTINCT VALUE c.value FROM c OFFSET @skip LIMIT @take", + vec![ + serde_json::json!({"name":"@skip","value":1}), + serde_json::json!({"name":"@take","value":2}), + ], + true, + 2, + ), + ( + "SELECT DISTINCT TOP 0 VALUE c.value FROM c", + vec![], + true, + 0, + ), + ( + "SELECT DISTINCT VALUE c.value FROM c OFFSET 0 LIMIT 0", + vec![], + true, + 0, + ), + ] { + let query = QuerySpec { + text: sql.into(), + parameters, + distinct_type: "Unordered".into(), + }; + let mut options = OperationOptions::default(); + options.allow_unbounded_queries = request; + for page_size in [1, 1000, u32::MAX] { + recorder.take(); + let result = Box::pin(driver.plan_operation( + query_operation(&container, &query, page_size), + &options, + None, + &PlanOptions::default().with_max_fan_out( + if !bound && !request.or(client).unwrap_or(false) { + 1 + } else { + partition_count + }, + ), + )) + .await; + assert!(recorder.take().is_empty(), "planning must not query items"); + if !bound && !request.or(client).unwrap_or(false) { + let error = result.err().expect("unbounded query must be denied"); + assert_eq!( + error.status(), + CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW + ); + let message = error.to_string(); + assert!(message.contains("unordered DISTINCT")); + assert!(message.contains("TOP or LIMIT")); + assert!(message.contains("allow_unbounded_queries")); + assert!(!message.contains(sql)); + continue; + } + let mut plan = result.unwrap(); + let mut values = Vec::new(); + while let Some(response) = driver + .execute_plan( + &mut plan, + Some(container.clone()), + OperationOptions::default(), + ) + .await + .unwrap() + { + values.extend(documents_of(response)); + } + let unique = sorted(values.clone()); + assert_eq!(values.len(), expected); + let mut deduped = unique.clone(); + deduped.dedup(); + assert_eq!(deduped, unique); + if !bound { + assert_eq!( + unique, + (0..6).map(|value| value.to_string()).collect::>() + ); + } + } + } + } + + // A complete logical key bypasses client deduplication even with an explicit denial. + let query = QuerySpec { + text: "SELECT DISTINCT VALUE c.value FROM c".into(), + parameters: vec![], + distinct_type: "Unordered".into(), + }; + let operation = CosmosOperation::query_items( + container.clone(), + Some(FeedRange::for_partition( + PartitionKey::from("pk-0"), + &PartitionKeyDefinition::new(vec!["/pk".into()]), + )), + ) + .with_body(query_body(&query)); + let options = OperationOptionsBuilder::new() + .with_allow_unbounded_queries(false) + .build(); + let mut plan = Box::pin(driver.plan_operation( + operation, + &options, + None, + &PlanOptions::default(), + )) + .await + .unwrap(); + let mut values = Vec::new(); + while let Some(response) = driver + .execute_plan(&mut plan, Some(container.clone()), options.clone()) + .await + .unwrap() + { + values.extend(documents_of(response)); + } + assert_eq!(values, vec![serde_json::json!(0)]); + } + } + } +} + async fn seed( driver: &CosmosDriver, container: &ContainerReference, @@ -419,12 +592,16 @@ async fn drain_all( query: &QuerySpec, page_size: u32, ) -> Vec { - let mut plan = Box::pin(driver.plan_operation( - query_operation(container, query, page_size), - &OperationOptions::default(), - None, - &PlanOptions::default(), - )) + let mut plan = Box::pin( + driver.plan_operation( + query_operation(container, query, page_size), + &OperationOptionsBuilder::new() + .with_allow_unbounded_queries(query.distinct_type == "Unordered") + .build(), + None, + &PlanOptions::default(), + ), + ) .await .expect("plan builds"); @@ -512,23 +689,31 @@ async fn catalog_emulator_error_scenarios_fail_as_expected() { let outcome = match expected.category.as_str() { // The unsupported-feature check happens while planning. - "clientUnsupportedQueryFeature" => Box::pin(driver.plan_operation( - query_operation(&container, &scenario.query, 10), - &OperationOptions::default(), - None, - &PlanOptions::default(), - )) + "clientUnsupportedQueryFeature" => Box::pin( + driver.plan_operation( + query_operation(&container, &scenario.query, 10), + &OperationOptionsBuilder::new() + .with_allow_unbounded_queries(true) + .build(), + None, + &PlanOptions::default(), + ), + ) .await .err() .map(|e| e.to_string()), // The continuation refusal happens when the caller mints a token. "clientDistinctContinuationUnsupported" => { - let mut plan = Box::pin(driver.plan_operation( - query_operation(&container, &scenario.query, 1), - &OperationOptions::default(), - None, - &PlanOptions::default(), - )) + let mut plan = Box::pin( + driver.plan_operation( + query_operation(&container, &scenario.query, 1), + &OperationOptionsBuilder::new() + .with_allow_unbounded_queries(true) + .build(), + None, + &PlanOptions::default(), + ), + ) .await .expect("an unordered DISTINCT query plans successfully"); let _ = driver @@ -673,12 +858,16 @@ async fn ordered_distinct_resume_matches_a_single_drain() { let mut resumed = Vec::new(); let mut token = None; loop { - let mut plan = Box::pin(driver.plan_operation( - query_operation(&container, &scenario.query, 1), - &OperationOptions::default(), - token.as_ref(), - &PlanOptions::default(), - )) + let mut plan = Box::pin( + driver.plan_operation( + query_operation(&container, &scenario.query, 1), + &OperationOptionsBuilder::new() + .with_allow_unbounded_queries(scenario.query.distinct_type == "Unordered") + .build(), + token.as_ref(), + &PlanOptions::default(), + ), + ) .await .expect("plan builds (fresh or resumed)"); @@ -750,12 +939,16 @@ async fn split_mid_drain_does_not_reemit_deduplicated_values() { .expect("container resolves"); seed(&driver, &container, &scenario.documents).await; - let mut plan = Box::pin(driver.plan_operation( - query_operation(&container, &scenario.query, 1), - &OperationOptions::default(), - None, - &PlanOptions::default(), - )) + let mut plan = Box::pin( + driver.plan_operation( + query_operation(&container, &scenario.query, 1), + &OperationOptionsBuilder::new() + .with_allow_unbounded_queries(scenario.query.distinct_type == "Unordered") + .build(), + None, + &PlanOptions::default(), + ), + ) .await .expect("plan builds"); @@ -845,6 +1038,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { }, ] { let text_options = OperationOptionsBuilder::new() + .with_allow_unbounded_queries(query.distinct_type == "Unordered") .with_binary_encoding(BinaryEncodingOptions::new().with_enabled(false)) .build(); let (text, text_formats) = drain_query_with_options( @@ -857,6 +1051,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { .await; let text_request_modes = recorder.take(); let binary_options = OperationOptionsBuilder::new() + .with_allow_unbounded_queries(query.distinct_type == "Unordered") .with_binary_encoding(BinaryEncodingOptions::new().with_enabled(true)) .build(); let (binary, binary_formats) = drain_query_with_options( @@ -869,6 +1064,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { .await; let binary_request_modes = recorder.take(); let binary_as_text_options = OperationOptionsBuilder::new() + .with_allow_unbounded_queries(query.distinct_type == "Unordered") .with_binary_encoding( BinaryEncodingOptions::new() .with_enabled(true) diff --git a/sdk/cosmos/docs/specs/0001-configuration-options.md b/sdk/cosmos/docs/specs/0001-configuration-options.md index 48764dc0b1..141e091886 100644 --- a/sdk/cosmos/docs/specs/0001-configuration-options.md +++ b/sdk/cosmos/docs/specs/0001-configuration-options.md @@ -203,6 +203,7 @@ pub struct OperationOptions { /* fields below */ } | Option | Type | Env Var | Notes | | ---------------------------------------------- | --------------------------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `allow_unbounded_queries` | `Option` | `AZURE_COSMOS_ALLOW_UNBOUNDED_QUERIES` | Allows client-buffered queries without a global finite TOP/LIMIT. `None` inherits; the final default is false. See section 5.3 for query usage. | | `read_consistency_strategy` | `Option` | `AZURE_COSMOS_READ_CONSISTENCY_STRATEGY` | Read consistency for the operation. Replaces the legacy `consistency_level` field. The SDK enforces weakening-only semantics relative to the account default. | | `excluded_regions` | `Option>` | `AZURE_COSMOS_EXCLUDED_REGIONS` | Regions to exclude from routing. `None` inherits from a lower layer; `Some(vec![])` explicitly clears exclusions. Env var is comma-separated (e.g. `"West US,East US"`). | | `content_response_on_write` | `Option` | `AZURE_COSMOS_CONTENT_RESPONSE_ON_WRITE` | Whether write operations return the resource body in the response. Only applicable to write operations; ignored by reads and queries. Cascades from runtime → account → operation, matching .NET/Java/Go behavior. | @@ -431,12 +432,43 @@ pub struct QueryOptions { | Option | Type | Notes | | ------------------------- | ---------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `operation` | `OperationOptions` | Layered group; `content_response_on_write` is ignored for queries. | +| `operation` | `OperationOptions` | Layered group, including `allow_unbounded_queries`; `content_response_on_write` is ignored for queries. | | `session_token` | `Option` | Session token for session-consistent queries. Operation-only. | | `enable_scan_if_no_index` | `Option` | If the query can't be served by indexes because the relevant paths are not indexed, setting this permits the query engine to perform a full container scan. Operation-only. | | `populate_index_metrics` | `Option` | If set to `true`, the response will contain metrics regarding indexes used. Operation-only. | | `populate_query_advice` | `Option` | If set to `true`, the response will include query optimization suggestions from the query advisor. Operation-only. | +#### Buffered-query opt-out + +`OperationOptions::allow_unbounded_queries: Option` participates in the +existing operation > account/client > runtime > environment precedence. +Unset inherits; the final default is false. Explicit false overrides inherited +true. `AZURE_COSMOS_ALLOW_UNBOUNDED_QUERIES=true` is a low-priority explicit +opt-out; there is no authoritative `_OVERRIDE` variant. + +The SDK convenience setter +`QueryOptions::with_allow_unbounded_queries(bool)` writes the nested operation +option, not a second setting. Prefer a global TOP or LIMIT where practical: + +```rust +use azure_data_cosmos::{ + CosmosClientBuilder, + options::{OperationOptionsBuilder, QueryOptions}, +}; + +// SQL alternative: SELECT DISTINCT TOP 100 VALUE c.category FROM c +let defaults = OperationOptionsBuilder::new() + .with_allow_unbounded_queries(true) + .build(); +let client_builder = CosmosClientBuilder::new() + .with_default_operation_options(defaults); +let query_options = QueryOptions::default().with_allow_unbounded_queries(false); +``` + +The Rust SDK and driver use the same operation-options hierarchy. This +setting only admits potentially unbounded client buffering; it supplies no +memory budget and does not change service or continuation restrictions. + ### 5.4 `TransactionalBatchOptions` Options for transactional batch operations. The batch as a whole carries cross-layer options via `OperationOptions`, plus batch-level operation-only fields. This follows the same pattern as `ItemWriteOptions`. diff --git a/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md b/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md index fc7b9234e7..6ab047e84e 100644 --- a/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md +++ b/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md @@ -4,6 +4,22 @@ This document describes the target retry behavior for the Azure Cosmos DB Rust d ## Design Philosophy +### Buffered-query input validation + +| Status | Symbol | Remedy | +| --- | --- | --- | +| 400/20126 | `CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW` | Add a finite global TOP/LIMIT to non-streaming ORDER BY (including buffered vector search) or unordered DISTINCT, or explicitly set `allow_unbounded_queries=true`. | + +The existing `CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW` constant +remains a compatibility alias. Both shapes report the symbolic name +`ClientBufferedQueryRequiresFiniteWindow`; no separate DISTINCT code is allocated. + +This is a non-retryable client input error. Messages identify the query shape +and remedies without SQL or parameter values. No fixed numeric ceiling applies. +400/20127 still reports unrepresentable non-streaming windows/candidate storage. +Continuation restrictions remain 400/20124 (unordered DISTINCT) and 400/20125 +(non-streaming ORDER BY), with or without an opt-out. + The Rust driver retries writes by default for retryable status codes. This is safe because Cosmos DB's write APIs are designed to be idempotent when used correctly: - **503 (Service Unavailable)**: Cosmos DB intentionally returns 503 when a write was **not processed** — it is always safe to retry. diff --git a/sdk/cosmos/docs/specs/0013-query-engine.md b/sdk/cosmos/docs/specs/0013-query-engine.md index 0fb3118727..f575f3ded4 100644 --- a/sdk/cosmos/docs/specs/0013-query-engine.md +++ b/sdk/cosmos/docs/specs/0013-query-engine.md @@ -16,6 +16,41 @@ The supported SDK query path now integrates the local planner as a pre-Gateway o ## Architecture +### Buffered-query admission + +Before allocating a client-buffering pipeline or issuing item-query requests, +the driver validates the normalized global plan from every provider. Metadata +and query-plan requests may precede this check. Non-streaming ORDER BY (including +buffered vector search) and unordered DISTINCT require a finite global TOP or +LIMIT, unless resolved `allow_unbounded_queries` is true. Ordering metadata does +not exempt an unordered DISTINCT stage. + +Any representable finite bound is accepted: there is no fixed numeric ceiling. +Zero is a bound, and TOP combined with LIMIT uses the smaller value. +OFFSET alone, page-size hints, fan-out, consumer-side take, nested subquery TOP, +and per-range rewritten bounds do not establish a global output bound. + +Complete logical-partition-key pass-through and provably empty local resolutions +are exempt. Partial hierarchical keys and explicit ranges still use the policy, +even if they currently touch only one physical partition. Ordinary streaming +queries, streaming ORDER BY without unordered DISTINCT, and ordered DISTINCT +retain their existing behavior. + +Bounded non-streaming execution retains its top-k heap and checked OFFSET + take +window. Explicit unbounded execution grows candidate storage incrementally, +sorts all candidates with the same key/ordinal comparison, applies OFFSET, and +paginates the remaining rows. Admission is fixed when the plan is built. +This is not a runtime memory budget or a guarantee that finite output bounds +bound all memory: OFFSET and page-level DISTINCT processing add retained work. + +Neither bounds nor opt-out enable unsupported query compositions or continuation +tokens (unordered DISTINCT: 400/20124; non-streaming ORDER BY: 400/20125). +Service validation remains authoritative. In particular, a service rejection of +a no-TOP vector query is not bypassed or replaced by a fabricated large TOP. +Live no-TOP vector support must be verified against an enabled account before +promising that service scenario; deterministic buffered-node tests cover the +client execution mode independently. + ```text SQL Text → Lexer (hand-crafted tokenizer) From d9a84c75cf23607a127d3c433875dfc946873e0f Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 14 Sep 2026 11:49:39 -0500 Subject: [PATCH 2/5] Link buffered query changelogs to PR 5301 Link the SDK and driver feature and breaking-change entries to the upstream pull request. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 4 ++-- sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 282abfa75c..68b648f21e 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -4,14 +4,14 @@ ### Features Added -- Added layered `OperationOptions::allow_unbounded_queries` and `QueryOptions::with_allow_unbounded_queries` to explicitly admit unbounded client-buffered queries. +- Added layered `OperationOptions::allow_unbounded_queries` and `QueryOptions::with_allow_unbounded_queries` to explicitly admit unbounded client-buffered queries. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) - Extended Cosmos binary JSON encoding to cross-partition `DISTINCT` query pages. ([#5070](https://github.com/Azure/azure-sdk-for-rust/pull/5070)) - Added `QueryPlanMode::{LocalPreferred, GatewayOnly}` to `OperationOptions`, allowing applications to force Gateway query planning globally or for an individual query as a livesite mitigation. ([#5181](https://github.com/Azure/azure-sdk-for-rust/pull/5181)) - Added finite cross-partition `ORDER BY VectorDistance(...)` queries with `TOP` or `OFFSET`/`LIMIT`. Results are fully buffered before the first page and cannot be resumed from continuation tokens. Hybrid/full-text vector ranking remains unsupported. ([#5130](https://github.com/Azure/azure-sdk-for-rust/pull/5130)) ### Breaking Changes -- Unordered cross-partition DISTINCT now requires a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`; non-streaming ORDER BY shares this admission policy, without a fixed numeric ceiling. +- Unordered cross-partition DISTINCT now requires a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`; non-streaming ORDER BY shares this admission policy, without a fixed numeric ceiling. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) ### Bugs Fixed diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 58718f3204..c4f9c6d396 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -4,14 +4,14 @@ ### Features Added -- Added layered `OperationOptions::allow_unbounded_queries` and incremental unbounded non-streaming ORDER BY execution, preserving finite top-k execution and continuation restrictions. +- Added layered `OperationOptions::allow_unbounded_queries` and incremental unbounded non-streaming ORDER BY execution, preserving finite top-k execution and continuation restrictions. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) - Extended Cosmos binary JSON query-page handling to cross-partition `DISTINCT`, including composition with streaming `ORDER BY` and `OFFSET`/`LIMIT`/`TOP`. ([#5070](https://github.com/Azure/azure-sdk-for-rust/pull/5070)) - Added local Rust query planning for supported cross-partition queries, avoiding Gateway query-plan requests while retaining native and Gateway fallbacks for advanced query shapes. Added layered `OperationOptions::query_plan_mode`, `QueryPlanMode::{LocalPreferred, GatewayOnly}`, and the authoritative `AZURE_COSMOS_QUERY_PLAN_MODE_OVERRIDE=gateway` break-glass setting to force Gateway planning globally or per operation. ([#5181](https://github.com/Azure/azure-sdk-for-rust/pull/5181)) - Added a fully buffered cross-partition merge for finite non-streaming `ORDER BY` plans, including `VectorDistance(...)`. Resumed, DISTINCT, and hybrid non-streaming plans are rejected with typed statuses. ([#5130](https://github.com/Azure/azure-sdk-for-rust/pull/5130)) ### Breaking Changes -- Unordered cross-partition DISTINCT and non-streaming ORDER BY require a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`, sharing 400/20126 (`CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`); the existing non-streaming status constant remains an alias. +- Unordered cross-partition DISTINCT and non-streaming ORDER BY require a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`, sharing 400/20126 (`CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`); the existing non-streaming status constant remains an alias. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) - `error::cosmos_status` is no longer a public module; `CosmosStatus` and `SubStatusCode` remain available as re-exports from `error`. The internal-only `query` module (gated behind the `__internal_testing` feature) is now `#[doc(hidden)]` so it no longer appears as an empty public module in generated API surfaces. ([#5205](https://github.com/Azure/azure-sdk-for-rust/pull/5205)) - Renamed several types for naming consistency: `diagnostics::PipelineType` is now `diagnostics::PipelineKind` (following the `Kind`-over-`Type` convention), `diagnostics::ProxyConfiguration` is now `diagnostics::ProxyConfig` (matching the `Config` naming used elsewhere), and `in_memory_emulator::RuChargingModel` is now `in_memory_emulator::RequestUnitChargingModel` (expanding the `RU` acronym). The unstable `testing` module (`__internal_mocking` feature) was renamed to `test`. ([#5203](https://github.com/Azure/azure-sdk-for-rust/pull/5203)) From dde537cc500ca894ca83dc2de5168cba2409df10 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Mon, 14 Sep 2026 13:43:12 -0500 Subject: [PATCH 3/5] Fix Cosmos live query and retry tests Cover default rejection and explicit opt-out for unbounded DISTINCT in the production query-plan comparison tests. Scope retry-count overrides to the injected read via the test client helper, keeping metadata setup retries intact. Inject a one-shot metadata 429 while preserving exact read attempt assertions. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../emulator_tests/driver_fault_injection.rs | 96 +++++++++++-------- .../tests/framework/test_client.rs | 19 +++- .../tests/gateway_query_plan_comparison.rs | 58 +++++++---- 3 files changed, 115 insertions(+), 58 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/emulator_tests/driver_fault_injection.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/emulator_tests/driver_fault_injection.rs index b9a133b3a6..2a39db5234 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/emulator_tests/driver_fault_injection.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/emulator_tests/driver_fault_injection.rs @@ -1119,11 +1119,25 @@ pub async fn fault_injection_429_honors_configurable_throttle_retry_count( .build(), ); - // Pin the throttle-retry budget at the runtime layer of the option - // view. A generous cumulative-wait budget keeps the attempt count the - // sole limiter for these small retry counts. No end-to-end latency - // policy is set, so the transport request carries no deadline and the - // forced-final retry is immediate. + let metadata_rule = Arc::new( + FaultInjectionRuleBuilder::new( + "setup-429", + FaultInjectionResultBuilder::new() + .with_error(FaultInjectionErrorType::TooManyRequests) + .with_probability(1.0) + .build(), + ) + .with_condition( + FaultInjectionConditionBuilder::new() + .with_operation_type(FaultOperationType::MetadataReadContainer) + .build(), + ) + .with_hit_limit(1) + .build(), + ); + + // Limit only the injected read; setup must retain normal metadata retries. + // A generous wait budget leaves retry count as the sole limiter. let operation_options = OperationOptionsBuilder::new() .with_throttling_retry_options( ThrottlingRetryOptionsBuilder::new() @@ -1134,45 +1148,49 @@ pub async fn fault_injection_429_honors_configurable_throttle_retry_count( .build(); let rule_for_assert = Arc::clone(&rule); - Box::pin( - DriverTestClient::run_with_unique_db_and_fault_injection_options( - vec![rule], - operation_options, - async move |context, database| { - let container_name = context.unique_container_name(); - let container = context - .create_container(&database, &container_name, "/pk") - .await?; - - // Seed the item with a write. The fault rule targets only - // ReadItem, so the seeding write is unaffected. - let item_json = br#"{"id": "item1", "pk": "pk1", "value": "test"}"#; - context - .create_item(&container, "item1", "pk1", item_json) - .await?; - - // The read always observes 429 and ultimately fails once - // the throttle budget is exhausted. - let read_result = context.read_item(&container, "item1", "pk1").await; - assert!( - read_result.is_err(), - "read must fail once the throttle budget is exhausted \ + Box::pin(DriverTestClient::run_with_unique_db_and_fault_injection( + vec![rule, Arc::clone(&metadata_rule)], + async move |context, database| { + let container_name = context.unique_container_name(); + let container = context + .create_container(&database, &container_name, "/pk") + .await?; + assert_eq!( + metadata_rule.hit_count(), + 1, + "setup must recover from metadata 429" + ); + + // Seed the item with a write. The fault rule targets only + // ReadItem, so the seeding write is unaffected. + let item_json = br#"{"id": "item1", "pk": "pk1", "value": "test"}"#; + context + .create_item(&container, "item1", "pk1", item_json) + .await?; + + // The read always observes 429 and ultimately fails once + // the throttle budget is exhausted. + let read_result = context + .read_item_with_options(&container, "item1", "pk1", operation_options) + .await; + assert!( + read_result.is_err(), + "read must fail once the throttle budget is exhausted \ (max_throttle_retry_count={max_throttle_retry_count})", - ); + ); - assert_eq!( - rule_for_assert.hit_count(), - expected_hits, - "max_throttle_retry_count={max_throttle_retry_count} must yield \ + assert_eq!( + rule_for_assert.hit_count(), + expected_hits, + "max_throttle_retry_count={max_throttle_retry_count} must yield \ {expected_hits} ReadItem attempts on the wire, but the 429 fault \ rule fired {} time(s)", - rule_for_assert.hit_count(), - ); + rule_for_assert.hit_count(), + ); - Ok(()) - }, - ), - ) + Ok(()) + }, + )) .await?; } diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/framework/test_client.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/framework/test_client.rs index 491734cc71..c039169cfa 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/framework/test_client.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/framework/test_client.rs @@ -1005,6 +1005,23 @@ impl DriverTestRunContext { container: &ContainerReference, item_id: &str, partition_key: impl Into, + ) -> Result> { + self.read_item_with_options( + container, + item_id, + partition_key, + OperationOptions::default(), + ) + .await + } + + /// Reads an item with options scoped to the read, not driver initialization. + pub async fn read_item_with_options( + &self, + container: &ContainerReference, + item_id: &str, + partition_key: impl Into, + options: OperationOptions, ) -> Result> { let driver = self .client @@ -1017,7 +1034,7 @@ impl DriverTestRunContext { let operation = CosmosOperation::read_item(item_ref); let result = driver - .execute_singleton_operation(operation, OperationOptions::default()) + .execute_singleton_operation(operation, options) .await?; Ok(result) diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs index 8f251f9ab4..7d0c504074 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/gateway_query_plan_comparison.rs @@ -32,8 +32,9 @@ use azure_data_cosmos_driver::driver::CosmosDriverRuntime; use azure_data_cosmos_driver::models::{ ContainerReference, CosmosOperation, FeedRange, PartitionKeyDefinition, }; -use azure_data_cosmos_driver::options::DriverOptions; -use azure_data_cosmos_driver::options::{OperationOptions, PlanOptions}; +use azure_data_cosmos_driver::options::{ + DriverOptions, OperationOptions, OperationOptionsBuilder, PlanOptions, +}; use azure_data_cosmos_driver::CosmosDriver; use framework::resolve_test_env; @@ -496,7 +497,7 @@ fn query_spec_body(sql: &str, parameters: &[(&str, serde_json::Value)]) -> Vec, ) { let (driver, container) = require_driver_and(get_driver().await, c_pk().await); let body = query_spec_body(sql, parameters); @@ -547,24 +548,15 @@ async fn validate_production_local_plan( "query ranges differ for '{sql}'" ); - if execute { + if let Some(options) = execution_options { let operation = CosmosOperation::query_items(container.clone(), Some(FeedRange::full())) .with_body(body); let mut plan = driver - .plan_operation( - operation, - &OperationOptions::default(), - None, - &PlanOptions::default(), - ) + .plan_operation(operation, &options, None, &PlanOptions::default()) .await .unwrap_or_else(|error| panic!("local plan failed for '{sql}': {error}")); while driver - .execute_plan( - &mut plan, - Some(container.clone()), - OperationOptions::default(), - ) + .execute_plan(&mut plan, Some(container.clone()), options.clone()) .await .unwrap_or_else(|error| panic!("local execution failed for '{sql}': {error}")) .is_some() @@ -1053,7 +1045,6 @@ async fn gw_production_local_plan_supported_surface() { "SELECT * FROM c ORDER BY c.name", "SELECT VALUE c.name FROM c ORDER BY c.name", "SELECT c.name, c.age AS years FROM c ORDER BY c.name", - "SELECT DISTINCT c.name FROM c", "SELECT DISTINCT TOP 5 c.name FROM c", "SELECT DISTINCT c.name FROM c OFFSET 2 LIMIT 3", "SELECT DISTINCT VALUE c.name FROM c ORDER BY c.name", @@ -1061,10 +1052,41 @@ async fn gw_production_local_plan_supported_surface() { "SELECT (SELECT VALUE 1) AS x FROM c", "SELECT * FROM c WHERE c.pk = 'production-local-plan'", ] { - validate_production_local_plan(sql, &[], true).await; + validate_production_local_plan(sql, &[], Some(OperationOptions::default())).await; } - validate_production_local_plan("SELECT VALUE udf.transform(c.data) FROM c", &[], false).await; + validate_production_local_plan("SELECT VALUE udf.transform(c.data) FROM c", &[], None).await; +} + +#[tokio::test] +#[cfg_attr( + not(test_category = "emulator"), + ignore = "requires test_category 'emulator'" +)] +async fn gw_production_local_plan_unbounded_distinct_requires_opt_out() { + let sql = "SELECT DISTINCT c.name FROM c"; + let (driver, container) = require_driver_and(get_driver().await, c_pk().await); + let operation = CosmosOperation::query_items(container.clone(), Some(FeedRange::full())) + .with_body(query_spec_body(sql, &[])); + let error = driver + .plan_operation( + operation, + &OperationOptions::default(), + None, + &PlanOptions::default(), + ) + .await + .err() + .expect("unbounded unordered DISTINCT must require an explicit opt-out"); + assert_eq!( + error.status(), + azure_data_cosmos_driver::error::CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW + ); + + let options = OperationOptionsBuilder::new() + .with_allow_unbounded_queries(true) + .build(); + validate_production_local_plan(sql, &[], Some(options)).await; } #[tokio::test] From 1014c1cc206dd08035755e9171666f6a22f61d9f Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Tue, 15 Sep 2026 13:42:33 -0400 Subject: [PATCH 4/5] Unify buffered query substatus codes Share continuation failures across buffered query shapes and compact the client error range to 20124-20126. Preserve existing constant names as aliases while updating symbolic diagnostics and regression coverage. Clarify partition-key exemptions in query specs, document numeric changes in both changelogs, and regenerate SDK and driver API artifacts. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 1 + sdk/cosmos/azure_data_cosmos/api/API.md | 5 +- .../azure_data_cosmos/api/API.metadata.yml | 2 +- .../azure_data_cosmos_driver/CHANGELOG.md | 3 +- .../azure_data_cosmos_driver/api/API.md | 20 +-- .../api/API.metadata.yml | 2 +- .../src/driver/dataflow/distinct.rs | 6 +- .../dataflow/non_streaming_ordered_merge.rs | 6 +- .../src/driver/dataflow/planner.rs | 14 ++- .../src/error/cosmos_status.rs | 114 +++++++++++++----- .../specs/0006-error-codes-and-retries.md | 13 +- .../0012-feed-operations-and-dataflow.md | 2 +- sdk/cosmos/docs/specs/0013-query-engine.md | 6 +- 13 files changed, 131 insertions(+), 63 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 68b648f21e..8f6f86e4dd 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -12,6 +12,7 @@ ### Breaking Changes - Unordered cross-partition DISTINCT now requires a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`; non-streaming ORDER BY shares this admission policy, without a fixed numeric ceiling. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) +- Unified client-buffered continuation errors under 400/20124, retaining the shape-specific constants as aliases; finite-window admission moved from 20126 to 20125 and non-streaming window/storage errors from 20127 to 20126. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) ### Bugs Fixed diff --git a/sdk/cosmos/azure_data_cosmos/api/API.md b/sdk/cosmos/azure_data_cosmos/api/API.md index 5d0abca648..6bd27f0ef6 100644 --- a/sdk/cosmos/azure_data_cosmos/api/API.md +++ b/sdk/cosmos/azure_data_cosmos/api/API.md @@ -1208,6 +1208,7 @@ pub mod models { impl CosmosStatus { const AUTHENTICATION_TOKEN_ACQUISITION_FAILED: CosmosStatus = _; const CLIENT_BAD_REQUEST: CosmosStatus = _; + const CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: CosmosStatus = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: CosmosStatus = _; @@ -1227,7 +1228,7 @@ pub mod models { const CLIENT_CROSS_PARTITION_FAN_OUT_EXCEEDED: CosmosStatus = _; const CLIENT_CROSS_PARTITION_QUERY_REQUIRES_CONTAINER_REF: CosmosStatus = _; const CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT: CosmosStatus = _; - const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = _; + const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: CosmosStatus = _; const CLIENT_DRIVER_NOT_INITIALIZED: CosmosStatus = _; const CLIENT_DUPLICATE_FAULT_INJECTION_RULE_ID: CosmosStatus = _; @@ -1241,7 +1242,7 @@ pub mod models { const CLIENT_INVALID_URL: CosmosStatus = _; const CLIENT_MIXED_NAME_RID_ADDRESSING: CosmosStatus = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: CosmosStatus = _; - const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; + const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: CosmosStatus = _; diff --git a/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml b/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml index a56dcee643..8b9457bd01 100644 --- a/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml +++ b/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml @@ -1,4 +1,4 @@ -apiMdSha256: 7e5716fe0862d00fbee0b0076792bd69837ca6c4df2a0833bab80506f170d928 +apiMdSha256: 438f72f9c3b42be5f0813f0a12e688b0a2fc9c1a55cf36b73d9767708e4873cc packageVersion: 0.39.0 parserVersion: 2.2.2 rustVersion: 1.97.0-nightly diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index c4f9c6d396..1813500fe0 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -11,7 +11,8 @@ ### Breaking Changes -- Unordered cross-partition DISTINCT and non-streaming ORDER BY require a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`, sharing 400/20126 (`CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`); the existing non-streaming status constant remains an alias. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) +- Unordered cross-partition DISTINCT and non-streaming ORDER BY require a global finite TOP/LIMIT or explicit `allow_unbounded_queries=true`, sharing 400/20125 (`CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`); the existing non-streaming status constant remains an alias. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) +- Unified client-buffered continuation errors under 400/20124 (`CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED`), retaining the shape-specific constants as aliases; finite-window admission moved from 20126 to 20125 and non-streaming window/storage errors from 20127 to 20126. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) - `error::cosmos_status` is no longer a public module; `CosmosStatus` and `SubStatusCode` remain available as re-exports from `error`. The internal-only `query` module (gated behind the `__internal_testing` feature) is now `#[doc(hidden)]` so it no longer appears as an empty public module in generated API surfaces. ([#5205](https://github.com/Azure/azure-sdk-for-rust/pull/5205)) - Renamed several types for naming consistency: `diagnostics::PipelineType` is now `diagnostics::PipelineKind` (following the `Kind`-over-`Type` convention), `diagnostics::ProxyConfiguration` is now `diagnostics::ProxyConfig` (matching the `Config` naming used elsewhere), and `in_memory_emulator::RuChargingModel` is now `in_memory_emulator::RequestUnitChargingModel` (expanding the `RU` acronym). The unstable `testing` module (`__internal_mocking` feature) was renamed to `test`. ([#5203](https://github.com/Azure/azure-sdk-for-rust/pull/5203)) diff --git a/sdk/cosmos/azure_data_cosmos_driver/api/API.md b/sdk/cosmos/azure_data_cosmos_driver/api/API.md index dc64080ac1..2537edecb6 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/api/API.md +++ b/sdk/cosmos/azure_data_cosmos_driver/api/API.md @@ -636,6 +636,7 @@ pub mod error { impl CosmosStatus { const AUTHENTICATION_TOKEN_ACQUISITION_FAILED: CosmosStatus = _; const CLIENT_BAD_REQUEST: CosmosStatus = _; + const CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: CosmosStatus = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: CosmosStatus = _; @@ -655,7 +656,7 @@ pub mod error { const CLIENT_CROSS_PARTITION_FAN_OUT_EXCEEDED: CosmosStatus = _; const CLIENT_CROSS_PARTITION_QUERY_REQUIRES_CONTAINER_REF: CosmosStatus = _; const CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT: CosmosStatus = _; - const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = _; + const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: CosmosStatus = _; const CLIENT_DRIVER_NOT_INITIALIZED: CosmosStatus = _; const CLIENT_DUPLICATE_FAULT_INJECTION_RULE_ID: CosmosStatus = _; @@ -669,7 +670,7 @@ pub mod error { const CLIENT_INVALID_URL: CosmosStatus = _; const CLIENT_MIXED_NAME_RID_ADDRESSING: CosmosStatus = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: CosmosStatus = _; - const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; + const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: CosmosStatus = _; @@ -793,6 +794,7 @@ pub mod error { const CANNOT_ACQUIRE_PKRANGE_LOCK: SubStatusCode = _; const CHANNEL_CLOSED: SubStatusCode = _; const CHECKPOINT_QUEUE_DEPTH_BACKPRESSURE: SubStatusCode = _; + const CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED: SubStatusCode = _; const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: SubStatusCode = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: SubStatusCode = _; @@ -813,7 +815,7 @@ pub mod error { const CLIENT_CROSS_PARTITION_FAN_OUT_EXCEEDED: SubStatusCode = _; const CLIENT_CROSS_PARTITION_QUERY_REQUIRES_CONTAINER_REF: SubStatusCode = _; const CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT: SubStatusCode = _; - const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: SubStatusCode = _; + const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: SubStatusCode = _; const CLIENT_DRIVER_NOT_INITIALIZED: SubStatusCode = _; const CLIENT_DUPLICATE_FAULT_INJECTION_RULE_ID: SubStatusCode = _; @@ -840,7 +842,7 @@ pub mod error { const CLIENT_INVALID_URL: SubStatusCode = _; const CLIENT_MIXED_NAME_RID_ADDRESSING: SubStatusCode = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: SubStatusCode = _; - const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = _; + const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: SubStatusCode = _; @@ -1798,6 +1800,7 @@ pub mod models { impl CosmosStatus { const AUTHENTICATION_TOKEN_ACQUISITION_FAILED: CosmosStatus = _; const CLIENT_BAD_REQUEST: CosmosStatus = _; + const CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: CosmosStatus = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: CosmosStatus = _; @@ -1817,7 +1820,7 @@ pub mod models { const CLIENT_CROSS_PARTITION_FAN_OUT_EXCEEDED: CosmosStatus = _; const CLIENT_CROSS_PARTITION_QUERY_REQUIRES_CONTAINER_REF: CosmosStatus = _; const CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT: CosmosStatus = _; - const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = _; + const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: CosmosStatus = _; const CLIENT_DRIVER_NOT_INITIALIZED: CosmosStatus = _; const CLIENT_DUPLICATE_FAULT_INJECTION_RULE_ID: CosmosStatus = _; @@ -1831,7 +1834,7 @@ pub mod models { const CLIENT_INVALID_URL: CosmosStatus = _; const CLIENT_MIXED_NAME_RID_ADDRESSING: CosmosStatus = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: CosmosStatus = _; - const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = _; + const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: CosmosStatus = _; @@ -2356,6 +2359,7 @@ pub mod models { const CANNOT_ACQUIRE_PKRANGE_LOCK: SubStatusCode = _; const CHANNEL_CLOSED: SubStatusCode = _; const CHECKPOINT_QUEUE_DEPTH_BACKPRESSURE: SubStatusCode = _; + const CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED: SubStatusCode = _; const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = _; const CLIENT_BUILD_RESPONSE_INVOKED_ON_FAILURE: SubStatusCode = _; const CLIENT_CHANGE_FEED_PIPELINE_UNEXPECTEDLY_DRAINED: SubStatusCode = _; @@ -2376,7 +2380,7 @@ pub mod models { const CLIENT_CROSS_PARTITION_FAN_OUT_EXCEEDED: SubStatusCode = _; const CLIENT_CROSS_PARTITION_QUERY_REQUIRES_CONTAINER_REF: SubStatusCode = _; const CLIENT_DISTINCT_CANNOT_FORWARD_SPLIT: SubStatusCode = _; - const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: SubStatusCode = _; + const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: SubStatusCode = _; const CLIENT_DRIVER_NOT_INITIALIZED: SubStatusCode = _; const CLIENT_DUPLICATE_FAULT_INJECTION_RULE_ID: SubStatusCode = _; @@ -2403,7 +2407,7 @@ pub mod models { const CLIENT_INVALID_URL: SubStatusCode = _; const CLIENT_MIXED_NAME_RID_ADDRESSING: SubStatusCode = _; const CLIENT_NON_MULTIHASH_PARTITION_KEY_ARITY_MISMATCH: SubStatusCode = _; - const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = _; + const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = _; const CLIENT_NO_OVERLAPPING_FEED_RANGES_FOR_SESSION_TOKEN: SubStatusCode = _; diff --git a/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml b/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml index 8a6a611d55..f33a944c94 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml +++ b/sdk/cosmos/azure_data_cosmos_driver/api/API.metadata.yml @@ -1,4 +1,4 @@ -apiMdSha256: 3da0427f1bd17b24b3a701d41d2bed9f2334fe496d8b0991f1b99ca482ffbf86 +apiMdSha256: 83b6120e0c55a436b45a6f7496a4871f9fbaedc437bebd0314043207ec9c018b packageVersion: 0.8.0 parserVersion: 2.2.2 rustVersion: 1.97.0-nightly diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs index bd292498a1..e7ef1e9195 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/distinct.rs @@ -33,7 +33,7 @@ //! Unordered `DISTINCT` is not. The set *is* the state, and serializing it //! would mean an unbounded token; truncating it would silently re-emit //! duplicates. [`Distinct::snapshot_state`] therefore fails with -//! [`CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED`], which surfaces +//! [`CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED`], which surfaces //! at `OperationPlan::to_continuation_token` time — while the caller still //! holds a live plan and can either keep draining in-process or rewrite the //! query with a matching `ORDER BY`. .NET refuses here too, with the same @@ -430,7 +430,7 @@ impl PipelineNode for Distinct { DistinctMap::Ordered { last_hash } => *last_hash, DistinctMap::Unordered { .. } => { return Err(crate::error::CosmosError::builder() - .with_status(CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED) + .with_status(CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED) .with_message(UNORDERED_CONTINUATION_MESSAGE) .build()); } @@ -914,7 +914,7 @@ mod tests { .expect_err("an unordered DISTINCT must not produce a resumable snapshot"); assert_eq!( err.status().sub_status(), - Some(crate::error::SubStatusCode::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED), + Some(crate::error::SubStatusCode::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED), ); assert!( err.to_string().contains("ORDER BY"), diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs index d1aea00571..1d826b175c 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/non_streaming_ordered_merge.rs @@ -235,7 +235,7 @@ impl PipelineNode for NonStreamingOrderedMerge { fn snapshot_state(&self) -> crate::error::Result { Err(CosmosError::builder() - .with_status(CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED) + .with_status(CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED) .with_message( "cross-partition non-streaming ORDER BY queries do not support continuation tokens", ) @@ -491,7 +491,7 @@ mod tests { let node = merge(Vec::new(), 1, 0, 1, None); assert_eq!( node.snapshot_state().unwrap_err().status(), - CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED + CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED ); } @@ -537,7 +537,7 @@ mod tests { assert_eq!(node.retained.capacity(), 0); assert_eq!( node.snapshot_state().unwrap_err().status(), - CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED + CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED ); let mut executor = NoopRequestExecutor; let mut topology = NoopTopologyProvider; diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs index c2c429262e..c35a7f22a8 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs @@ -393,7 +393,7 @@ pub(crate) async fn build_non_streaming_ordered_merge( if resume.is_some() { return Err(crate::error::CosmosError::builder() .with_status( - crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED, + crate::error::CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED, ) .with_message( "cross-partition non-streaming ORDER BY queries cannot be resumed from a continuation token", @@ -1635,7 +1635,7 @@ fn peel_distinct_resume( // checkpoint. return Err(crate::error::CosmosError::builder() .with_status( - crate::error::CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED, + crate::error::CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED, ) .with_message( "continuation token carries unordered DISTINCT state, which cannot be \ @@ -4079,8 +4079,10 @@ mod tests { assert_eq!(actual, expected); assert_eq!(charge, 5.0); assert!(pipeline.next_page(&mut context).await.unwrap().is_none()); - assert_eq!(pipeline.snapshot_state().unwrap_err().status(), - crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED); + assert_eq!( + pipeline.snapshot_state().unwrap_err().status(), + crate::error::CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED + ); assert_eq!( executor.continuation_calls, vec![None, Some("a-next".into()), None, Some("b-next".into())] @@ -4158,7 +4160,7 @@ mod tests { .unwrap_err(); assert_eq!( err.status(), - crate::error::CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED + crate::error::CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED ); let mut streaming_plan = non_streaming_order_by_plan(); @@ -4317,7 +4319,7 @@ mod tests { .expect_err("an unordered DISTINCT token is never resumable"); assert_eq!( err.status().sub_status(), - Some(crate::error::SubStatusCode::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED) + Some(crate::error::SubStatusCode::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED) ); assert!(err.to_string().contains("ORDER BY")); } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs index 8f1edc50da..b1a7b32c5e 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/error/cosmos_status.rs @@ -501,10 +501,9 @@ impl SubStatusCode { 20121 => Some("ClientMixedNameRidAddressing"), 20122 => Some("ClientQueryRewriteBodyInvalid"), 20123 => Some("ClientDistinctValueTooDeeplyNested"), - 20124 => Some("ClientDistinctContinuationUnsupported"), - 20125 => Some("ClientNonStreamingOrderByContinuationUnsupported"), - 20126 => Some("ClientBufferedQueryRequiresFiniteWindow"), - 20127 => Some("ClientNonStreamingOrderByWindowTooLarge"), + 20124 => Some("ClientBufferedQueryContinuationUnsupported"), + 20125 => Some("ClientBufferedQueryRequiresFiniteWindow"), + 20126 => Some("ClientNonStreamingOrderByWindowTooLarge"), 20150 => Some("ClientDuplicateFaultInjectionRuleId"), 20151 => Some("ClientThroughputControlGroupRegistrationFailed"), 20152 => Some("ClientThroughputControlGroupNotRegistered"), @@ -1394,29 +1393,29 @@ impl SubStatusCode { /// this indicates a hand-crafted or corrupt payload. pub const CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED: SubStatusCode = SubStatusCode(20123); - /// A continuation token was requested for an unordered `DISTINCT` query - /// (20124). Resuming would require carrying the entire set of seen values, - /// so the token is refused rather than silently re-emitting duplicates. - /// Adding a matching `ORDER BY` makes the query resumable. - pub const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: SubStatusCode = SubStatusCode(20124); + /// A cross-partition client-buffered query cannot use continuation tokens + /// (20124); its buffered state must be drained in-process. + pub const CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED: SubStatusCode = SubStatusCode(20124); - /// A continuation token was supplied or requested for a non-streaming - /// `ORDER BY` query (20125). Resuming would require serializing the buffered - /// result set, so the operation must be drained in-process. + /// Compatibility alias for [`Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED`]. + pub const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: SubStatusCode = + Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; + + /// Compatibility alias for [`Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED`]. pub const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: SubStatusCode = - SubStatusCode(20125); + Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; /// A buffered query requires a finite global TOP/LIMIT or explicit opt-out - /// (20126), including non-streaming ORDER BY and unordered DISTINCT. - pub const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = SubStatusCode(20126); + /// (20125), including non-streaming ORDER BY and unordered DISTINCT. + pub const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = SubStatusCode(20125); /// Compatibility alias for [`Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`]. pub const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: SubStatusCode = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; /// A non-streaming `ORDER BY` query's candidate window cannot be represented - /// by the current process (20127). - pub const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = SubStatusCode(20127); + /// by the current process (20126). + pub const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = SubStatusCode(20126); // ----- 20150-20199: SDK configuration / setup errors ----- @@ -2383,21 +2382,22 @@ impl CosmosStatus { sub_status: Some(SubStatusCode::CLIENT_DISTINCT_VALUE_TOO_DEEPLY_NESTED), }; - /// 400 / 20124 — a continuation token was requested for an unordered - /// `DISTINCT` query, which cannot be resumed safely. - pub const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = CosmosStatus { + /// 400 / 20124 — continuation tokens are unsupported by the cross-partition + /// client-buffering stage. + pub const CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED: CosmosStatus = CosmosStatus { status_code: StatusCode::BadRequest, - sub_status: Some(SubStatusCode::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED), + sub_status: Some(SubStatusCode::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED), }; - /// 400 / 20125 — continuation tokens are not supported by non-streaming - /// `ORDER BY`. - pub const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = CosmosStatus { - status_code: StatusCode::BadRequest, - sub_status: Some(SubStatusCode::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED), - }; + /// Compatibility alias for [`Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED`]. + pub const CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED: CosmosStatus = + Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; + + /// Compatibility alias for [`Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED`]. + pub const CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED: CosmosStatus = + Self::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; - /// 400 / 20126 — a buffered query requires a finite global TOP/LIMIT or + /// 400 / 20125 — a buffered query requires a finite global TOP/LIMIT or /// explicit opt-out, including non-streaming ORDER BY and unordered DISTINCT. pub const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: CosmosStatus = CosmosStatus { status_code: StatusCode::BadRequest, @@ -2408,7 +2408,7 @@ impl CosmosStatus { pub const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = Self::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW; - /// 400 / 20127 — the non-streaming `ORDER BY` candidate window cannot be + /// 400 / 20126 — the non-streaming `ORDER BY` candidate window cannot be /// represented by the current process. pub const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = CosmosStatus { status_code: StatusCode::BadRequest, @@ -2776,8 +2776,60 @@ mod tests { } #[test] - fn buffered_query_status_preserves_existing_code() { - let status = CosmosStatus::new(StatusCode::BadRequest).with_sub_status(20126); + fn buffered_query_status_codes_and_names() { + for (code, expected, name) in [ + ( + 20124, + CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED, + "ClientBufferedQueryContinuationUnsupported", + ), + ( + 20125, + CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW, + "ClientBufferedQueryRequiresFiniteWindow", + ), + ( + 20126, + CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE, + "ClientNonStreamingOrderByWindowTooLarge", + ), + ] { + let status = CosmosStatus::new(StatusCode::BadRequest).with_sub_status(code); + assert_eq!(status, expected); + assert_eq!(status.name(), Some(name)); + } + assert_eq!( + CosmosStatus::new(StatusCode::BadRequest) + .with_sub_status(20127) + .name(), + None + ); + } + + #[test] + fn buffered_query_continuation_aliases() { + let status = CosmosStatus::CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED; + assert_eq!( + status, + CosmosStatus::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED + ); + assert_eq!( + status, + CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED + ); + assert_eq!( + status.sub_status(), + Some(SubStatusCode::CLIENT_DISTINCT_CONTINUATION_UNSUPPORTED) + ); + assert_eq!( + status.sub_status(), + Some(SubStatusCode::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED) + ); + } + + #[test] + fn buffered_query_admission_aliases() { + let status = CosmosStatus::new(StatusCode::BadRequest).with_sub_status(20125); assert_eq!( status, CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW diff --git a/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md b/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md index af0db2459e..8af721ee4f 100644 --- a/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md +++ b/sdk/cosmos/docs/specs/0006-error-codes-and-retries.md @@ -8,7 +8,9 @@ This document describes the implemented retry behavior for the Azure Cosmos DB R | Status | Symbol | Remedy | | --- | --- | --- | -| 400/20126 | `CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW` | Add a finite global TOP/LIMIT to non-streaming ORDER BY (including buffered vector search) or unordered DISTINCT, or explicitly set `allow_unbounded_queries=true`. | +| 400/20124 | `CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED` | Drain the cross-partition client-buffered query in-process rather than using continuation tokens. | +| 400/20125 | `CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW` | Add a finite global TOP/LIMIT to non-streaming ORDER BY (including buffered vector search) or unordered DISTINCT, or explicitly set `allow_unbounded_queries=true`. | +| 400/20126 | `CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE` | Reduce the non-streaming candidate window or required storage. | The existing `CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW` constant remains a compatibility alias. Both shapes report the symbolic name @@ -16,9 +18,12 @@ remains a compatibility alias. Both shapes report the symbolic name This is a non-retryable client input error. Messages identify the query shape and remedies without SQL or parameter values. No fixed numeric ceiling applies. -400/20127 still reports unrepresentable non-streaming windows/candidate storage. -Continuation restrictions remain 400/20124 (unordered DISTINCT) and 400/20125 -(non-streaming ORDER BY), with or without an opt-out. +400/20126 reports unrepresentable non-streaming windows/candidate storage. +Cross-partition plans using client-side unordered DISTINCT or non-streaming +ORDER BY stages share continuation restriction 400/20124, with or without an +opt-out. Complete logical-partition-key queries bypass these stages and their +client-side continuation restrictions. The existing shape-specific continuation +constants remain aliases of `CLIENT_BUFFERED_QUERY_CONTINUATION_UNSUPPORTED`. The Rust driver retries writes by default for retryable status codes. This is safe because Cosmos DB's write APIs are designed to be idempotent when used correctly: diff --git a/sdk/cosmos/docs/specs/0012-feed-operations-and-dataflow.md b/sdk/cosmos/docs/specs/0012-feed-operations-and-dataflow.md index cf6e271d90..0c32634399 100644 --- a/sdk/cosmos/docs/specs/0012-feed-operations-and-dataflow.md +++ b/sdk/cosmos/docs/specs/0012-feed-operations-and-dataflow.md @@ -171,7 +171,7 @@ Within a resumed full-key tie, RID filtering follows each backend page's `x-ms-c Ordered `DISTINCT` — only `SELECT DISTINCT VALUE … ORDER BY `, the exact shape the service reports as `Ordered` (see below); list projections and multi-column `ORDER BY` stay unordered even when every projected path is covered — deduplicates by adjacency, so it keeps one hash, runs in O(1) memory, and resumes from the 16 bytes `PipelineNodeState::Distinct` persists: a value the stage has moved past can never reappear. This complements rather than duplicates the merge's own resume trim, which is positional (`_rid` + `skipCount`); `last_hash` catches a *different* `_rid` carrying the *same* projected value — two documents that are one `DISTINCT` row but two `ORDER BY` rows. -Unordered `DISTINCT` retains every hash seen (unbounded, ~16 bytes per distinct value) and is **not** resumable: the set *is* the state, serializing it would produce an unbounded token, and truncating it would silently re-emit duplicates. `Distinct::snapshot_state` fails with `400 / 20124 ClientDistinctContinuationUnsupported`, so `OperationPlan::to_continuation_token` errors at mint time — while the caller still holds a live plan and can keep draining in process or rewrite with a matching `ORDER BY`. In-process paging is fully supported. Once drained there is no state left to lose, so the stage snapshots as `Drained` like any other finished node. +Unordered `DISTINCT` retains every hash seen (unbounded, ~16 bytes per distinct value) and is **not** resumable: the set *is* the state, serializing it would produce an unbounded token, and truncating it would silently re-emit duplicates. `Distinct::snapshot_state` fails with `400 / 20124 ClientBufferedQueryContinuationUnsupported`, so `OperationPlan::to_continuation_token` errors at mint time — while the caller still holds a live plan and can keep draining in process or rewrite with a matching `ORDER BY`. In-process paging is fully supported. Once drained there is no state left to lose, so the stage snapshots as `Drained` like any other finished node. The driver executes whatever `distinctType` the plan reports and never upgrades `Unordered` to `Ordered`. The local plan generator (`query::plan`, backing the in-memory emulator) is deliberately stricter than the service planner — see `plan::distinct_is_ordered` — because misclassifying a stream as adjacency-safe drops rows, while the reverse only costs resumability. diff --git a/sdk/cosmos/docs/specs/0013-query-engine.md b/sdk/cosmos/docs/specs/0013-query-engine.md index f575f3ded4..915866bb67 100644 --- a/sdk/cosmos/docs/specs/0013-query-engine.md +++ b/sdk/cosmos/docs/specs/0013-query-engine.md @@ -43,8 +43,10 @@ paginates the remaining rows. Admission is fixed when the plan is built. This is not a runtime memory budget or a guarantee that finite output bounds bound all memory: OFFSET and page-level DISTINCT processing add retained work. -Neither bounds nor opt-out enable unsupported query compositions or continuation -tokens (unordered DISTINCT: 400/20124; non-streaming ORDER BY: 400/20125). +Neither bounds nor opt-out enable unsupported query compositions. Cross-partition +plans using client-side unordered DISTINCT or non-streaming ORDER BY stages +reject continuation tokens with 400/20124. Complete logical-partition-key queries +bypass these stages and their client-side continuation restrictions. Service validation remains authoritative. In particular, a service rejection of a no-TOP vector query is not bypassed or replaced by a fabricated large TOP. Live no-TOP vector support must be verified against an enabled account before From d69c58c9459128be53e27a1c18ea3f65c5757729 Mon Sep 17 00:00:00 2001 From: tvaron3 Date: Wed, 16 Sep 2026 13:58:34 -0400 Subject: [PATCH 5/5] Clarify DISTINCT memory and format query tests Describe DISTINCT state growth under the finite-window policy and explain why configurable state still makes continuation tokens impractical. Expand HPK, vector, and in-memory query-test formatting without changing behavior. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../tests/emulator_tests/cosmos_hpk.rs | 15 ++-- .../emulator_tests/cosmos_vector_query.rs | 79 +++++++++++++------ .../query_comparison.rs | 11 ++- .../src/driver/dataflow/distinct.rs | 4 +- .../0012-feed-operations-and-dataflow.md | 4 +- 5 files changed, 76 insertions(+), 37 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs index 18951377e2..d3dfced2f0 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_hpk.rs @@ -834,12 +834,15 @@ pub async fn hpk_query_cross_partition_advanced_not_servable() -> Result<(), Box // Servable: DISTINCT has a client-side stage, and it must // deduplicate correctly across the container's physical partitions. - let mut countries = container.query_items::( - "SELECT DISTINCT TOP 1000 VALUE c.country FROM c", FeedScope::full_container(), - None, - ) - .await? - .try_collect::>().await?; + let mut countries = container + .query_items::( + "SELECT DISTINCT TOP 1000 VALUE c.country FROM c", + FeedScope::full_container(), + None, + ) + .await? + .try_collect::>() + .await?; countries.sort(); assert_eq!( countries, diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs index a6125aad08..33844cc6e3 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_vector_query.rs @@ -678,48 +678,76 @@ pub async fn finite_vector_query_admission_and_execution() -> Result<(), Box( - query.clone(), FeedScope::full_container(), None, - ).await; + ) + .with_parameter("@queryVector", QUERY_VECTOR.as_slice())?; + let denied = container + .query_items::(query.clone(), FeedScope::full_container(), None) + .await; let error = match denied { Err(error) => error, Ok(_) => panic!("missing global bound must be rejected"), }; // A service rejection is not evidence of client admission. - assert_eq!(error.status(), CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW, - "the service must supply a no-TOP vector plan to validate client admission: {error}"); + assert_eq!( + error.status(), + CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW, + "the service must supply a no-TOP vector plan to validate client admission: {error}" + ); let bounded = Query::from( "SELECT TOP 6 c.id, VectorDistance(c.embedding, @queryVector, true) AS score \ FROM c WHERE c.active = true \ ORDER BY VectorDistance(c.embedding, @queryVector, true)", - ).with_parameter("@queryVector", QUERY_VECTOR.as_slice())?; - let mut pages = container.query_items::( - bounded, FeedScope::full_container(), - Some(QueryOptions::default().with_max_buffered_query_window(6).with_max_item_count( - MaxItemCountHint::Limit(NonZeroU32::new(2).unwrap()), - )), - ).await?.into_pages(); - assert_eq!(pages.to_continuation_token().unwrap_err().status(), - CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED); + ) + .with_parameter("@queryVector", QUERY_VECTOR.as_slice())?; + let mut pages = container + .query_items::( + bounded, + FeedScope::full_container(), + Some( + QueryOptions::default() + .with_max_buffered_query_window(6) + .with_max_item_count(MaxItemCountHint::Limit( + NonZeroU32::new(2).unwrap(), + )), + ), + ) + .await? + .into_pages(); + assert_eq!( + pages.to_continuation_token().unwrap_err().status(), + CosmosStatus::CLIENT_NON_STREAMING_ORDER_BY_CONTINUATION_UNSUPPORTED + ); let mut ids = Vec::new(); while let Some(page) = pages.next().await { ids.extend(page?.into_items().into_iter().map(|item| item.id)); } - assert_eq!(ids, ["origin", "other-partition-origin", "near", - "other-partition-near", "far", "farthest"]); - let mut pages = container.query_items::( - query, FeedScope::partition(SEARCH_PARTITION), - Some(QueryOptions::default().with_max_buffered_query_window(0)), - ).await?.into_pages(); + assert_eq!( + ids, + [ + "origin", + "other-partition-origin", + "near", + "other-partition-near", + "far", + "farthest", + ] + ); + let mut pages = container + .query_items::( + query, + FeedScope::partition(SEARCH_PARTITION), + Some(QueryOptions::default().with_max_buffered_query_window(0)), + ) + .await? + .into_pages(); let mut ids = Vec::new(); while let Some(page) = pages.next().await { ids.extend(page?.into_items().into_iter().map(|item| item.id)); @@ -728,7 +756,8 @@ pub async fn finite_vector_query_admission_and_execution() -> Result<(), Box … ORDER BY `, the exact shape the service reports as `Ordered` (see below); list projections and multi-column `ORDER BY` stay unordered even when every projected path is covered — deduplicates by adjacency, so it keeps one hash, runs in O(1) memory, and resumes from the 16 bytes `PipelineNodeState::Distinct` persists: a value the stage has moved past can never reappear. This complements rather than duplicates the merge's own resume trim, which is positional (`_rid` + `skipCount`); `last_hash` catches a *different* `_rid` carrying the *same* projected value — two documents that are one `DISTINCT` row but two `ORDER BY` rows. -Unordered `DISTINCT` retains every hash seen (unbounded, ~16 bytes per distinct value) and is **not** resumable: the set *is* the state, serializing it would produce an unbounded token, and truncating it would silently re-emit duplicates. `Distinct::snapshot_state` fails with `400 / 20124 ClientBufferedQueryContinuationUnsupported`, so `OperationPlan::to_continuation_token` errors at mint time — while the caller still holds a live plan and can keep draining in process or rewrite with a matching `ORDER BY`. In-process paging is fully supported. Once drained there is no state left to lose, so the stage snapshots as `Drained` like any other finished node. +Unordered `DISTINCT` retains every distinct hash seen (16 bytes per hash, plus hash-set overhead). The admission policy caps global OFFSET plus effective take, and the outer `SkipTake` stops after that many distinct rows. Retained state therefore grows with the admitted window plus page-level processing overhead: `Distinct` processes a whole page before `SkipTake` trims its output. This is not a byte-level memory budget. + +Unordered `DISTINCT` is **not** resumable: the set *is* the state, serializing that configurable set can produce impractically large tokens, and truncating it would silently re-emit duplicates. `Distinct::snapshot_state` fails with `400 / 20124 ClientBufferedQueryContinuationUnsupported`, so `OperationPlan::to_continuation_token` errors at mint time — while the caller still holds a live plan and can keep draining in process or rewrite with a matching `ORDER BY`. In-process paging is fully supported. Once drained there is no state left to lose, so the stage snapshots as `Drained` like any other finished node. The driver executes whatever `distinctType` the plan reports and never upgrades `Unordered` to `Ordered`. The local plan generator (`query::plan`, backing the in-memory emulator) is deliberately stricter than the service planner — see `plan::distinct_is_ordered` — because misclassifying a stream as adjacency-safe drops rows, while the reverse only costs resumability.