From a7bc4f420a1871700745e36c76352f8a95311e65 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 9 Jul 2026 10:49:38 -0400 Subject: [PATCH 1/5] Fix HPK prefix and cross-partition query bugs Both bugs share a single root cause: the driver emitted the point key-type `EffectivePartitionKey` in `x-ms-read-key-type` while sending `x-ms-start-epk`/`x-ms-end-epk`, which describe an EPK *range*. The gateway rejects that combination with `400 "One of the input values is invalid"`. The correct value is `EffectivePartitionKeyRange` (validated against a live Cosmos account for both query and change-feed paths; no api-version bump needed). - azure_data_cosmos_driver/operation_pipeline.rs: emit `EffectivePartitionKeyRange` when EPK bounds are present; update the apply_headers unit test. - azure_data_cosmos_driver/planner.rs: prefix EPK-range clipping (prefix_epk_range / clip_to_prefix) so a `FeedScope::partition` prefix fans out over `[prefix_epk, prefix_epk + "FF")` instead of scanning whole physical partitions. - azure_data_cosmos/feed/query.rs: document prefix support on `FeedScope::partition`. - azure_data_cosmos in-memory emulator tests: end-to-end HPK prefix and cross-partition coverage through the public SDK. - CHANGELOG entries in both crates. Fixes #4680 Fixes #4681 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 1 + .../azure_data_cosmos/src/feed/query.rs | 9 +- .../tests/in_memory_emulator_tests/hpk.rs | 247 ++++++++++++++++++ .../tests/in_memory_emulator_tests/mod.rs | 1 + .../azure_data_cosmos_driver/CHANGELOG.md | 1 + .../src/driver/dataflow/planner.rs | 207 ++++++++++++++- .../src/driver/pipeline/operation_pipeline.rs | 9 +- 7 files changed, 469 insertions(+), 6 deletions(-) create mode 100644 sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index e6861e92720..3509bc7c423 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -12,6 +12,7 @@ ### Bugs Fixed +- Fixed hierarchical-partition-key (HPK) queries. A `FeedScope::partition` scope with only a *prefix* of the key hierarchy (e.g. `("USA", "CA")` on a `/country/state/city` container) now filters to that prefix instead of returning every item in the physical partition (issue [#4680](https://github.com/Azure/azure-sdk-for-rust/issues/4680)), and cross-partition queries over an HPK container no longer fail with `400 Bad Request` (issue [#4681](https://github.com/Azure/azure-sdk-for-rust/issues/4681)). - 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/src/feed/query.rs b/sdk/cosmos/azure_data_cosmos/src/feed/query.rs index a56623de8c0..b2df2c7bada 100644 --- a/sdk/cosmos/azure_data_cosmos/src/feed/query.rs +++ b/sdk/cosmos/azure_data_cosmos/src/feed/query.rs @@ -22,8 +22,13 @@ pub enum FeedScope { impl FeedScope { /// Returns a [`FeedScope`] that represents the given partition key, which is used for targeting a specific partition in the container. /// - /// The provided [`PartitionKey`] MUST specify all levels of the hierarchy (e.g. in a multi-level hierarchical partition key, you must provide values for all levels, not just a prefix). - /// Use [`range()`](FeedScope::range) with a [`FeedRange`] that covers the desired partition(s) to specify anything beyond a single logical partition. + /// The provided [`PartitionKey`] may specify all levels of a hierarchical + /// (multi-level) partition key to target a single logical partition, or only + /// a *prefix* of the levels (e.g. `("USA", "CA")` on a `/country/state/city` + /// container) to target every logical partition that shares that prefix. A + /// prefix scope is executed as a filtered cross-partition query bounded to the + /// prefix's effective-partition-key range, so it only returns items under the + /// prefix rather than scanning whole physical partitions. pub fn partition(pk: impl Into) -> Self { Self::Partition(pk.into()) } diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs new file mode 100644 index 00000000000..68c08e77603 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs @@ -0,0 +1,247 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +//! Hierarchical-partition-key (HPK / MultiHash) query coverage driven through +//! the public SDK against the in-memory emulator. +//! +//! These tests exercise the fixes for: +//! - #4680 — prefix scopes (`FeedScope::partition` with fewer components than +//! the container's partition-key hierarchy) must filter to the prefix instead +//! of returning every item in the physical partition. +//! - #4681 — a full-container cross-partition query over an HPK container must +//! fan out successfully instead of failing. +//! +//! The in-memory emulator honors `x-ms-start-epk`/`x-ms-end-epk` filtering for +//! MultiHash containers, so the prefix counts below are deterministic without a +//! live account. + +use azure_core::credentials::Secret; +use azure_core::http::Url; +use azure_data_cosmos::{ + options::Region, AccountEndpoint, AccountReference, ContainerClient, CosmosClientBuilder, + CosmosRuntimeBuilder, FeedScope, PartitionKey, Query, RoutingStrategy, +}; +use azure_data_cosmos_driver::in_memory_emulator::{ + ConsistencyLevel, ContainerConfig, InMemoryEmulatorHttpClient, VirtualAccountConfig, + VirtualRegion, +}; +use futures::TryStreamExt; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +const EMULATOR_GATEWAY_URL: &str = "https://eastus.emulator.local"; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +struct GeoItem { + id: String, + country: String, + state: String, + city: String, +} + +impl GeoItem { + fn new(country: &str, state: &str, city: &str) -> Self { + Self { + id: format!("{country}-{state}-{city}"), + country: country.to_string(), + state: state.to_string(), + city: city.to_string(), + } + } +} + +/// The 9-row dataset shared by every HPK test: +/// - USA/CA: 5 rows +/// - USA/WA: 2 rows +/// - USA/NY: 1 row → USA total = 8 +/// - CANADA/ON: 1 row → grand total = 9 +fn dataset() -> Vec { + vec![ + GeoItem::new("USA", "CA", "SanFrancisco"), + GeoItem::new("USA", "CA", "LosAngeles"), + GeoItem::new("USA", "CA", "SanDiego"), + GeoItem::new("USA", "CA", "SanJose"), + GeoItem::new("USA", "CA", "Oakland"), + GeoItem::new("USA", "WA", "Seattle"), + GeoItem::new("USA", "WA", "Tacoma"), + GeoItem::new("USA", "NY", "NewYork"), + GeoItem::new("CANADA", "ON", "Toronto"), + ] +} + +/// Provisions a `/country/state/city` MultiHash container (4 physical +/// partitions) seeded with [`dataset`], and returns a ready SDK +/// [`ContainerClient`] wired to the in-memory emulator. +async fn setup_hpk_container() -> ContainerClient { + let run_id = Uuid::new_v4().to_string()[..8].to_string(); + + let config = VirtualAccountConfig::new(vec![VirtualRegion::new( + "East US", + Url::parse(EMULATOR_GATEWAY_URL).unwrap(), + )]) + .unwrap() + .with_consistency(ConsistencyLevel::Session); + + let emulator = std::sync::Arc::new(InMemoryEmulatorHttpClient::new(config)); + let emulator_store = emulator.store(); + + let db_name = format!("hpk-{run_id}"); + emulator_store.create_database(&db_name); + emulator_store.create_container_with_config( + &db_name, + "geo", + serde_json::from_value(serde_json::json!({ + "paths": ["/country", "/state", "/city"], + "kind": "MultiHash", + "version": 2 + })) + .unwrap(), + ContainerConfig::new() + .with_partition_count(4) + .build() + .unwrap(), + ); + + let emulator_account = AccountReference::with_authentication_key( + EMULATOR_GATEWAY_URL.parse::().unwrap(), + Secret::new("dGVzdGtleQ=="), + ); + let client = CosmosClientBuilder::new() + .with_runtime( + CosmosRuntimeBuilder::from(emulator.runtime_builder()) + .build() + .await + .unwrap(), + ) + .build( + emulator_account, + RoutingStrategy::ProximityTo(Region::EAST_US), + ) + .await + .unwrap(); + + let container = client + .database_client(&db_name) + .container_client("geo") + .await + .unwrap(); + + for item in dataset() { + container + .create_item( + PartitionKey::from((item.country.clone(), item.state.clone(), item.city.clone())), + &item.id, + &item, + None, + ) + .await + .unwrap(); + } + + container +} + +async fn query_scope(container: &ContainerClient, scope: FeedScope) -> Vec { + container + .query_items(Query::from("SELECT * FROM c"), scope, None) + .await + .unwrap() + .try_collect() + .await + .unwrap() +} + +/// #4680 — a one-level prefix `(USA,)` returns only the 8 USA rows, not CANADA. +#[tokio::test] +async fn hpk_query_prefix_level1() { + let container = setup_hpk_container().await; + let items = query_scope(&container, FeedScope::partition(PartitionKey::from("USA"))).await; + + assert_eq!(items.len(), 8, "prefix (USA,) should return 8 items"); + assert!( + items.iter().all(|i| i.country == "USA"), + "prefix (USA,) leaked non-USA rows: {items:?}" + ); +} + +/// #4680 — a two-level prefix `(USA, CA)` returns only the 5 CA rows. +#[tokio::test] +async fn hpk_query_prefix_level2() { + let container = setup_hpk_container().await; + let items = query_scope( + &container, + FeedScope::partition(PartitionKey::from(("USA", "CA"))), + ) + .await; + + assert_eq!(items.len(), 5, "prefix (USA, CA) should return 5 items"); + assert!( + items.iter().all(|i| i.country == "USA" && i.state == "CA"), + "prefix (USA, CA) leaked non-CA rows: {items:?}" + ); +} + +/// #4680 — a prefix with no matching rows `(USA, TX)` returns nothing. +#[tokio::test] +async fn hpk_query_prefix_no_match_returns_empty() { + let container = setup_hpk_container().await; + let items = query_scope( + &container, + FeedScope::partition(PartitionKey::from(("USA", "TX"))), + ) + .await; + + assert!(items.is_empty(), "prefix (USA, TX) should return 0 items"); +} + +/// #4680 — explicit anti-leak guard: the `(USA, CA)` result must exclude every +/// WA, NY, and CANADA row. +#[tokio::test] +async fn hpk_query_prefix_correctness_guard() { + let container = setup_hpk_container().await; + let items = query_scope( + &container, + FeedScope::partition(PartitionKey::from(("USA", "CA"))), + ) + .await; + + for item in &items { + assert_eq!(item.country, "USA", "leaked foreign country: {item:?}"); + assert_eq!(item.state, "CA", "leaked foreign state: {item:?}"); + } + assert!( + !items + .iter() + .any(|i| i.state == "WA" || i.state == "NY" || i.country == "CANADA"), + "prefix (USA, CA) leaked WA/NY/CANADA rows: {items:?}" + ); +} + +/// #4681 — a full-container cross-partition query over an HPK container fans out +/// successfully and returns every row. +#[tokio::test] +async fn hpk_query_cross_partition_full_container() { + let container = setup_hpk_container().await; + let items = query_scope(&container, FeedScope::full_container()).await; + + assert_eq!( + items.len(), + 9, + "full-container query over an HPK container should return all 9 items" + ); +} + +/// Sanity: a complete key still targets a single logical partition and returns +/// exactly its one row. +#[tokio::test] +async fn hpk_query_full_key_single_partition() { + let container = setup_hpk_container().await; + let items = query_scope( + &container, + FeedScope::partition(PartitionKey::from(("USA", "CA", "SanFrancisco"))), + ) + .await; + + assert_eq!(items.len(), 1, "complete key should return exactly 1 item"); + assert_eq!(items[0].city, "SanFrancisco"); +} diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs index 894188df569..44800b3eb3b 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/mod.rs @@ -13,6 +13,7 @@ pub mod dtx_live_comparison; pub mod dtx_sdk_validation; pub mod dual_backend; pub mod end_to_end; +pub mod hpk; pub mod session_token; pub mod user_agent; pub mod validation; diff --git a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index c61851bfdbd..eaabc9311a0 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 hierarchical-partition-key (HPK / MultiHash) queries emitting the wrong `x-ms-read-key-type` for effective-partition-key range-scoped requests. The driver sent the point value `EffectivePartitionKey` alongside `x-ms-start-epk`/`x-ms-end-epk`, which the gateway rejects with `400 "One of the input values is invalid"`; it now sends `EffectivePartitionKeyRange`. This enables filtered partition-key *prefix* queries (issue [#4680](https://github.com/Azure/azure-sdk-for-rust/issues/4680)) and fixes cross-partition queries over HPK containers failing with `400 Bad Request` (issue [#4681](https://github.com/Azure/azure-sdk-for-rust/issues/4681)). - 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..c31e08e4249 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 @@ -173,10 +173,30 @@ pub(crate) async fn build_sequential_drain( } }; + // When the operation targets a hierarchical-partition-key *prefix* (an + // incomplete logical partition key), the fan-out must be scoped to the EPK + // range that covers every completion of that prefix. Without this clip the + // query scans whole physical partitions and returns rows outside the prefix + // (issue #4680). + let prefix_clip = prefix_epk_range(operation)?; + let request_nodes = if let Some(saved) = saved_snapshot.as_ref() { - plan_resume_from_saved_snapshot(query_plan, topology_provider, operation, saved).await? + plan_resume_from_saved_snapshot( + query_plan, + topology_provider, + operation, + saved, + prefix_clip.as_ref(), + ) + .await? } else { - plan_fresh(query_plan, topology_provider, operation).await? + plan_fresh( + query_plan, + topology_provider, + operation, + prefix_clip.as_ref(), + ) + .await? }; // TODO: enforce max fan-out (default 100, configurable). See FEED_OPERATIONS_REQS.md §3. @@ -372,14 +392,24 @@ fn push_change_feed_leaf( } /// Builds the request leaves for a fresh (non-resumed) cross-partition plan. +/// +/// When `prefix_clip` is `Some`, each query-plan range is first intersected +/// with it so a hierarchical-partition-key prefix target only fans out over — +/// and is EPK-scoped to — the physical partitions its completions can live in +/// (issue #4680). Query-plan ranges that fall entirely outside the prefix are +/// skipped. async fn plan_fresh( query_plan: &QueryPlan, topology_provider: &mut dyn TopologyProvider, operation: &Arc, + prefix_clip: Option<&FeedRange>, ) -> crate::error::Result>> { let mut nodes: Vec> = Vec::new(); for query_range in &query_plan.query_ranges { let feed_range = query_range_to_feed_range(query_range)?; + let Some(feed_range) = clip_to_prefix(feed_range, prefix_clip) else { + continue; + }; let resolved = topology_provider .resolve_ranges(&feed_range, PartitionRoutingRefresh::UseCached) .await?; @@ -423,12 +453,16 @@ async fn plan_resume_from_saved_snapshot( topology_provider: &mut dyn TopologyProvider, operation: &Arc, saved: &SavedSnapshot, + prefix_clip: Option<&FeedRange>, ) -> crate::error::Result>> { let mut nodes: Vec> = Vec::new(); let mut coverage: Vec> = vec![Vec::new(); saved.active_tokens.len()]; for query_range in &query_plan.query_ranges { let feed_range = query_range_to_feed_range(query_range)?; + let Some(feed_range) = clip_to_prefix(feed_range, prefix_clip) else { + continue; + }; let resolved = topology_provider .resolve_ranges(&feed_range, PartitionRoutingRefresh::UseCached) .await?; @@ -572,6 +606,50 @@ fn query_range_to_feed_range( FeedRange::new(min, max) } +/// Computes the EPK range a hierarchical-partition-key *prefix* target must be +/// clipped to, or `None` when the operation does not target a partition-key +/// prefix. +/// +/// A cross-partition query whose target is a *logical partition* only reaches +/// the fan-out planner when the key is an incomplete HPK prefix — a complete +/// key is trivial and handled by [`build_trivial_pipeline`]. Such a prefix +/// collapses to a single EPK point in its [`FeedRange`] +/// (`min_inclusive() == max_exclusive()`), so on its own it cannot scope the +/// fan-out. Instead the fan-out is scoped to `[prefix_epk, prefix_epk + "FF")` +/// from [`EffectivePartitionKey::compute_range`]; otherwise the query scans +/// every physical partition and returns rows outside the prefix (issue #4680). +fn prefix_epk_range(operation: &CosmosOperation) -> crate::error::Result> { + let Some(target) = operation.target() else { + return Ok(None); + }; + let Some(partition_key) = target.partition_key() else { + return Ok(None); + }; + let Some(container) = operation.container() else { + return Ok(None); + }; + let pk_def = container.partition_key_definition(); + if pk_def.is_complete(partition_key) { + // A complete key hashes to a single EPK point and should already have + // been planned as a trivial single-partition request; nothing to clip. + return Ok(None); + } + let range = EffectivePartitionKey::compute_range(partition_key.values(), pk_def)?; + Ok(Some(FeedRange::new(range.start, range.end)?)) +} + +/// Intersects a query-plan `feed_range` with the operation's prefix clip range. +/// +/// Returns the unmodified `feed_range` when there is no prefix clip, the +/// intersection when the two overlap, or `None` when the query-plan range falls +/// entirely outside the prefix (in which case the caller skips it). +fn clip_to_prefix(feed_range: FeedRange, prefix_clip: Option<&FeedRange>) -> Option { + match prefix_clip { + None => Some(feed_range), + Some(clip) => intersect_feed_ranges(&feed_range, clip), + } +} + /// Returns true if the union of `pieces` covers `range` end-to-end. /// /// Assumes pieces are subsets of `range`. The check sorts pieces by @@ -1180,6 +1258,131 @@ mod tests { assert_drain_requests_with_partitions(pipeline, &[("20", "80", "pkrange-wide", "", "FF")]); } + // --- hierarchical partition key prefix clipping (issue #4680) --- + + fn multihash_pk_def() -> PartitionKeyDefinition { + serde_json::from_str(r#"{"paths":["/country","/state"],"kind":"MultiHash","version":2}"#) + .unwrap() + } + + fn multihash_container() -> ContainerReference { + let props = ContainerProperties { + id: Cow::Owned("coll".into()), + partition_key: multihash_pk_def(), + system_properties: SystemProperties::default(), + }; + ContainerReference::new(test_account(), "db", "db_rid", "coll", "coll_rid", &props) + } + + /// Builds a cross-partition query whose target is a one-level prefix + /// (`(country,)`) of a two-level `/country/state` MultiHash container. + fn prefix_query_operation(prefix: &PartitionKey) -> CosmosOperation { + let target = FeedRange::for_partition(prefix.clone(), &multihash_pk_def()); + CosmosOperation::query_items(multihash_container(), Some(target)) + .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec()) + } + + #[tokio::test] + async fn clips_prefix_fan_out_to_epk_range() { + // A prefix scope over an HPK container must be EPK-scoped to + // [prefix_epk, prefix_epk+"FF") inside its physical partition, not scan + // the whole partition (issue #4680). + let prefix = PartitionKey::from("USA"); + let range = EffectivePartitionKey::compute_range(prefix.values(), &multihash_pk_def()) + .expect("prefix range"); + let (start, end) = ( + range.start.as_str().to_owned(), + range.end.as_str().to_owned(), + ); + assert_ne!( + start, end, + "a one-level prefix must be a real range, not a point" + ); + + // Query plan covers the whole space; a single physical partition owns it. + let plan = plan_with_ranges(vec![qr("", "FF")]); + let op = prefix_query_operation(&prefix); + 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(); + + assert_drain_requests_with_partitions( + pipeline, + &[(start.as_str(), end.as_str(), "pkrange-0", "", "FF")], + ); + } + + #[tokio::test] + async fn prefix_fan_out_spans_multiple_partitions() { + // When the prefix EPK range straddles a physical partition boundary the + // fan-out must emit one EPK-scoped request per overlapping partition. + let prefix = PartitionKey::from("USA"); + let range = EffectivePartitionKey::compute_range(prefix.values(), &multihash_pk_def()) + .expect("prefix range"); + let (start, end) = ( + range.start.as_str().to_owned(), + range.end.as_str().to_owned(), + ); + // A split point strictly inside [start, end): start is a prefix of both, + // and "80" < "FF" so start < mid < end. + let mid = format!("{start}80"); + + let plan = plan_with_ranges(vec![qr("", "FF")]); + let op = prefix_query_operation(&prefix); + let mut topology = MockTopologyProvider::new(vec![Ok(vec![ + rr("", &mid, "pkrange-left"), + rr(&mid, "FF", "pkrange-right"), + ])]); + + let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None) + .await + .unwrap(); + + assert_drain_requests_with_partitions( + pipeline, + &[ + ( + start.as_str(), + mid.as_str(), + "pkrange-left", + "", + mid.as_str(), + ), + ( + mid.as_str(), + end.as_str(), + "pkrange-right", + mid.as_str(), + "FF", + ), + ], + ); + } + + #[tokio::test] + async fn full_container_query_is_not_prefix_clipped() { + // A full-container cross-partition query over an HPK container carries no + // prefix target, so every physical partition is queried unscoped. + let plan = plan_with_ranges(vec![qr("", "FF")]); + let op = CosmosOperation::query_items(multihash_container(), Some(FeedRange::full())) + .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec()); + 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 rejects_query_plan_with_top() { let plan = QueryPlan { 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..7bcdf9882cf 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 @@ -105,9 +105,14 @@ impl OperationOverrides { HeaderValue::from(feed_range.max_exclusive().as_str().to_owned()), ); } + // `x-ms-start-epk`/`x-ms-end-epk` describe an effective-partition-key + // *range*, so the key type must be `EffectivePartitionKeyRange`. The + // point value `EffectivePartitionKey` is rejected by the gateway with + // `400 "One of the input values is invalid"` for range-scoped requests + // (issues #4680 and #4681). headers.insert( HeaderName::from_static(request_header_names::READ_FEED_KEY_TYPE), - HeaderValue::from_static("EffectivePartitionKey"), + HeaderValue::from_static("EffectivePartitionKeyRange"), ); } @@ -3673,7 +3678,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 From 505ba37274f029f8ac48d0dc2f0a8bda824fee02 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:38:45 -0400 Subject: [PATCH 2/5] Add PR link to HPK fix changelog entries Reference PR #4729 in the azure_data_cosmos and azure_data_cosmos_driver CHANGELOG entries for the HPK prefix and cross-partition query fixes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sdk/cosmos/azure_data_cosmos/CHANGELOG.md | 2 +- sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index 3509bc7c423..3a0c56cb449 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -12,7 +12,7 @@ ### Bugs Fixed -- Fixed hierarchical-partition-key (HPK) queries. A `FeedScope::partition` scope with only a *prefix* of the key hierarchy (e.g. `("USA", "CA")` on a `/country/state/city` container) now filters to that prefix instead of returning every item in the physical partition (issue [#4680](https://github.com/Azure/azure-sdk-for-rust/issues/4680)), and cross-partition queries over an HPK container no longer fail with `400 Bad Request` (issue [#4681](https://github.com/Azure/azure-sdk-for-rust/issues/4681)). +- Fixed hierarchical-partition-key (HPK) queries. A `FeedScope::partition` scope with only a *prefix* of the key hierarchy (e.g. `("USA", "CA")` on a `/country/state/city` container) now filters to that prefix instead of returning every item in the physical partition (issue [#4680](https://github.com/Azure/azure-sdk-for-rust/issues/4680)), and cross-partition queries over an HPK container no longer fail with `400 Bad Request` (issue [#4681](https://github.com/Azure/azure-sdk-for-rust/issues/4681)). ([#4729](https://github.com/Azure/azure-sdk-for-rust/pull/4729)) - 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_driver/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md index eaabc9311a0..d4ad859f30c 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -12,7 +12,7 @@ ### Bugs Fixed -- Fixed hierarchical-partition-key (HPK / MultiHash) queries emitting the wrong `x-ms-read-key-type` for effective-partition-key range-scoped requests. The driver sent the point value `EffectivePartitionKey` alongside `x-ms-start-epk`/`x-ms-end-epk`, which the gateway rejects with `400 "One of the input values is invalid"`; it now sends `EffectivePartitionKeyRange`. This enables filtered partition-key *prefix* queries (issue [#4680](https://github.com/Azure/azure-sdk-for-rust/issues/4680)) and fixes cross-partition queries over HPK containers failing with `400 Bad Request` (issue [#4681](https://github.com/Azure/azure-sdk-for-rust/issues/4681)). +- Fixed hierarchical-partition-key (HPK / MultiHash) queries emitting the wrong `x-ms-read-key-type` for effective-partition-key range-scoped requests. The driver sent the point value `EffectivePartitionKey` alongside `x-ms-start-epk`/`x-ms-end-epk`, which the gateway rejects with `400 "One of the input values is invalid"`; it now sends `EffectivePartitionKeyRange`. This enables filtered partition-key *prefix* queries (issue [#4680](https://github.com/Azure/azure-sdk-for-rust/issues/4680)) and fixes cross-partition queries over HPK containers failing with `400 Bad Request` (issue [#4681](https://github.com/Azure/azure-sdk-for-rust/issues/4681)). ([#4729](https://github.com/Azure/azure-sdk-for-rust/pull/4729)) - 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)) From 08dd1db943199cb17dce828d782a030daa456fab Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 9 Jul 2026 13:03:32 -0400 Subject: [PATCH 3/5] Harden HPK prefix query handling per PR review Address review feedback on the HPK query fixes (#4680, #4681), aligning edge-case behavior with the .NET SDK (which routes a prefix scope directly to its effective-partition-key range and lets the server filter, rather than intersecting with query-plan ranges). Behavioral changes (planner.rs): - A fresh prefix-scoped fan-out whose query ranges fall entirely outside the prefix EPK range now drains to an empty page instead of raising CLIENT_QUERY_PLAN_PRODUCED_EMPTY_RANGES, matching .NET. - On resume, a saved continuation range that lies outside the current prefix scope is treated as already satisfied (clipped to the prefix before the coverage guard) instead of hard-failing with an unhonored- range error. - prefix_epk_range debug_asserts the incomplete-key invariant while still degrading gracefully in release. Regression hardening: - The in-memory emulator now rejects an EPK-range-scoped read (x-ms-start-epk/x-ms-end-epk) that declares the point key type instead of EffectivePartitionKeyRange with 400 "One of the input values is invalid", so the emulator-backed HPK tests fail if the driver's key-type value regresses. - Strengthen hpk_query_prefix_correctness_guard with an explicit count so an empty result cannot pass the anti-leak loop. - Add planner tests for a three-level prefix, a disjoint fresh prefix, a query range skipped outside the prefix, and resume outside the prefix; assert the split point in the multi-partition prefix test; guard the FF/0x3F EPK sentinel invariant. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/in_memory_emulator_tests/hpk.rs | 7 + .../src/driver/dataflow/planner.rs | 191 +++++++++++++++++- .../src/in_memory_emulator/dispatch.rs | 27 +++ .../src/in_memory_emulator/operations.rs | 14 ++ .../src/models/effective_partition_key.rs | 34 ++++ 5 files changed, 265 insertions(+), 8 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs index 68c08e77603..4fc4cc47c93 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs @@ -205,6 +205,13 @@ async fn hpk_query_prefix_correctness_guard() { ) .await; + // A green result on an *empty* set would silently pass the loop-based + // anti-leak checks below, so pin the expected count first (issue #4680). + assert_eq!( + items.len(), + 5, + "prefix (USA, CA) should return exactly 5 items" + ); for item in &items { assert_eq!(item.country, "USA", "leaked foreign country: {item:?}"); assert_eq!(item.state, "CA", "leaked foreign state: {item:?}"); 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 c31e08e4249..335ceacfdab 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 @@ -202,10 +202,14 @@ pub(crate) async fn build_sequential_drain( // TODO: enforce max fan-out (default 100, configurable). See FEED_OPERATIONS_REQS.md §3. if request_nodes.is_empty() { - // Resumed past every range that still has work: the pipeline is - // fully drained. Otherwise the plan / topology yielded nothing to - // query — that's a service contract violation. - if saved_snapshot.is_some() { + // Resumed past every range that still has work: the pipeline is fully + // drained. A prefix-scoped fan-out whose query ranges all fall outside + // the prefix EPK range is likewise empty — the query simply matches no + // partition, which the reference SDKs (.NET routes directly to the + // prefix's effective ranges and lets the server filter) surface as an + // empty page rather than an error. Only an *unscoped* plan that yields + // no ranges is a genuine service-contract violation. + if saved_snapshot.is_some() || prefix_clip.is_some() { return Ok(Pipeline::new(Box::new(DrainedLeaf))); } return Err(crate::error::CosmosError::builder() @@ -553,7 +557,17 @@ async fn plan_resume_from_saved_snapshot( // saved continuation without risking duplicate emission or data loss — // fail loudly. for (idx, entry) in saved.active_tokens.iter().enumerate() { - if !range_fully_covered(&entry.range, &coverage[idx]) { + // A saved token range can extend outside the current prefix scope when + // resuming a continuation captured before the prefix clip existed (for + // example a token issued by an older SDK that over-scanned the whole + // physical partition). Only the portion of the saved range that falls + // inside the prefix must be covered by the current topology; the rest + // is out of scope and treated as already satisfied, matching the .NET + // route-to-prefix model which never queries outside the prefix. + let Some(effective_range) = clip_to_prefix(entry.range.clone(), prefix_clip) else { + continue; + }; + if !range_fully_covered(&effective_range, &coverage[idx]) { const MAX_COVERAGE_PIECES_RENDERED: usize = 8; let coverage_summary = if coverage[idx].is_empty() { "(no overlapping topology ranges)".to_string() @@ -586,8 +600,8 @@ async fn plan_resume_from_saved_snapshot( "continuation token active range [{}, {}) could not be fully covered \ by the current topology above the cursor (covered: {}); the query \ cannot be safely resumed", - entry.range.min_inclusive().as_str(), - entry.range.max_exclusive().as_str(), + effective_range.min_inclusive().as_str(), + effective_range.max_exclusive().as_str(), coverage_summary, )) .build()); @@ -631,7 +645,14 @@ fn prefix_epk_range(operation: &CosmosOperation) -> crate::error::Result PartitionKeyDefinition { + serde_json::from_str( + r#"{"paths":["/country","/state","/city"],"kind":"MultiHash","version":2}"#, + ) + .unwrap() + } + + fn multihash_container_three_level() -> ContainerReference { + let props = ContainerProperties { + id: Cow::Owned("coll".into()), + partition_key: multihash_pk_def_three_level(), + system_properties: SystemProperties::default(), + }; + ContainerReference::new(test_account(), "db", "db_rid", "coll", "coll_rid", &props) + } + + #[tokio::test] + async fn three_level_two_component_prefix_clips_fan_out_to_epk_range() { + // A two-component prefix (country, state) of a three-level container is + // EPK-scoped the same way a one-level prefix is: the prefix mechanism is + // independent of hierarchy depth (issue #4680). + let prefix = PartitionKey::from(("USA", "CA")); + let range = + EffectivePartitionKey::compute_range(prefix.values(), &multihash_pk_def_three_level()) + .expect("prefix range"); + let (start, end) = ( + range.start.as_str().to_owned(), + range.end.as_str().to_owned(), + ); + assert_ne!(start, end, "a two-of-three prefix must be a real range"); + assert_eq!(start.len(), 64, "two hashed components => 64 hex chars"); + + let plan = plan_with_ranges(vec![qr("", "FF")]); + let target = FeedRange::for_partition(prefix.clone(), &multihash_pk_def_three_level()); + let op = CosmosOperation::query_items(multihash_container_three_level(), Some(target)) + .with_body(br#"{"query":"SELECT * FROM c"}"#.to_vec()); + 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(); + + assert_drain_requests_with_partitions( + pipeline, + &[(start.as_str(), end.as_str(), "pkrange-0", "", "FF")], + ); + } + + #[tokio::test] + async fn fresh_prefix_disjoint_from_query_ranges_yields_empty_pipeline() { + // A predicated prefix query whose plan ranges fall entirely outside the + // prefix EPK range matches no partition. Like the .NET SDK — which routes + // directly to the prefix's effective ranges and lets the server filter — + // this is an empty page, not an error. The clip happens before topology + // resolution, so the (empty) topology script is never consulted. + let prefix = PartitionKey::from("USA"); + let range = EffectivePartitionKey::compute_range(prefix.values(), &multihash_pk_def()) + .expect("prefix range"); + // A query range that starts at the prefix's exclusive max cannot overlap + // the half-open prefix range [start, end). + let end = range.end.as_str().to_owned(); + let plan = plan_with_ranges(vec![qr(&end, "FF")]); + let op = prefix_query_operation(&prefix); + let mut topology = MockTopologyProvider::new(vec![]); + + let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None) + .await + .unwrap(); + + assert!( + matches!( + pipeline.snapshot_state().unwrap(), + PipelineNodeState::Drained + ), + "a fully-disjoint prefix query must drain to an empty page, not error" + ); + } + + #[tokio::test] + async fn prefix_skips_query_ranges_outside_the_prefix() { + // With multiple query ranges, only the ranges overlapping the prefix + // survive the clip; ranges entirely outside are dropped without + // consulting the topology for them. + let prefix = PartitionKey::from("USA"); + let range = EffectivePartitionKey::compute_range(prefix.values(), &multihash_pk_def()) + .expect("prefix range"); + let (start, end) = ( + range.start.as_str().to_owned(), + range.end.as_str().to_owned(), + ); + + // First range is exactly the prefix interior (survives); second range is + // entirely above the prefix (dropped). Only one topology result is + // scripted, proving the second range never reaches resolution. + let plan = plan_with_ranges(vec![qr(&start, &end), qr(&end, "FF")]); + let op = prefix_query_operation(&prefix); + let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr(&start, &end, "pkrange-0")])]); + + let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), None) + .await + .unwrap(); + + assert_drain_requests_with_partitions( + pipeline, + &[( + start.as_str(), + end.as_str(), + "pkrange-0", + start.as_str(), + end.as_str(), + )], + ); + } + + #[tokio::test] + async fn resume_saved_range_outside_prefix_is_dropped() { + // Resuming a continuation whose saved token range lies entirely outside + // the current prefix scope (for example a token captured by an older SDK + // that over-scanned the whole physical partition) must not hard-fail the + // coverage guard: the out-of-prefix portion is treated as already + // satisfied. Here the saved token covers [end, FF) — above the prefix — + // so the resume drains to an empty page. + let prefix = PartitionKey::from("USA"); + let range = EffectivePartitionKey::compute_range(prefix.values(), &multihash_pk_def()) + .expect("prefix range"); + let end = range.end.as_str().to_owned(); + + let plan = plan_with_ranges(vec![qr("", "FF")]); + let op = prefix_query_operation(&prefix); + let mut topology = MockTopologyProvider::new(vec![Ok(vec![rr("", "FF", "pkrange-0")])]); + + let resume = saved_drain(vec![( + end.as_str(), + "FF", + saved_request(Some("server-token-abc")), + )]); + + let pipeline = build_sequential_drain(&plan, &mut topology, &Arc::new(op), Some(resume)) + .await + .unwrap(); + + assert!( + matches!( + pipeline.snapshot_state().unwrap(), + PipelineNodeState::Drained + ), + "a saved range outside the prefix must drain, not raise an unhonored-range error" + ); + } + #[tokio::test] async fn rejects_query_plan_with_top() { let plan = QueryPlan { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs index 6c27ec9e53f..3562dfedf30 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs @@ -43,6 +43,11 @@ pub(crate) enum OperationType { /// for these (not 501) — keep the variant separate so the handler can /// emit the right status and substatus. BadRequestPath(String), + /// A request whose headers are internally inconsistent — for example an + /// EPK-range-scoped read (`x-ms-start-epk`/`x-ms-end-epk`) that declares + /// the point key type instead of `EffectivePartitionKeyRange`. Real Cosmos + /// rejects these with `400 "One of the input values is invalid"`. + InvalidInput(String), } /// Parsed request data extracted from an HTTP request. @@ -104,6 +109,7 @@ static PARTITION_KEY_RANGE_ID: HeaderName = HeaderName::from_static("x-ms-documentdb-partitionkeyrangeid"); static START_EPK: HeaderName = HeaderName::from_static("x-ms-start-epk"); static END_EPK: HeaderName = HeaderName::from_static("x-ms-end-epk"); +static READ_FEED_KEY_TYPE: HeaderName = HeaderName::from_static("x-ms-read-key-type"); static IS_BATCH_REQUEST: HeaderName = HeaderName::from_static("x-ms-cosmos-is-batch-request"); static OFFER_THROUGHPUT: HeaderName = HeaderName::from_static("x-ms-offer-throughput"); static OFFER_AUTOPILOT_SETTINGS: HeaderName = @@ -158,6 +164,9 @@ pub(crate) fn parse_request(request: &Request) -> ParsedRequest { .map(|s| s.to_string()); let start_epk = headers.get_optional_str(&START_EPK).map(|s| s.to_string()); let end_epk = headers.get_optional_str(&END_EPK).map(|s| s.to_string()); + let read_key_type = headers + .get_optional_str(&READ_FEED_KEY_TYPE) + .map(|s| s.to_string()); // Parse `x-ms-offer-throughput` (RU/s) from the request headers. Invalid / // non-numeric values are treated as absent; the container creation handler // then uses `ContainerConfig::default()`. A failing parse is intentionally @@ -193,6 +202,24 @@ pub(crate) fn parse_request(request: &Request) -> ParsedRequest { ) }; + // A read scoped to an effective-partition-key *range* + // (`x-ms-start-epk`/`x-ms-end-epk`) must declare the range key type + // `EffectivePartitionKeyRange`. The point key type `EffectivePartitionKey` + // is rejected by the real gateway with `400 "One of the input values is + // invalid"`; mirror that here so a regression that reverts the driver's + // key-type value is caught by the emulator-backed tests (issues #4680 and + // #4681). + let operation = if (start_epk.is_some() || end_epk.is_some()) + && read_key_type.as_deref() != Some("EffectivePartitionKeyRange") + { + OperationType::InvalidInput(format!( + "x-ms-read-key-type must be 'EffectivePartitionKeyRange' when \ + x-ms-start-epk/x-ms-end-epk are present, got {read_key_type:?}" + )) + } else { + operation + }; + // Index by *position*, not by keyword search. Cosmos URLs are // `/dbs/{db}/colls/{coll}/docs/{doc}/...`, so the keyword always // appears at an even index and the value follows it. Searching by diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs index e59ab7eb331..33fdafcf8f9 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/operations.rs @@ -272,6 +272,7 @@ pub(crate) async fn handle_operation( .await } OperationType::BadRequestPath(desc) => bad_request_path_response(desc, start), + OperationType::InvalidInput(desc) => invalid_input_response(desc, start), OperationType::Unsupported(desc) => unsupported_response(desc, start), }; @@ -4831,6 +4832,19 @@ fn bad_request_path_response(path: &str, start: Instant) -> AsyncRawResponse { .build() } +fn invalid_input_response(message: &str, start: Instant) -> AsyncRawResponse { + error_response( + StatusCode::BadRequest, + None, + "BadRequest", + &format!("One of the input values is invalid. {message}"), + 0.0, + "", + start, + ) + .build() +} + fn unsupported_response(operation: &str, start: Instant) -> AsyncRawResponse { error_response( StatusCode::NotImplemented, diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/effective_partition_key.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/effective_partition_key.rs index 7b1de01a0ef..f008fb0da22 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/effective_partition_key.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/effective_partition_key.rs @@ -676,6 +676,40 @@ mod tests { assert!(range.end.as_str().ends_with("FF")); } + /// The `[epk, epk+"FF")` prefix range relies on every hash component's first + /// hex digit being masked into `[0-3]` (`hash_bytes[0] & 0x3F`), so the "FF" + /// sentinel is always lexicographically greater than any real completion of + /// the prefix. Guard that coupling so `compute_range` and the planner's + /// prefix clip cannot silently break. + #[test] + fn compute_range_first_nibble_masked_keeps_ff_sentinel_valid() { + let pk_def = PartitionKeyDefinition::from(("/tenantId", "/userId", "/sessionId")); + for values in [ + vec![PartitionKeyValue::from("tenant1".to_string())], + vec![ + PartitionKeyValue::from("tenant1".to_string()), + PartitionKeyValue::from("user1".to_string()), + ], + ] { + let range = EffectivePartitionKey::compute_range(&values, &pk_def).unwrap(); + let start = range.start.as_str(); + // Each hash component is 32 hex chars; the first digit of every + // component must be in [0-3] for the "FF" upper bound to dominate. + for component_start in (0..start.len()).step_by(32) { + let first = start.as_bytes()[component_start]; + assert!( + (b'0'..=b'3').contains(&first), + "component at {component_start} starts with '{}', expected [0-3]", + first as char + ); + } + assert!( + range.end.as_str() > start, + "the FF sentinel must exceed the prefix EPK" + ); + } + } + /// compute_range returns a point for single-hash (non-MultiHash) keys /// when the component count matches the definition path count. #[test] From c570d5bf12b39db83b4af7af13b51bc07b24a8bc Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:10:21 -0400 Subject: [PATCH 4/5] Fix clippy large-futures and hoist EPK-range key-type constant Resolve the `clippy::large-futures` CI failures in the in-memory HPK tests by boxing the query pipeline futures inside the `query_scope` helper, so callers awaiting it stay under the lint's stack-size threshold (one fix instead of six `Box::pin` call sites). Address a review nit (Fabian): hoist the `EffectivePartitionKeyRange` `x-ms-read-key-type` value into a shared `request_header_names::READ_FEED_KEY_TYPE_EPK_RANGE` constant and use it in the driver's `apply_headers`, its unit test, and the in-memory emulator dispatch guard instead of repeating the string literal. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../tests/in_memory_emulator_tests/hpk.rs | 11 +++++------ .../src/driver/pipeline/operation_pipeline.rs | 4 ++-- .../src/in_memory_emulator/dispatch.rs | 6 ++++-- .../src/models/cosmos_headers.rs | 5 +++++ 4 files changed, 16 insertions(+), 10 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs index 4fc4cc47c93..4e767003933 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs @@ -142,13 +142,12 @@ async fn setup_hpk_container() -> ContainerClient { } async fn query_scope(container: &ContainerClient, scope: FeedScope) -> Vec { - container - .query_items(Query::from("SELECT * FROM c"), scope, None) - .await - .unwrap() - .try_collect() + // Box the query pipeline futures so callers awaiting `query_scope` don't trip + // `clippy::large-futures` (the `query_items` future is ~16 KB on the stack). + let iter = Box::pin(container.query_items(Query::from("SELECT * FROM c"), scope, None)) .await - .unwrap() + .unwrap(); + Box::pin(iter.try_collect()).await.unwrap() } /// #4680 — a one-level prefix `(USA,)` returns only the 8 USA rows, not CANADA. 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 fc190060c22..41843d8044a 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 @@ -123,7 +123,7 @@ impl OperationOverrides { // (issues #4680 and #4681). headers.insert( HeaderName::from_static(request_header_names::READ_FEED_KEY_TYPE), - HeaderValue::from_static("EffectivePartitionKeyRange"), + HeaderValue::from_static(request_header_names::READ_FEED_KEY_TYPE_EPK_RANGE), ); } @@ -4057,7 +4057,7 @@ mod tests { request_header_names::READ_FEED_KEY_TYPE )) .map(|s| s.to_string()), - Some("EffectivePartitionKeyRange".to_string()) + Some(request_header_names::READ_FEED_KEY_TYPE_EPK_RANGE.to_string()) ); assert_eq!( headers diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs index 3562dfedf30..6e98df40b1a 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/in_memory_emulator/dispatch.rs @@ -210,8 +210,10 @@ pub(crate) fn parse_request(request: &Request) -> ParsedRequest { // key-type value is caught by the emulator-backed tests (issues #4680 and // #4681). let operation = if (start_epk.is_some() || end_epk.is_some()) - && read_key_type.as_deref() != Some("EffectivePartitionKeyRange") - { + && read_key_type.as_deref() + != Some( + crate::models::cosmos_headers::request_header_names::READ_FEED_KEY_TYPE_EPK_RANGE, + ) { OperationType::InvalidInput(format!( "x-ms-read-key-type must be 'EffectivePartitionKeyRange' when \ x-ms-start-epk/x-ms-end-epk are present, got {read_key_type:?}" diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs index 863431e59ab..34e5235e18a 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs @@ -67,6 +67,11 @@ pub(crate) mod request_header_names { pub const START_EPK: &str = "x-ms-start-epk"; pub const END_EPK: &str = "x-ms-end-epk"; pub const READ_FEED_KEY_TYPE: &str = "x-ms-read-key-type"; + /// The [`READ_FEED_KEY_TYPE`] value for an effective-partition-key *range* + /// (`x-ms-start-epk`/`x-ms-end-epk`)-scoped request. The point value + /// `EffectivePartitionKey` is rejected by the gateway with `400 "One of the + /// input values is invalid"` for range-scoped requests (issues #4680, #4681). + pub const READ_FEED_KEY_TYPE_EPK_RANGE: &str = "EffectivePartitionKeyRange"; /// Internal-only headers carrying the physical pkrange's full EPK bounds /// to the GW_V2 dispatcher. The thin-client (Gateway 2.0) proxy requires /// `StartEpkHash`/`EndEpkHash` RNTBD body tokens on every Query frame; in From 36870c8f1e2195164560bc51e82c386363004b20 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 9 Jul 2026 14:44:16 -0400 Subject: [PATCH 5/5] Add feed-range scope-bug regression test Fold Ananth's standard-gateway (single-hash `/pk`) end-to-end regression test into cosmos_feed_ranges.rs rather than adding a separate file. It seeds one document per partition key, computes each doc's effective partition key locally, and asserts that FeedScope::range(interior_window) returns exactly the windowed subset across wide, tight, and empty-gap windows. This exercises a path not covered by the in-memory HPK tests or the GW2.0 query tests: it fails on pre-fix main with "SCOPE BUG" (full-scan fallback) and passes once both the planner scope-clip and the x-ms-read-key-type: EffectivePartitionKeyRange header fix are in. drain_ids drains page-by-page via into_pages() instead of item-by-item so the query future stays under the clippy::large-futures threshold. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../emulator_tests/cosmos_feed_ranges.rs | 177 +++++++++++++++++- 1 file changed, 176 insertions(+), 1 deletion(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_feed_ranges.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_feed_ranges.rs index 521679260d3..d817da7cc5f 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_feed_ranges.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_feed_ranges.rs @@ -3,11 +3,20 @@ use super::framework; +use std::borrow::Cow; +use std::collections::BTreeSet; use std::error::Error; -use azure_data_cosmos::models::{ContainerProperties, ThroughputProperties}; +use azure_data_cosmos::clients::ContainerClient; +use azure_data_cosmos::feed::{FeedRange, FeedScope}; +use azure_data_cosmos::models::{ + ContainerProperties, EffectivePartitionKey, PartitionKeyDefinition, ThroughputProperties, +}; use azure_data_cosmos::options::CreateContainerOptions; +use azure_data_cosmos::{PartitionKey, Query}; use base64::Engine; +use futures::StreamExt; +use serde::{Deserialize, Serialize}; use framework::{TestClient, TestOptions}; @@ -278,3 +287,169 @@ pub async fn feed_range_from_partition_key_single_hash_full_key() -> Result<(), ) .await } + +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. +/// +/// Drains the result set page-by-page (via [`into_pages`]) rather than +/// item-by-item; the paged future is small enough to satisfy +/// `clippy::large-futures`, which the item-level stream would otherwise trip. +async fn drain_ids( + container: &ContainerClient, + scope: FeedScope, +) -> Result, Box> { + let mut pages = container + .query_items::(Query::from("SELECT * FROM c"), scope, None) + .await? + .into_pages(); + let mut got = BTreeSet::new(); + while let Some(page) = pages.next().await { + for doc in page?.into_items() { + got.insert(doc.id); + } + } + Ok(got) +} + +/// 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 seeds one document per partition key, computes each document's +/// effective partition key locally, and issues cross-partition queries scoped +/// to interior EPK windows, asserting the result set is exactly the windowed +/// subset. Without the planner fix these fail with `SCOPE BUG`; with the fix — +/// plus the `x-ms-read-key-type: EffectivePartitionKeyRange` header correction +/// that lets the gateway accept the emitted interior EPK window — they pass. +#[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().cmp(b.1.min_inclusive())); + 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. Appending a whole byte (`80`) to + // `X_mid`'s hex keeps `X_mid` as a byte prefix, so the result is + // strictly greater than `X_mid` and — because adjacent single-hash + // EPKs differ within their leading bytes — strictly less than + // `X_{mid+1}`. So `[X_mid || 0x80, X_{mid+1})` contains no doc's EPK. + // (A single hex nibble would be dropped by the byte-wise hex parser, + // so a full byte is required.) This guards against an off-by-one + // where the lower bound is treated as inclusive of `X_mid`. + let gap_start = EffectivePartitionKey::from(format!( + "{}80", + points[mid].1.min_inclusive().to_hex() + )); + 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 +}