cosmos: scope cross-partition queries to their requested feed range

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### Bugs Fixed
 
+- Fixed cross-partition queries scoped to a sub-partition feed range (`ContainerClient::query_items` with `FeedScope::range(..)`) ignoring the scope and scanning the whole container, returning documents outside the requested effective-partition-key window. The planner now clips each query-plan range to the operation's scope, and the SDK tags the resulting interior `x-ms-start-epk`/`x-ms-end-epk` window with `x-ms-read-key-type: EffectivePartitionKeyRange` (the value .NET and Java pair with an EPK window) instead of the point value `EffectivePartitionKey`, which the gateway rejected with HTTP 400. ([#4726](https://github.com/Azure/azure-sdk-for-rust/pull/4726))
 - Fixed the `AZURE_COSMOS_PPCB_*` environment variables (including `AZURE_COSMOS_PPCB_ENABLED` and the `AZURE_COSMOS_PPCB_ENABLED_OVERRIDE` kill switch) being ignored when a `CosmosClient` was built without calling `CosmosClientBuilder::with_partition_failover_options`. The per-partition circuit breaker (PPCB) stayed enabled even with `AZURE_COSMOS_PPCB_ENABLED=false`. The client's driver now resolves these options from the environment when they are not supplied explicitly. ([#4655](https://github.com/Azure/azure-sdk-for-rust/pull/4655))
 
 ### Other Changes

sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query_feed_range_scope.rs

