Cosmos: Cut over query methods to the driver (#4174)

Migrates query_items, query_databases, and query_containers from the
gateway HTTP pipeline to the Cosmos driver pipeline, aligning query
operations with the item CRUD operations that were already cut over.

- QueryExecutor now takes an Arc<CosmosDriver> and an operation factory
closure instead of Arc<GatewayPipeline> and a ResourceLink,
executing queries via driver.execute_operation().
- Query-specific headers (content type, continuation, custom headers)
are passed through DriverOperationOptions instead of being set
directly on the HTTP request.
- QueryOptions::apply_headers is now #[cfg(test)]-only since production
code no longer uses it.
- query_items_cross_partition simplified to delegate to query_items with
PartitionKey::EMPTY.
- Adds emulator tests for single-partition and cross-partition query
pagination.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: simorenoh

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -22,6 +22,7 @@
 ### Other Changes
 
 - Database and container CRUD operations (`create_database`, `read`, `create_container`, `delete`) now route through the Cosmos driver pipeline. Throughput provisioning uses typed request headers via the driver. ([#4147](https://github.com/Azure/azure-sdk-for-rust/pull/4147))
+- Query operations (`query_items`, `query_databases`, `query_containers`) now route through the Cosmos driver pipeline, gaining driver-level transport, routing, and retry capabilities. ([#4174](https://github.com/Azure/azure-sdk-for-rust/pull/4174))
 
 ## 0.32.0 (2026-04-09)
 

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

@@ -21,7 +21,6 @@ use crate::cosmos_request::CosmosRequest;
 use crate::handler::container_connection::ContainerConnection;
 use crate::operation_context::OperationType;
 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::{
     effective_partition_key::EffectivePartitionKey as DriverEpk, ContainerReference,
@@ -698,23 +697,20 @@ impl ContainerClient {
         options: Option<QueryOptions>,
     ) -> azure_core::Result<FeedItemIterator<T>> {
         let options = options.unwrap_or_default();
-        let partition_key = partition_key.into();
+        let partition_key: PartitionKey = partition_key.into();
         let query = query.into();
 
-        let mut headers = azure_core::http::headers::Headers::new();
-
-        // Convert PartitionKey and query options into headers.
-        for (name, value) in partition_key.as_headers()? {
-            headers.insert(name, value);
-        }
-        options.apply_headers(&mut headers);
+        let driver_pk = partition_key.into_driver_partition_key();
+        let container_ref = self.container_ref.clone();
+        let factory =
+            move || CosmosOperation::query_items(container_ref.clone(), driver_pk.clone());
 
         crate::query::executor::QueryExecutor::new(
-            self.context.pipeline.clone(),
-            self.items_link.clone(),
-            Context::default(),
+            self.context.driver.clone(),
+            factory,
             query,
-            headers,
+            options.operation,
+            options.session_token,
         )
         .into_stream()
     }

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

@@ -4,10 +4,9 @@
 use crate::{
     clients::{ClientContext, DatabaseClient},
     models::{DatabaseProperties, ResourceResponse},
-    resource_context::ResourceLink,
     CreateDatabaseOptions, FeedItemIterator, Query, QueryDatabasesOptions,
 };
-use azure_core::http::{Context, Url};
+use azure_core::http::Url;
 use azure_data_cosmos_driver::models::CosmosOperation;
 use azure_data_cosmos_driver::options::OperationOptions;
 use serde::Serialize;
@@ -62,7 +61,6 @@ pub use super::cosmos_client_builder::CosmosClientBuilder;
 /// ```
 #[derive(Debug, Clone)]
 pub struct CosmosClient {
-    pub(crate) databases_link: ResourceLink,
     pub(crate) context: ClientContext,
 }
 
@@ -131,12 +129,15 @@ impl CosmosClient {
         query: impl Into<Query>,
         _options: Option<QueryDatabasesOptions>,
     ) -> azure_core::Result<FeedItemIterator<DatabaseProperties>> {
+        let account = self.context.driver.account().clone();
+        let factory = move || CosmosOperation::query_databases(account.clone());
+
         crate::query::executor::QueryExecutor::new(
-            self.context.pipeline.clone(),
-            self.databases_link.clone(),
-            Context::default(),
+            self.context.driver.clone(),
+            factory,
             query.into(),
-            azure_core::http::headers::Headers::new(),
+            Default::default(),
+            None,
         )
         .into_stream()
     }

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

@@ -7,7 +7,6 @@ use crate::{
     clients::ClientContext,
     options::ThroughputControlGroupOptions,
     pipeline::{AuthorizationPolicy, CosmosHeadersPolicy, GatewayPipeline},
-    resource_context::{ResourceLink, ResourceType},
     CosmosAccountReference, CosmosClient, CosmosClientOptions, CosmosCredential, RoutingStrategy,
 };
 
@@ -449,7 +448,6 @@ impl CosmosClientBuilder {
             .await?;
 
         Ok(CosmosClient {
-            databases_link: ResourceLink::root(ResourceType::Databases),
             context: ClientContext {
                 pipeline,
                 driver,

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

@@ -9,7 +9,6 @@ use crate::{
     CreateContainerOptions, DeleteDatabaseOptions, FeedItemIterator, Query, QueryContainersOptions,
     ThroughputOptions,
 };
-use azure_core::http::Context;
 use azure_data_cosmos_driver::models::{CosmosOperation, DatabaseReference};
 use azure_data_cosmos_driver::options::OperationOptions;
 
@@ -20,7 +19,6 @@ use super::ThroughputPoller;
 /// You can get a `DatabaseClient` by calling [`CosmosClient::database_client()`](crate::CosmosClient::database_client()).
 pub struct DatabaseClient {
     link: ResourceLink,
-    containers_link: ResourceLink,
     database_id: String,
     context: ClientContext,
     database_ref: DatabaseReference,
@@ -30,13 +28,11 @@ impl DatabaseClient {
     pub(crate) fn new(context: ClientContext, database_id: &str) -> Self {
         let database_id = database_id.to_string();
         let link = ResourceLink::root(ResourceType::Databases).item(&database_id);
-        let containers_link = link.feed(ResourceType::Containers);
         let database_ref =
             DatabaseReference::from_name(context.driver.account().clone(), database_id.clone());
 
         Self {
             link,
-            containers_link,
             database_id,
             context,
             database_ref,
@@ -128,12 +124,18 @@ impl DatabaseClient {
         query: impl Into<Query>,
         options: Option<QueryContainersOptions>,
     ) -> azure_core::Result<FeedItemIterator<ContainerProperties>> {
+        let db_ref = DatabaseReference::from_name(
+            self.context.driver.account().clone(),
+            self.database_id.clone(),
+        );
+        let factory = move || CosmosOperation::query_containers(db_ref.clone());
+
         crate::query::executor::QueryExecutor::new(
-            self.context.pipeline.clone(),
-            self.containers_link.clone(),
-            Context::default(),
+            self.context.driver.clone(),
+            factory,
             query.into(),
-            azure_core::http::headers::Headers::new(),
+            Default::default(),
+            None,
         )
         .into_stream()
     }

sdk/cosmos/azure_data_cosmos/src/driver_bridge.rs

@@ -24,8 +24,11 @@ use crate::{
 /// Converts a driver [`DriverResponse`] into the SDK's typed [`CosmosResponse<T>`].
 ///
 /// This reconstructs an `azure_core::Response<T>` from the driver's raw bytes,
-/// status code, and headers, then wraps it in the SDK's response type using
-/// the pre-parsed headers from the driver to avoid a redundant parse.
+/// status code, and headers, then wraps it in the SDK's response type.
+///
+/// The driver's pre-parsed [`CosmosResponseHeaders`] are passed directly to
+/// avoid double-parsing. Some headers (e.g., `index_metrics`) are base64-decoded
+/// by the driver; re-parsing from raw headers would fail on already-decoded values.
 pub(crate) fn driver_response_to_cosmos_response<T>(
     driver_response: DriverResponse,
 ) -> CosmosResponse<T> {
@@ -266,4 +269,60 @@ mod tests {
         assert_eq!(headers.get_optional_str(&SUB_STATUS), None);
         assert_eq!(headers.get_optional_str(&OFFER_REPLACE_PENDING), None);
     }
+
+    /// Regression test: index_metrics (base64-decoded by the driver) must survive
+    /// the driver→SDK bridge without double-decoding.
+    ///
+    /// Exercises `CosmosResponse::from_driver_response` which accepts pre-parsed
+    /// headers, ensuring that already-decoded index_metrics are preserved rather
+    /// than being base64-decoded a second time (which would silently return None).
+    #[test]
+    fn driver_response_preserves_index_metrics() {
+        use crate::feed::{FeedBody, QueryFeedPage};
+        use crate::models::CosmosResponse;
+
+        let mut cosmos_headers = CosmosResponseHeaders::new();
+        cosmos_headers.index_metrics = Some(r#"{"UtilizedSingleIndexes":[]}"#.to_string());
+        cosmos_headers.query_metrics =
+            Some("totalExecutionTimeInMs=1.23;queryCompileTimeInMs=0.01".to_string());
+
+        // Build a minimal raw response with synthesized headers (as the bridge does).
+        let raw_headers = driver_response_headers_to_headers(&cosmos_headers);
+        let raw_response = azure_core::http::RawResponse::from_bytes(
+            StatusCode::Ok,
+            raw_headers,
+            Bytes::from_static(br#"{"Documents":[]}"#),
+        );
+        let typed_response: azure_core::http::response::Response<FeedBody<serde_json::Value>> =
+            raw_response.into();
+
+        // This is the code path used by driver_response_to_cosmos_response:
+        // pre-parsed headers are passed directly, skipping re-parsing.
+        let cosmos_response = CosmosResponse::from_driver_response(typed_response, cosmos_headers);
+
+        assert_eq!(
+            cosmos_response.cosmos_headers().index_metrics.as_deref(),
+            Some(r#"{"UtilizedSingleIndexes":[]}"#),
+            "index_metrics should survive the driver bridge without double base64-decoding"
+        );
+        assert_eq!(
+            cosmos_response.cosmos_headers().query_metrics.as_deref(),
+            Some("totalExecutionTimeInMs=1.23;queryCompileTimeInMs=0.01"),
+        );
+
+        let rt = tokio::runtime::Runtime::new().unwrap();
+        let page = rt
+            .block_on(QueryFeedPage::<serde_json::Value>::from_response(
+                cosmos_response,
+            ))
+            .unwrap();
+        assert_eq!(
+            page.index_metrics(),
+            Some(r#"{"UtilizedSingleIndexes":[]}"#)
+        );
+        assert_eq!(
+            page.query_metrics(),
+            Some("totalExecutionTimeInMs=1.23;queryCompileTimeInMs=0.01")
+        );
+    }
 }

sdk/cosmos/azure_data_cosmos/src/models/cosmos_response.rs

@@ -50,10 +50,12 @@ impl<T> CosmosResponse<T> {
         }
     }
 
-    /// Creates a `CosmosResponse` from a typed response and pre-parsed headers.
+    /// Creates a `CosmosResponse` from a typed response and pre-parsed driver headers.
     ///
-    /// Used by the driver bridge to avoid re-parsing headers that were already
-    /// parsed by the driver pipeline.
+    /// Used by the driver bridge to avoid double-parsing response headers.
+    /// The driver already decodes headers (e.g., base64 for index metrics),
+    /// so re-parsing from raw headers would fail on values that are no longer
+    /// in their wire format.
     pub(crate) fn from_driver_response(
         response: Response<T>,
         cosmos_headers: CosmosResponseHeaders,

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

@@ -359,24 +359,6 @@ impl QueryOptions {
     }
 }
 
-impl QueryOptions {
-    // Temporary: applies option values as HTTP headers for the SDK pipeline.
-    // Will be removed when query 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());
-        }
-    }
-}
-
 /// Options to be passed to [`ContainerClient::read()`](crate::clients::ContainerClient::read()).
 #[derive(Clone, Default)]
 #[non_exhaustive]
@@ -515,36 +497,6 @@ mod tests {
         );
     }
 
-    #[test]
-    fn query_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 query_options = QueryOptions {
-            operation,
-            ..Default::default()
-        }
-        .with_session_token("QuerySessionToken".to_string());
-
-        let mut headers_result = Headers::new();
-        query_options.apply_headers(&mut headers_result);
-
-        let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![
-            ("x-custom-header".into(), "custom_value".into()),
-            (constants::SESSION_TOKEN, "QuerySessionToken".into()),
-        ];
-
-        assert_eq!(
-            headers_to_map(headers_result),
-            headers_to_map(headers_expected)
-        );
-    }
-
     #[test]
     fn item_write_options_default_as_headers() {
         let options = ItemWriteOptions::default();