Cosmos: Cut over replace_item, upsert_item, delete_item to driver

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,
@@ -322,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
@@ -417,20 +413,31 @@ impl ContainerClient {
         item: T,
         options: Option<ItemWriteOptions>,
     ) -> azure_core::Result<ItemResponse<()>> {
-        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 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
+            .context
+            .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 +513,26 @@ impl ContainerClient {
         item: T,
         options: Option<ItemWriteOptions>,
     ) -> azure_core::Result<ItemResponse<()>> {
-        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 operation =
+            CosmosOperation::upsert_item(self.container_ref.clone(), driver_pk).with_body(body);
+        let operation = apply_item_options(operation, options.session_token, options.precondition);
+
+        // Execute through the driver.
+        let driver_response = self
+            .context
+            .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.
@@ -567,15 +579,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
@@ -616,19 +621,30 @@ impl ContainerClient {
         item_id: &str,
         options: Option<ItemWriteOptions>,
     ) -> azure_core::Result<ItemResponse<()>> {
-        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 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
+            .context
+            .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.
@@ -906,6 +922,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<SessionToken>,
+    precondition: Option<Precondition>,
+) -> 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::*;

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

@@ -359,7 +359,7 @@ impl CosmosClientBuilder {
 
         let global_endpoint_manager = GlobalEndpointManager::new(
             endpoint.clone(),
-            preferred_regions,
+            preferred_regions.clone(),
             Vec::new(),
             pipeline_core.clone(),
         );
@@ -443,8 +443,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 {

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);