diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index e6861e92720..1ac160d3134 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/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 diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query_feed_range_scope.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_query_feed_range_scope.rs new file mode 100644 index 00000000000..e5b9a66707b --- /dev/null +++ b/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, Box> { + let mut iter = container + .query_items::(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> { + 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 = 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 = 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 = + 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::>() + ); + + // 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 = 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::>() + ); + + // 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::>() + ); + + Ok(()) + }, + Some(TestOptions::for_emulator()), + ) + .await +} diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/mod.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/mod.rs index 202f9475fe8..0654227dcf8 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/mod.rs +++ b/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"] diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index c61851bfdbd..9d0f1f51a4f 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/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)) 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 b2ea49f11fe..5e170d04984 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 @@ -380,6 +380,15 @@ async fn plan_fresh( let mut nodes: Vec> = 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 { + 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 diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs index f9b401b5b5c..2e194e9998b 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs @@ -93,22 +93,35 @@ impl OperationOverrides { continuation_as_if_none_match: bool, ) -> crate::error::Result<()> { if let Some(feed_range) = &self.feed_range { + let mut emitted_epk_bound = false; if feed_range.min_inclusive() != &EffectivePartitionKey::MIN { headers.insert( HeaderName::from_static(request_header_names::START_EPK), HeaderValue::from(feed_range.min_inclusive().as_str().to_owned()), ); + emitted_epk_bound = true; } if feed_range.max_exclusive() != &EffectivePartitionKey::MAX { headers.insert( HeaderName::from_static(request_header_names::END_EPK), HeaderValue::from(feed_range.max_exclusive().as_str().to_owned()), ); + emitted_epk_bound = true; + } + // An EPK *window* (start/end-epk) must be tagged as an + // `EffectivePartitionKeyRange` read — this is the value both the + // .NET (`RequestInvokerHandler`) and Java (`FeedRangeEpkImpl`) SDKs + // pair with start/end-epk. Sending the point value + // (`EffectivePartitionKey`) alongside a range window is an invalid, + // self-contradictory header set that the gateway rejects with a bare + // HTTP 400 (no sub-status). Only emit it when a bound was actually + // written; a whole-partition feed range carries neither. + if emitted_epk_bound { + headers.insert( + HeaderName::from_static(request_header_names::READ_FEED_KEY_TYPE), + HeaderValue::from_static("EffectivePartitionKeyRange"), + ); } - headers.insert( - HeaderName::from_static(request_header_names::READ_FEED_KEY_TYPE), - HeaderValue::from_static("EffectivePartitionKey"), - ); } if let Some(pk_range_id) = &self.partition_key_range_id { @@ -3673,7 +3686,7 @@ mod tests { request_header_names::READ_FEED_KEY_TYPE )) .map(|s| s.to_string()), - Some("EffectivePartitionKey".to_string()) + Some("EffectivePartitionKeyRange".to_string()) ); assert_eq!( headers @@ -3689,6 +3702,41 @@ mod tests { ); } + #[test] + fn apply_headers_whole_partition_feed_range_omits_read_key_type_and_epk_bounds() { + // A whole-partition feed range (`["", "FF")`) writes neither start-epk + // nor end-epk, so it must NOT emit `x-ms-read-key-type` either. The + // read-key-type tag is only meaningful alongside an interior EPK window; + // emitting the point value here is what produced the original HTTP 400 + // once scoped queries started routing through this path. + let feed_range = + FeedRange::new(EffectivePartitionKey::MIN, EffectivePartitionKey::MAX).unwrap(); + let overrides = OperationOverrides { + partition_key_range_id: Some("pkrange".to_string()), + feed_range: Some(feed_range), + ..Default::default() + }; + let mut headers = azure_core::http::headers::Headers::new(); + overrides + .apply_headers(&mut headers, false) + .expect("apply_headers should succeed"); + + assert!( + headers + .get_optional_str(&HeaderName::from_static( + request_header_names::READ_FEED_KEY_TYPE + )) + .is_none(), + "whole-partition feed range must not emit x-ms-read-key-type" + ); + assert!(headers + .get_optional_str(&HeaderName::from_static(request_header_names::START_EPK)) + .is_none()); + assert!(headers + .get_optional_str(&HeaderName::from_static(request_header_names::END_EPK)) + .is_none()); + } + #[test] fn build_transport_request_feed_path_is_resolved() { let operation = CosmosOperation::read_all_databases(test_account());