Fix HPK prefix and cross-partition query bugs

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)). ([#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

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<PartitionKey>) -> Self {
         Self::Partition(pk.into())
     }

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<GeoItem> {
+    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::<AccountEndpoint>().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<GeoItem> {
+    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");
+}

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;

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)). ([#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))