Cosmos: Regression tests for feed-range-scoped query fan-out

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### Bugs Fixed
 
+- Fixed cross-partition queries scoped to an explicit feed range (`ContainerClient::query_items` with `FeedScope::range`) fanning out to partitions outside the requested range. The query planner now intersects the query plan's ranges with the requested feed-range target, so only the partitions overlapping the requested range are queried. This prevents returning documents outside the requested feed range and the extra scans / RU usage that occurred when enumerating `read_feed_ranges()` and querying each range. ([#4705](https://github.com/Azure/azure-sdk-for-rust/pull/4705))
 - 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_driver/CHANGELOG.md

@@ -12,6 +12,7 @@
 
 ### Bugs Fixed
 
+- Fixed cross-partition query fan-out ignoring an operation's explicit feed-range target. When a query was scoped to an explicit EPK range (`CosmosOperation::query_items` with a `FeedRange` target, surfaced as `FeedScope::range`), the query planner built request leaves from the query plan's ranges — which are derived from the query text alone — without intersecting them against the requested target. Both fan-out paths (`plan_fresh` and the resume path) now clip each query-plan range to the operation's target, so only partitions overlapping the requested range are queried. Absent, full, and logical-partition/prefix targets are unaffected. ([#4705](https://github.com/Azure/azure-sdk-for-rust/pull/4705))
 - 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/mocks.rs

@@ -193,6 +193,45 @@ impl TopologyProvider for MockTopologyProvider {
     }
 }
 
+/// A range-aware topology provider backed by a fixed physical partition layout.
+///
+/// Unlike [`MockTopologyProvider`] (which ignores the requested range and
+/// replays a queue), this resolves *against* the requested range:
+/// `resolve_ranges(range)` returns exactly the configured physical partitions
+/// whose EPK span overlaps `range`, mirroring real routing. This lets fan-out
+/// tests observe when the planner resolves partitions outside the requested
+/// scope — a plain replay mock cannot, because it returns the same partitions
+/// regardless of the range it is asked about.
+pub(crate) struct PhysicalTopologyProvider {
+    partitions: Vec<ResolvedRange>,
+}
+
+impl PhysicalTopologyProvider {
+    pub fn new(partitions: Vec<ResolvedRange>) -> Self {
+        Self { partitions }
+    }
+}
+
+impl TopologyProvider for PhysicalTopologyProvider {
+    fn resolve_ranges<'a>(
+        &'a mut self,
+        range: &'a FeedRange,
+        _refresh: PartitionRoutingRefresh,
+    ) -> BoxFuture<'a, crate::error::Result<Vec<ResolvedRange>>> {
+        let resolved: Vec<ResolvedRange> = self
+            .partitions
+            .iter()
+            .filter(|p| {
+                // Half-open overlap: [a, b) intersects [c, d) iff a < d && c < b.
+                p.range.min_inclusive() < range.max_exclusive()
+                    && range.min_inclusive() < p.range.max_exclusive()
+            })
+            .cloned()
+            .collect();
+        Box::pin(async move { Ok(resolved) })
+    }
+}
+
 // ── Test helpers ────────────────────────────────────────────────────────────
 
 /// Extracts the `CosmosResponse` from a `PageResult::Page`, panicking otherwise.

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

@@ -380,6 +380,11 @@ 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)?;
+        // Restrict the query-plan range to the operation's requested scope.
+        // Ranges entirely outside the target contribute no request leaves.
+        let Some(feed_range) = clip_to_target(feed_range, operation) else {
+            continue;
+        };
         let resolved = topology_provider
             .resolve_ranges(&feed_range, PartitionRoutingRefresh::UseCached)
             .await?;
@@ -429,6 +434,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)?;
+        // Restrict the query-plan range to the operation's requested scope,
+        // mirroring the fresh path. Ranges outside the target are skipped.
+        let Some(feed_range) = clip_to_target(feed_range, operation) else {
+            continue;
+        };
         let resolved = topology_provider
             .resolve_ranges(&feed_range, PartitionRoutingRefresh::UseCached)
             .await?;
