diff --git a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md index b1d4a397534..8610ab17f8d 100644 --- a/sdk/cosmos/azure_data_cosmos/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos/CHANGELOG.md @@ -15,6 +15,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)). ([#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/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/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 +} 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..4e767003933 --- /dev/null +++ b/sdk/cosmos/azure_data_cosmos/tests/in_memory_emulator_tests/hpk.rs @@ -0,0 +1,253 @@ +// 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 { + // 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(); + Box::pin(iter.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; + + // 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:?}"); + } + 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 37046b7cb92..4e1924e7260 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md +++ b/sdk/cosmos/azure_data_cosmos_driver/CHANGELOG.md @@ -15,6 +15,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)). ([#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)) 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 ec065110067..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 @@ -116,9 +116,14 @@ impl OperationOverrides { HeaderName::from_static(request_header_names::END_EPK), HeaderValue::from(feed_range.max_exclusive().to_hex()), ); + // `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(request_header_names::READ_FEED_KEY_TYPE_EPK_RANGE), ); } @@ -4052,7 +4057,7 @@ mod tests { request_header_names::READ_FEED_KEY_TYPE )) .map(|s| s.to_string()), - Some("EffectivePartitionKey".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 6c27ec9e53f..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 @@ -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,26 @@ 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( + 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:?}" + )) + } 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 77cfbeb4e68..917825e852a 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), }; @@ -2576,6 +2577,16 @@ fn collect_item_documents( let requested_epk = match parsed.partition_key_header.as_deref() { Some(header) => match parse_partition_key_header(header) { Ok(components) if components.is_empty() => None, + // A partial hierarchical partition key (fewer components than the + // container's PK paths) targets a *prefix* of logical partitions. + // Real Cosmos scopes such reads via the `x-ms-start-epk`/ + // `x-ms-end-epk` range (below) rather than an exact point EPK, so + // don't compute a point to exact-match here — that would compare a + // 2-component prefix EPK against 3-component item EPKs and drop + // every row. + Ok(components) if components.len() < state.metadata.partition_key.paths().len() => { + None + } Ok(components) => Some(compute_epk( &components, state.metadata.partition_key.kind(), @@ -4831,6 +4842,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/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 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 c188b2494c7..5cfffad07f2 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 @@ -741,6 +741,40 @@ mod tests { assert!(range.end.to_hex().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.to_hex(); + // 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.to_hex() > 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]