@@ -0,0 +1,185 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+//! End-to-end regression test for the cross-partition query **scope bug**: the
+//! planner used to ignore the caller's `FeedScope::range(..)` feed range.
+//!
+//! Before the fix, `plan_fresh` built its request ranges purely from the query
+//! plan (`query_plan.query_ranges`) and never intersected them with the
+//! operation's scope feed range (`operation.target()`). For a plain `SELECT *`
+//! the query plan reports the whole container, so a `FeedScope::range([X, Y))`
+//! window was silently dropped and the query did a **full scan** — returning
+//! documents outside the requested window.
+//!
+//! This test seeds one document per partition key, computes each document's
+//! effective partition key locally, and issues cross-partition queries scoped
+//! to interior EPK windows. It asserts the result set is exactly the windowed
+//! subset. Without the planner fix these assertions fail with
+//! `SCOPE BUG` (too many documents returned); with the fix — plus the
+//! `x-ms-read-key-type: EffectivePartitionKeyRange` header correction that lets
+//! the gateway accept the now-emitted interior EPK window — they pass.
+//!
+//! Run against the emulator:
+//!   RUSTFLAGS='--cfg test_category="emulator"' \
+//!     cargo test -p azure_data_cosmos --test emulator \
+//!     --features "key_auth fault_injection" feed_range_scope
+
+use super::framework;
+
+use std::borrow::Cow;
+use std::collections::BTreeSet;
+use std::error::Error;
+
+use azure_data_cosmos::clients::ContainerClient;
+use azure_data_cosmos::feed::{FeedRange, FeedScope};
+use azure_data_cosmos::models::{
+    ContainerProperties, EffectivePartitionKey, PartitionKeyDefinition,
+};
+use azure_data_cosmos::{PartitionKey, Query};
+use futures::TryStreamExt;
+use serde::{Deserialize, Serialize};
+
+use framework::{TestClient, TestOptions};
+
+const DOC_COUNT: usize = 40;
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+struct Doc {
+    id: String,
+    pk: String,
+}
+
+fn pk_definition() -> PartitionKeyDefinition {
+    PartitionKeyDefinition::new(vec![Cow::Borrowed("/pk")])
+}
+
+/// Runs a feed-range-scoped `SELECT *` and returns the set of returned ids.
+async fn drain_ids(
+    container: &ContainerClient,
+    scope: FeedScope,
+) -> Result<BTreeSet<String>, Box<dyn Error>> {
+    let mut iter = container
+        .query_items::<Doc>(Query::from("SELECT * FROM c"), scope, None)
+        .await?;
+    let mut got = BTreeSet::new();
+    while let Some(doc) = iter.try_next().await? {
+        got.insert(doc.id);
+    }
+    Ok(got)
+}
+
+#[tokio::test]
+#[cfg_attr(
+    not(any(test_category = "emulator", test_category = "emulator_vnext")),
+    ignore = "requires test_category 'emulator' or 'emulator_vnext'"
+)]
+pub async fn feed_range_scope_restricts_cross_partition_query() -> Result<(), Box<dyn Error>> {
+    TestClient::run_with_unique_db(
+        async |run_context, db_client| {
+            let properties = ContainerProperties::new("ScopeBugContainer", pk_definition());
+            let container = run_context
+                .create_container(db_client, properties, None)
+                .await?;
+
+            // Seed DOC_COUNT documents, each on its own partition key.
+            let mut docs: Vec<Doc> = Vec::with_capacity(DOC_COUNT);
+            for i in 0..DOC_COUNT {
+                let id = format!("doc-{i:03}");
+                let doc = Doc {
+                    id: id.clone(),
+                    pk: id.clone(),
+                };
+                container
+                    .create_item(PartitionKey::from(doc.pk.clone()), &id, doc.clone(), None)
+                    .await?;
+                docs.push(doc);
+            }
+
+            // Compute each document's effective partition key locally (same hash
+            // the service uses), then sort ascending by EPK.
+            let pk_def = pk_definition();
+            let mut points: Vec<(String, FeedRange)> = docs
+                .iter()
+                .map(|d| {
+                    let fr = FeedRange::for_partition(PartitionKey::from(d.pk.clone()), &pk_def);
+                    (d.id.clone(), fr)
+                })
+                .collect();
+            points.sort_by(|a, b| {
+                a.1.min_inclusive()
+                    .as_str()
+                    .cmp(b.1.min_inclusive().as_str())
+            });
+            let k = points.len();
+
+            // Control: a full-container scan must return every seeded document.
+            let all: BTreeSet<String> = points.iter().map(|(id, _)| id.clone()).collect();
+            let control = drain_ids(&container, FeedScope::full_container()).await?;
+            assert_eq!(
+                control, all,
+                "control full-container scan should return all {DOC_COUNT} docs"
+            );
+
+            // Test B: a WIDE interior window [X1, X_{k-1}) must exclude the
+            // globally smallest and largest EPKs => expect k-2 documents.
+            let window_b = FeedRange::new(
+                points[1].1.min_inclusive().clone(),
+                points[k - 1].1.min_inclusive().clone(),
+            )?;
+            let expected_b: BTreeSet<String> =
+                points[1..k - 1].iter().map(|(id, _)| id.clone()).collect();
+            let got_b = drain_ids(&container, FeedScope::range(window_b)).await?;
+            assert_eq!(
+                got_b,
+                expected_b,
+                "SCOPE BUG (wide window): FeedScope::range was ignored and the \
+                 query fell back to a full scan. got {} docs, expected {}. \
+                 unexpected(outside window)={:?}",
+                got_b.len(),
+                expected_b.len(),
+                got_b.difference(&expected_b).collect::<Vec<_>>()
+            );
+
+            // Test C: the TIGHTEST interior window [X_mid, X_{mid+1}) must
+            // return exactly one document.
+            let mid = k / 2;
+            let window_c = FeedRange::new(
+                points[mid].1.min_inclusive().clone(),
+                points[mid + 1].1.min_inclusive().clone(),
+            )?;
+            let expected_c: BTreeSet<String> = std::iter::once(points[mid].0.clone()).collect();
+            let got_c = drain_ids(&container, FeedScope::range(window_c)).await?;
+            assert_eq!(
+                got_c,
+                expected_c,
+                "SCOPE BUG (tight window): FeedScope::range was ignored. \
+                 got {} docs, expected 1. unexpected(outside window)={:?}",
+                got_c.len(),
+                got_c.difference(&expected_c).collect::<Vec<_>>()
+            );
+
+            // Test D: a window that lies ENTIRELY in the gap between two adjacent
+            // EPKs must return zero documents. `X_mid + "8"` is strictly greater
+            // than `X_mid` (it keeps `X_mid` as a prefix) and, because adjacent
+            // 32-hex EPKs differ within those 32 chars, strictly less than
+            // `X_{mid+1}` — so `[X_mid + "8", X_{mid+1})` contains no doc's EPK.
+            // This guards against an off-by-one where the lower bound is treated
+            // as inclusive of `X_mid`.
+            let gap_start =
+                EffectivePartitionKey::from(format!("{}8", points[mid].1.min_inclusive().as_str()));
+            let window_d = FeedRange::new(gap_start, points[mid + 1].1.min_inclusive().clone())?;
+            let got_d = drain_ids(&container, FeedScope::range(window_d)).await?;
+            assert!(
+                got_d.is_empty(),
+                "SCOPE BUG (empty gap window): a window between adjacent EPKs \
+                 returned {} docs, expected 0. unexpected={:?}",
+                got_d.len(),
+                got_d.iter().collect::<Vec<_>>()
+            );
+
+            Ok(())
+        },
+        Some(TestOptions::for_emulator()),
+    )
+    .await
+}

