Gate unbounded buffered Cosmos queries

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<Duration>`; `TransactionalBatchOperationResult::retry_after_milliseconds()` and `DistributedTransactionResponse::retry_after_ms()` are similarly replaced by `retry_after() -> Option<Duration>`. `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

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"

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<crate::options::QueryPlanMode>,
         pub patch_strategy: Option<crate::options::PatchStrategy>,
         pub read_consistency_strategy: Option<crate::options::ReadConsistencyStrategy>,
         pub excluded_regions: Option<crate::options::ExcludedRegions>,
@@ -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<OperationOptions>>, runtime: Option<::std::sync::Arc<OperationOptions>>, account: Option<::std::sync::Arc<OperationOptions>>, operation: Option<&'a OperationOptions>) -> Self;
         pub fn new_with_override(env_override: Option<::std::sync::Arc<OperationOptions>>, env: Option<::std::sync::Arc<OperationOptions>>, runtime: Option<::std::sync::Arc<OperationOptions>>, account: Option<::std::sync::Arc<OperationOptions>>, 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<azure_data_cosmos_driver::models::SessionToken>,
@@ -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<impl Into<SessionToken>: Into<SessionToken>>(self, session_token: impl Into<SessionToken>) -> Self;
     }
+    impl Default for QueryOptions {
+        fn default() -> Self;
+    }
     #[derive(Clone, Default)]
     #[non_exhaustive]
     pub struct ReadContainerOptions {

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

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,

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]

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<bool>,
 }
 
+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<SessionToken>) -> 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() {

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<T: DeserializeOwned + Send + 'static>(
     sql: &str,
     run_id: &str,
     context: &str,
+    max_buffered_query_window: u64,
 ) -> Result<Vec<T>, Box<dyn Error>> {
     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<dyn Error>> {
             }
         }
 
+        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<dyn Error>> {
                     .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<dyn Error>> {
                 "SELECT VALUE {\"int\": 7} FROM c WHERE c.fuzzRun = @run",
                 &run_id,
                 &context,
+                query_window,
             )
             .await?;
             assert!(

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::<serde_json::Value>(
-                &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::<Vec<_>>();
+            let mut countries = container
+                .query_items::<String>(
+                    "SELECT DISTINCT TOP 1000 VALUE c.country FROM c",
+                    FeedScope::full_container(),
+                    None,
+                )
+                .await?
+                .try_collect::<Vec<_>>()
+                .await?;
             countries.sort();
             assert_eq!(
                 countries,