Skip to content

Commit b50fb4c

Browse files
Copilotsimorenoh
andauthored
Cosmos: Cut over replace_item, upsert_item, delete_item to driver (#4128)
Routes the remaining item-level operations (`replace_item`, `upsert_item`, `delete_item`) through the `CosmosDriver`, completing the item operation cutover from the legacy gateway pipeline. ### Driver - **`CosmosOperation::upsert_item`** now takes `(ContainerReference, PartitionKey)` instead of `ItemReference`. Upsert is a `POST` to the collection feed (same as create), not a `PUT` to a document URL. ### SDK - **`replace_item`**, **`upsert_item`**, **`delete_item`** in `ContainerClient` follow the same pattern as `create_item`/`read_item`: build `CosmosOperation` → wire session token/precondition → `driver.execute_operation()` → bridge response. ### Cleanup - Removed `ItemWriteOptions::apply_headers()` and `apply_precondition_headers()` — dead code now that all item writes use the driver. - Removed associated unit tests for the removed shims. ### Updated state | Method | Pipeline | |---|---| | `read_item` | Driver | | `create_item` | Driver | | `replace_item` | **Driver** ← | | `upsert_item` | **Driver** ← | | `delete_item` | **Driver** ← | | `query_items` | Gateway | | Database/Container CRUD | Gateway | | Batch | Gateway | > [!WARNING] > > <details> > <summary>Firewall rules blocked me from connecting to one or more addresses (expand for details)</summary> > > #### I tried to connect to the following addresses, but was blocked by firewall rules: > > - `test.documents.azure.com` > - Triggering command: `/home/REDACTED/work/azure-sdk-for-rust/azure-sdk-for-rust/target/debug/deps/azure_data_cosmos-6ba76bd8313db34b /home/REDACTED/work/azure-sdk-for-rust/azure-sdk-for-rust/target/debug/deps/azure_data_cosmos-6ba76bd8313db34b /home/REDACTED/work/azure-sdk-for-rust/azure-sdk-for-rust/target/debug/build/azure_core_test-3ead228e4bac545f/build_script_build-3ead228e4bac545f.28rxmzcc53ieso0vhgnjhp0mb.1g0umtu.rcgu.o -Wl,--as-needed -Wl,-Bstatic get/debug/deps/rustcVORl3H/list /home/REDACTED/.rustup/toolchains/stable-x86_64-REDACTED-linux-gnu/lib/rustlib/x86_64-REDACTED-linux-gnu/lib/libpanic_unwind-e462f106b2b26a06.rlib /home/REDACTED/.rustup/toolchains/stable-x86_64-REDACTED-linux-gnu/lib/rustlib/x86_64-REDACTED-linux-gnu/lib/libobject-926daa94a00ee327.rlib ORl3H/symbols.o core_test_macros-a1ab7b6aba03e8ee.00uyiyx5pyx2v1ngfia649863.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.0wxt1muwk7ybp3n7mh0wllbni.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.10a2tbvi9x7z0h5lmbonpdh71.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.16kd4stujzqt4xf2jl2sbtd59.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.1kccwp8eolfs8k1p8t95mpyuj.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.1lmiqx49m2x5aacy6jphmlzgy.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.1zbl7otmqptz39v95knjdsqes.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.2f42o1vbvljajd8rp0o10wval.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.2zwx85bqwtramffcrv2xdiq6i.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.31h1vvrpiuwu05d0sfhyv10f4.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.34eawspeaufijc3hxiybm1m5s.1iuy4et.rcgu.o core_test_macros-a1ab7b6aba03e8ee.3oxva6e4m60m4raphd0wdn5rb.1iuy4et.rcgu.o` (dns block) > - Triggering command: `/home/REDACTED/work/azure-sdk-for-rust/azure-sdk-for-rust/target/debug/deps/azure_data_cosmos-6ba76bd8313db34b /home/REDACTED/work/azure-sdk-for-rust/azure-sdk-for-rust/target/debug/deps/azure_data_cosmos-6ba76bd8313db34b _data_cosmos-ef5244a14879f580/build_script_build-ef5244a14879f580.5hz7v9l3n3w76s28d4�� _data_cosmos-ef5244a14879f580/build_script_build-ef5244a14879f580.b3xwjkkkf035nfaknghx3bnrq.1v2nf23.rcgu.o _data_cosmos-ef5244a14879f580/build_script_build-ef5244a14879f580.c2fzye5iqbx61t8txhcw5znvn.1v2nf23.rcgu.o _data_cosmos-ef5244a14879f580/build_script_build-ef5244a14879f580.d6uyvyxfwkk1tmp94qt99wief.1v2nf23.rcgu.o rqab2.rcgu.o rqab2.rcgu.o 64-REDACTED-linux-gnu/lib/libstd-46d936097e8c5b85.rlib 64-REDACTED-linux-gnu/lib/libpanic_unwind-e462f106b2b26a06.rlib 64-REDACTED-linux-gnu/lib/libobject-926daa94a00ee327.rlib 64-REDACTED-linux-gnu/lib/libmemchr-48d5b0db80402653.rlib 64-REDACTED-linux-gnu/lib/libaddr2line-3367f26bd486b29d.rlib 64-u�� 64-REDACTED-linux-gnu/lib/libcfg_if-a5addfdc94c3bad3.rlib 64-REDACTED-linux-gnu/lib/librustc_demangle-789fb9c0cb1c7158.rlib 4et.rcgu.o 4et.rcgu.o 4et.rcgu.o 4et.rcgu.o 4et.rcgu.o` (dns block) > > If you need me to access, download, or install something from one of these locations, you can either: > > - Configure [Actions setup steps](https://gh.io/copilot/actions-setup-steps) to set up my environment, which run before the firewall is enabled > - Add the appropriate URLs or hosts to the custom allowlist in this repository's [Copilot coding agent settings](https://github.com/Azure/azure-sdk-for-rust/settings/copilot/coding_agent) (admins only) > > </details> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: simorenoh <30335873+simorenoh@users.noreply.github.com>
1 parent ab3ca55 commit b50fb4c

16 files changed

Lines changed: 532 additions & 426 deletions

File tree

sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs

Lines changed: 92 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,10 @@ use crate::{
88
BatchResponse, ContainerProperties, CosmosResponse, ItemResponse, ResourceResponse,
99
ThroughputProperties,
1010
},
11-
options::{BatchOptions, QueryOptions, ReadContainerOptions, ReadFeedRangesOptions},
11+
options::{
12+
BatchOptions, Precondition, QueryOptions, ReadContainerOptions, ReadFeedRangesOptions,
13+
SessionToken,
14+
},
1215
resource_context::{ResourceLink, ResourceType},
1316
transactional_batch::TransactionalBatch,
1417
DeleteContainerOptions, FeedItemIterator, ItemReadOptions, ItemWriteOptions, PartitionKey,
@@ -322,16 +325,9 @@ impl ContainerClient {
322325
let driver_pk = partition_key.into().into_driver_partition_key();
323326

324327
// Create the driver operation and apply ItemWriteOptions fields.
325-
let mut operation =
328+
let operation =
326329
CosmosOperation::create_item(self.container_ref.clone(), driver_pk).with_body(body);
327-
328-
// Wire session token and precondition from SDK options onto the operation.
329-
if let Some(session_token) = options.session_token {
330-
operation = operation.with_session_token(session_token);
331-
}
332-
if let Some(precondition) = options.precondition {
333-
operation = operation.with_precondition(precondition);
334-
}
330+
let operation = apply_item_options(operation, options.session_token, options.precondition);
335331

336332
// Execute through the driver.
337333
let driver_response = self
@@ -417,20 +413,31 @@ impl ContainerClient {
417413
item: T,
418414
options: Option<ItemWriteOptions>,
419415
) -> azure_core::Result<ItemResponse<()>> {
420-
let link = self.items_link.item(item_id);
421-
let options = options.clone().unwrap_or_default();
422-
let excluded_regions = options.operation.excluded_regions.clone();
423-
let mut cosmos_request = CosmosRequest::builder(OperationType::Replace, link)
424-
.json(&item)
425-
.partition_key(partition_key.into())
426-
.excluded_regions(excluded_regions)
427-
.build()?;
428-
options.apply_headers(&mut cosmos_request.headers);
416+
let options = options.unwrap_or_default();
417+
let body = serde_json::to_vec(&item)?;
429418

430-
self.container_connection
431-
.send(cosmos_request, Context::default())
432-
.await
433-
.map(ItemResponse::new)
419+
// Build the driver's item reference from our stored container metadata.
420+
let item_ref = ItemReference::from_name(
421+
&self.container_ref,
422+
partition_key.into().into_driver_partition_key(),
423+
item_id.to_owned(),
424+
);
425+
426+
// Create the driver operation and apply ItemWriteOptions fields.
427+
let operation = CosmosOperation::replace_item(item_ref).with_body(body);
428+
let operation = apply_item_options(operation, options.session_token, options.precondition);
429+
430+
// Execute through the driver.
431+
let driver_response = self
432+
.context
433+
.driver
434+
.execute_operation(operation, options.operation)
435+
.await?;
436+
437+
// Bridge the driver response to the SDK response type.
438+
Ok(ItemResponse::new(
439+
crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
440+
))
434441
}
435442

436443
/// Creates or replaces an item in the container.
@@ -506,21 +513,26 @@ impl ContainerClient {
506513
item: T,
507514
options: Option<ItemWriteOptions>,
508515
) -> azure_core::Result<ItemResponse<()>> {
509-
let options = options.clone().unwrap_or_default();
510-
let excluded_regions = options.operation.excluded_regions.clone();
511-
let mut cosmos_request =
512-
CosmosRequest::builder(OperationType::Upsert, self.items_link.clone())
513-
.json(&item)
514-
.partition_key(partition_key.into())
515-
.excluded_regions(excluded_regions)
516-
.build()?;
517-
options.apply_headers(&mut cosmos_request.headers);
516+
let options = options.unwrap_or_default();
517+
let body = serde_json::to_vec(&item)?;
518+
let driver_pk = partition_key.into().into_driver_partition_key();
518519

519-
return self
520-
.container_connection
521-
.send(cosmos_request, Context::default())
522-
.await
523-
.map(ItemResponse::new);
520+
// Create the driver operation and apply ItemWriteOptions fields.
521+
let operation =
522+
CosmosOperation::upsert_item(self.container_ref.clone(), driver_pk).with_body(body);
523+
let operation = apply_item_options(operation, options.session_token, options.precondition);
524+
525+
// Execute through the driver.
526+
let driver_response = self
527+
.context
528+
.driver
529+
.execute_operation(operation, options.operation)
530+
.await?;
531+
532+
// Bridge the driver response to the SDK response type.
533+
Ok(ItemResponse::new(
534+
crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
535+
))
524536
}
525537

526538
/// Reads a specific item from the container.
@@ -567,15 +579,8 @@ impl ContainerClient {
567579
);
568580

569581
// Create the driver operation.
570-
let mut operation = CosmosOperation::read_item(item_ref);
571-
572-
// Wire session token and precondition from SDK options onto the operation.
573-
if let Some(session_token) = options.session_token {
574-
operation = operation.with_session_token(session_token);
575-
}
576-
if let Some(precondition) = options.precondition {
577-
operation = operation.with_precondition(precondition);
578-
}
582+
let operation = CosmosOperation::read_item(item_ref);
583+
let operation = apply_item_options(operation, options.session_token, options.precondition);
579584

580585
// Execute through the driver.
581586
let driver_response = self
@@ -616,19 +621,30 @@ impl ContainerClient {
616621
item_id: &str,
617622
options: Option<ItemWriteOptions>,
618623
) -> azure_core::Result<ItemResponse<()>> {
619-
let link = self.items_link.item(item_id);
620-
let options = options.clone().unwrap_or_default();
621-
let excluded_regions = options.operation.excluded_regions.clone();
622-
let mut cosmos_request = CosmosRequest::builder(OperationType::Delete, link)
623-
.partition_key(partition_key.into())
624-
.excluded_regions(excluded_regions)
625-
.build()?;
626-
options.apply_headers(&mut cosmos_request.headers);
624+
let options = options.unwrap_or_default();
627625

628-
self.container_connection
629-
.send(cosmos_request, Context::default())
630-
.await
631-
.map(ItemResponse::new)
626+
// Build the driver's item reference from our stored container metadata.
627+
let item_ref = ItemReference::from_name(
628+
&self.container_ref,
629+
partition_key.into().into_driver_partition_key(),
630+
item_id.to_owned(),
631+
);
632+
633+
// Create the driver operation (no body for delete).
634+
let operation = CosmosOperation::delete_item(item_ref);
635+
let operation = apply_item_options(operation, options.session_token, options.precondition);
636+
637+
// Execute through the driver.
638+
let driver_response = self
639+
.context
640+
.driver
641+
.execute_operation(operation, options.operation)
642+
.await?;
643+
644+
// Bridge the driver response to the SDK response type.
645+
Ok(ItemResponse::new(
646+
crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
647+
))
632648
}
633649

634650
/// Executes a single-partition query against items in the container.
@@ -906,6 +922,24 @@ impl ContainerClient {
906922
}
907923
}
908924

925+
/// Applies optional `session_token` and `precondition` to a [`CosmosOperation`].
926+
///
927+
/// Both [`ItemReadOptions`] and [`ItemWriteOptions`] carry these fields;
928+
/// this helper avoids duplicating the wiring logic in every item operation.
929+
fn apply_item_options(
930+
mut operation: CosmosOperation,
931+
session_token: Option<SessionToken>,
932+
precondition: Option<Precondition>,
933+
) -> CosmosOperation {
934+
if let Some(session_token) = session_token {
935+
operation = operation.with_session_token(session_token);
936+
}
937+
if let Some(precondition) = precondition {
938+
operation = operation.with_precondition(precondition);
939+
}
940+
operation
941+
}
942+
909943
#[cfg(test)]
910944
mod tests {
911945
use super::*;

sdk/cosmos/azure_data_cosmos/src/clients/cosmos_client_builder.rs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,7 @@ impl CosmosClientBuilder {
359359

360360
let global_endpoint_manager = GlobalEndpointManager::new(
361361
endpoint.clone(),
362-
preferred_regions,
362+
preferred_regions.clone(),
363363
Vec::new(),
364364
pipeline_core.clone(),
365365
);
@@ -443,8 +443,12 @@ impl CosmosClientBuilder {
443443
})?;
444444
}
445445
let driver_runtime = driver_runtime_builder.build().await?;
446+
let driver_options =
447+
azure_data_cosmos_driver::options::DriverOptions::builder(driver_account)
448+
.with_preferred_regions(preferred_regions)
449+
.build();
446450
let driver = driver_runtime
447-
.get_or_create_driver(driver_account, None)
451+
.get_or_create_driver(driver_options.account().clone(), Some(driver_options))
448452
.await?;
449453

450454
Ok(CosmosClient {

sdk/cosmos/azure_data_cosmos/src/cosmos_request.rs

Lines changed: 1 addition & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -292,10 +292,7 @@ mod tests {
292292
use super::*;
293293
use crate::operation_context::OperationType;
294294
use crate::resource_context::ResourceType;
295-
use crate::{
296-
constants, CosmosClientOptions, ItemWriteOptions, OperationOptions, PartitionKey,
297-
SessionToken,
298-
};
295+
use crate::{constants, PartitionKey};
299296

300297
fn make_base_request(op: OperationType) -> CosmosRequest {
301298
let req = CosmosRequest::builder(op, ResourceLink::root(ResourceType::Documents))
@@ -402,59 +399,6 @@ mod tests {
402399
assert!(has_upsert);
403400
}
404401

405-
#[test]
406-
fn prioritize_request_headers_over_client_headers() {
407-
let mut request_custom_headers = std::collections::HashMap::new();
408-
request_custom_headers.insert(
409-
HeaderName::from_static("x-custom-header"),
410-
HeaderValue::from_static("custom_value"),
411-
);
412-
413-
let operation = OperationOptions::default().with_custom_headers(request_custom_headers);
414-
let item_options = ItemWriteOptions {
415-
operation,
416-
..Default::default()
417-
}
418-
.with_session_token(SessionToken::from("RequestSession"));
419-
let mut req = CosmosRequest::builder(
420-
OperationType::Create,
421-
ResourceLink::root(ResourceType::Documents),
422-
)
423-
.build()
424-
.unwrap();
425-
item_options.apply_headers(&mut req.headers);
426-
427-
req.request_context.location_endpoint_to_route =
428-
Some("https://example.com/".parse().unwrap());
429-
430-
let mut client_custom_headers = std::collections::HashMap::new();
431-
client_custom_headers.insert(
432-
HeaderName::from_static("x-custom-header"),
433-
HeaderValue::from_static("custom_value-2"),
434-
);
435-
436-
let client_operation =
437-
OperationOptions::default().with_custom_headers(client_custom_headers);
438-
let client_options =
439-
CosmosClientOptions::default().with_operation_options(client_operation);
440-
client_options.apply_headers(&mut req.headers);
441-
442-
let raw = req.into_raw_request();
443-
let get_header = |name: &HeaderName| {
444-
raw.headers()
445-
.iter()
446-
.find(|(n, _)| n == &name)
447-
.map(|(_, v)| v.as_str())
448-
.unwrap()
449-
};
450-
451-
assert_eq!(get_header(&constants::SESSION_TOKEN), "RequestSession");
452-
assert_eq!(
453-
get_header(&HeaderName::from_static("x-custom-header")),
454-
"custom_value"
455-
);
456-
}
457-
458402
#[test]
459403
fn to_raw_request_batch_sets_batch_headers() {
460404
let req = make_base_request(OperationType::Batch);

0 commit comments

Comments
 (0)