@@ -572,6 +582,28 @@ fn query_range_to_feed_range(
     FeedRange::new(min, max)
 }
 
+/// Clips a query-plan feed range to the operation's requested scope
+/// ([`CosmosOperation::target`]).
+///
+/// The backend query plan's ranges are derived from the query text alone and
+/// never reflect the caller's requested [`FeedRange`] target, so the planner
+/// must intersect them here. Returns `None` when the query range lies entirely
+/// outside the target, meaning it contributes no request leaves.
+///
+/// When the operation has no target, or targets a logical partition / prefix,
+/// the range is returned unchanged: a logical-partition [`FeedRange`] collapses
+/// its `[min, max)` bounds to a single effective partition key and therefore
+/// cannot be intersected as an EPK range. Prefix-scoped fan-out is handled
+/// elsewhere; clipping it here would incorrectly drop every range.
+fn clip_to_target(feed_range: FeedRange, operation: &CosmosOperation) -> Option<FeedRange> {
+    match operation.target() {
+        Some(target) if !target.is_logical_partition() => {
+            intersect_feed_ranges(&feed_range, target)
+        }
+        _ => Some(feed_range),
+    }
+}
+
 /// Returns true if the union of `pieces` covers `range` end-to-end.
 ///
 /// Assumes pieces are subsets of `range`. The check sorts pieces by
@@ -864,6 +896,18 @@ mod tests {
             .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec())
     }
 
+    /// A cross-partition query scoped to an explicit EPK feed-range target,
+    /// e.g. `FeedScope::range([min, max))`.
+    fn query_operation_with_target(min: &str, max: &str) -> CosmosOperation {
+        let target = FeedRange::new(
+            EffectivePartitionKey::from(min),
+            EffectivePartitionKey::from(max),
+        )
+        .unwrap();
+        CosmosOperation::query_items(test_container(), Some(target))
+            .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec())
+    }
+
     // --- build_trivial_pipeline tests ---
 
     #[test]
@@ -1180,6 +1224,145 @@ mod tests {
         assert_drain_requests_with_partitions(pipeline, &[("20", "80", "pkrange-wide", "", "FF")]);
     }
 
