diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 504f2e1f1a..0d34ef8244 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -4,13 +4,17 @@ ### Features Added +- Added per-query `QueryOptions::max_buffered_query_window` and `with_max_buffered_query_window` to configure the maximum global OFFSET plus effective take for client-buffered queries (default 1000, no opt-out). ([#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 `QueryPlanMode::{LocalPreferred, GatewayOnly}`, allowing applications to force Gateway query planning 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)) - Extended Cosmos binary JSON encoding to the thin client (Gateway 2.0) transport. ([#5284](https://github.com/Azure/azure-sdk-for-rust/pull/5284)) ### Breaking Changes +- Moved `query_plan_mode` from `OperationOptions` to per-query `QueryOptions`; removed client/runtime defaults and environment settings, including the query-plan-mode override. The default remains `LocalPreferred`. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) +- Unordered cross-partition DISTINCT and non-streaming ORDER BY require finite global TOP/LIMIT with OFFSET plus effective take within the configured maximum; missing bounds, excess windows, and overflow fail with 400/20125. ([#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)) - Reviewed public API type consistency for time durations and integer sizes. `ResponseHeaders::server_duration_ms()` and `retry_after_ms()` are replaced by `server_duration()` and `retry_after()`, both returning `Option`; `TransactionalBatchOperationResult::retry_after_milliseconds()` and `DistributedTransactionResponse::retry_after_ms()` are similarly replaced by `retry_after() -> Option`. `ThroughputProperties::manual`, `autoscale`, `throughput()`, and `autoscale_maximum()` now use `u64` instead of the platform-dependent `usize`, allowing RU/s values above 4 billion; `autoscale_increment()` now uses `u32` because it returns a percentage. ([#5204](https://github.com/Azure/azure-sdk-for-rust/pull/5204)) ### Bugs Fixed diff --git a/sdk/cosmos/azure_data_cosmos/Cargo.toml b/sdk/cosmos/azure_data_cosmos/Cargo.toml index 863b6312a5..1ca8e057a3 100644 --- a/sdk/cosmos/azure_data_cosmos/Cargo.toml +++ b/sdk/cosmos/azure_data_cosmos/Cargo.toml @@ -140,6 +140,11 @@ name = "emulator" path = "tests/emulator.rs" required-features = ["key_auth", "control_plane", "fault_injection"] +[[test]] +name = "live" +path = "tests/live.rs" +required-features = ["key_auth", "control_plane", "fault_injection"] + [[test]] name = "multi_write" path = "tests/multi_write.rs" diff --git a/sdk/cosmos/azure_data_cosmos/api/API.md b/sdk/cosmos/azure_data_cosmos/api/API.md index 21373405e0..53b828be13 100644 --- a/sdk/cosmos/azure_data_cosmos/api/API.md +++ b/sdk/cosmos/azure_data_cosmos/api/API.md @@ -1208,6 +1208,8 @@ 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 = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: CosmosStatus = _; @@ -1226,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 = _; @@ -1240,8 +1242,8 @@ 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_REQUIRES_FINITE_WINDOW: 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 = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: CosmosStatus = _; @@ -2456,7 +2458,6 @@ pub mod options { #[derive(Clone, Debug, Default)] #[non_exhaustive] pub struct OperationOptions { - pub query_plan_mode: Option, pub patch_strategy: Option, pub read_consistency_strategy: Option, pub excluded_regions: Option, @@ -2500,7 +2501,6 @@ pub mod options { pub fn with_max_failover_retry_count(self, value: u32) -> Self; pub fn with_max_session_retry_count(self, value: u32) -> Self; pub fn with_patch_strategy(self, value: PatchStrategy) -> Self; - pub fn with_query_plan_mode(self, value: QueryPlanMode) -> Self; pub fn with_read_consistency_strategy(self, value: ReadConsistencyStrategy) -> Self; pub fn with_session_capturing_disabled(self, value: bool) -> Self; pub fn with_throttling_retry_options(self, value: ThrottlingRetryOptions) -> Self; @@ -2526,7 +2526,6 @@ pub mod options { pub fn new(env: Option<::std::sync::Arc>, runtime: Option<::std::sync::Arc>, account: Option<::std::sync::Arc>, operation: Option<&'a OperationOptions>) -> Self; pub fn new_with_override(env_override: Option<::std::sync::Arc>, env: Option<::std::sync::Arc>, runtime: Option<::std::sync::Arc>, account: Option<::std::sync::Arc>, operation: Option<&'a OperationOptions>) -> Self; pub fn patch_strategy(&self) -> Option<&PatchStrategy>; - pub fn query_plan_mode(&self) -> Option<&QueryPlanMode>; pub fn read_consistency_strategy(&self) -> Option<&ReadConsistencyStrategy>; pub fn session_capturing_disabled(&self) -> Option<&bool>; pub fn throttling_retry_options(&self) -> ThrottlingRetryOptionsView<'_>; @@ -2613,9 +2612,11 @@ pub mod options { impl QueryDatabasesOptions { pub fn with_operation_options(self, operation: OperationOptions) -> Self; } - #[derive(Clone, Default)] + #[derive(Clone)] #[non_exhaustive] pub struct QueryOptions { + pub max_buffered_query_window: u64, + pub query_plan_mode: crate::options::QueryPlanMode, pub operation: azure_data_cosmos_driver::options::OperationOptions, pub feed: FeedOptions, pub session_token: Option, @@ -2625,12 +2626,17 @@ pub mod options { impl QueryOptions { pub fn with_continuation_token(self, continuation_token: ContinuationToken) -> Self; pub fn with_feed_options(self, feed: FeedOptions) -> Self; + pub fn with_max_buffered_query_window(self, max_buffered_query_window: u64) -> Self; pub fn with_max_item_count(self, max_item_count: MaxItemCountHint) -> Self; pub fn with_operation_options(self, operation: OperationOptions) -> Self; pub fn with_populate_index_metrics(self, enable: bool) -> Self; pub fn with_populate_query_metrics(self, enable: bool) -> Self; + pub fn with_query_plan_mode(self, mode: QueryPlanMode) -> Self; pub fn with_session_token: Into>(self, session_token: impl Into) -> Self; } + impl Default for QueryOptions { + fn default() -> Self; + } #[derive(Clone, Default)] #[non_exhaustive] pub struct ReadContainerOptions { diff --git a/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml b/sdk/cosmos/azure_data_cosmos/api/API.metadata.yml index 2238e6b882..b6a492ce02 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: 9cc6b0b348bf227b6e61e7211e3c762630fbcb14a45985bededb0ee3e81e673c +apiMdSha256: 5bbb1f703332631ab748f09624547b3fd27afa80b66ab76a200f553bf81767bc packageVersion: 0.39.0 parserVersion: 2.2.2 rustVersion: 1.97.0-nightly diff --git a/sdk/cosmos/azure_data_cosmos/build.rs b/sdk/cosmos/azure_data_cosmos/build.rs index 60b4f930f4..133a31f232 100644 --- a/sdk/cosmos/azure_data_cosmos/build.rs +++ b/sdk/cosmos/azure_data_cosmos/build.rs @@ -10,7 +10,7 @@ fn main() { // Allow `#[cfg_attr(not(test_category = "..."), ignore)]` in `tests/*.rs`. println!( - "cargo:rustc-check-cfg=cfg(test_category, values(\"emulator\", \"emulator_vnext\", \"emulator_inmemory\", \"emulator_inmemory_gateway_v2\", \"multi_write\", \"split\", \"merge\", \"binary_encoding\", \"gateway_v2\", \"gateway_v2_multi_region\"))" + "cargo:rustc-check-cfg=cfg(test_category, values(\"live\", \"emulator\", \"emulator_vnext\", \"emulator_inmemory\", \"emulator_inmemory_gateway_v2\", \"multi_write\", \"split\", \"merge\", \"binary_encoding\", \"gateway_v2\", \"gateway_v2_multi_region\"))" ); // Marker cfg set by test setups where the target Cosmos account is provisioned // for AAD data-plane access (local emulator started with /enableaadauthentication, diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs index 34174e2c49..9943104547 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs @@ -442,7 +442,7 @@ mod tests { use super::*; use crate::{ - options::{PartitionFailoverOptions, QueryPlanMode, Region, UserAgentSuffix}, + options::{PartitionFailoverOptions, Region, UserAgentSuffix}, RoutingStrategy, }; @@ -579,18 +579,6 @@ mod tests { assert_eq!(opts.preferred_regions(), input.as_slice()); } - #[test] - fn query_plan_mode_flows_to_driver_options() { - let mut input = test_driver_options_input(RoutingStrategy::PreferredRegions(Vec::new())); - input.operation_options.query_plan_mode = Some(QueryPlanMode::GatewayOnly); - let opts = input.build().expect("driver options should build"); - - assert_eq!( - opts.operation_options().query_plan_mode, - Some(QueryPlanMode::GatewayOnly) - ); - } - /// The user-agent suffix must flow through to the per-driver options so /// the driver builds a User-Agent that overrides the runtime default. #[test] diff --git a/sdk/cosmos/azure_data_cosmos/src/options/feed.rs b/sdk/cosmos/azure_data_cosmos/src/options/feed.rs index 7bca2495fd..80cabaa141 100644 --- a/sdk/cosmos/azure_data_cosmos/src/options/feed.rs +++ b/sdk/cosmos/azure_data_cosmos/src/options/feed.rs @@ -4,9 +4,11 @@ //! Feed/query options: paging, query metrics, and continuation tokens. use azure_data_cosmos_driver::models::{MaxItemCountHint, SessionToken}; -use azure_data_cosmos_driver::options::{OperationOptions, PlanOptions, DEFAULT_MAX_FAN_OUT}; +use azure_data_cosmos_driver::options::{ + OperationOptions, PlanOptions, DEFAULT_MAX_BUFFERED_QUERY_WINDOW, DEFAULT_MAX_FAN_OUT, +}; -use crate::feed::ContinuationToken; +use crate::{feed::ContinuationToken, options::QueryPlanMode}; /// Options that apply to feed-style operations (paged reads, queries, etc.). /// @@ -119,9 +121,20 @@ impl FeedOptions { /// [`with_max_item_count`](Self::with_max_item_count) and /// [`with_continuation_token`](Self::with_continuation_token) delegate to the inner /// [`FeedOptions`]. -#[derive(Clone, Default)] +#[derive(Clone)] #[non_exhaustive] pub struct QueryOptions { + /// Maximum global OFFSET plus effective take for client-buffered queries. + /// + /// Requires a finite TOP or LIMIT; when both exist, the smaller is used. + /// Defaults to 1000. Zero is a valid limit. + pub max_buffered_query_window: u64, + + /// Query-plan provider selection for this query. + /// + /// Defaults to [`QueryPlanMode::LocalPreferred`]. + pub query_plan_mode: QueryPlanMode, + /// General-purpose options that apply to this request. /// See [`OperationOptions`] for available settings and layered resolution behavior. pub operation: OperationOptions, @@ -144,7 +157,33 @@ pub struct QueryOptions { pub populate_query_metrics: Option, } +impl Default for QueryOptions { + fn default() -> Self { + Self { + max_buffered_query_window: DEFAULT_MAX_BUFFERED_QUERY_WINDOW, + query_plan_mode: QueryPlanMode::default(), + operation: OperationOptions::default(), + feed: FeedOptions::default(), + session_token: None, + populate_index_metrics: None, + populate_query_metrics: None, + } + } +} + impl QueryOptions { + /// Sets the maximum global OFFSET plus effective take for client-buffered queries. + pub fn with_max_buffered_query_window(mut self, max_buffered_query_window: u64) -> Self { + self.max_buffered_query_window = max_buffered_query_window; + self + } + + /// Sets the query-plan provider selection for this query. + pub fn with_query_plan_mode(mut self, mode: QueryPlanMode) -> Self { + self.query_plan_mode = mode; + 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()); @@ -196,13 +235,43 @@ impl QueryOptions { } pub(crate) fn to_plan_options(&self) -> PlanOptions { - self.feed.to_plan_options() + self.feed + .to_plan_options() + .with_max_buffered_query_window(self.max_buffered_query_window) + .with_query_plan_mode(self.query_plan_mode) } } #[cfg(test)] mod tests { - use super::*; + use super::{FeedOptions, QueryOptions, QueryPlanMode, DEFAULT_MAX_FAN_OUT}; + + #[test] + fn query_plan_options_preserve_defaults() { + let plan = QueryOptions::default().to_plan_options(); + assert_eq!(plan.max_buffered_query_window, 1000); + assert_eq!(plan.query_plan_mode, QueryPlanMode::LocalPreferred); + assert_eq!(plan.max_fan_out, DEFAULT_MAX_FAN_OUT); + } + + #[test] + fn query_plan_options_map_query_fields_and_feed_limits() { + for maximum in [0, 999, 1001, u64::MAX] { + for mode in [QueryPlanMode::LocalPreferred, QueryPlanMode::GatewayOnly] { + for (max_fan_out, expected) in [(0, DEFAULT_MAX_FAN_OUT), (250, 250)] { + let options = QueryOptions::default() + .with_max_buffered_query_window(100) + .with_max_buffered_query_window(maximum) + .with_query_plan_mode(mode) + .with_feed_options(FeedOptions::default().with_max_fan_out(max_fan_out)); + let plan = options.to_plan_options(); + assert_eq!(plan.max_buffered_query_window, maximum); + assert_eq!(plan.query_plan_mode, mode); + assert_eq!(plan.max_fan_out, expected); + } + } + } + } #[test] fn plan_options_uses_default_fan_out_when_unset() { 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..50254d8c37 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,21 @@ async fn query_values( sql: &str, run_id: &str, context: &str, + max_buffered_query_window: u64, ) -> 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_max_buffered_query_window(max_buffered_query_window), + ), + ) .await { Ok(iterator) => Box::pin(iterator.try_collect()).await, @@ -2383,17 +2391,17 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { } } + let query_window = u64::try_from(clients.len())?; + let distinct_query = format!( + "SELECT DISTINCT TOP {query_window} VALUE c._sampler.int FROM c WHERE c.fuzzRun = @run" + ); let query_cases = [ ( "SELECT * FROM c WHERE c.fuzzRun = @run", false, "select-all", ), - ( - "SELECT DISTINCT VALUE c._sampler.int FROM c WHERE c.fuzzRun = @run", - false, - "distinct", - ), + (distinct_query.as_str(), false, "distinct"), ( "SELECT DISTINCT VALUE c._sampler.int FROM c WHERE c.fuzzRun = @run \ ORDER BY c._sampler.int", @@ -2410,7 +2418,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, query_window).await?, ordered, ); if let Some(expected) = &expected { @@ -2440,6 +2448,7 @@ async fn binary_encoding_roundtrip_fuzz() -> Result<(), Box> { "SELECT VALUE {\"int\": 7} FROM c WHERE c.fuzzRun = @run", &run_id, &context, + query_window, ) .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..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,15 +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 = collect_query::( - &container, - "SELECT DISTINCT VALUE c.country FROM c", - FeedScope::full_container(), - ) - .await? - .into_iter() - .map(|v| v.as_str().unwrap_or_default().to_owned()) - .collect::>(); + 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_query.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query.rs index 047912f996..c94583d54a 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 @@ -404,7 +404,7 @@ pub async fn cross_partition_query_with_unordered_distinct() -> Result<(), Box( - "select distinct value c.partitionKey from c", + "select distinct top 1000 value c.partitionKey from c", FeedScope::full_container(), Some( QueryOptions::default().with_max_item_count(MaxItemCountHint::Limit( @@ -515,7 +515,7 @@ pub async fn unordered_distinct_refuses_a_continuation_token() -> Result<(), Box let mut pages = container_client .query_items::( - "select distinct value c.partitionKey from c", + "select distinct top 1000 value c.partitionKey from c", FeedScope::full_container(), Some( QueryOptions::default().with_max_item_count(MaxItemCountHint::Limit( @@ -1116,7 +1116,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, + None, + ) .await? .into_pages(); let mut n = 0; @@ -1131,7 +1135,7 @@ pub async fn distinct_projection_shapes() -> Result<(), Box> { assert_eq!( count( &container, - "select distinct * from c", + "select distinct top 1000 * from c", FeedScope::full_container() ) .await?, @@ -1146,7 +1150,7 @@ pub async fn distinct_projection_shapes() -> Result<(), Box> { assert_eq!( count( &container, - "select distinct value 1 from c", + "select distinct top 1000 value 1 from c", FeedScope::full_container() ) .await?, @@ -1159,7 +1163,7 @@ pub async fn distinct_projection_shapes() -> Result<(), Box> { count( &container, Query::from( - "select distinct value c.partitionKey from c where c.mergeOrder >= @m" + "select distinct top 1000 value c.partitionKey from c where c.mergeOrder >= @m" ) .with_parameter("@m", 0)?, FeedScope::full_container() @@ -1222,13 +1226,17 @@ pub async fn distinct_combined_with_unsupported_stages_is_rejected() -> Result<( // longer here — `SkipTake` composes above `DISTINCT`, so those // shapes are servable and are asserted positively below. let unsupported = [ - "select distinct c.partitionKey, count(1) as n from c group by c.partitionKey", - "select distinct value max(c.mergeOrder) from c", + "select distinct top 1000 c.partitionKey, count(1) as n from c group by c.partitionKey", + "select distinct top 1000 value max(c.mergeOrder) from c", ]; for query in unsupported { let outcome = container - .query_items::(query, FeedScope::full_container(), None) + .query_items::( + query, + FeedScope::full_container(), + None, + ) .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 6134e97ce2..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 @@ -666,6 +666,100 @@ 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 finite_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. + 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 + ); + 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(); + 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/binary_round_trip.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/binary_round_trip.rs index dda44ba794..bc3f0e835a 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/binary_round_trip.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/binary_round_trip.rs @@ -16,7 +16,7 @@ use azure_data_cosmos::{ options::{ BinaryEncodingOptions, ContentResponseOnWrite, ItemWriteOptions, OperationOptions, - OperationOptionsBuilder, QueryPlanMode, Region, RoutingStrategy, + QueryOptions, QueryPlanMode, Region, RoutingStrategy, }, AccountEndpoint, AccountReference, ContainerClient, CosmosClientBuilder, CosmosRuntimeBuilder, FeedScope, Query, @@ -144,18 +144,12 @@ async fn build_multi_partition_container_with_recorder( azure_core::credentials::Secret::new("dGVzdGtleQ=="), ); // `None` leaves the binary option unset, exercising the resolved default. - let mut builder = CosmosClientBuilder::new() - .with_runtime( - CosmosRuntimeBuilder::from(emulator.runtime_builder()) - .build() - .await - .unwrap(), - ) - .with_default_operation_options( - OperationOptionsBuilder::new() - .with_query_plan_mode(QueryPlanMode::GatewayOnly) - .build(), - ); + let mut builder = CosmosClientBuilder::new().with_runtime( + CosmosRuntimeBuilder::from(emulator.runtime_builder()) + .build() + .await + .unwrap(), + ); if let Some(binary) = binary { builder = builder.with_binary_encoding_options(BinaryEncodingOptions::new().with_enabled(binary)); @@ -624,7 +618,7 @@ async fn binary_cross_partition_query_round_trips() { let iter = Box::pin(container.query_items( Query::from("SELECT * FROM c"), FeedScope::full_container(), - None, + Some(QueryOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly)), )) .await .unwrap(); @@ -674,7 +668,7 @@ async fn binary_cross_partition_order_by_merges_and_round_trips() { let iter = Box::pin(container.query_items( Query::from("SELECT * FROM c ORDER BY c.value"), FeedScope::full_container(), - None, + Some(QueryOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly)), )) .await .unwrap(); @@ -729,7 +723,7 @@ async fn binary_cross_partition_skip_take_round_trips() { let offset_limit = Box::pin(container.query_items::( Query::from("SELECT * FROM c OFFSET 2 LIMIT 3"), FeedScope::full_container(), - None, + Some(QueryOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly)), )) .await .unwrap(); @@ -743,7 +737,7 @@ async fn binary_cross_partition_skip_take_round_trips() { let topped = Box::pin(container.query_items::( Query::from("SELECT TOP 4 * FROM c"), FeedScope::full_container(), - None, + Some(QueryOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly)), )) .await .unwrap(); @@ -835,7 +829,7 @@ async fn disabled_binary_query_advertises_no_format() { let iter = Box::pin(container.query_items::( Query::from("SELECT * FROM c"), FeedScope::full_container(), - None, + Some(QueryOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly)), )) .await .unwrap(); @@ -881,7 +875,7 @@ async fn default_client_negotiates_binary_without_any_option() { let iter = Box::pin(container.query_items::( Query::from("SELECT * FROM c"), FeedScope::full_container(), - None, + Some(QueryOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly)), )) .await .unwrap(); 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..fb45b4028f 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 @@ -19,7 +19,7 @@ use azure_data_cosmos::{ models::{ContainerProperties, PartitionKeyDefinition, PartitionKeyVersion}, options::{ ConnectionPoolOptions, ExcludedRegions, MaxItemCountHint, OperationOptions, QueryOptions, - Region, ServerCertificateValidation, + QueryPlanMode, Region, ServerCertificateValidation, }, AccountEndpoint, AccountReference, ContainerClient, CosmosClient, CosmosClientBuilder, CosmosRuntimeBuilder, FeedScope, PartitionKey, Query, RoutingStrategy, @@ -310,6 +310,110 @@ enum FixtureKind { Hpk, } +#[tokio::test] +async fn buffered_query_policy_per_query_options_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 mode in [QueryPlanMode::LocalPreferred, QueryPlanMode::GatewayOnly] { + let client = CosmosClientBuilder::new() + .with_runtime( + CosmosRuntimeBuilder::from(harness.emulator_http.runtime_builder()) + .build() + .await?, + ) + .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 (sql, maximum, denied) in [ + (query, None, true), + (query, Some(u64::MAX), true), + ("SELECT DISTINCT TOP 1000 VALUE c.value FROM c", None, false), + ("SELECT DISTINCT TOP 1001 VALUE c.value FROM c", None, true), + ( + "SELECT DISTINCT TOP 1001 VALUE c.value FROM c", + Some(1001), + false, + ), + ( + "SELECT DISTINCT TOP 1000 VALUE c.value FROM c", + Some(999), + true, + ), + ] { + let mut options = QueryOptions::default() + .with_query_plan_mode(mode) + .with_max_item_count(MaxItemCountHint::Limit(NonZeroU32::new(1).unwrap())); + if let Some(maximum) = maximum { + options = options.with_max_buffered_query_window(maximum); + } + let result = container + .query_items::(sql, scope.clone(), Some(options)) + .await; + if buffered && denied { + 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/live.rs b/sdk/cosmos/azure_data_cosmos/tests/live.rs new file mode 100644 index 0000000000..14b1ccbbec --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/tests/live.rs @@ -0,0 +1,9 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +// Integration tests compose large Cosmos operation futures on tokio test threads. +#![allow(clippy::large_futures)] + +#[path = "live_tests/cosmos_query.rs"] +mod cosmos_query; +mod framework; diff --git a/sdk/cosmos/azure_data_cosmos/tests/live_tests/cosmos_query.rs b/sdk/cosmos/azure_data_cosmos/tests/live_tests/cosmos_query.rs new file mode 100644 index 0000000000..0ee1d3d262 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/tests/live_tests/cosmos_query.rs @@ -0,0 +1,160 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +use crate::framework::{self, test_client::TEST_MODE_ENV_VAR, test_data, TestClient, TestOptions}; +use azure_data_cosmos::{ + feed::FeedScope, + options::{ + BinaryEncodingOptions, MaxItemCountHint, OperationOptions, QueryOptions, QueryPlanMode, + Region, + }, + AccountReference, CosmosClient, CosmosStatus, Query, RoutingStrategy, +}; +use futures::StreamExt; +use std::error::Error; + +#[tokio::test] +#[cfg_attr(not(test_category = "live"), ignore = "requires live account")] +async fn live_distinct_admission_and_per_query_options() -> Result<(), Box> { + let connection = framework::resolve_connection_string() + .ok_or("live tests require a valid AZURE_COSMOS_CONNECTION_STRING")?; + assert!( + !framework::targets_emulator(), + "live DISTINCT admission coverage requires a live account, not an emulator" + ); + assert!( + !std::env::var(TEST_MODE_ENV_VAR).is_ok_and(|mode| mode.eq_ignore_ascii_case("skipped")), + "explicitly selected live tests cannot use AZURE_COSMOS_TEST_MODE=skipped" + ); + 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 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] { + let client = CosmosClient::builder() + .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_query_plan_mode(mode) + .with_operation_options(operation) + .with_max_item_count(MaxItemCountHint::Limit( + std::num::NonZeroU32::new(1).unwrap(), + )); + for (query, maximum, denied) in [ + (unbounded, None, true), + (unbounded, Some(u64::MAX), true), + ( + "SELECT DISTINCT TOP 1001 VALUE c.partitionKey FROM c", + None, + true, + ), + ( + "SELECT DISTINCT TOP 1000 VALUE c.partitionKey FROM c", + None, + false, + ), + ( + "SELECT DISTINCT TOP 1001 VALUE c.partitionKey FROM c", + Some(1001), + false, + ), + ( + "SELECT DISTINCT TOP 4 VALUE c.partitionKey FROM c", + Some(3), + true, + ), + ( + "SELECT DISTINCT TOP 4 VALUE c.partitionKey FROM c", + Some(4), + false, + ), + ] { + let mut options = options.clone(); + if let Some(maximum) = maximum { + options = options.with_max_buffered_query_window(maximum); + } + let result = container + .query_items::( + query, + FeedScope::full_container(), + Some(options), + ) + .await; + if denied { + let error = match result { + Err(error) => error, + Ok(_) => { + panic!("DISTINCT outside finite window must fail at admission") + } + }; + assert_eq!( + error.status(), + CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW, + "{mode:?}, maximum={maximum:?}, 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_max_buffered_query_window(3)), + ) + .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 +} 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..c9155c7d22 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 @@ -51,7 +51,7 @@ const PK_COUNT: usize = 40; const GROUP_COUNT: usize = 8; const PAGE_SIZE: u32 = 5; -const UNORDERED_QUERY: &str = "SELECT DISTINCT VALUE c.groupKey FROM c"; +const UNORDERED_QUERY: &str = "SELECT DISTINCT TOP 1000 VALUE c.groupKey FROM c"; const ORDERED_QUERY: &str = "SELECT DISTINCT VALUE c.groupKey FROM c ORDER BY c.groupKey"; /// `DISTINCT` composed under a global row window. `GROUP_COUNT` distinct keys /// exist, so skipping one and taking two must yield exactly two — the window diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index 1db9fa2e81..2e49297dcd 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -4,13 +4,18 @@ ### Features Added +- Added per-query `PlanOptions::max_buffered_query_window`, `with_max_buffered_query_window`, and `DEFAULT_MAX_BUFFERED_QUERY_WINDOW` (1000) to cap global OFFSET plus effective take for client-buffered queries, with no opt-out. ([#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(...)`. 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 local Rust query planning for supported cross-partition queries, avoiding Gateway query-plan requests while retaining native and Gateway fallbacks for advanced query shapes, and `QueryPlanMode::{LocalPreferred, GatewayOnly}` to select providers per query. ([#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)) - Added Cosmos binary JSON encoding support over the thin client (Gateway 2.0) path by forwarding the `x-ms-cosmos-supported-serialization-formats` header as the RNTBD `SupportedSerializationFormats` token. ([#5284](https://github.com/Azure/azure-sdk-for-rust/pull/5284)) ### Breaking Changes +- Moved `query_plan_mode` from `OperationOptions` to `PlanOptions`; removed account/runtime defaults and environment settings, including the query-plan-mode override. The default remains `LocalPreferred`. ([#5301](https://github.com/Azure/azure-sdk-for-rust/pull/5301)) +- Unordered cross-partition DISTINCT and non-streaming ORDER BY require finite global TOP/LIMIT with OFFSET plus effective take within the configured maximum; missing bounds, excess windows, and overflow share 400/20125 (`CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW`). ([#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)) - `CosmosRequestHeaders::offer_throughput`, `OfferAutoscaleSettings::max_throughput`, `OfferAutoscaleSettings::new`, `OfferAutoscaleSettings::with_increment_percent`, and `AutoscaleThroughputPolicy::increment_percent` now use `u32` instead of the platform-dependent `usize`, matching the RU/s values Cosmos DB actually returns. ([#5204](https://github.com/Azure/azure-sdk-for-rust/pull/5204)) - 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 be4bc759ef..8266973056 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/api/API.md +++ b/sdk/cosmos/azure_data_cosmos_driver/api/API.md @@ -636,6 +636,8 @@ 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 = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: CosmosStatus = _; @@ -654,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 = _; @@ -668,8 +670,8 @@ 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_REQUIRES_FINITE_WINDOW: 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 = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: CosmosStatus = _; @@ -792,6 +794,8 @@ 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 = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: SubStatusCode = _; @@ -811,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 = _; @@ -838,8 +842,8 @@ 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_REQUIRES_FINITE_WINDOW: 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 = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: SubStatusCode = _; @@ -1796,6 +1800,8 @@ 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 = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: CosmosStatus = _; @@ -1814,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 = _; @@ -1828,8 +1834,8 @@ 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_REQUIRES_FINITE_WINDOW: 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 = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: CosmosStatus = _; @@ -2353,6 +2359,8 @@ 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 = _; const CLIENT_COMPUTE_RANGE_INVOKED_WITH_EMPTY_PARTITION_KEY: SubStatusCode = _; @@ -2372,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 = _; @@ -2399,8 +2407,8 @@ 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_REQUIRES_FINITE_WINDOW: 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 = _; const CLIENT_NO_THROUGHPUT_OFFER_FOR_RESOURCE: SubStatusCode = _; @@ -3261,7 +3269,6 @@ pub mod options { #[derive(Clone, Debug, Default)] #[non_exhaustive] pub struct OperationOptions { - pub query_plan_mode: Option, pub patch_strategy: Option, pub read_consistency_strategy: Option, pub excluded_regions: Option, @@ -3302,7 +3309,6 @@ pub mod options { pub fn with_max_failover_retry_count(self, value: u32) -> Self; pub fn with_max_session_retry_count(self, value: u32) -> Self; pub fn with_patch_strategy(self, value: PatchStrategy) -> Self; - pub fn with_query_plan_mode(self, value: QueryPlanMode) -> Self; pub fn with_read_consistency_strategy(self, value: ReadConsistencyStrategy) -> Self; pub fn with_session_capturing_disabled(self, value: bool) -> Self; pub fn with_throttling_retry_options(self, value: ThrottlingRetryOptions) -> Self; @@ -3326,7 +3332,6 @@ pub mod options { pub fn new(env: Option<::std::sync::Arc>, runtime: Option<::std::sync::Arc>, account: Option<::std::sync::Arc>, operation: Option<&'a OperationOptions>) -> Self; pub fn new_with_override(env_override: Option<::std::sync::Arc>, env: Option<::std::sync::Arc>, runtime: Option<::std::sync::Arc>, account: Option<::std::sync::Arc>, operation: Option<&'a OperationOptions>) -> Self; pub fn patch_strategy(&self) -> Option<&PatchStrategy>; - pub fn query_plan_mode(&self) -> Option<&QueryPlanMode>; pub fn read_consistency_strategy(&self) -> Option<&ReadConsistencyStrategy>; pub fn session_capturing_disabled(&self) -> Option<&bool>; pub fn throttling_retry_options(&self) -> ThrottlingRetryOptionsView<'_>; @@ -3367,10 +3372,14 @@ pub mod options { #[derive(Clone, Debug)] #[non_exhaustive] pub struct PlanOptions { + pub max_buffered_query_window: u64, + pub query_plan_mode: crate::options::QueryPlanMode, pub max_fan_out: u32, } impl PlanOptions { + pub fn with_max_buffered_query_window(self, max_buffered_query_window: u64) -> Self; pub fn with_max_fan_out(self, max_fan_out: u32) -> Self; + pub fn with_query_plan_mode(self, mode: QueryPlanMode) -> Self; } impl Default for PlanOptions { fn default() -> Self; @@ -3747,6 +3756,7 @@ pub mod options { #[default] Rustls, } + pub const DEFAULT_MAX_BUFFERED_QUERY_WINDOW: u64 = 1000; pub const DEFAULT_MAX_CONCURRENT_METADATA_ATTEMPTS: usize = 32; pub const DEFAULT_MAX_FAN_OUT: u32 = 100; } 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 a68fd1ee0f..17b21e1fe0 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: 382404332bc1317f824c07ead0a648902fcbd379f00e49edbcda5ffffdfcb775 +apiMdSha256: 5d073cd5b049c3b73a81766792dfc182d279773df423cf96a96f19efb9c8680a packageVersion: 0.8.0 parserVersion: 2.2.2 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 5a782e98b0..c53f26f6ff 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 @@ -42,7 +42,7 @@ use crate::{ }, options::{ ConnectionPoolOptions, DriverOptions, OperationOptions, OperationOptionsView, PlanOptions, - QueryPlanMode, ResolvedThroughputControl, ThroughputControlGroupSnapshot, + ResolvedThroughputControl, ThroughputControlGroupSnapshot, }, ActivityId, CosmosResponse, DiagnosticsContext, }; @@ -2169,13 +2169,6 @@ impl CosmosDriver { ) } - fn effective_query_plan_mode(&self, options: &OperationOptions) -> QueryPlanMode { - self.operation_options_view(options) - .query_plan_mode() - .copied() - .unwrap_or_default() - } - /// Computes the effective throughput-control header values for an operation. /// /// Resolves the per-request `x-ms-cosmos-throughput-bucket` and @@ -4157,7 +4150,7 @@ impl CosmosDriver { operation_type = ?operation.operation_type(), resource_type = ?operation.resource_type(), resource_reference = ?operation.resource_reference(), - query_plan_mode = ?self.effective_query_plan_mode(options), + query_plan_mode = ?plan_options.query_plan_mode, "planning operation" ); @@ -4264,7 +4257,9 @@ impl CosmosDriver { Err(error) => { if matches!( query_planning::try_resolve_without_topology( - self, container, &operation, options, + container, + &operation, + plan_options, ), Some(ResolvedQueryPlan::Empty) ) { @@ -4278,7 +4273,11 @@ impl CosmosDriver { // `Box::pin` keeps `plan_operation`'s future small. Inlined, it grows to // 17,288 bytes and trips `clippy::large_futures` at five caller sites. let resolved = Box::pin(query_planning::resolve_query_plan( - self, container, &operation, options, + self, + container, + &operation, + options, + plan_options, )) .await?; @@ -4292,6 +4291,8 @@ impl CosmosDriver { ResolvedQueryPlan::Plan(plan) => *plan, }; + planner::validate_buffered_query(&query_plan, plan_options.max_buffered_query_window)?; + // 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/cosmos_driver/query_planning.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver/query_planning.rs index f7e187779c..1344ca7dd0 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver/query_planning.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver/query_planning.rs @@ -9,7 +9,7 @@ use crate::{ pipeline::operation_pipeline::OperationOverrides, }, models::{ContainerReference, CosmosOperation}, - options::{OperationOptions, QueryPlanMode}, + options::{OperationOptions, PlanOptions, QueryPlanMode}, }; use super::CosmosDriver; @@ -33,12 +33,11 @@ impl From for ResolvedQuer /// Returns an empty resolution when local planning proves topology is unnecessary. pub(super) fn try_resolve_without_topology( - driver: &CosmosDriver, container: &ContainerReference, operation: &CosmosOperation, - options: &OperationOptions, + plan_options: &PlanOptions, ) -> Option { - if driver.effective_query_plan_mode(options) == QueryPlanMode::GatewayOnly { + if plan_options.query_plan_mode == QueryPlanMode::GatewayOnly { return None; } @@ -58,8 +57,9 @@ pub(super) async fn resolve_query_plan( container: &ContainerReference, operation: &CosmosOperation, options: &OperationOptions, + plan_options: &PlanOptions, ) -> crate::error::Result { - let mode = driver.effective_query_plan_mode(options); + let mode = plan_options.query_plan_mode; if mode != QueryPlanMode::GatewayOnly { if let Some(plan) = try_plan_query_using_native_planner(driver, container, operation, mode).await 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..133f1b1528 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 @@ -31,9 +31,9 @@ //! node needs, because a value it has moved past can never reappear. //! //! 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 +//! can produce impractically large tokens even with a finite admission window; +//! truncating it would silently re-emit duplicates. [`Distinct::snapshot_state`] fails with +//! [`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 @@ -80,9 +80,8 @@ enum DistinctMap { /// Global deduplication over an unordered stream. /// - /// Unbounded by design, matching .NET's `UnorderedDistinctMap` and Java's - /// `UnorderedDistinctMap`: ~16 bytes per *distinct* value seen. Since the - /// query cannot be resumed anyway, the set only has to survive one drain. + /// Retained for one drain; admission requires a finite global window and + /// the enclosing SkipTake stops pulling after that window. Unordered { seen: HashSet }, } @@ -430,7 +429,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 +913,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 8b37b5a9e2..59965d88bd 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,7 +25,7 @@ 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]>, @@ -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)), @@ -86,6 +86,9 @@ impl NonStreamingOrderedMerge { } fn retain(&mut self, row: EnvelopeRow) -> crate::error::Result<()> { + if self.take == 0 { + return Ok(()); + } let ordinal = self.next_ordinal; self.next_ordinal = self.next_ordinal.checked_add(1).ok_or_else(|| { CosmosError::builder() @@ -94,12 +97,14 @@ impl NonStreamingOrderedMerge { .build() })?; - if self.retention_limit == 0 { - return Ok(()); - } - let candidate = RetainedRow { row, ordinal }; if self.retained.len() < self.retention_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() + })?; 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 +125,16 @@ 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 = results.take(self.take).collect(); self.session_token = self .aggregator .as_ref() @@ -220,7 +224,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", ) @@ -402,6 +406,34 @@ mod tests { assert_eq!(response.headers().request_charge.unwrap().value(), 3.5); } + #[tokio::test] + async fn zero_take_retains_no_candidates_even_with_offset() { + for skip in [0, 1000] { + let mut node = merge( + vec![page(&[("a", 1.0, "a"), ("b", 2.0, "b")], 3.5, true)], + skip, + skip, + 0, + None, + ); + let mut executor = NoopRequestExecutor; + let mut topology = NoopTopologyProvider; + let mut context = context(&mut executor, &mut topology); + let PageResult::Page { + response, + is_terminal, + } = node.next_page(&mut context).await.unwrap() + else { + panic!("expected charged empty page"); + }; + assert!(is_terminal); + assert!(ids(&response).is_empty()); + assert_eq!(response.headers().request_charge.unwrap().value(), 3.5); + assert_eq!(node.next_ordinal, 0); + assert_eq!(node.retained.capacity(), 0); + } + } + #[tokio::test] async fn emits_binary_items_when_negotiated() { let mut node = NonStreamingOrderedMerge::new( @@ -476,10 +508,132 @@ 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 ); } + #[tokio::test] + async fn finite_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], + skip + rows.len(), + skip, + rows.len(), + 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_BUFFERED_QUERY_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 [0, 2, 1000] { + 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 { 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 13e4309a72..f53fdf84b4 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,41 @@ 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, + max_buffered_query_window: u64, +) -> crate::error::Result<()> { + let Some(info) = query_plan.query_info.as_ref() else { + 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(()); + }; + buffered_query_window(info, max_buffered_query_window, shape).map(|_| ()) +} + +fn buffered_query_window(info: &QueryInfo, maximum: u64, shape: &str) -> crate::error::Result { + if let Some(window) = + combine_take(info).and_then(|take| info.offset.unwrap_or(0).checked_add(take)) + { + if window <= maximum { + return Ok(window); + } + } + 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 and OFFSET plus effective take at most max_buffered_query_window ({maximum})" + )) + .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, @@ -367,7 +401,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", @@ -430,23 +464,9 @@ 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(|_| { + // Admission has checked the per-query policy; enforce a finite representation here too. + let window = buffered_query_window(info, u64::MAX, "non-streaming ORDER BY")?; + let retention_limit = usize::try_from(window).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") @@ -458,12 +478,7 @@ pub(crate) async fn build_non_streaming_ordered_merge( .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 = retention_limit - skip; let effective_operation = rewritten_operation(operation, query_plan)?; let request_nodes = plan_fresh(query_plan, topology_provider, &effective_operation).await?; @@ -1643,7 +1658,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 \ @@ -1861,6 +1876,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::{ @@ -1869,8 +1885,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, }, }; @@ -3878,19 +3895,265 @@ 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_always_requires_finite_window() { 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, u64::MAX).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 + ); + } + + #[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, offset, maximum, accepted) in [ + (None, None, 0, u64::MAX, false), + (Some(0), None, 0, 0, true), + (None, Some(0), 1, 0, false), + (Some(0), None, u64::MAX, u64::MAX, true), + (Some(1), None, 0, 0, false), + (Some(1000), None, 0, 1000, true), + (None, Some(1001), 0, 1000, false), + (Some(1001), None, 0, 1001, true), + (Some(3), None, 997, 1000, true), + (None, Some(3), 998, 1000, false), + (Some(10), Some(3), 997, 1000, true), + (Some(3), Some(10), 997, 1000, true), + (Some(u64::MAX), Some(0), 0, 0, true), + (Some(u64::MAX), None, 1, u64::MAX, false), + (Some(1), None, u64::MAX, u64::MAX, false), + ] { + let plan = QueryPlan { + query_info: Some(QueryInfo { + top, + limit, + offset: Some(offset), + 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 = !accepted + && (non_streaming || distinct_type == DistinctType::Unordered); + let result = validate_buffered_query(&plan, maximum); + 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("max_buffered_query_window")); + 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(1000), Some(1001)] { + 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, 1000).is_ok(), + top.is_some_and(|take| take <= 1000) + ); + assert_eq!( + validate_buffered_query(&plan, u64::MAX).is_ok(), + top.is_some() + ); + } + } + } + + #[tokio::test] + async fn admitted_finite_non_streaming_plan_preserves_partition_rewrite() { + for input_binary in [false, true] { + for output_binary in [false, true] { + for (offset, limit, expected) in [ + (0, Some(6), vec!["a", "b", "c", "d", "e", "f"]), + (1, Some(6), vec!["b", "c", "d", "e", "f"]), + (1, Some(3), vec!["b", "c", "d"]), + (20, Some(6), 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, 1000).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_BUFFERED_QUERY_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_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW ); } @@ -3902,11 +4165,12 @@ mod tests { info.top = None; info.offset = Some(50_000); info.limit = Some(3); + validate_buffered_query(&plan, 50_003).unwrap(); let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr("", "FF", "pk-range")])]); let pipeline = build_non_streaming_ordered_merge(&plan, &mut topology, &operation, None) .await - .expect("finite windows are not capped by the client"); + .expect("explicitly admitted finite window builds"); assert!(pipeline .into_root() .downcast::() @@ -3928,7 +4192,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(); @@ -4087,7 +4351,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 9d4e090af9..3698a2e0da 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("ClientNonStreamingOrderByRequiresFiniteWindow"), - 20127 => Some("ClientNonStreamingOrderByWindowTooLarge"), + 20124 => Some("ClientBufferedQueryContinuationUnsupported"), + 20125 => Some("ClientBufferedQueryRequiresFiniteWindow"), + 20126 => Some("ClientNonStreamingOrderByWindowTooLarge"), 20150 => Some("ClientDuplicateFaultInjectionRuleId"), 20151 => Some("ClientThroughputControlGroupRegistrationFailed"), 20152 => Some("ClientThroughputControlGroupNotRegistered"), @@ -1394,26 +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 finite global TOP/LIMIT and OFFSET plus take + /// within its configured maximum (20125). + pub const CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW: SubStatusCode = SubStatusCode(20125); - /// A non-streaming `ORDER BY` query did not contain a finite `TOP` or - /// `OFFSET`/`LIMIT` window (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). - pub const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = SubStatusCode(20127); + /// A non-streaming `ORDER BY` query's candidate storage cannot be represented + /// or allocated by the current process (20126). + pub const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: SubStatusCode = SubStatusCode(20126); // ----- 20150-20199: SDK configuration / setup errors ----- @@ -2380,28 +2382,34 @@ 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 — non-streaming `ORDER BY` requires a finite result window. - pub const CLIENT_NON_STREAMING_ORDER_BY_REQUIRES_FINITE_WINDOW: CosmosStatus = CosmosStatus { + /// 400 / 20125 — a buffered query requires finite global TOP/LIMIT and + /// OFFSET plus take within its configured maximum. + 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), }; - /// 400 / 20127 — the non-streaming `ORDER BY` candidate window cannot be - /// represented by the current process. + /// 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 / 20126 — the non-streaming `ORDER BY` candidate window cannot be + /// represented or allocated by the current process. pub const CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE: CosmosStatus = CosmosStatus { status_code: StatusCode::BadRequest, sub_status: Some(SubStatusCode::CLIENT_NON_STREAMING_ORDER_BY_WINDOW_TOO_LARGE), @@ -2767,6 +2775,79 @@ mod tests { ); } + #[test] + 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 + ); + 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/driver_options.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/driver_options.rs index 202261157a..bc91ece3d9 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/driver_options.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/driver_options.rs @@ -491,12 +491,6 @@ mod tests { assert_eq!(options.preferred_regions(), ®ions); } - #[test] - fn query_plan_mode_is_unset_by_default() { - let options = DriverOptionsBuilder::new(test_account()).build_from_env(&|_| None); - assert_eq!(options.operation_options().query_plan_mode, None); - } - // ── Partition-failover / PPCB end-to-end env resolution ───────────────── // // These guard the customer-reported bug: when the caller omits diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs index 84cfc06781..76dc17c093 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/mod.rs @@ -56,7 +56,7 @@ pub use operation_options::{ }; pub use partition_failover::{PartitionFailoverOptions, PartitionFailoverOptionsBuilder}; pub use patch_strategy::PatchStrategy; -pub use plan_options::{PlanOptions, DEFAULT_MAX_FAN_OUT}; +pub use plan_options::{PlanOptions, DEFAULT_MAX_BUFFERED_QUERY_WINDOW, DEFAULT_MAX_FAN_OUT}; pub use policies::{ ContentResponseOnWrite, EndToEndOperationLatencyPolicy, ExcludedRegions, ServerCertificateValidation, TlsBackend, 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..2f0ba82894 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 @@ -14,7 +14,7 @@ use crate::{ options::{ AvailabilityStrategy, BinaryEncodingOptions, ContentResponseOnWrite, EndToEndOperationLatencyPolicy, ExcludedRegions, PatchStrategy, PriorityLevel, - QueryPlanMode, ReadConsistencyStrategy, + ReadConsistencyStrategy, }, }; @@ -40,15 +40,6 @@ use crate::{ #[options(layers(runtime, account, operation))] #[non_exhaustive] pub struct OperationOptions { - /// Query-plan provider selection for query operations. - /// - /// `None` inherits from a lower layer (default: - /// [`QueryPlanMode::LocalPreferred`]). The - /// `AZURE_COSMOS_QUERY_PLAN_MODE_OVERRIDE` environment variable takes - /// precedence over every programmatic layer as a livesite kill switch. - #[option(env = "AZURE_COSMOS_QUERY_PLAN_MODE", overridable)] - pub query_plan_mode: Option, - /// How PATCH operations are executed. /// /// `None` inherits from a lower layer (default: [`PatchStrategy::Auto`]). @@ -725,94 +716,6 @@ mod tests { assert_eq!(view.hedging_enabled(), Some(&true)); } - #[test] - fn query_plan_mode_resolves_across_all_layers() { - let env = std::sync::Arc::new(OperationOptions { - query_plan_mode: Some(QueryPlanMode::LocalPreferred), - ..Default::default() - }); - let runtime = std::sync::Arc::new(OperationOptions { - query_plan_mode: Some(QueryPlanMode::GatewayOnly), - ..Default::default() - }); - let account = std::sync::Arc::new(OperationOptions { - query_plan_mode: Some(QueryPlanMode::LocalPreferred), - ..Default::default() - }); - let operation = OperationOptions { - query_plan_mode: Some(QueryPlanMode::GatewayOnly), - ..Default::default() - }; - - let view = - OperationOptionsView::new(Some(env), Some(runtime), Some(account), Some(&operation)); - - assert_eq!(view.query_plan_mode(), Some(&QueryPlanMode::GatewayOnly)); - } - - #[test] - fn query_plan_mode_environment_override_is_authoritative() { - let env_override = std::sync::Arc::new(OperationOptions { - query_plan_mode: Some(QueryPlanMode::GatewayOnly), - ..Default::default() - }); - let operation = OperationOptions { - query_plan_mode: Some(QueryPlanMode::LocalPreferred), - ..Default::default() - }; - - let view = OperationOptionsView::new_with_override( - Some(env_override), - None, - None, - None, - Some(&operation), - ); - - assert_eq!(view.query_plan_mode(), Some(&QueryPlanMode::GatewayOnly)); - } - - #[test] - fn query_plan_mode_environment_variables_are_parsed() { - let base = OperationOptions::from_env_vars(|key| match key { - "AZURE_COSMOS_QUERY_PLAN_MODE" => Ok("LocalPreferred".to_string()), - _ => Err(std::env::VarError::NotPresent), - }); - let override_options = OperationOptions::from_env_override_vars(|key| match key { - "AZURE_COSMOS_QUERY_PLAN_MODE_OVERRIDE" => Ok("gateway".to_string()), - _ => Err(std::env::VarError::NotPresent), - }); - - assert_eq!(base.query_plan_mode, Some(QueryPlanMode::LocalPreferred)); - assert_eq!( - override_options.query_plan_mode, - Some(QueryPlanMode::GatewayOnly) - ); - } - - #[test] - fn invalid_query_plan_mode_override_falls_through() { - let env_override = - std::sync::Arc::new(OperationOptions::from_env_override_vars(|key| match key { - "AZURE_COSMOS_QUERY_PLAN_MODE_OVERRIDE" => Ok("invalid".to_string()), - _ => Err(std::env::VarError::NotPresent), - })); - let operation = OperationOptions { - query_plan_mode: Some(QueryPlanMode::GatewayOnly), - ..Default::default() - }; - - let view = OperationOptionsView::new_with_override( - Some(env_override), - None, - None, - None, - Some(&operation), - ); - - assert_eq!(view.query_plan_mode(), Some(&QueryPlanMode::GatewayOnly)); - } - /// `from_env_override_vars` populates only the `overridable` fields from /// their `{ENV}_OVERRIDE` variants and leaves every other env field /// `None` (the base `from_env_vars` path is unaffected). diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/plan_options.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/plan_options.rs index 83171c9ed3..e542e94d9b 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/plan_options.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/plan_options.rs @@ -3,12 +3,17 @@ //! Plan-time options for [`CosmosDriver::plan_operation`](crate::driver::CosmosDriver::plan_operation). +use crate::options::QueryPlanMode; + /// Default maximum fan-out for a fresh cross-partition operation. /// /// A plan that would fan out to more than this many leaf request nodes is /// rejected unless the caller raises [`PlanOptions::max_fan_out`]. pub const DEFAULT_MAX_FAN_OUT: u32 = 100; +/// Default maximum global OFFSET plus take for client-buffered queries. +pub const DEFAULT_MAX_BUFFERED_QUERY_WINDOW: u64 = 1000; + /// Options that shape how an operation is planned into a dataflow pipeline. /// /// Unlike [`OperationOptions`](crate::options::OperationOptions), which controls @@ -18,6 +23,17 @@ pub const DEFAULT_MAX_FAN_OUT: u32 = 100; #[derive(Clone, Debug)] #[non_exhaustive] pub struct PlanOptions { + /// Maximum global OFFSET plus effective take for client-buffered queries. + /// + /// Requires a finite TOP or LIMIT; when both exist, the smaller is used. + /// Defaults to [`DEFAULT_MAX_BUFFERED_QUERY_WINDOW`]. Zero is a valid limit. + pub max_buffered_query_window: u64, + + /// Query-plan provider selection for this query. + /// + /// Defaults to [`QueryPlanMode::LocalPreferred`]. + pub query_plan_mode: QueryPlanMode, + /// Maximum number of leaf request nodes a fresh cross-partition plan may /// fan out to. /// @@ -40,15 +56,79 @@ pub struct PlanOptions { impl Default for PlanOptions { fn default() -> Self { Self { + max_buffered_query_window: DEFAULT_MAX_BUFFERED_QUERY_WINDOW, + query_plan_mode: QueryPlanMode::LocalPreferred, max_fan_out: DEFAULT_MAX_FAN_OUT, } } } impl PlanOptions { + /// Sets the maximum global OFFSET plus effective take for client-buffered queries. + pub fn with_max_buffered_query_window(mut self, max_buffered_query_window: u64) -> Self { + self.max_buffered_query_window = max_buffered_query_window; + self + } + + /// Sets the query-plan provider selection for this query. + pub fn with_query_plan_mode(mut self, mode: QueryPlanMode) -> Self { + self.query_plan_mode = mode; + self + } + /// Sets the maximum fan-out for a fresh cross-partition plan. pub fn with_max_fan_out(mut self, max_fan_out: u32) -> Self { self.max_fan_out = max_fan_out; self } } + +#[cfg(test)] +mod tests { + use super::PlanOptions; + use crate::{ + driver::dataflow::{ + planner::validate_buffered_query, + query_plan::{DistinctType, QueryInfo, QueryPlan}, + }, + error::CosmosStatus, + }; + + #[test] + fn buffered_admission_uses_plan_default_and_explicit_values() { + let mut query = QueryPlan { + query_info: Some(QueryInfo { + distinct_type: DistinctType::Unordered, + top: Some(1000), + ..Default::default() + }), + ..Default::default() + }; + for (maximum, take, accepted) in [ + (None, Some(1000), true), + (None, Some(1001), false), + (Some(1001), Some(1001), true), + (Some(999), Some(1000), false), + (Some(0), Some(0), true), + (Some(0), Some(1), false), + (Some(u64::MAX), None, false), + ] { + let options = maximum.map_or_else(PlanOptions::default, |maximum| { + PlanOptions::default().with_max_buffered_query_window(maximum) + }); + query.query_info.as_mut().unwrap().top = take; + let result = validate_buffered_query(&query, options.max_buffered_query_window); + assert_eq!( + result.is_ok(), + accepted, + "maximum={maximum:?}, take={take:?}" + ); + if let Err(error) = result { + assert_eq!( + error.status(), + CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW + ); + } + } + } +} 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 3a539aa2b9..3d4a16c7cb 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 @@ -1033,6 +1033,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 @@ -1045,7 +1062,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..69a8727c16 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,7 @@ 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, PlanOptions}; use azure_data_cosmos_driver::CosmosDriver; use framework::resolve_test_env; @@ -496,7 +495,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 +546,16 @@ async fn validate_production_local_plan( "query ranges differ for '{sql}'" ); - if execute { + if let Some(plan_options) = plan_options { + let options = OperationOptions::default(); 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, &plan_options) .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 +1044,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 +1051,43 @@ 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(PlanOptions::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("unordered DISTINCT must require a finite global window"); + assert_eq!( + error.status(), + azure_data_cosmos_driver::error::CosmosStatus::CLIENT_BUFFERED_QUERY_REQUIRES_FINITE_WINDOW + ); + + validate_production_local_plan( + "SELECT DISTINCT TOP 1000 VALUE c.city FROM c", + &[], + Some(PlanOptions::default()), + ) + .await; } #[tokio::test] 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..96ce090001 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"; @@ -78,7 +79,21 @@ struct ExpectedError { } fn catalog() -> Catalog { - serde_json::from_str(CATALOG_JSON).expect("catalog must parse") + let mut catalog: Catalog = serde_json::from_str(CATALOG_JSON).expect("catalog must parse"); + // Keep the shared stage fixtures, but bound their end-to-end execution. + for scenario in &mut catalog.scenarios { + if scenario.query.distinct_type == "Unordered" + && !scenario.query.text.contains(" TOP ") + && !scenario.query.text.contains(" LIMIT ") + { + scenario.query.text = + scenario + .query + .text + .replacen("SELECT DISTINCT ", "SELECT DISTINCT TOP 1000 ", 1); + } + } + catalog } #[derive(Debug, Default)] @@ -122,6 +137,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 +162,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 +182,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 +203,183 @@ async fn setup_with_query_recorder() -> ( (emulator, driver, recorder) } +#[tokio::test] +async fn buffered_admission_precedes_items_and_uses_per_plan_options() { + use azure_data_cosmos_driver::error::CosmosStatus; + + for partition_count in [1, 2] { + for mode in [QueryPlanMode::LocalPreferred, QueryPlanMode::GatewayOnly] { + let recorder = Arc::new(QueryRequestRecorder::default()); + let (_, driver) = setup_with_policy( + Some(recorder.clone()), + OperationOptions::default(), + 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(0), Some(2), Some(3), Some(1001), Some(u64::MAX)] { + for (sql, parameters, window, expected) in [ + ("SELECT DISTINCT VALUE c.value FROM c", vec![], None, 6), + ( + "SELECT DISTINCT TOP @take VALUE c.value FROM c", + vec![serde_json::json!({"name":"@take","value":3})], + Some(3), + 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}), + ], + Some(3), + 2, + ), + ( + "SELECT DISTINCT TOP 0 VALUE c.value FROM c", + vec![], + Some(0), + 0, + ), + ( + "SELECT DISTINCT VALUE c.value FROM c OFFSET 0 LIMIT 0", + vec![], + Some(0), + 0, + ), + ( + "SELECT DISTINCT TOP 1000 VALUE c.value FROM c", + vec![], + Some(1000), + 6, + ), + ( + "SELECT DISTINCT TOP 1001 VALUE c.value FROM c", + vec![], + Some(1001), + 6, + ), + ] { + let query = QuerySpec { + text: sql.into(), + parameters, + distinct_type: "Unordered".into(), + }; + let mut options = PlanOptions::default().with_query_plan_mode(mode); + if let Some(maximum) = request { + options = options.with_max_buffered_query_window(maximum); + } + let denied = + window.is_none_or(|window| window > options.max_buffered_query_window); + for page_size in [1, 1000, u32::MAX] { + recorder.take(); + let result = Box::pin(driver.plan_operation( + query_operation(&container, &query, page_size), + &OperationOptions::default(), + None, + &options.clone().with_max_fan_out(if denied { + 1 + } else { + partition_count + }), + )) + .await; + assert!(recorder.take().is_empty(), "planning must not query items"); + if denied { + let error = result + .err() + .expect("query outside finite policy 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("max_buffered_query_window")); + 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 expected == 6 { + assert_eq!( + unique, + (0..6).map(|value| value.to_string()).collect::>() + ); + } + } + } + } + + // A complete logical key bypasses client buffering even with a zero window. + 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 = OperationOptions::default(); + let mut plan = Box::pin( + driver.plan_operation( + operation, + &options, + None, + &PlanOptions::default() + .with_query_plan_mode(mode) + .with_max_buffered_query_window(0), + ), + ) + .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, @@ -834,7 +1038,7 @@ async fn text_and_binary_query_pages_have_pipeline_parity() { distinct_type: "None".to_owned(), }, QuerySpec { - text: "SELECT DISTINCT VALUE c.value FROM c".to_owned(), + text: "SELECT DISTINCT TOP 1000 VALUE c.value FROM c".to_owned(), parameters: Vec::new(), distinct_type: "Unordered".to_owned(), }, diff --git a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/local_query_planning.rs b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/local_query_planning.rs index 15ef81a65c..256940649f 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/local_query_planning.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/tests/in_memory_emulator_tests/local_query_planning.rs @@ -13,9 +13,7 @@ use azure_data_cosmos_driver::{ VirtualRegion, }, models::{AccountReference, CosmosOperation, FeedRange, PartitionKeyDefinition}, - options::{ - DriverOptions, OperationOptions, OperationOptionsBuilder, PlanOptions, QueryPlanMode, - }, + options::{DriverOptions, OperationOptions, PlanOptions, QueryPlanMode}, }; use super::{host_recorder::HostRecorder, GATEWAY_URL}; @@ -25,21 +23,10 @@ async fn setup() -> ( Arc, Arc, ) { - setup_with_query_plan_mode(QueryPlanMode::LocalPreferred).await -} - -async fn setup_with_query_plan_mode( - mode: QueryPlanMode, -) -> ( - Arc, - Arc, - Arc, -) { - setup_with_driver_options(mode, true).await + setup_with_driver_options(true).await } async fn setup_with_driver_options( - mode: QueryPlanMode, partition_key_range_cache_enabled: bool, ) -> ( Arc, @@ -72,11 +59,6 @@ async fn setup_with_driver_options( let driver = runtime .create_driver( DriverOptions::builder(account) - .with_operation_options( - OperationOptionsBuilder::new() - .with_query_plan_mode(mode) - .build(), - ) .with_partition_key_range_cache_enabled(partition_key_range_cache_enabled) .build(), ) @@ -93,16 +75,14 @@ async fn per_request_gateway_only_mode_bypasses_local_query_planning() { .await .unwrap(); recorder.clear(); - let options = OperationOptionsBuilder::new() - .with_query_plan_mode(QueryPlanMode::GatewayOnly) - .build(); + let options = PlanOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly); driver .plan_operation( query(&container, "SELECT * FROM c WHERE c.pk = 'a'"), - &options, + &OperationOptions::default(), None, - &PlanOptions::default(), + &options, ) .await .unwrap(); @@ -111,22 +91,28 @@ async fn per_request_gateway_only_mode_bypasses_local_query_planning() { #[cfg(not(feature = "__internal_native_query_plan"))] #[tokio::test] -async fn per_request_local_preferred_overrides_gateway_only_client_default() { - let (_emulator, recorder, driver) = - setup_with_query_plan_mode(QueryPlanMode::GatewayOnly).await; +async fn per_plan_mode_does_not_leak_to_subsequent_queries() { + let (_emulator, recorder, driver) = setup().await; let container = driver .resolve_container("testdb", "testcoll", OperationOptions::default()) .await .unwrap(); recorder.clear(); - let options = OperationOptionsBuilder::new() - .with_query_plan_mode(QueryPlanMode::LocalPreferred) - .build(); - driver .plan_operation( query(&container, "SELECT * FROM c WHERE c.pk = 'a'"), - &options, + &OperationOptions::default(), + None, + &PlanOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly), + ) + .await + .unwrap(); + assert_eq!(recorder.query_plan_count(), 1); + recorder.clear(); + driver + .plan_operation( + query(&container, "SELECT * FROM c WHERE c.pk = 'a'"), + &OperationOptions::default(), None, &PlanOptions::default(), ) @@ -227,10 +213,41 @@ async fn contradictory_query_short_circuits_all_query_io() { assert_eq!(recorder.document_query_count(), 0); } +#[tokio::test] +async fn contradictory_buffered_query_is_exempt_from_zero_window() { + for topology_enabled in [false, true] { + let (_emulator, recorder, driver) = setup_with_driver_options(topology_enabled).await; + let container = driver + .resolve_container("testdb", "testcoll", OperationOptions::default()) + .await + .unwrap(); + recorder.clear(); + let mut plan = driver + .plan_operation( + query( + &container, + "SELECT DISTINCT VALUE c.value FROM c WHERE c.pk = 'a' AND c.pk = 'b'", + ), + &OperationOptions::default(), + None, + &PlanOptions::default().with_max_buffered_query_window(0), + ) + .await + .unwrap(); + assert!(driver + .execute_plan(&mut plan, Some(container), OperationOptions::default()) + .await + .unwrap() + .is_none()); + assert_eq!(recorder.query_plan_count(), 0); + assert_eq!(recorder.routing_metadata_count(), 0); + assert_eq!(recorder.document_query_count(), 0); + } +} + #[tokio::test] async fn gateway_only_contradiction_without_partition_topology_still_fails() { - let (_emulator, recorder, driver) = - setup_with_driver_options(QueryPlanMode::GatewayOnly, false).await; + let (_emulator, recorder, driver) = setup_with_driver_options(false).await; let container = driver .resolve_container("testdb", "testcoll", OperationOptions::default()) .await @@ -245,7 +262,7 @@ async fn gateway_only_contradiction_without_partition_topology_still_fails() { ), &OperationOptions::default(), None, - &PlanOptions::default(), + &PlanOptions::default().with_query_plan_mode(QueryPlanMode::GatewayOnly), ) .await .err() @@ -262,8 +279,7 @@ async fn gateway_only_contradiction_without_partition_topology_still_fails() { #[tokio::test] async fn contradictory_query_does_not_require_partition_topology() { - let (_emulator, recorder, driver) = - setup_with_driver_options(QueryPlanMode::LocalPreferred, false).await; + let (_emulator, recorder, driver) = setup_with_driver_options(false).await; let container = driver .resolve_container("testdb", "testcoll", OperationOptions::default()) .await @@ -295,8 +311,7 @@ async fn contradictory_query_does_not_require_partition_topology() { #[tokio::test] async fn nonempty_query_without_partition_topology_fails_before_gateway() { - let (_emulator, recorder, driver) = - setup_with_driver_options(QueryPlanMode::LocalPreferred, false).await; + let (_emulator, recorder, driver) = setup_with_driver_options(false).await; let container = driver .resolve_container("testdb", "testcoll", OperationOptions::default()) .await diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/README.md b/sdk/cosmos/azure_data_cosmos_driver_native/README.md index 3b1950000c..9ec73f6cd0 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/README.md +++ b/sdk/cosmos/azure_data_cosmos_driver_native/README.md @@ -174,10 +174,10 @@ below for the production-shape guidance. > > Query-plan selection is per operation through > `cosmos_operation_options_t.query_plan_mode`. Leave it -> `COSMOS_QUERY_PLAN_MODE_UNSET` to inherit, or set +> `COSMOS_QUERY_PLAN_MODE_UNSET` for the `LocalPreferred` default, or set > `COSMOS_QUERY_PLAN_MODE_LOCAL_PREFERRED` or > `COSMOS_QUERY_PLAN_MODE_GATEWAY_ONLY` for an individual query. The -> environment override remains authoritative over this field. +> setting is per-query only; no client/runtime or environment defaults apply. > > The v1 functions take `(driver, const cosmos_operation_request_t *request, queue, > user_data, out_pre_error)` and return a `cosmos_operation_handle_t *`. diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h b/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h index c79e384e47..831e809ab7 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h +++ b/sdk/cosmos/azure_data_cosmos_driver_native/include/azurecosmosdriver.h @@ -1453,8 +1453,8 @@ typedef struct cosmos_operation_options_t { int8_t binary_encoding_request_text_response; /** * Query-plan mode encoded as a [`CosmosQueryPlanMode`] discriminant. - * `0` (`Unset`) inherits. Stored as a raw `i32` so invalid host values can - * be rejected before materializing the enum. + * `0` (`Unset`) uses LocalPreferred. Raw `i32` storage allows invalid host + * values to be rejected before materializing the enum. */ int32_t query_plan_mode; } cosmos_operation_options_t; diff --git a/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs b/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs index 05fc204ab5..0f3e300feb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs +++ b/sdk/cosmos/azure_data_cosmos_driver_native/src/op_request.rs @@ -219,11 +219,11 @@ impl CosmosPatchStrategy { } /// Tri-state mirror of [`QueryPlanMode`] for the flat options struct. -/// `0` (`Unset`) means "inherit from a lower-priority layer". +/// `0` (`Unset`) means "use the LocalPreferred default". #[repr(i32)] #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum CosmosQueryPlanMode { - /// Inherit from account / runtime / environment. + /// Use the LocalPreferred default. CosmosQueryPlanModeUnset = 0, /// Prefer local planning, falling back to the Gateway before execution. CosmosQueryPlanModeLocalPreferred = 1, @@ -478,12 +478,19 @@ pub struct CosmosOperationOptions { /// [`binary_encoding_enabled`](Self::binary_encoding_enabled) to `1`. pub binary_encoding_request_text_response: i8, /// Query-plan mode encoded as a [`CosmosQueryPlanMode`] discriminant. - /// `0` (`Unset`) inherits. Stored as a raw `i32` so invalid host values can - /// be rejected before materializing the enum. + /// `0` (`Unset`) uses LocalPreferred. Raw `i32` storage allows invalid host + /// values to be rejected before materializing the enum. pub query_plan_mode: i32, } impl CosmosOperationOptions { + fn to_plan_options(&self, max_fan_out: u32) -> Result { + let mode = CosmosQueryPlanMode::from_i32(self.query_plan_mode)? + .to_driver() + .unwrap_or_default(); + Ok(plan_options_from_max_fan_out(max_fan_out).with_query_plan_mode(mode)) + } + /// Builds the driver [`OperationOptions`] from this flat struct. /// /// # Safety @@ -499,7 +506,7 @@ impl CosmosOperationOptions { CosmosContentResponseOnWriteOpt::from_i32(self.content_response_on_write)? .to_driver()?; opts.patch_strategy = CosmosPatchStrategy::from_i32(self.patch_strategy)?.to_driver(); - opts.query_plan_mode = CosmosQueryPlanMode::from_i32(self.query_plan_mode)?.to_driver(); + CosmosQueryPlanMode::from_i32(self.query_plan_mode)?; opts.session_capturing_disabled = decode_tristate_bool(self.session_capturing_disabled)?; opts.max_failover_retry_count = decode_opt_u32(self.max_failover_retry_count); @@ -932,7 +939,12 @@ pub(crate) unsafe fn build_request( // A `max_fan_out` of 0 means "unset": fall back to the driver default. A // non-zero value opts into a broader (or narrower) fan-out. - let plan_options = plan_options_from_max_fan_out(req.max_fan_out); + let plan_options = if req.options.is_null() { + plan_options_from_max_fan_out(req.max_fan_out) + } else { + // SAFETY: non-NULL checked; caller guarantees a valid struct. + unsafe { (*req.options).to_plan_options(req.max_fan_out)? } + }; Ok(BuiltRequest { operation, @@ -1614,21 +1626,22 @@ mod tests { fn query_plan_mode_maps_to_driver() { use CosmosQueryPlanMode as M; for (mode, expected) in [ - (M::CosmosQueryPlanModeUnset, None), + (M::CosmosQueryPlanModeUnset, QueryPlanMode::LocalPreferred), ( M::CosmosQueryPlanModeLocalPreferred, - Some(QueryPlanMode::LocalPreferred), + QueryPlanMode::LocalPreferred, ), ( M::CosmosQueryPlanModeGatewayOnly, - Some(QueryPlanMode::GatewayOnly), + QueryPlanMode::GatewayOnly, ), ] { let mut options = cosmos_operation_options_default(); options.query_plan_mode = mode as i32; - // SAFETY: all pointer fields are NULL / len 0. - let driver = unsafe { options.to_driver() }.expect("options convert"); + let driver = options.to_plan_options(250).expect("options convert"); assert_eq!(driver.query_plan_mode, expected); + assert_eq!(driver.max_fan_out, 250); + assert_eq!(driver.max_buffered_query_window, 1000); } } @@ -1782,7 +1795,10 @@ mod tests { assert_eq!(driver.read_consistency_strategy, None); assert_eq!(driver.content_response_on_write, None); assert_eq!(driver.patch_strategy, None); - assert_eq!(driver.query_plan_mode, None); + assert_eq!( + o.to_plan_options(0).unwrap().query_plan_mode, + QueryPlanMode::LocalPreferred + ); assert_eq!(driver.session_capturing_disabled, None); assert_eq!(driver.max_failover_retry_count, None); assert_eq!(driver.max_session_retry_count, None); diff --git a/sdk/cosmos/docs/specs/0001-configuration-options.md b/sdk/cosmos/docs/specs/0001-configuration-options.md index 48764dc0b1..0b4ce8abc6 100644 --- a/sdk/cosmos/docs/specs/0001-configuration-options.md +++ b/sdk/cosmos/docs/specs/0001-configuration-options.md @@ -413,11 +413,14 @@ pub struct ItemWriteOptions { ### 5.3 `QueryOptions` Options for query operations (`query_items`, `query_items_single_partition`). +The manual `Default` implementation sets `max_buffered_query_window` to 1000. ```rust -#[derive(Clone, Default)] +#[derive(Clone)] #[non_exhaustive] pub struct QueryOptions { + pub max_buffered_query_window: u64, + pub query_plan_mode: QueryPlanMode, // Layered option group pub operation: OperationOptions, @@ -429,13 +432,43 @@ pub struct QueryOptions { } ``` -| Option | Type | Notes | -| ------------------------- | ---------------------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| `operation` | `OperationOptions` | Layered group; `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. | +| Option | Type | Notes | +| --- | --- | --- | +| `operation` | `OperationOptions` | Layered group; `content_response_on_write` is ignored for queries. | +| `max_buffered_query_window` | `u64` | Maximum global OFFSET plus effective take for client-buffered queries. Defaults to 1000; zero is valid. | +| `query_plan_mode` | `QueryPlanMode` | Per-query provider selection. Defaults to `LocalPreferred`; `GatewayOnly` bypasses local planning. | +| `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 finite window + +`QueryOptions::max_buffered_query_window: u64` applies only to this query and +defaults to 1000. The driver receives it through `PlanOptions`, alongside +`query_plan_mode` (default `LocalPreferred`). Neither setting participates in +operation/account/runtime layering or reads environment variables. + +The SDK convenience setter +`QueryOptions::with_max_buffered_query_window(u64)` sets the per-query field. +A global TOP or LIMIT is always required for affected buffered shapes: + +```rust +use azure_data_cosmos::options::{QueryOptions, QueryPlanMode}; + +// SQL: SELECT DISTINCT TOP 2000 VALUE c.category FROM c +let query_options = QueryOptions::default() + .with_max_buffered_query_window(2000) + .with_query_plan_mode(QueryPlanMode::GatewayOnly); +``` + +Admission requires checked `OFFSET + min(TOP, LIMIT) <= max_buffered_query_window`, +using whichever finite clause is present. Zero is valid; OFFSET still counts +when take is zero. Missing bounds, arithmetic overflow, and excess windows fail +with 400/20125, even at `u64::MAX`. The driver exports +`DEFAULT_MAX_BUFFERED_QUERY_WINDOW = 1000`, reused internally by the SDK. +There is no opt-out or C ABI option. This is a row-window admission policy, +not a byte budget, and does not change service or continuation restrictions. ### 5.4 `TransactionalBatchOptions` 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 c695179951..e8fe1deb17 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,28 @@ This document describes the implemented retry behavior for the Azure Cosmos DB R ## Design Philosophy +### Buffered-query input validation + +| Status | Symbol | Remedy | +| --- | --- | --- | +| 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` | Supply a finite global TOP/LIMIT and keep OFFSET plus effective take within per-query `max_buffered_query_window` (default 1000), or raise that finite maximum. Overflow is rejected. | +| 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 +`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. The maximum is configurable per +query, with no opt-out; `u64::MAX` still requires a finite bound. +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 regardless of the +configured maximum. 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: - **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/0012-feed-operations-and-dataflow.md b/sdk/cosmos/docs/specs/0012-feed-operations-and-dataflow.md index cf6e271d90..0a5056e5b5 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,9 @@ 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 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. diff --git a/sdk/cosmos/docs/specs/0013-query-engine.md b/sdk/cosmos/docs/specs/0013-query-engine.md index 0fb3118727..5a012f06e0 100644 --- a/sdk/cosmos/docs/specs/0013-query-engine.md +++ b/sdk/cosmos/docs/specs/0013-query-engine.md @@ -16,6 +16,44 @@ 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, with OFFSET plus effective take at most the per-query +`max_buffered_query_window: u64` (default 1000). There is no opt-out, including +when the configured maximum is `u64::MAX`. Ordering metadata does not exempt +an unordered DISTINCT stage. + +Zero is a valid maximum and bound; TOP combined with LIMIT uses the smaller +value. OFFSET still counts when take is zero. Checked addition rejects overflow +as an admission error (400/20125), as it does missing bounds or excess windows. +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. + +Non-streaming execution uses a bounded top-k heap with a checked OFFSET + take +window, then sorts retained candidates, applies OFFSET, and paginates the result. +There is no unbounded execution representation. 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: document sizes and page-level DISTINCT processing add retained work. + +Raising the maximum does not 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. +Deterministic buffered-node tests cover finite-window execution independently +of live service availability. + ```text SQL Text → Lexer (hand-crafted tokenizer) @@ -158,18 +196,11 @@ Normal builds resolve cross-partition query plans in this order: 1. Local Rust planner. 2. Gateway query-plan endpoint. -Set `QueryPlanMode::GatewayOnly` on `OperationOptions` to bypass local planning. -Use `DriverOptionsBuilder::with_operation_options` or -`CosmosClientBuilder::with_default_operation_options` to change the default for -a driver or client, and the operation-specific `with_operation_options` setter -to override it for one query. - -Query-plan mode follows the standard operation → account → runtime → -environment option hierarchy. `AZURE_COSMOS_QUERY_PLAN_MODE` supplies the -lowest-priority process default. For livesite mitigation, -`AZURE_COSMOS_QUERY_PLAN_MODE_OVERRIDE=gateway` authoritatively forces Gateway -planning above every programmatic layer. Environment settings are captured -when the runtime is constructed. +Set `PlanOptions::query_plan_mode` to `QueryPlanMode::GatewayOnly` to bypass +local planning, or use `QueryOptions::with_query_plan_mode` in the SDK. +Query-plan mode and `max_buffered_query_window` are per-query only, with defaults +of `LocalPreferred` and 1000. They have no client, account, runtime, or +environment defaults and no authoritative environment override. When `__internal_native_query_plan` is enabled, the existing native-first behavior is preserved: diff --git a/sdk/cosmos/eng/pipelines/README.md b/sdk/cosmos/eng/pipelines/README.md index fb92b88228..a4b7254cd3 100644 --- a/sdk/cosmos/eng/pipelines/README.md +++ b/sdk/cosmos/eng/pipelines/README.md @@ -90,7 +90,7 @@ pwsh sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.tests.ps1 ``` You can also invoke the resolver directly against the sample JSON. Dot-source -it (note the leading `. `) so the resolved variables land directly in your +it (note the leading `.` followed by a space) so the resolved variables land directly in your current shell instead of just being printed: ```powershell @@ -100,6 +100,35 @@ $env:COSMOS_TEST_ACCOUNTS_JSON = Get-Content -Raw sdk/cosmos/eng/pipelines/live- . ./sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.ps1 ``` +### Live-only query tests + +The SDK's `live` test target shares the integration-test framework but keeps +live-only DISTINCT coverage separate from emulator-compatible query tests. +It requires the `key_auth`, `control_plane`, and `fault_injection` features +and is ignored by default. + +To run it locally, set `AZURE_COSMOS_CONNECTION_STRING` to a real account's +connection string, then run from the repository root: + +```powershell +$env:AZURE_COSMOS_TEST_MODE = 'required' +$env:RUSTFLAGS = '--cfg=test_category="live"' +cargo test -p azure_data_cosmos --all-features --test live live_distinct_admission_and_per_query_options +``` + +The test creates and cleans up a unique database, so the account key must +permit database and container management. Explicit selection (including +`--ignored`) fails on missing/invalid connection strings, local emulator +endpoints, or `AZURE_COSMOS_TEST_MODE=skipped`; it never reports a skipped +live scenario as a pass. + +The fixed-account resolver appends `test_category="live"` alongside each +account's existing category. Test setup propagates these flags before its +preconfigured-connection early return, so `Cosmos_live_test` runs this target +without changing the existing suites. Ordinary emulator, vnext, in-memory, +and codec-fuzz jobs do not use that resolver and do not enable the live gate. +Legacy ARM/AAD jobs retain their existing categories. + ## Adding or rotating an account See `account-provisioning/README.md`. diff --git a/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.ps1 b/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.ps1 index fb6bbb9db0..80f1478e63 100644 --- a/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.ps1 +++ b/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.ps1 @@ -89,7 +89,8 @@ if ([string]::IsNullOrWhiteSpace($consistency)) { Fail "Account '$AccountSelecto if ([string]::IsNullOrWhiteSpace($testCategory)) { Fail "Account '$AccountSelector' is missing required 'testCategory'." } $connectionString = "AccountEndpoint=$endpoint;AccountKey=$key;" -$rustFlags = "--cfg=test_category=`"$testCategory`"" +# Only fixed live-account jobs use this resolver; emulator setup never adds "live". +$rustFlags = "--cfg=test_category=`"$testCategory`" --cfg=test_category=`"live`"" function Emit-Public([string]$name, [string]$value) { if ($Local) { diff --git a/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.tests.ps1 b/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.tests.ps1 index c722757f2e..8186911912 100644 --- a/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.tests.ps1 +++ b/sdk/cosmos/eng/pipelines/resolve-cosmos-test-account.tests.ps1 @@ -38,7 +38,7 @@ Write-Host "Test 1: resolves a valid selector and exports connection string + ru $result = Invoke-Resolver 'session-multiwrite' $sampleJson if ($result.ExitCode -eq 0 -and $result.Output -match 'AZURE_COSMOS_CONNECTION_STRING=AccountEndpoint=https://REPLACE-session-multiwrite' -and - $result.Output -match 'COSMOS_RUSTFLAGS=--cfg=test_category="multi_write"' -and + $result.Output -match 'COSMOS_RUSTFLAGS=--cfg=test_category="multi_write" --cfg=test_category="live"' -and $result.Output -match 'AZURE_COSMOS_DEFAULT_CONSISTENCY=Session') { Test-Ok "resolved connection string + rustflags + consistency" } @@ -83,6 +83,20 @@ Write-Host "Test 9: missing testCategory fails" $result = Invoke-Resolver 'x' '{"version":1,"accounts":{"x":{"endpoint":"https://x","key":"k","consistency":"Session"}}}' if ($result.ExitCode -ne 0) { Test-Ok "missing testCategory rejected" } else { Test-Fail "missing testCategory should fail" } +Write-Host "" +Write-Host "Test 10: every fixed live account enables the dedicated live target" +$accounts = ($sampleJson | ConvertFrom-Json).accounts +foreach ($property in $accounts.PSObject.Properties) { + $result = Invoke-Resolver $property.Name $sampleJson + $expected = 'COSMOS_RUSTFLAGS=--cfg=test_category="{0}" --cfg=test_category="live"' -f $property.Value.testCategory + if ($result.ExitCode -eq 0 -and $result.Output.Contains($expected)) { + Test-Ok "$($property.Name) preserves its category and enables live coverage" + } + else { + Test-Fail "$($property.Name) live selection" "(rc=$($result.ExitCode)): $($result.Output)" + } +} + Write-Host "" Write-Host "Results: $script:pass passed, $script:fail failed" if ($script:fail -ne 0) { exit 1 }