sdk/cosmos/azure_data_cosmos/tests/emulator_tests/mod.rs

@@ -13,6 +13,7 @@ mod cosmos_offers;
 mod cosmos_patch;
 mod cosmos_proxy;
 mod cosmos_query;
+mod cosmos_query_feed_range_scope;
 mod cosmos_response_metadata;
 
 #[path = "../framework/mod.rs"]

sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### Bugs Fixed
 
+- Fixed the cross-partition query planner ignoring an operation's scope feed range. `plan_fresh` and `plan_resume_from_saved_snapshot` built request ranges solely from the backend query plan's `query_ranges` (which cover the whole container for an unfiltered `SELECT`) and never intersected them with `CosmosOperation::target()`, so a query scoped to a sub-partition effective-partition-key window fanned out over the entire container instead of the requested slice. The planner now clips each query-plan range to the operation's scope (dropping ranges that fall entirely outside it; a full-key-space scope imposes no restriction). Additionally, an emitted `x-ms-start-epk`/`x-ms-end-epk` window is now tagged `x-ms-read-key-type: EffectivePartitionKeyRange` (matching .NET/Java) and only when a bound is actually written; previously the point value `EffectivePartitionKey` was sent alongside the range window, which the gateway rejected with a bare HTTP 400. ([#4726](https://github.com/Azure/azure-sdk-for-rust/pull/4726))
 - Fixed session-token parsing rejecting the `-1` region-progress sentinel that multi-region accounts emit for regions with no local progress. Region LSNs in the version vector are now parsed as signed values, with `-1` (and any negative value) modeled as "no progress" instead of failing the whole token as malformed. This applies to all Session-consistency traffic (not just distributed transactions) and prevents otherwise-valid multi-region session tokens from being discarded. ([#4702](https://github.com/Azure/azure-sdk-for-rust/pull/4702))
 - Fixed `AZURE_COSMOS_PPCB_*` environment variables (including the `AZURE_COSMOS_PPCB_ENABLED` master switch and the `AZURE_COSMOS_PPCB_ENABLED_OVERRIDE` kill switch) being silently ignored when a caller built `DriverOptions` without calling `DriverOptionsBuilder::with_partition_failover_options`. The environment is read only by `PartitionFailoverOptionsBuilder::build`, but the omitted-options path used a bare `PartitionFailoverOptions::default()` (which hard-codes PPCB enabled and reads no environment), so PPCB stayed on even with `AZURE_COSMOS_PPCB_ENABLED=false`. `DriverOptionsBuilder::build` now resolves the partition-failover options from the environment when the caller does not supply them (falling back to defaults, fail-soft, if an environment value is out of bounds). An explicitly supplied `PartitionFailoverOptions` continues to take precedence. ([#4655](https://github.com/Azure/azure-sdk-for-rust/pull/4655))
 

sdk/cosmos/azure_data_cosmos_driver/src/driver/dataflow/planner.rs

@@ -380,6 +380,15 @@ async fn plan_fresh(
     let mut nodes: Vec<Box<dyn PipelineNode>> = Vec::new();
     for query_range in &query_plan.query_ranges {
         let feed_range = query_range_to_feed_range(query_range)?;
+        // Clip the query-plan range to the operation's requested scope (e.g. a
+        // sub-partition EPK window from `FeedScope::range`). The query plan's
+        // ranges describe what the query *text* touches (often the whole
+        // container for a `SELECT`), so they must be narrowed to the caller's
+        // scope; otherwise a scoped query fans out over the whole container. A
+        // range that falls entirely outside the scope contributes nothing.
+        let Some(feed_range) = clip_to_operation_scope(operation, feed_range) else {
+            continue;
+        };
         let resolved = topology_provider
             .resolve_ranges(&feed_range, PartitionRoutingRefresh::UseCached)
             .await?;
@@ -429,6 +438,11 @@ async fn plan_resume_from_saved_snapshot(
 
     for query_range in &query_plan.query_ranges {
         let feed_range = query_range_to_feed_range(query_range)?;
+        // Clip the query-plan range to the operation's requested scope, mirroring
+        // `plan_fresh` so a resumed sub-partition-scoped query stays scoped.
+        let Some(feed_range) = clip_to_operation_scope(operation, feed_range) else {
+            continue;
+        };
         let resolved = topology_provider
             .resolve_ranges(&feed_range, PartitionRoutingRefresh::UseCached)
             .await?;
@@ -572,6 +586,36 @@ fn query_range_to_feed_range(
     FeedRange::new(min, max)
 }
 
+/// Returns true when `scope` covers the entire EPK key space (`""..FF`), i.e.
+/// imposes no restriction on a cross-partition query.
+fn scope_covers_full_key_space(scope: &FeedRange) -> bool {
+    let full = FeedRange::full();
+    scope.min_inclusive() == full.min_inclusive() && scope.max_exclusive() == full.max_exclusive()
+}
+
+/// Clips a query-plan range to the operation's requested scope feed range.
+///
+/// A cross-partition query may be scoped to a user-supplied [`FeedRange`] (e.g.
+/// a sub-partition EPK window from `FeedScope::range`). The query plan's ranges
+/// describe what the query *text* touches (often the whole container for a
+/// `SELECT`), so the ranges we actually query must be clipped to the requested
+/// scope. A scope covering the whole key space imposes no restriction.
+///
+/// Returns `None` when the query range lies entirely outside the scope — that
+/// range contributes nothing and the caller skips it.
+fn clip_to_operation_scope(
+    operation: &CosmosOperation,
+    feed_range: FeedRange,
+) -> Option<FeedRange> {
+    let Some(scope) = operation
+        .target()
+        .filter(|s| !scope_covers_full_key_space(s))
+    else {
+        return Some(feed_range);
+    };
+    intersect_feed_ranges(&feed_range, scope)
+}
+
 /// Returns true if the union of `pieces` covers `range` end-to-end.
 ///
 /// Assumes pieces are subsets of `range`. The check sorts pieces by
@@ -1072,6 +1116,59 @@ mod tests {
         assert_drain_requests(pipeline, &[("", "FF", "pkrange-0")]);
     }
 
+    #[tokio::test]
+    async fn scopes_cross_partition_query_to_requested_feed_range() {
+        // Regression: a cross-partition `SELECT` scoped to a sub-partition EPK
+        // window (`FeedScope::range`) must clip the query-plan range to that
+        // window and emit an interior `start`/`end-epk` slice — not fan out
+        // over the whole container. Before the fix, the operation's scope feed
+        // range was ignored entirely and the query did a full scan.
+        let plan = plan_with_ranges(vec![qr("", "FF")]);
+        let op = CosmosOperation::query_items(
+            test_container(),
+            Some(
+                FeedRange::new(
+                    EffectivePartitionKey::from("20"),
+                    EffectivePartitionKey::from("80"),
+                )
+                .unwrap(),
+            ),
+        )
+        .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec());
+        // The planner resolves the CLIPPED window; the mock returns the single
+        // owning partition (the whole `["", "FF")`).
+        let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr("", "FF", "pkrange-0")])]);
+
+        let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None)
+            .await
+            .unwrap();
+
+        // Emitted EPK slice is clipped to the scope window, inside the whole
+        // partition (partition range stays `["", "FF")`).
+        assert_drain_requests_with_partitions(pipeline, &[("20", "80", "pkrange-0", "", "FF")]);
+    }
+
+    #[tokio::test]
+    async fn full_container_scope_does_not_restrict_cross_partition_query() {
+        // A `FeedScope::full_container()` scope (`FeedRange::full()`) must
+        // impose no restriction: the query fans out over every partition just
+        // as an unscoped cross-partition query does.
+        let plan = plan_with_ranges(vec![qr("", "FF")]);
+        let op = cross_partition_query_operation();
+        let mut topology = MockTopologyProvider::new(vec![Ok(vec![
+            rr("", "80", "pkrange-left"),
+            rr("80", "FF", "pkrange-right"),
+        ])]);
+
+        let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None)
+            .await
+            .unwrap();
+        assert_drain_requests(
+            pipeline,
+            &[("", "80", "pkrange-left"), ("80", "FF", "pkrange-right")],
+        );
+    }
+
     #[tokio::test]
     async fn builds_sequential_drain_for_multiple_partitions() {
         // Query targets full range, topology has two partitions split at "80".
@@ -1469,6 +1566,43 @@ mod tests {
         assert_drain_requests(pipeline, &[("55", "AA", "pk-b"), ("AA", "FF", "pk-c")]);
     }
 
+    #[tokio::test]
+    async fn resume_scopes_cross_partition_query_to_requested_feed_range() {
+        // Regression companion to `scopes_cross_partition_query_to_requested_feed_range`
+        // for the RESUME path: `plan_resume_from_saved_snapshot` must clip the
+        // query-plan range to the operation's `FeedScope::range` window too.
+        //
+        // The saved token covers exactly the scoped window `[20, 80)`. Without
+        // the clip, the query range stays `["", "FF")`, so the leaf above the
+        // cursor would additionally emit a fresh-start `[80, FF)` sub-leaf. With
+        // the clip, only the single `[20, 80)` leaf (carrying the saved token) is
+        // emitted.
+        let plan = plan_with_ranges(vec![qr("", "FF")]);
+        let op = CosmosOperation::query_items(
+            test_container(),
+            Some(
+                FeedRange::new(
+                    EffectivePartitionKey::from("20"),
+                    EffectivePartitionKey::from("80"),
+                )
+                .unwrap(),
+            ),
+        )
+        .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec());
+        let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr("", "FF", "pkrange-0")])]);
+
+        let resume = saved_drain(vec![("20", "80", saved_request(Some("server-token-xyz")))]);
+
+        let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), Some(resume))
+            .await
+            .unwrap();
+
+        assert_drain_requests_with_partitions_and_continuation(
+            pipeline,
+            &[("20", "80", "pkrange-0", "", "FF", Some("server-token-xyz"))],
+        );
+    }
+
     #[tokio::test]
     async fn resume_propagates_server_continuation_to_every_surviving_leaf_after_split() {
         // The saved `[55, AA)` child held a server continuation. Between