From e1477dfb9aeac00ab74087b30441fc6a8dfb0154 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Apr 2026 20:10:18 +0000 Subject: [PATCH 01/13] Initial plan From f6f0b18b470bd1d1312aee3139c60dbd4cbc2445 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 9 Apr 2026 20:39:06 +0000 Subject: [PATCH 02/13] Cut over replace_item, upsert_item, delete_item to driver Route replace_item, upsert_item, and delete_item through the driver instead of the legacy gateway pipeline, following the same pattern established by create_item and read_item. - Change CosmosOperation::upsert_item to take (ContainerReference, PartitionKey) instead of ItemReference, since upsert is a POST to the collection feed (same as create). - Remove ItemWriteOptions::apply_headers() shim and the apply_precondition_headers helper, no longer needed now that all item operations use the driver. - Remove tests for the removed gateway shim code. Agent-Logs-Url: https://github.com/Azure/azure-sdk-for-rust/sessions/9397eee6-2f94-4a70-ac3e-253e5fff4f90 Co-authored-by: simorenoh <30335873+simorenoh@users.noreply.github.com> --- .../src/clients/container_client.rs | 123 +++++++++++----- .../azure_data_cosmos/src/cosmos_request.rs | 58 +------- .../azure_data_cosmos/src/options/mod.rs | 136 ------------------ .../src/models/cosmos_operation.rs | 13 +- 4 files changed, 93 insertions(+), 237 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs index 8d5ee036d8e..2f9272a892d 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs @@ -417,20 +417,37 @@ impl ContainerClient { item: T, options: Option, ) -> azure_core::Result> { - let link = self.items_link.item(item_id); - let options = options.clone().unwrap_or_default(); - let excluded_regions = options.operation.excluded_regions.clone(); - let mut cosmos_request = CosmosRequest::builder(OperationType::Replace, link) - .json(&item) - .partition_key(partition_key.into()) - .excluded_regions(excluded_regions) - .build()?; - options.apply_headers(&mut cosmos_request.headers); + let options = options.unwrap_or_default(); + let body = serde_json::to_vec(&item)?; - self.container_connection - .send(cosmos_request, Context::default()) - .await - .map(ItemResponse::new) + // Build the driver's item reference from our stored container metadata. + let item_ref = ItemReference::from_name( + &self.container_ref, + partition_key.into().into_driver_partition_key(), + item_id.to_owned(), + ); + + // Create the driver operation and apply ItemWriteOptions fields. + let mut operation = CosmosOperation::replace_item(item_ref).with_body(body); + + // Wire session token and precondition from SDK options onto the operation. + if let Some(session_token) = options.session_token { + operation = operation.with_session_token(session_token); + } + if let Some(precondition) = options.precondition { + operation = operation.with_precondition(precondition); + } + + // Execute through the driver. + let driver_response = self + .driver + .execute_operation(operation, options.operation) + .await?; + + // Bridge the driver response to the SDK response type. + Ok(ItemResponse::new( + crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + )) } /// Creates or replaces an item in the container. @@ -506,21 +523,32 @@ impl ContainerClient { item: T, options: Option, ) -> azure_core::Result> { - let options = options.clone().unwrap_or_default(); - let excluded_regions = options.operation.excluded_regions.clone(); - let mut cosmos_request = - CosmosRequest::builder(OperationType::Upsert, self.items_link.clone()) - .json(&item) - .partition_key(partition_key.into()) - .excluded_regions(excluded_regions) - .build()?; - options.apply_headers(&mut cosmos_request.headers); + let options = options.unwrap_or_default(); + let body = serde_json::to_vec(&item)?; + let driver_pk = partition_key.into().into_driver_partition_key(); - return self - .container_connection - .send(cosmos_request, Context::default()) - .await - .map(ItemResponse::new); + // Create the driver operation and apply ItemWriteOptions fields. + let mut operation = + CosmosOperation::upsert_item(self.container_ref.clone(), driver_pk).with_body(body); + + // Wire session token and precondition from SDK options onto the operation. + if let Some(session_token) = options.session_token { + operation = operation.with_session_token(session_token); + } + if let Some(precondition) = options.precondition { + operation = operation.with_precondition(precondition); + } + + // Execute through the driver. + let driver_response = self + .driver + .execute_operation(operation, options.operation) + .await?; + + // Bridge the driver response to the SDK response type. + Ok(ItemResponse::new( + crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + )) } /// Reads a specific item from the container. @@ -615,19 +643,36 @@ impl ContainerClient { item_id: &str, options: Option, ) -> azure_core::Result> { - let link = self.items_link.item(item_id); - let options = options.clone().unwrap_or_default(); - let excluded_regions = options.operation.excluded_regions.clone(); - let mut cosmos_request = CosmosRequest::builder(OperationType::Delete, link) - .partition_key(partition_key.into()) - .excluded_regions(excluded_regions) - .build()?; - options.apply_headers(&mut cosmos_request.headers); + let options = options.unwrap_or_default(); - self.container_connection - .send(cosmos_request, Context::default()) - .await - .map(ItemResponse::new) + // Build the driver's item reference from our stored container metadata. + let item_ref = ItemReference::from_name( + &self.container_ref, + partition_key.into().into_driver_partition_key(), + item_id.to_owned(), + ); + + // Create the driver operation (no body for delete). + let mut operation = CosmosOperation::delete_item(item_ref); + + // Wire session token and precondition from SDK options onto the operation. + if let Some(session_token) = options.session_token { + operation = operation.with_session_token(session_token); + } + if let Some(precondition) = options.precondition { + operation = operation.with_precondition(precondition); + } + + // Execute through the driver. + let driver_response = self + .driver + .execute_operation(operation, options.operation) + .await?; + + // Bridge the driver response to the SDK response type. + Ok(ItemResponse::new( + crate::driver_bridge::driver_response_to_cosmos_response(driver_response), + )) } /// Executes a single-partition query against items in the container. diff --git a/sdk/cosmos/azure_data_cosmos/src/cosmos_request.rs b/sdk/cosmos/azure_data_cosmos/src/cosmos_request.rs index f8b3a11b4c1..4ee5e0f1a59 100644 --- a/sdk/cosmos/azure_data_cosmos/src/cosmos_request.rs +++ b/sdk/cosmos/azure_data_cosmos/src/cosmos_request.rs @@ -292,10 +292,7 @@ mod tests { use super::*; use crate::operation_context::OperationType; use crate::resource_context::ResourceType; - use crate::{ - constants, CosmosClientOptions, ItemWriteOptions, OperationOptions, PartitionKey, - SessionToken, - }; + use crate::{constants, PartitionKey}; fn make_base_request(op: OperationType) -> CosmosRequest { let req = CosmosRequest::builder(op, ResourceLink::root(ResourceType::Documents)) @@ -402,59 +399,6 @@ mod tests { assert!(has_upsert); } - #[test] - fn prioritize_request_headers_over_client_headers() { - let mut request_custom_headers = std::collections::HashMap::new(); - request_custom_headers.insert( - HeaderName::from_static("x-custom-header"), - HeaderValue::from_static("custom_value"), - ); - - let operation = OperationOptions::default().with_custom_headers(request_custom_headers); - let item_options = ItemWriteOptions { - operation, - ..Default::default() - } - .with_session_token(SessionToken::from("RequestSession")); - let mut req = CosmosRequest::builder( - OperationType::Create, - ResourceLink::root(ResourceType::Documents), - ) - .build() - .unwrap(); - item_options.apply_headers(&mut req.headers); - - req.request_context.location_endpoint_to_route = - Some("https://example.com/".parse().unwrap()); - - let mut client_custom_headers = std::collections::HashMap::new(); - client_custom_headers.insert( - HeaderName::from_static("x-custom-header"), - HeaderValue::from_static("custom_value-2"), - ); - - let client_operation = - OperationOptions::default().with_custom_headers(client_custom_headers); - let client_options = - CosmosClientOptions::default().with_operation_options(client_operation); - client_options.apply_headers(&mut req.headers); - - let raw = req.into_raw_request(); - let get_header = |name: &HeaderName| { - raw.headers() - .iter() - .find(|(n, _)| n == &name) - .map(|(_, v)| v.as_str()) - .unwrap() - }; - - assert_eq!(get_header(&constants::SESSION_TOKEN), "RequestSession"); - assert_eq!( - get_header(&HeaderName::from_static("x-custom-header")), - "custom_value" - ); - } - #[test] fn to_raw_request_batch_sets_batch_headers() { let req = make_base_request(OperationType::Batch); diff --git a/sdk/cosmos/azure_data_cosmos/src/options/mod.rs b/sdk/cosmos/azure_data_cosmos/src/options/mod.rs index 3b66772c6d0..7e90df48c83 100644 --- a/sdk/cosmos/azure_data_cosmos/src/options/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/options/mod.rs @@ -16,21 +16,6 @@ pub use azure_data_cosmos_driver::options::{ OperationOptionsBuilder, OperationOptionsView, ReadConsistencyStrategy, Region, }; -// Temporary: these helpers allow the SDK pipeline to apply OperationOptions values -// as HTTP headers. They will be removed when individual operations use the internal -// pipeline directly. -fn apply_precondition_headers(precondition: &Precondition, headers: &mut Headers) { - match precondition { - Precondition::IfMatch(etag) => { - headers.insert(headers::IF_MATCH, etag.to_string()); - } - Precondition::IfNoneMatch(etag) => { - headers.insert(constants::IF_NONE_MATCH, etag.to_string()); - } - _ => {} - } -} - // Temporary: applies the prefer header based on the content_response_on_write option. // Will be removed when write operations use the internal pipeline directly. fn apply_content_response_on_write_header( @@ -243,31 +228,6 @@ impl ItemWriteOptions { } } -impl ItemWriteOptions { - // Temporary: applies option values as HTTP headers for the SDK pipeline. - // Will be removed when write operations use the internal pipeline directly. - pub(crate) fn apply_headers(&self, headers: &mut Headers) { - if let Some(custom_headers) = self.operation.custom_headers() { - for (name, value) in custom_headers { - // Only insert if not already set — SDK/request headers take priority. - if headers.get_optional_str(name).is_none() { - headers.insert(name.clone(), value.clone()); - } - } - } - if let Some(session_token) = &self.session_token { - headers.insert(constants::SESSION_TOKEN, session_token.to_string()); - } - if let Some(precondition) = &self.precondition { - apply_precondition_headers(precondition, headers); - } - apply_content_response_on_write_header( - self.operation.content_response_on_write.as_ref(), - headers, - ); - } -} - /// Options for transactional batch operations. /// /// Used by [`ContainerClient::execute_transactional_batch()`](crate::clients::ContainerClient::execute_transactional_batch()). @@ -410,69 +370,6 @@ mod tests { headers.into_iter().collect() } - #[test] - fn item_write_options_as_headers() { - let mut custom_headers = HashMap::new(); - custom_headers.insert( - HeaderName::from_static("x-custom-header"), - HeaderValue::from_static("custom_value"), - ); - - let operation = OperationOptions::default().with_custom_headers(custom_headers); - - let options = ItemWriteOptions { - operation, - ..Default::default() - } - .with_session_token("SessionToken".to_string()) - .with_precondition(Precondition::IfMatch(ETag::from("etag_value"))); - - let mut headers_result = Headers::new(); - options.apply_headers(&mut headers_result); - - let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![ - ("x-custom-header".into(), "custom_value".into()), - (constants::SESSION_TOKEN, "SessionToken".into()), - (headers::IF_MATCH, "etag_value".into()), - (headers::PREFER, constants::PREFER_MINIMAL), - ]; - - assert_eq!( - headers_to_map(headers_result), - headers_to_map(headers_expected) - ); - } - - #[test] - fn custom_headers_should_not_override_sdk_set_headers() { - let mut custom_headers = HashMap::new(); - custom_headers.insert( - constants::SESSION_TOKEN, - HeaderValue::from_static("CustomSession"), - ); - - let operation = OperationOptions::default().with_custom_headers(custom_headers); - - let options = ItemWriteOptions { - operation, - ..Default::default() - } - .with_session_token("RealSessionToken".to_string()); - - let mut headers_result = Headers::new(); - options.apply_headers(&mut headers_result); - - let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![ - (constants::SESSION_TOKEN, "RealSessionToken".into()), - (headers::PREFER, constants::PREFER_MINIMAL), - ]; - - assert_eq!( - headers_to_map(headers_result), - headers_to_map(headers_expected) - ); - } - #[test] fn client_options_as_headers() { let mut custom_headers = HashMap::new(); @@ -530,39 +427,6 @@ mod tests { ); } - #[test] - fn item_write_options_default_as_headers() { - let options = ItemWriteOptions::default(); - - let mut headers_result = Headers::new(); - options.apply_headers(&mut headers_result); - let headers_result: Vec<(HeaderName, HeaderValue)> = headers_result.into_iter().collect(); - - let headers_expected: Vec<(HeaderName, HeaderValue)> = - vec![(headers::PREFER, constants::PREFER_MINIMAL)]; - - assert_eq!(headers_result, headers_expected); - } - - #[test] - fn item_write_options_with_content_response_enabled() { - let mut operation = OperationOptions::default(); - operation.content_response_on_write = Some(ContentResponseOnWrite::Enabled); - - let options = ItemWriteOptions { - operation, - ..Default::default() - }; - - let mut headers_result = Headers::new(); - options.apply_headers(&mut headers_result); - let headers_result: Vec<(HeaderName, HeaderValue)> = headers_result.into_iter().collect(); - - let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![]; - - assert_eq!(headers_result, headers_expected); - } - #[test] fn batch_options_as_headers() { let mut custom_headers = HashMap::new(); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs index 755472759f4..2ee2b9a4acf 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs @@ -465,13 +465,16 @@ impl CosmosOperation { /// Upserts (creates or replaces) an item (document) in a container. /// - /// The `ItemReference` contains the container, partition key, and item identifier, - /// providing all the information needed for the operation. + /// The Cosmos DB REST API treats upsert as a `POST` to the collection feed + /// (same as create), so this takes a `ContainerReference` and `PartitionKey` + /// rather than an `ItemReference`. /// Use `with_body()` to provide the document JSON. /// If an item with the same ID exists, it will be replaced; otherwise, a new item is created. - pub fn upsert_item(item: ItemReference) -> Self { - let partition_key = item.partition_key().clone(); - Self::new(OperationType::Upsert, item).with_partition_key(partition_key) + pub fn upsert_item(container: ContainerReference, partition_key: PartitionKey) -> Self { + let resource_ref: CosmosResourceReference = CosmosResourceReference::from(container) + .with_resource_type(ResourceType::Document) + .into_feed_reference(); + Self::new(OperationType::Upsert, resource_ref).with_partition_key(partition_key) } /// Replaces an existing item (document) in a container. From 595181545b5417d669398eba75afa40bbe4a4bd8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Fri, 10 Apr 2026 16:53:35 +0000 Subject: [PATCH 03/13] Extract apply_item_options helper to deduplicate session_token/precondition wiring Consolidates the repeated session_token and precondition wiring into a shared apply_item_options() function used by all five item operations (read, create, replace, upsert, delete). Agent-Logs-Url: https://github.com/Azure/azure-sdk-for-rust/sessions/1e363c06-e4bc-4fca-88d5-80527c2f5d92 Co-authored-by: simorenoh <30335873+simorenoh@users.noreply.github.com> --- .../src/clients/container_client.rs | 77 ++++++++----------- 1 file changed, 31 insertions(+), 46 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs index 2f9272a892d..37f1a856e58 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs @@ -24,7 +24,9 @@ use crate::routing::global_partition_endpoint_manager::GlobalPartitionEndpointMa use crate::routing::partition_key_range_cache::PartitionKeyRangeCache; use azure_core::http::headers::AsHeaders; use azure_core::http::Context; -use azure_data_cosmos_driver::models::{ContainerReference, CosmosOperation, ItemReference}; +use azure_data_cosmos_driver::models::{ + ContainerReference, CosmosOperation, ItemReference, Precondition, SessionToken, +}; use azure_data_cosmos_driver::CosmosDriver; use serde::{de::DeserializeOwned, Serialize}; @@ -323,16 +325,9 @@ impl ContainerClient { let driver_pk = partition_key.into().into_driver_partition_key(); // Create the driver operation and apply ItemWriteOptions fields. - let mut operation = + let operation = CosmosOperation::create_item(self.container_ref.clone(), driver_pk).with_body(body); - - // Wire session token and precondition from SDK options onto the operation. - if let Some(session_token) = options.session_token { - operation = operation.with_session_token(session_token); - } - if let Some(precondition) = options.precondition { - operation = operation.with_precondition(precondition); - } + let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. let driver_response = self @@ -428,15 +423,8 @@ impl ContainerClient { ); // Create the driver operation and apply ItemWriteOptions fields. - let mut operation = CosmosOperation::replace_item(item_ref).with_body(body); - - // Wire session token and precondition from SDK options onto the operation. - if let Some(session_token) = options.session_token { - operation = operation.with_session_token(session_token); - } - if let Some(precondition) = options.precondition { - operation = operation.with_precondition(precondition); - } + let operation = CosmosOperation::replace_item(item_ref).with_body(body); + let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. let driver_response = self @@ -528,16 +516,9 @@ impl ContainerClient { let driver_pk = partition_key.into().into_driver_partition_key(); // Create the driver operation and apply ItemWriteOptions fields. - let mut operation = + let operation = CosmosOperation::upsert_item(self.container_ref.clone(), driver_pk).with_body(body); - - // Wire session token and precondition from SDK options onto the operation. - if let Some(session_token) = options.session_token { - operation = operation.with_session_token(session_token); - } - if let Some(precondition) = options.precondition { - operation = operation.with_precondition(precondition); - } + let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. let driver_response = self @@ -595,15 +576,8 @@ impl ContainerClient { ); // Create the driver operation. - let mut operation = CosmosOperation::read_item(item_ref); - - // Wire session token and precondition from SDK options onto the operation. - if let Some(session_token) = options.session_token { - operation = operation.with_session_token(session_token); - } - if let Some(precondition) = options.precondition { - operation = operation.with_precondition(precondition); - } + let operation = CosmosOperation::read_item(item_ref); + let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. let driver_response = self @@ -653,15 +627,8 @@ impl ContainerClient { ); // Create the driver operation (no body for delete). - let mut operation = CosmosOperation::delete_item(item_ref); - - // Wire session token and precondition from SDK options onto the operation. - if let Some(session_token) = options.session_token { - operation = operation.with_session_token(session_token); - } - if let Some(precondition) = options.precondition { - operation = operation.with_precondition(precondition); - } + let operation = CosmosOperation::delete_item(item_ref); + let operation = apply_item_options(operation, options.session_token, options.precondition); // Execute through the driver. let driver_response = self @@ -825,6 +792,24 @@ impl ContainerClient { } } +/// Applies optional `session_token` and `precondition` to a [`CosmosOperation`]. +/// +/// Both [`ItemReadOptions`] and [`ItemWriteOptions`] carry these fields; +/// this helper avoids duplicating the wiring logic in every item operation. +fn apply_item_options( + mut operation: CosmosOperation, + session_token: Option, + precondition: Option, +) -> CosmosOperation { + if let Some(session_token) = session_token { + operation = operation.with_session_token(session_token); + } + if let Some(precondition) = precondition { + operation = operation.with_precondition(precondition); + } + operation +} + #[cfg(test)] mod tests { use super::*; From 50b20a75659a14aa225cddbfe007ef0126cacfac Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Mon, 13 Apr 2026 12:09:57 -0400 Subject: [PATCH 04/13] missing bits from initial copilot work --- .../emulator_tests/cosmos_fault_injection.rs | 4 +- .../src/driver/pipeline/operation_pipeline.rs | 60 ++++++++++++++++++- .../src/models/cosmos_headers.rs | 1 + 3 files changed, 61 insertions(+), 4 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs index b9f4a4ad76e..44e000de38b 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs @@ -1051,9 +1051,7 @@ pub async fn fault_injection_pkrange_readfeed_succeeds() -> Result<(), Box Date: Mon, 13 Apr 2026 14:55:14 -0400 Subject: [PATCH 05/13] remove test now that pkrange cache is moving to driver --- .../emulator_tests/cosmos_fault_injection.rs | 81 ------------------- 1 file changed, 81 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs index 44e000de38b..ebb2908a68c 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/emulator_tests/cosmos_fault_injection.rs @@ -996,84 +996,3 @@ pub async fn fault_injection_enable_disable_rule() -> Result<(), Box> ) .await } - -/// Verifies that the pkranges request uses the correct URL and returns 200. -/// -/// Uses a "spy" fault injection rule with no error or custom response. The rule -/// matches `MetadataPartitionKeyRanges` requests but lets them pass through to the -/// real service, recording the response status code. This proves: -/// 1. The pkranges code path is exercised (`hit_count >= 1`). -/// 2. The service returns 200 (the URL uses name-based addressing, not mixed RID). -/// 3. The item operation succeeds end-to-end with correct partition key routing. -#[tokio::test] -#[cfg_attr( - not(test_category = "emulator"), - ignore = "requires test_category 'emulator'" -)] -pub async fn fault_injection_pkrange_readfeed_succeeds() -> Result<(), Box> { - // Spy rule: no error, no custom response → passthrough with status recording. - let spy_result = FaultInjectionResultBuilder::new().build(); - - let condition = FaultInjectionConditionBuilder::new() - .with_operation_type(FaultOperationType::MetadataPartitionKeyRanges) - .build(); - - let rule = Arc::new( - FaultInjectionRuleBuilder::new("pkrange-spy", spy_result) - .with_condition(condition) - .build(), - ); - - let rule_handle = Arc::clone(&rule); - - let fault_builder = FaultInjectionClientBuilder::new().with_rule(rule); - - TestClient::run_with_unique_db( - async move |run_context, db_client| { - let container_id = format!("PkRangeTest-{}", Uuid::new_v4()); - let _container_client = run_context - .create_container_with_throughput( - db_client, - ContainerProperties::new(container_id.clone(), "/partition_key".into()), - ThroughputProperties::manual(400), - ) - .await?; - - let unique_id = Uuid::new_v4().to_string(); - let item = create_test_item(&unique_id); - let pk = format!("Partition-{}", unique_id); - - // Use the fault client so the spy rule is in the HTTP pipeline. - let fault_client = run_context - .fault_client() - .expect("fault client should be available"); - let fault_db_client = fault_client.database_client(db_client.id()); - let fault_container_client = fault_db_client.container_client(&container_id).await?; - - // Upsert an item — this triggers partition key range resolution - // through the driver's transport pipeline where the spy rule can observe. - fault_container_client - .upsert_item(&pk, item, None) - .await?; - - // The spy rule should have been hit during the pkrange fetch. - assert!( - rule_handle.hit_count() >= 1, - "Expected the MetadataPartitionKeyRanges spy rule to be hit at least once, but hit_count was {}", - rule_handle.hit_count() - ); - - // Verify the pkranges request returned 200 from the service. - let statuses = rule_handle.passthrough_statuses(); - assert!( - statuses.contains(&StatusCode::Ok), - "Expected pkranges request to return 200, got: {:?}", - statuses - ); - - Ok(()) - }, - Some(TestOptions::new().with_fault_injection_builder(fault_builder)), - ) - .await -} From e3eaded460edabb99c0a5e9a59b6cf049d051f29 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Tue, 14 Apr 2026 15:53:52 -0400 Subject: [PATCH 06/13] Update cosmos_multi_write.rs --- .../multi_write_tests/cosmos_multi_write.rs | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write.rs b/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write.rs index 67a9323f32b..86afec06471 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write.rs @@ -53,13 +53,14 @@ impl tracing::field::Visit for StringVisitor { } } -/// Finds log lines containing the expected operation and returns them -fn find_upsert_document_logs(logs: &[String]) -> Vec { +/// Finds log lines from the driver's operation pipeline containing routing decisions. +/// +/// The driver emits `tracing::debug!(routing_decision = %routing, "routing decision made")` +/// which includes the resolved region and endpoint URL. +fn find_upsert_routing_logs(logs: &[String]) -> Vec { logs.iter() .filter(|line| { - line.contains("azure_data_cosmos::retry_handler") - && line.contains("Upsert") - && line.contains("Document") + line.contains("azure_data_cosmos_driver") && line.contains("routing decision made") }) .cloned() .collect() @@ -86,7 +87,7 @@ fn create_container_and_write_item<'a>( .create_container_with_throughput(db_client, properties, throughput) .await?; - // This upsert operation should be logged by the retry_handler + // This upsert operation triggers a routing decision log in the driver container_client .upsert_item( "item1", @@ -142,7 +143,7 @@ pub async fn multi_write_preferred_locations() -> Result<(), Box> { // Check that the upsert went to the hub region { let logs = log_buffer.lock().unwrap(); - let upsert_logs = find_upsert_document_logs(&logs); + let upsert_logs = find_upsert_routing_logs(&logs); println!("Hub region upsert logs: {:?}", upsert_logs); assert!( @@ -184,7 +185,7 @@ pub async fn multi_write_preferred_locations() -> Result<(), Box> { // Check that the upsert went to the satellite region { let logs = log_buffer.lock().unwrap(); - let upsert_logs = find_upsert_document_logs(&logs); + let upsert_logs = find_upsert_routing_logs(&logs); println!("Satellite region upsert logs: {:?}", upsert_logs); assert!( From b6ba1f14939db4fda291b20530d16477fade5796 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Mon, 20 Apr 2026 17:41:45 -0400 Subject: [PATCH 07/13] Update operation_pipeline.rs --- .../src/driver/pipeline/operation_pipeline.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 a0ba9359b82..85cd6fdec02 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 @@ -1293,7 +1293,8 @@ mod tests { .get_optional_str(&HeaderName::from_static("x-ms-documentdb-is-upsert")) .is_none(), "is-upsert header should not be set for create" - } + ); + } #[test] fn build_transport_request_sets_priority_level_header() { From 223430424d5a3f245339727af5e367e17b7ce625 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Mon, 20 Apr 2026 18:08:18 -0400 Subject: [PATCH 08/13] formatting --- .../src/clients/container_client.rs | 9 +++- .../azure_data_cosmos/src/options/mod.rs | 17 ------- .../src/driver/pipeline/operation_pipeline.rs | 49 ++++++++++--------- .../src/models/cosmos_headers.rs | 1 + 4 files changed, 34 insertions(+), 42 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs index e9fe9856965..8859de1f886 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs @@ -8,7 +8,10 @@ use crate::{ BatchResponse, ContainerProperties, CosmosResponse, ItemResponse, ResourceResponse, ThroughputProperties, }, - options::{BatchOptions, QueryOptions, ReadContainerOptions, ReadFeedRangesOptions}, + options::{ + BatchOptions, Precondition, QueryOptions, ReadContainerOptions, ReadFeedRangesOptions, + SessionToken, + }, resource_context::{ResourceLink, ResourceType}, transactional_batch::TransactionalBatch, DeleteContainerOptions, FeedItemIterator, ItemReadOptions, ItemWriteOptions, PartitionKey, @@ -27,7 +30,6 @@ use azure_data_cosmos_driver::models::{ effective_partition_key::EffectivePartitionKey as DriverEpk, ContainerReference, CosmosOperation, ItemReference, PartitionKeyKind, }; -use azure_data_cosmos_driver::CosmosDriver; use serde::{de::DeserializeOwned, Serialize}; /// A client for working with a specific container in a Cosmos DB account. @@ -428,6 +430,7 @@ impl ContainerClient { // Execute through the driver. let driver_response = self + .context .driver .execute_operation(operation, options.operation) .await?; @@ -522,6 +525,7 @@ impl ContainerClient { // Execute through the driver. let driver_response = self + .context .driver .execute_operation(operation, options.operation) .await?; @@ -633,6 +637,7 @@ impl ContainerClient { // Execute through the driver. let driver_response = self + .context .driver .execute_operation(operation, options.operation) .await?; diff --git a/sdk/cosmos/azure_data_cosmos/src/options/mod.rs b/sdk/cosmos/azure_data_cosmos/src/options/mod.rs index 1b5ed4a0d70..b3f3fb230f3 100644 --- a/sdk/cosmos/azure_data_cosmos/src/options/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/options/mod.rs @@ -543,21 +543,4 @@ mod tests { assert_eq!(headers_result, headers_expected); } - - #[test] - fn no_throughput_control_headers_from_apply_headers_alone() { - // apply_headers() does not set throughput control headers — those are - // applied by the driver pipeline. Verify that priority and bucket - // headers are absent after apply_headers() only. - let options = ItemWriteOptions::default(); - let mut headers = Headers::new(); - options.apply_headers(&mut headers); - - assert!(headers - .get_optional_str(&constants::PRIORITY_LEVEL) - .is_none()); - assert!(headers - .get_optional_str(&constants::THROUGHPUT_BUCKET) - .is_none()); - } } 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 85cd6fdec02..bf9df4bcc20 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 @@ -20,8 +20,7 @@ use crate::{ }, models::{ request_header_names, AccountEndpoint, ActivityId, CosmosOperation, CosmosResponse, - CosmosResponseHeaders, Credential, DefaultConsistencyLevel, OperationType, SessionToken, - SubStatusCode, + Credential, DefaultConsistencyLevel, OperationType, SessionToken, SubStatusCode, }, options::{OperationOptionsView, ReadConsistencyStrategy, ThroughputControlGroupSnapshot}, }; @@ -500,7 +499,7 @@ fn build_transport_request( // distinguishes them via this header. if operation.operation_type() == OperationType::Upsert { headers.insert( - request_header_names::IS_UPSERT.clone(), + HeaderName::from_static(request_header_names::IS_UPSERT), HeaderValue::from_static("true"), ); } @@ -1253,16 +1252,18 @@ mod tests { let operation = CosmosOperation::upsert_item(test_container(), PartitionKey::from("pk1")) .with_body(b"{}".to_vec()); - let request = build_transport_request( - &operation, - None, - &test_routing(), - &ActivityId::from_string("default-activity".to_string()), - ExecutionContext::Initial, - None, - None, - ) - .expect("request should build"); + let routing = test_routing(); + let activity_id = ActivityId::from_string("default-activity".to_string()); + let ctx = TransportRequestContext { + routing: &routing, + activity_id: &activity_id, + execution_context: ExecutionContext::Initial, + deadline: None, + resolved_session_token: None, + throughput_control: None, + }; + let request = + build_transport_request(&operation, None, &ctx).expect("request should build"); let is_upsert = request .headers @@ -1276,16 +1277,18 @@ mod tests { let operation = CosmosOperation::create_item(test_container(), PartitionKey::from("pk1")) .with_body(b"{}".to_vec()); - let request = build_transport_request( - &operation, - None, - &test_routing(), - &ActivityId::from_string("default-activity".to_string()), - ExecutionContext::Initial, - None, - None, - ) - .expect("request should build"); + let routing = test_routing(); + let activity_id = ActivityId::from_string("default-activity".to_string()); + let ctx = TransportRequestContext { + routing: &routing, + activity_id: &activity_id, + execution_context: ExecutionContext::Initial, + deadline: None, + resolved_session_token: None, + throughput_control: None, + }; + let request = + build_transport_request(&operation, None, &ctx).expect("request should build"); assert!( request 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 9e3dfb96e37..87c8c99f8e7 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 @@ -19,6 +19,7 @@ pub(crate) mod request_header_names { pub const IF_MATCH: &str = "if-match"; pub const IF_NONE_MATCH: &str = "if-none-match"; pub const PREFER: &str = "prefer"; + pub const IS_UPSERT: &str = "x-ms-documentdb-is-upsert"; pub const PRIORITY_LEVEL: &str = "x-ms-cosmos-priority-level"; pub const THROUGHPUT_BUCKET: &str = "x-ms-cosmos-throughput-bucket"; } From 3594a0401c514686010c76493b6df9d768bbc8fa Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Thu, 23 Apr 2026 15:14:58 -0400 Subject: [PATCH 09/13] Update retry_evaluation.rs --- .../src/driver/pipeline/retry_evaluation.rs | 104 +++++++++++++++--- 1 file changed, 88 insertions(+), 16 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/retry_evaluation.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/retry_evaluation.rs index b034a19dc06..8184d0068cb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/retry_evaluation.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/retry_evaluation.rs @@ -10,8 +10,9 @@ //! - Transport error (Sent/Unknown, non-idempotent) → Abort //! - 403/3 WriteForbidden → FailoverRetry + refresh + mark unavailable //! - 404/1002 ReadSessionNotAvailable → SessionRetry (advances region) -//! - 503, 429/3092, 410 → FailoverRetry + mark partition/endpoint unavailable -//! - 500 (reads only) → FailoverRetry +//! - 429/3092, 410 → FailoverRetry + mark partition/endpoint unavailable (idempotency-gated) +//! - ≥500, 408 → FailoverRetry for all ops (reads + writes) +//! - 503 additionally marks partition unavailable //! - Other HTTP errors → Abort use azure_core::http::headers::Headers; @@ -87,15 +88,12 @@ pub(crate) fn evaluate_transport_result( ); } + // 429/3092 (SystemResourceUnavailable) and 410 (Gone). let is_system_resource_unavailable = status.is_throttled() && status.sub_status() == Some(SubStatusCode::SYSTEM_RESOURCE_UNAVAILABLE); - let is_service_unavailable = - status.status_code() == azure_core::http::StatusCode::ServiceUnavailable; let is_gone = status.is_gone(); - if (is_system_resource_unavailable || is_service_unavailable || is_gone) - && retry_state.can_retry_failover() - { + if (is_system_resource_unavailable || is_gone) && retry_state.can_retry_failover() { if request_definitely_not_sent { return ( OperationAction::FailoverRetry { @@ -138,19 +136,36 @@ pub(crate) fn evaluate_transport_result( ); } - if status.status_code() == azure_core::http::StatusCode::InternalServerError - && operation.is_read_only() + // Server errors (≥500) and request timeouts (408). + // Retries ALL operations (reads + writes). + // Cross-region retry will be attempted, if not possible in-region retry will be attempted. + let status_code = status.status_code(); + if (status_code.is_server_error() + || status_code == azure_core::http::StatusCode::RequestTimeout) && retry_state.can_retry_failover() { + let mut effects = vec![LocationEffect::MarkEndpointUnavailable { + endpoint: endpoint.clone(), + reason: UnavailableReason::InternalServerError, + }]; + if status_code == azure_core::http::StatusCode::ServiceUnavailable { + effects.push(LocationEffect::MarkPartitionUnavailable( + UnavailablePartition { + // TODO(partition-routing): Wire the actual partition key range ID from + // TransportResult or CosmosOperation once partition-level + // routing is implemented. + partition_key_range_id: String::new(), + region: endpoint.region().cloned(), + is_read: operation.is_read_only(), + }, + )); + } return ( OperationAction::FailoverRetry { new_state: retry_state.clone().advance_failover(), delay: None, }, - vec![LocationEffect::MarkEndpointUnavailable { - endpoint: endpoint.clone(), - reason: UnavailableReason::InternalServerError, - }], + effects, ); } @@ -543,7 +558,7 @@ mod tests { } #[test] - fn service_unavailable_sent_non_idempotent_aborts() { + fn service_unavailable_write_retries_and_marks_partition() { let op = make_create_operation(); let result = make_http_error(StatusCode::ServiceUnavailable); let state = OperationRetryState::initial(0, false, Vec::new(), 3, 1); @@ -552,8 +567,13 @@ mod tests { ); let (action, effects) = evaluate_transport_result(&op, &endpoint, result, &state); - assert!(matches!(action, OperationAction::Abort { .. })); - assert!(effects.is_empty()); + assert!(matches!(action, OperationAction::FailoverRetry { .. })); + assert!(effects + .iter() + .any(|e| matches!(e, LocationEffect::MarkEndpointUnavailable { .. }))); + assert!(effects + .iter() + .any(|e| matches!(e, LocationEffect::MarkPartitionUnavailable(_)))); } #[test] @@ -599,4 +619,56 @@ mod tests { .iter() .any(|e| matches!(e, LocationEffect::MarkEndpointUnavailable { .. }))); } + + #[test] + fn internal_server_error_on_write_retries() { + let op = make_create_operation(); + let result = make_http_error(StatusCode::InternalServerError); + let state = OperationRetryState::initial(0, false, Vec::new(), 3, 1); + let endpoint = CosmosEndpoint::global( + url::Url::parse("https://test.documents.azure.com:443/").unwrap(), + ); + + let (action, effects) = evaluate_transport_result(&op, &endpoint, result, &state); + assert!(matches!(action, OperationAction::FailoverRetry { .. })); + assert!(effects + .iter() + .any(|e| matches!(e, LocationEffect::MarkEndpointUnavailable { .. }))); + // 500 should NOT mark partition unavailable (only 503 does) + assert!(!effects + .iter() + .any(|e| matches!(e, LocationEffect::MarkPartitionUnavailable(_)))); + } + + #[test] + fn request_timeout_read_retries() { + let op = make_read_operation(); + let result = make_http_error(StatusCode::RequestTimeout); + let state = OperationRetryState::initial(0, false, Vec::new(), 3, 1); + let endpoint = CosmosEndpoint::global( + url::Url::parse("https://test.documents.azure.com:443/").unwrap(), + ); + + let (action, effects) = evaluate_transport_result(&op, &endpoint, result, &state); + assert!(matches!(action, OperationAction::FailoverRetry { .. })); + assert!(effects + .iter() + .any(|e| matches!(e, LocationEffect::MarkEndpointUnavailable { .. }))); + } + + #[test] + fn request_timeout_write_retries() { + let op = make_create_operation(); + let result = make_http_error(StatusCode::RequestTimeout); + let state = OperationRetryState::initial(0, false, Vec::new(), 3, 1); + let endpoint = CosmosEndpoint::global( + url::Url::parse("https://test.documents.azure.com:443/").unwrap(), + ); + + let (action, effects) = evaluate_transport_result(&op, &endpoint, result, &state); + assert!(matches!(action, OperationAction::FailoverRetry { .. })); + assert!(effects + .iter() + .any(|e| matches!(e, LocationEffect::MarkEndpointUnavailable { .. }))); + } } From 1c57cfeb0ab64078fdbbc96cfe0725d5e7835b42 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Fri, 24 Apr 2026 10:41:41 -0400 Subject: [PATCH 10/13] Update cosmos_multi_write_fault_injection.rs --- .../cosmos_multi_write_fault_injection.rs | 29 +++++++++---------- 1 file changed, 13 insertions(+), 16 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs b/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs index b26c4adae9f..ad151e362c1 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/multi_write_tests/cosmos_multi_write_fault_injection.rs @@ -348,13 +348,19 @@ pub async fn fault_injection_read_region_retry_503() -> Result<(), Box Result<(), Box> { +pub async fn fault_injection_transport_generated_503_write_aborts() -> Result<(), Box> { let server_error = FaultInjectionResultBuilder::new() .with_error(FaultInjectionErrorType::ServiceUnavailable) .build(); @@ -364,7 +370,7 @@ pub async fn fault_injection_write_region_retry_503() -> Result<(), Box Result<(), Box Date: Fri, 24 Apr 2026 13:02:51 -0400 Subject: [PATCH 11/13] add preferred regions logic to driver for regional routing --- .../src/clients/cosmos_client_builder.rs | 8 +- .../src/driver/cosmos_driver.rs | 1 + .../driver/routing/location_state_store.rs | 8 + .../src/driver/routing/routing_systems.rs | 189 +++++++++++++++++- .../src/options/driver_options.rs | 41 +++- 5 files changed, 234 insertions(+), 13 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs index 24336f9d175..697f0370fcc 100644 --- a/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs +++ b/sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs @@ -360,7 +360,7 @@ impl CosmosClientBuilder { let global_endpoint_manager = GlobalEndpointManager::new( endpoint.clone(), - preferred_regions, + preferred_regions.clone(), Vec::new(), pipeline_core.clone(), ); @@ -444,8 +444,12 @@ impl CosmosClientBuilder { })?; } let driver_runtime = driver_runtime_builder.build().await?; + let driver_options = + azure_data_cosmos_driver::options::DriverOptions::builder(driver_account) + .with_preferred_regions(preferred_regions) + .build(); let driver = driver_runtime - .get_or_create_driver(driver_account, None) + .get_or_create_driver(driver_options.account().clone(), Some(driver_options)) .await?; Ok(CosmosClient { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs index 9993ae10510..b427b978f3d 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/cosmos_driver.rs @@ -752,6 +752,7 @@ impl CosmosDriver { refresh_callback, runtime.connection_pool().is_gateway20_allowed(), endpoint_unavailability_ttl, + options.preferred_regions().to_vec(), )); Self { diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs index 69d7bd5d5a6..44c48cbee18 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/location_state_store.rs @@ -17,6 +17,7 @@ use futures::future::BoxFuture; use crate::{ driver::cache::{AccountMetadataCache, AccountProperties}, models::AccountEndpoint, + options::Region, }; use super::{ @@ -62,6 +63,7 @@ pub(crate) struct LocationStateStore { account_endpoint: AccountEndpoint, account_refresh_fn: AccountRefreshFn, default_endpoint: CosmosEndpoint, + preferred_regions: Vec, gateway20_enabled: bool, endpoint_unavailability_ttl: Duration, refresh_interval: Duration, @@ -126,6 +128,7 @@ impl LocationStateStore { account_refresh_fn: AccountRefreshFn, gateway20_enabled: bool, endpoint_unavailability_ttl: Duration, + preferred_regions: Vec, ) -> Self { let account_state = AccountEndpointState::single(default_endpoint.clone()); @@ -141,6 +144,7 @@ impl LocationStateStore { account_endpoint, account_refresh_fn, default_endpoint, + preferred_regions, gateway20_enabled, endpoint_unavailability_ttl, // TODO(refresh-config): Make refresh interval configurable. @@ -342,6 +346,7 @@ impl LocationStateStore { default_endpoint.clone(), Some(current.generation), self.gateway20_enabled, + &self.preferred_regions, ); // Carry forward unavailability marks from the current state, // filtering out entries that have expired past the configured TTL. @@ -437,6 +442,7 @@ mod tests { refresh, false, Duration::from_secs(60), + Vec::new(), ); store @@ -473,6 +479,7 @@ mod tests { refresh, false, Duration::from_secs(60), + Vec::new(), ); store @@ -503,6 +510,7 @@ mod tests { refresh, false, Duration::from_secs(60), + Vec::new(), ); let properties = Arc::new(test_refresh_payload()); diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/routing_systems.rs b/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/routing_systems.rs index 1a934086e9a..8ab9d704bcb 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/routing_systems.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/driver/routing/routing_systems.rs @@ -12,22 +12,22 @@ use std::{ use tracing::warn; use crate::driver::cache::AccountProperties; +use crate::options::Region; use super::{AccountEndpointState, CosmosEndpoint, UnavailableReason}; /// Builds account endpoint state from account metadata. /// -/// TODO: Accept `preferred_locations: &[Region]` to reorder endpoint lists -/// based on user configuration (derived from `application_region` via -/// `generate_preferred_region_list` in `azure_data_cosmos`). Wire this when -/// operations in `azure_data_cosmos` are migrated to the driver's -/// `execute_operation` API — that cross-crate change is the natural point -/// to thread preferred regions through `DriverOptions` → `LocationStateStore`. +/// When `preferred_regions` is non-empty, read and write endpoint lists are +/// reordered so that endpoints matching the preferred regions appear first +/// (in preference order). Endpoints whose region is not in the preferred +/// list are appended in their original account-metadata order. pub(crate) fn build_account_endpoint_state( properties: &AccountProperties, default_endpoint: CosmosEndpoint, previous_generation: Option, gateway20_enabled: bool, + preferred_regions: &[Region], ) -> AccountEndpointState { let generation = previous_generation.map_or(0, |g| g.saturating_add(1)); @@ -43,6 +43,13 @@ pub(crate) fn build_account_endpoint_state( gateway20_enabled, ); + if !preferred_regions.is_empty() { + preferred_read_endpoints = + reorder_by_preferred_regions(preferred_read_endpoints, preferred_regions); + preferred_write_endpoints = + reorder_by_preferred_regions(preferred_write_endpoints, preferred_regions); + } + if preferred_read_endpoints.is_empty() { preferred_read_endpoints.push(default_endpoint.clone()); } @@ -128,6 +135,36 @@ fn parse_thin_client_locations( urls } +/// Reorders endpoints so that those matching `preferred_regions` appear first, +/// in the same order as the preference list. Endpoints whose region is not in +/// the preference list (or that have no region, e.g. global endpoints) are +/// appended in their original order. +fn reorder_by_preferred_regions( + endpoints: Vec, + preferred_regions: &[Region], +) -> Vec { + let mut ordered = Vec::with_capacity(endpoints.len()); + let mut remaining: Vec> = endpoints.into_iter().map(Some).collect(); + + for region in preferred_regions { + for slot in remaining.iter_mut() { + if let Some(ep) = slot { + if ep.region().is_some_and(|r| r == region) { + ordered.push(slot.take().unwrap()); + break; + } + } + } + } + + // Append any endpoints not matched by the preferred list. + for ep in remaining.into_iter().flatten() { + ordered.push(ep); + } + + ordered +} + /// Returns a new state with an endpoint marked unavailable. pub(crate) fn mark_endpoint_unavailable( state: &AccountEndpointState, @@ -176,6 +213,7 @@ pub(crate) fn expire_unavailable_endpoints( mod tests { use super::*; use crate::driver::cache::AccountProperties; + use crate::options::Region; fn default_endpoint() -> CosmosEndpoint { CosmosEndpoint::global(url::Url::parse("https://test.documents.azure.com:443/").unwrap()) @@ -204,7 +242,7 @@ mod tests { #[test] fn build_state_uses_metadata_locations() { let state = - build_account_endpoint_state(&test_properties(), default_endpoint(), None, false); + build_account_endpoint_state(&test_properties(), default_endpoint(), None, false, &[]); assert_eq!(state.generation, 0); assert_eq!(state.preferred_write_endpoints.len(), 1); assert_eq!(state.preferred_read_endpoints.len(), 1); @@ -232,7 +270,7 @@ mod tests { })) .unwrap(); - let state = build_account_endpoint_state(&properties, default_endpoint(), None, true); + let state = build_account_endpoint_state(&properties, default_endpoint(), None, true, &[]); assert!(state.preferred_read_endpoints[0].gateway20_url().is_some()); assert!(state.preferred_write_endpoints[0].gateway20_url().is_none()); @@ -260,7 +298,7 @@ mod tests { })) .unwrap(); - let state = build_account_endpoint_state(&properties, default_endpoint(), None, true); + let state = build_account_endpoint_state(&properties, default_endpoint(), None, true, &[]); assert!(state.preferred_read_endpoints[0].gateway20_url().is_some()); assert!(state.preferred_write_endpoints[0].gateway20_url().is_some()); @@ -269,7 +307,7 @@ mod tests { #[test] fn mark_and_expire_unavailable_endpoint() { let state = - build_account_endpoint_state(&test_properties(), default_endpoint(), None, false); + build_account_endpoint_state(&test_properties(), default_endpoint(), None, false, &[]); let endpoint = state.preferred_read_endpoints[0].clone(); let marked = mark_endpoint_unavailable(&state, &endpoint, UnavailableReason::TransportError); @@ -282,4 +320,135 @@ mod tests { ); assert!(expired.unavailable_endpoints.is_empty()); } + + fn multi_region_properties() -> AccountProperties { + serde_json::from_value(serde_json::json!({ + "_self": "", + "id": "test", + "_rid": "test.documents.azure.com", + "media": "//media/", + "addresses": "//addresses/", + "_dbs": "//dbs/", + "writableLocations": [ + { "name": "eastus", "databaseAccountEndpoint": "https://test-eastus.documents.azure.com:443/" }, + { "name": "westus2", "databaseAccountEndpoint": "https://test-westus2.documents.azure.com:443/" }, + { "name": "westus3", "databaseAccountEndpoint": "https://test-westus3.documents.azure.com:443/" } + ], + "readableLocations": [ + { "name": "eastus", "databaseAccountEndpoint": "https://test-eastus.documents.azure.com:443/" }, + { "name": "westus2", "databaseAccountEndpoint": "https://test-westus2.documents.azure.com:443/" }, + { "name": "westus3", "databaseAccountEndpoint": "https://test-westus3.documents.azure.com:443/" } + ], + "enableMultipleWriteLocations": true, + "userReplicationPolicy": { "minReplicaSetSize": 3, "maxReplicasetSize": 4 }, + "userConsistencyPolicy": { "defaultConsistencyLevel": "Session" }, + "systemReplicationPolicy": { "minReplicaSetSize": 3, "maxReplicasetSize": 4 }, + "readPolicy": { "primaryReadCoefficient": 1, "secondaryReadCoefficient": 1 }, + "queryEngineConfiguration": "{}" + })) + .unwrap() + } + + #[test] + fn preferred_regions_reorder_write_endpoints() { + let preferred = vec![Region::WEST_US_3, Region::EAST_US]; + let state = build_account_endpoint_state( + &multi_region_properties(), + default_endpoint(), + None, + false, + &preferred, + ); + + assert_eq!(state.preferred_write_endpoints.len(), 3); + assert_eq!( + state.preferred_write_endpoints[0].region().unwrap(), + &Region::WEST_US_3 + ); + assert_eq!( + state.preferred_write_endpoints[1].region().unwrap(), + &Region::EAST_US + ); + assert_eq!( + state.preferred_write_endpoints[2].region().unwrap(), + &Region::WEST_US_2 + ); + } + + #[test] + fn preferred_regions_reorder_read_endpoints() { + let preferred = vec![Region::WEST_US_2, Region::WEST_US_3]; + let state = build_account_endpoint_state( + &multi_region_properties(), + default_endpoint(), + None, + false, + &preferred, + ); + + assert_eq!(state.preferred_read_endpoints.len(), 3); + assert_eq!( + state.preferred_read_endpoints[0].region().unwrap(), + &Region::WEST_US_2 + ); + assert_eq!( + state.preferred_read_endpoints[1].region().unwrap(), + &Region::WEST_US_3 + ); + assert_eq!( + state.preferred_read_endpoints[2].region().unwrap(), + &Region::EAST_US + ); + } + + #[test] + fn preferred_regions_unknown_regions_are_skipped() { + let preferred = vec![Region::new("nonexistent"), Region::WEST_US_3]; + let state = build_account_endpoint_state( + &multi_region_properties(), + default_endpoint(), + None, + false, + &preferred, + ); + + // westus3 should be first; the nonexistent region is skipped. + assert_eq!( + state.preferred_write_endpoints[0].region().unwrap(), + &Region::WEST_US_3 + ); + assert_eq!( + state.preferred_write_endpoints[1].region().unwrap(), + &Region::EAST_US + ); + assert_eq!( + state.preferred_write_endpoints[2].region().unwrap(), + &Region::WEST_US_2 + ); + } + + #[test] + fn empty_preferred_regions_preserves_original_order() { + let state = build_account_endpoint_state( + &multi_region_properties(), + default_endpoint(), + None, + false, + &[], + ); + + // Original account-metadata order: eastus, westus2, westus3. + assert_eq!( + state.preferred_write_endpoints[0].region().unwrap(), + &Region::EAST_US + ); + assert_eq!( + state.preferred_write_endpoints[1].region().unwrap(), + &Region::WEST_US_2 + ); + assert_eq!( + state.preferred_write_endpoints[2].region().unwrap(), + &Region::WEST_US_3 + ); + } } diff --git a/sdk/cosmos/azure_data_cosmos_driver/src/options/driver_options.rs b/sdk/cosmos/azure_data_cosmos_driver/src/options/driver_options.rs index cd5af0d6928..de8a3aeb2ea 100644 --- a/sdk/cosmos/azure_data_cosmos_driver/src/options/driver_options.rs +++ b/sdk/cosmos/azure_data_cosmos_driver/src/options/driver_options.rs @@ -5,7 +5,10 @@ use std::sync::Arc; -use crate::{models::AccountReference, options::OperationOptions}; +use crate::{ + models::AccountReference, + options::{OperationOptions, Region}, +}; /// Configuration options for a Cosmos DB driver instance. /// @@ -43,6 +46,12 @@ pub struct DriverOptions { account: AccountReference, /// Driver-level operation options (e.g., consistency, excluded regions, failover, session retry). operation_options: Arc, + /// Preferred regions for routing, ordered by proximity to the application. + /// + /// When non-empty, read and write endpoint lists are reordered so that + /// endpoints matching these regions appear first. Regions that don't match + /// any account endpoint are silently skipped. + preferred_regions: Vec, } impl DriverOptions { @@ -62,6 +71,11 @@ impl DriverOptions { pub fn operation_options(&self) -> &Arc { &self.operation_options } + + /// Returns the preferred regions for routing. + pub fn preferred_regions(&self) -> &[Region] { + &self.preferred_regions + } } /// Builder for creating [`DriverOptions`]. @@ -73,6 +87,7 @@ impl DriverOptions { pub struct DriverOptionsBuilder { account: AccountReference, operation_options: Option, + preferred_regions: Vec, } impl DriverOptionsBuilder { @@ -81,6 +96,7 @@ impl DriverOptionsBuilder { Self { account, operation_options: None, + preferred_regions: Vec::new(), } } @@ -90,11 +106,22 @@ impl DriverOptionsBuilder { self } + /// Sets the preferred regions for routing. + /// + /// Regions should be ordered by proximity to the application (closest first). + /// The driver reorders endpoint lists to prefer these regions for both reads + /// and writes. Regions not present in the account are silently skipped. + pub fn with_preferred_regions(mut self, regions: Vec) -> Self { + self.preferred_regions = regions; + self + } + /// Builds the [`DriverOptions`]. pub fn build(self) -> DriverOptions { DriverOptions { account: self.account, operation_options: Arc::new(self.operation_options.unwrap_or_default()), + preferred_regions: self.preferred_regions, } } } @@ -170,5 +197,17 @@ mod tests { .operation_options() .read_consistency_strategy .is_none()); + assert!(options.preferred_regions().is_empty()); + } + + #[test] + fn builder_sets_preferred_regions() { + let regions = vec![Region::WEST_US_2, Region::EAST_US]; + + let options = DriverOptionsBuilder::new(test_account()) + .with_preferred_regions(regions.clone()) + .build(); + + assert_eq!(options.preferred_regions(), ®ions); } } From e48a139d8b5a11f54a65f158b684faf6bf1aca9d Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Fri, 24 Apr 2026 13:55:18 -0400 Subject: [PATCH 12/13] adjust test infra --- .../tests/framework/test_client.rs | 40 +++++++++++-------- 1 file changed, 24 insertions(+), 16 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/tests/framework/test_client.rs b/sdk/cosmos/azure_data_cosmos/tests/framework/test_client.rs index f1261bb020e..ff193a2eb0d 100644 --- a/sdk/cosmos/azure_data_cosmos/tests/framework/test_client.rs +++ b/sdk/cosmos/azure_data_cosmos/tests/framework/test_client.rs @@ -746,15 +746,20 @@ impl TestRunContext { let container_id = &created_properties.id; - // Wait for hub region client to successfully read the container + // Wait for hub region client to successfully resolve and read the container. + // Both `container_client()` (which resolves metadata via the driver) and + // `read()` can fail with 404 while the container replicates. loop { - match hub_client - .database_client(db_client.id()) - .container_client(container_id) - .await? - .read(None) - .await - { + let result = async { + hub_client + .database_client(db_client.id()) + .container_client(container_id) + .await? + .read(None) + .await + } + .await; + match result { Ok(_) => break, Err(e) => { println!( @@ -767,15 +772,18 @@ impl TestRunContext { } } - // Wait for satellite region client to successfully read the container + // Wait for satellite region client to successfully resolve and read the container. loop { - match satellite_client - .database_client(db_client.id()) - .container_client(container_id) - .await? - .read(None) - .await - { + let result = async { + satellite_client + .database_client(db_client.id()) + .container_client(container_id) + .await? + .read(None) + .await + } + .await; + match result { Ok(_) => break, Err(e) => { println!( From 66d00461bf7f0e3f2b516dc5bbb54013113a4792 Mon Sep 17 00:00:00 2001 From: Simon Moreno <30335873+simorenoh@users.noreply.github.com> Date: Fri, 24 Apr 2026 15:56:20 -0400 Subject: [PATCH 13/13] Update mod.rs --- .../azure_data_cosmos/src/options/mod.rs | 33 ------------------- 1 file changed, 33 deletions(-) diff --git a/sdk/cosmos/azure_data_cosmos/src/options/mod.rs b/sdk/cosmos/azure_data_cosmos/src/options/mod.rs index d4fae1e1817..6486acc4d6c 100644 --- a/sdk/cosmos/azure_data_cosmos/src/options/mod.rs +++ b/sdk/cosmos/azure_data_cosmos/src/options/mod.rs @@ -394,39 +394,6 @@ mod tests { ); } - #[test] - fn item_write_options_default_as_headers() { - let options = ItemWriteOptions::default(); - - let mut headers_result = Headers::new(); - options.apply_headers(&mut headers_result); - let headers_result: Vec<(HeaderName, HeaderValue)> = headers_result.into_iter().collect(); - - let headers_expected: Vec<(HeaderName, HeaderValue)> = - vec![(headers::PREFER, constants::PREFER_MINIMAL)]; - - assert_eq!(headers_result, headers_expected); - } - - #[test] - fn item_write_options_with_content_response_enabled() { - let mut operation = OperationOptions::default(); - operation.content_response_on_write = Some(ContentResponseOnWrite::Enabled); - - let options = ItemWriteOptions { - operation, - ..Default::default() - }; - - let mut headers_result = Headers::new(); - options.apply_headers(&mut headers_result); - let headers_result: Vec<(HeaderName, HeaderValue)> = headers_result.into_iter().collect(); - - let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![]; - - assert_eq!(headers_result, headers_expected); - } - #[test] fn batch_options_as_headers() { let mut custom_headers = HashMap::new();