Cosmos: Cutover transactional batch operation to Driver

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

@@ -37,7 +37,6 @@ use serde::{de::DeserializeOwned, Serialize};
 #[derive(Clone)]
 pub struct ContainerClient {
     link: ResourceLink,
-    items_link: ResourceLink,
     container_connection: Arc<ContainerConnection>,
     container_ref: ContainerReference,
     context: ClientContext,
@@ -53,7 +52,6 @@ impl ContainerClient {
         let link = database_link
             .feed(ResourceType::Containers)
             .item(container_id);
-        let items_link = link.feed(ResourceType::Documents);
 
         // Eagerly resolve immutable container metadata from the driver.
         let container_ref = context
@@ -80,7 +78,6 @@ impl ContainerClient {
 
         Ok(Self {
             link,
-            items_link,
             container_connection,
             container_ref,
             context,
@@ -792,19 +789,22 @@ impl ContainerClient {
         options: Option<BatchOptions>,
     ) -> azure_core::Result<BatchResponse> {
         let options = options.unwrap_or_default();
-        let partition_key = batch.partition_key().clone();
+        let body = serde_json::to_vec(batch.operations())?;
+        let driver_pk = batch.partition_key().clone().into_driver_partition_key();
 
-        let mut cosmos_request =
-            CosmosRequest::builder(OperationType::Batch, self.items_link.clone())
-                .partition_key(partition_key)
-                .json(batch.operations())
-                .build()?;
-        options.apply_headers(&mut cosmos_request.headers);
+        let operation =
+            CosmosOperation::batch(self.container_ref.clone(), driver_pk).with_body(body);
+        let operation = apply_batch_options(operation, &options);
 
-        self.container_connection
-            .send(cosmos_request, Context::default())
-            .await
-            .map(BatchResponse::new)
+        let driver_response = self
+            .context
+            .driver
+            .execute_operation(operation, options.operation)
+            .await?;
+
+        Ok(BatchResponse::new(
+            crate::driver_bridge::driver_response_to_cosmos_response(driver_response),
+        ))
     }
 
     /// Gets the feed ranges for this container.
@@ -1001,6 +1001,17 @@ fn apply_item_options(
     operation
 }
 
+/// Applies [`BatchOptions`] fields to a [`CosmosOperation`].
+///
+/// [`BatchOptions`] carries a session token but no precondition (ETag-based
+/// conditions are specified per-operation within the batch itself).
+fn apply_batch_options(mut operation: CosmosOperation, options: &BatchOptions) -> CosmosOperation {
+    if let Some(session_token) = &options.session_token {
+        operation = operation.with_session_token(session_token.clone());
+    }
+    operation
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

sdk/cosmos/azure_data_cosmos/src/constants.rs

@@ -6,11 +6,7 @@
 
 //! Constants defining HTTP headers and other values relevant to Azure Cosmos DB APIs.
 
-use azure_core::http::{
-    headers::{HeaderName, HeaderValue},
-    request::options::ContentType,
-    StatusCode,
-};
+use azure_core::http::{headers::HeaderName, request::options::ContentType, StatusCode};
 
 /// Macro to define Cosmos DB header constants and the allowed headers list in one place.
 macro_rules! cosmos_headers {
@@ -201,8 +197,6 @@ cosmos_headers! {
 
 pub const QUERY_CONTENT_TYPE: ContentType = ContentType::from_static("application/query+json");
 
-pub(crate) const PREFER_MINIMAL: HeaderValue = HeaderValue::from_static("return=minimal");
-
 pub const ACCOUNT_PROPERTIES_KEY: &str = "account_properties_key";
 
 /// The Cosmos DB-specific 449 Retry With status code.

sdk/cosmos/azure_data_cosmos/src/options/mod.rs

@@ -1,9 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-use crate::constants;
 use crate::models::ThroughputProperties;
-use azure_core::http::headers::{self, Headers};
+use azure_core::http::headers::Headers;
 use std::fmt;
 use std::fmt::Display;
 
@@ -19,20 +18,6 @@ pub use azure_data_cosmos_driver::options::{
     ThroughputControlGroupOptions,
 };
 
-// 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(
-    content_response_on_write: Option<&ContentResponseOnWrite>,
-    headers: &mut Headers,
-) {
-    match content_response_on_write {
-        Some(ContentResponseOnWrite::Enabled) => {}
-        _ => {
-            headers.insert(headers::PREFER, constants::PREFER_MINIMAL);
-        }
-    }
-}
-
 /// Options used when creating a [`CosmosClient`](crate::CosmosClient).
 ///
 /// This struct is used internally by [`CosmosClientBuilder`](crate::CosmosClientBuilder).
@@ -256,28 +241,6 @@ impl BatchOptions {
     }
 }
 
-impl BatchOptions {
-    // Temporary: applies option values as HTTP headers for the SDK pipeline.
-    // Will be removed when batch 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());
-        }
-        apply_content_response_on_write_header(
-            self.operation.content_response_on_write.as_ref(),
-            headers,
-        );
-    }
-}
-
 /// Options to be passed to [`DatabaseClient::query_containers()`](crate::clients::DatabaseClient::query_containers()).
 #[derive(Clone, Default)]
 #[non_exhaustive]
@@ -393,98 +356,4 @@ mod tests {
             headers_to_map(headers_expected)
         );
     }
-
-    #[test]
-    fn batch_options_as_headers() {
-        let mut custom_headers = HashMap::new();
-        custom_headers.insert(
-            HeaderName::from_static("x-custom-header"),
-            HeaderValue::from_static("custom_value"),
-        );
-
-        let mut operation = OperationOptions::default().with_custom_headers(custom_headers);
-        operation.content_response_on_write = Some(ContentResponseOnWrite::Enabled);
-
-        let batch_options = BatchOptions {
-            operation,
-            ..Default::default()
-        }
-        .with_session_token("BatchSessionToken".to_string());
-
-        let mut headers_result = Headers::new();
-        batch_options.apply_headers(&mut headers_result);
-
-        let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![
-            ("x-custom-header".into(), "custom_value".into()),
-            (constants::SESSION_TOKEN, "BatchSessionToken".into()),
-        ];
-
-        assert_eq!(
-            headers_to_map(headers_result),
-            headers_to_map(headers_expected)
-        );
-    }
-
-    #[test]
-    fn batch_options_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 batch_options = BatchOptions {
-            operation,
-            ..Default::default()
-        }
-        .with_session_token("RealSessionToken".to_string());
-
-        let mut headers_result = Headers::new();
-        batch_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 batch_options_default_as_headers() {
-        let batch_options = BatchOptions::default();
-
-        let mut headers_result = Headers::new();
-        batch_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 batch_options_with_content_response_enabled() {
-        let mut operation = OperationOptions::default();
-        operation.content_response_on_write = Some(ContentResponseOnWrite::Enabled);
-
-        let batch_options = BatchOptions {
-            operation,
-            ..Default::default()
-        };
-
-        let mut headers_result = Headers::new();
-        batch_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);
-    }
 }

sdk/cosmos/azure_data_cosmos_driver/src/driver/pipeline/operation_pipeline.rs

@@ -502,6 +502,23 @@ fn build_transport_request(
         );
     }
 
+    // Cosmos DB uses POST for batch (same endpoint as create/upsert);
+    // the service requires these headers to process the request as a batch.
+    if operation.operation_type() == OperationType::Batch {
+        headers.insert(
+            HeaderName::from_static(request_header_names::IS_BATCH_REQUEST),
+            HeaderValue::from_static("True"),
+        );
+        headers.insert(
+            HeaderName::from_static(request_header_names::BATCH_ATOMIC),
+            HeaderValue::from_static("True"),
+        );
+        headers.insert(
+            HeaderName::from_static(request_header_names::BATCH_CONTINUE_ON_ERROR),
+            HeaderValue::from_static("False"),
+        );
+    }
+
     // Add operation type header for fault injection rule matching
     #[cfg(feature = "fault_injection")]
     {
@@ -1313,6 +1330,79 @@ mod tests {
         );
     }
 
+    #[test]
+    fn build_transport_request_sets_batch_headers() {
+        let operation = CosmosOperation::batch(test_container(), PartitionKey::from("pk1"))
+            .with_body(b"[]".to_vec());
+
+        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_eq!(
+            request
+                .headers
+                .get_optional_str(&HeaderName::from_static("x-ms-cosmos-is-batch-request")),
+            Some("True"),
+            "is-batch-request header should be set"
+        );
+        assert_eq!(
+            request
+                .headers
+                .get_optional_str(&HeaderName::from_static("x-ms-cosmos-batch-atomic")),
+            Some("True"),
+            "batch-atomic header should be set"
+        );
+        assert_eq!(
+            request.headers.get_optional_str(&HeaderName::from_static(
+                "x-ms-cosmos-batch-continue-on-error"
+            )),
+            Some("False"),
+            "batch-continue-on-error header should be set"
+        );
+    }
+
+    #[test]
+    fn build_transport_request_omits_batch_headers_for_create() {
+        let container = test_container();
+        let operation = CosmosOperation::create_item(ItemReference::from_name(
+            &container,
+            PartitionKey::from("pk1"),
+            "doc1",
+        ))
+        .with_body(b"{}".to_vec());
+
+        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
+                .headers
+                .get_optional_str(&HeaderName::from_static("x-ms-cosmos-is-batch-request"))
+                .is_none(),
+            "batch headers should not be set for create"
+        );
+    }
+
     #[test]
     fn build_transport_request_sets_priority_level_header() {
         let container = test_container();

sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_headers.rs

@@ -21,6 +21,9 @@ pub(crate) mod request_header_names {
     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 IS_BATCH_REQUEST: &str = "x-ms-cosmos-is-batch-request";
+    pub const BATCH_ATOMIC: &str = "x-ms-cosmos-batch-atomic";
+    pub const BATCH_CONTINUE_ON_ERROR: &str = "x-ms-cosmos-batch-continue-on-error";
     pub const OFFER_THROUGHPUT: &str = "x-ms-offer-throughput";
     pub const OFFER_AUTOPILOT_SETTINGS: &str = "x-ms-cosmos-offer-autopilot-settings";
     pub const PRIORITY_LEVEL: &str = "x-ms-cosmos-priority-level";

sdk/cosmos/azure_data_cosmos_driver/src/models/cosmos_operation.rs

@@ -481,6 +481,18 @@ impl CosmosOperation {
         Self::new(OperationType::Delete, item).with_partition_key(partition_key)
     }
 
+    /// Executes a transactional batch of operations against a single partition.
+    ///
+    /// All operations in the batch target the same `partition_key` and are
+    /// committed atomically. Use `with_body()` to provide the JSON-encoded
+    /// array of batch operations.
+    pub fn batch(container: ContainerReference, partition_key: PartitionKey) -> Self {
+        let resource_ref: CosmosResourceReference = CosmosResourceReference::from(container)
+            .with_resource_type(ResourceType::Document)
+            .into_feed_reference();
+        Self::new(OperationType::Batch, resource_ref).with_partition_key(partition_key)
+    }
+
     /// Upserts (creates or replaces) an item (document) in a container.
     ///
     /// The `ItemReference` contains the container, partition key, and item identifier,