+    #[tokio::test]
+    async fn restricts_fanout_to_explicit_target_range() {
+        // The reported bug: query plan spans the whole space `[, FF)` but the
+        // caller scoped the query to `[00, 80)` via `FeedScope::range`. Only
+        // the requested slice must be queried, not the neighbouring `[80, FF)`
+        // partition. The physical topology actually contains both partitions,
+        // so a planner that ignores the target would resolve and emit a leaf
+        // for `[80, FF)` as well.
+        let plan = plan_with_ranges(vec![qr("", "FF")]);
+        let op = query_operation_with_target("00", "80");
+        let mut topology = PhysicalTopologyProvider::new(vec![
+            rr("00", "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, &[("00", "80", "pkrange-left")]);
+    }
+
+    #[tokio::test]
+    async fn drops_query_ranges_outside_target() {
+        // Two disjoint query-plan ranges; the target only overlaps the first.
+        // The second range must contribute no request leaves (and its topology
+        // must never be resolved).
+        let plan = plan_with_ranges(vec![qr("", "40"), qr("80", "FF")]);
+        let op = query_operation_with_target("", "40");
+        let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr("", "40", "pkrange-A")])]);
+
+        let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None)
+            .await
+            .unwrap();
+        assert_drain_requests(pipeline, &[("", "40", "pkrange-A")]);
+    }
+
+    #[tokio::test]
+    async fn clips_query_range_to_partial_target_overlap() {
+        // The target `[20, 60)` partially overlaps a single query-plan range
+        // spanning several partitions. Only partitions within the target,
+        // clipped to its bounds, may be queried.
+        let plan = plan_with_ranges(vec![qr("", "FF")]);
+        let op = query_operation_with_target("20", "60");
+        let mut topology = MockTopologyProvider::new(vec![Ok(vec![
+            rr("00", "40", "pkrange-1"),
+            rr("40", "80", "pkrange-2"),
+        ])]);
+
+        let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None)
+            .await
+            .unwrap();
+        assert_drain_requests_with_partitions(
+            pipeline,
+            &[
+                ("20", "40", "pkrange-1", "00", "40"),
+                ("40", "60", "pkrange-2", "40", "80"),
+            ],
+        );
+    }
+
+    #[tokio::test]
+    async fn drops_query_range_touching_target_boundary() {
+        // A query-plan range that only *touches* the target's exclusive upper
+        // bound (target `[, 40)`, range `[40, FF)`) shares no EPKs with the
+        // target and must be dropped, not queried. Exercises the exact
+        // boundary case where `intersect_feed_ranges` collapses to empty.
+        let plan = plan_with_ranges(vec![qr("", "40"), qr("40", "FF")]);
+        let op = query_operation_with_target("", "40");
+        let mut topology = PhysicalTopologyProvider::new(vec![
+            rr("", "40", "pkrange-A"),
+            rr("40", "FF", "pkrange-B"),
+        ]);
+
+        let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None)
+            .await
+            .unwrap();
+        assert_drain_requests(pipeline, &[("", "40", "pkrange-A")]);
+    }
+
+    #[tokio::test]
+    async fn logical_partition_target_is_not_clipped() {
+        // A logical-partition (prefix) target must NOT be clipped: its
+        // `[min, max)` bounds collapse to a single EPK, so intersecting the
+        // query-plan ranges against it would drop everything. The guard in
+        // `clip_to_target` leaves such ranges untouched, letting the query
+        // fan out across the resolved topology. This also documents that
+        // prefix-scoped over-scan is handled elsewhere (see follow-up).
+        let plan = plan_with_ranges(vec![qr("", "FF")]);
+        let target =
+            FeedRange::for_partition(PartitionKey::from("pk1"), &test_partition_key_definition());
+        assert!(target.is_logical_partition());
+        let op = CosmosOperation::query_items(test_container(), Some(target))
+            .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec());
+        let mut topology = PhysicalTopologyProvider::new(vec![
+            rr("00", "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,
+            &[("00", "80", "pkrange-left"), ("80", "FF", "pkrange-right")],
+        );
+    }
+
+    #[tokio::test]
+    async fn resume_restricts_fanout_to_target_range() {
+        // Resuming a target-scoped query must also honour the target: the
+        // `[80, FF)` partition lies outside `[00, 80)` and must not be queried
+        // even though the query plan spans `[, FF)` and that partition exists
+        // in the physical topology. An unclipped resume would emit a second,
+        // fresh-start leaf for `[80, FF)`.
+        let plan = plan_with_ranges(vec![qr("", "FF")]);
+        let op = query_operation_with_target("00", "80");
+        let mut topology = PhysicalTopologyProvider::new(vec![
+            rr("00", "80", "pkrange-left"),
+            rr("80", "FF", "pkrange-right"),
+        ]);
+
+        let resume = saved_drain(vec![("00", "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,
+            &[(
+                "00",
+                "80",
+                "pkrange-left",
+                "00",
+                "80",
+                Some("server-token-xyz"),
+            )],
+        );
+    }
+
     #[tokio::test]
     async fn rejects_query_plan_with_top() {
         let plan = QueryPlan {
@@ -1346,6 +1529,27 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn rejects_target_disjoint_from_query_ranges() {
+        // A target that shares no EPKs with any query-plan range clips every
+        // range away, leaving zero request leaves. On the fresh path this is
+        // reported as the same hard error as an empty query plan — the clip
+        // does not silently swallow the query. `NoopTopologyProvider` asserts
+        // no partition is ever resolved (the clip skips before resolution).
+        let plan = plan_with_ranges(vec![qr("80", "FF")]);
+        let op = query_operation_with_target("00", "40");
+        let mut topology = NoopTopologyProvider;
+
+        let err = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None)
+            .await
+            .unwrap_err();
+        let rendered = err.to_string();
+        assert!(
+            rendered.ends_with("query plan produced no partition ranges to query"),
+            "unexpected: {rendered}"
+        );
+    }
+
     #[tokio::test]
     async fn propagates_topology_resolution_error() {
         let plan = plan_with_ranges(vec![qr("", "FF")]);