[Cosmos] Cross-Partition queries, with optional external Query Engine integration (#2577)

* begin work on query engine

* update non-query-engine tests

* implement query-via-query-engine

* add 'deque' to rust cspell dictionary

* remove raw_value and tracing features from dependencies

* update assets.json for cosmos

* address clippy lint

* rename Context::into_borrowed to Context::to_borrowed

* a few docs updates

* pr feedback

* re-enable tracing feature on azure_data_cosmos tests

Author: analogrelay

eng/dict/rust-custom.txt

@@ -2,6 +2,7 @@ bindgen
 cfg
 cfgs
 consts
+deque
 impls
 newtype
 oneshot

sdk/core/azure_core/src/http/pager.rs

@@ -120,6 +120,18 @@ impl<T> PageStream<T> {
             stream: Box::pin(stream),
         }
     }
+
+    /// Creates a [`Pager<T>`] from a raw stream of [`Result<T>`](typespec::Result<T>) values.
+    ///
+    /// This constructor is used when you are implementing a completely custom stream and want to use it as a pager.
+    pub fn from_stream<S>(stream: S) -> Self
+    where
+        S: Stream<Item = Result<T, Error>> + Send + 'static,
+    {
+        Self {
+            stream: Box::pin(stream),
+        }
+    }
 }
 
 impl<T> futures::Stream for PageStream<T> {

sdk/cosmos/assets.json

@@ -1,6 +1,6 @@
 {
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "rust",
-  "Tag": "rust/azure_data_cosmos_43a7f435b6",
+  "Tag": "rust/azure_data_cosmos_ff23846344",
   "TagPrefix": "rust/azure_data_cosmos"
 }

sdk/cosmos/azure_data_cosmos/Cargo.toml

@@ -25,7 +25,7 @@ url.workspace = true
 
 [dev-dependencies]
 azure_identity.workspace = true
-azure_core_test.workspace = true
+azure_core_test = { workspace = true, features = ["tracing"] }
 clap.workspace = true
 reqwest.workspace = true
 time.workspace = true
@@ -38,6 +38,7 @@ workspace = true
 [features]
 default = ["hmac_rust"]
 key_auth = [] # Enables support for key-based authentication (Primary Keys and Resource Tokens)
+preview_query_engine = [] # Enables support for the PREVIEW external query engine
 hmac_rust = ["azure_core/hmac_rust"]
 hmac_openssl = ["azure_core/hmac_openssl"]
 

sdk/cosmos/azure_data_cosmos/examples/cosmos/query.rs

@@ -25,7 +25,7 @@ enum Subcommands {
 
         /// The partition key to use when querying the container. Currently this only supports a single string partition key.
         #[arg(long, short)]
-        partition_key: String,
+        partition_key: Option<String>,
     },
     Databases {
         /// The query to execute.
@@ -52,7 +52,11 @@ impl QueryCommand {
                 let db_client = client.database_client(&database);
                 let container_client = db_client.container_client(&container);
 
-                let pk = PartitionKey::from(&partition_key);
+                let pk = match partition_key {
+                    Some(pk) => PartitionKey::from(pk),
+                    None => PartitionKey::EMPTY,
+                };
+
                 let mut items =
                     container_client.query_items::<serde_json::Value>(&query, pk, None)?;
 

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

@@ -7,16 +7,22 @@ use crate::{
     options::{QueryOptions, ReadContainerOptions},
     pipeline::CosmosPipeline,
     resource_context::{ResourceLink, ResourceType},
-    DeleteContainerOptions, FeedPager, ItemOptions, PartitionKey, Query, QueryPartitionStrategy,
-    ReplaceContainerOptions, ThroughputOptions,
+    DeleteContainerOptions, FeedPager, ItemOptions, PartitionKey, Query, ReplaceContainerOptions,
+    ThroughputOptions,
 };
 
-use azure_core::http::{headers, request::Request, response::Response, Method};
+use azure_core::http::{
+    headers,
+    request::{options::ContentType, Request},
+    response::Response,
+    Method,
+};
 use serde::{de::DeserializeOwned, Serialize};
 
 /// A client for working with a specific container in a Cosmos DB account.
 ///
 /// You can get a `Container` by calling [`DatabaseClient::container_client()`](crate::clients::DatabaseClient::container_client()).
+#[derive(Clone)]
 pub struct ContainerClient {
     link: ResourceLink,
     items_link: ResourceLink,
@@ -109,6 +115,7 @@ impl ContainerClient {
         let options = options.unwrap_or_default();
         let url = self.pipeline.url(&self.link);
         let mut req = Request::new(url, Method::Put);
+        req.insert_headers(&ContentType::APPLICATION_JSON)?;
         req.set_json(&properties)?;
         self.pipeline
             .send(options.method_options.context, &mut req, self.link.clone())
@@ -260,6 +267,7 @@ impl ContainerClient {
             req.insert_header(headers::PREFER, constants::PREFER_MINIMAL);
         }
         req.insert_headers(&partition_key.into())?;
+        req.insert_headers(&ContentType::APPLICATION_JSON)?;
         req.set_json(&item)?;
         self.pipeline
             .send(
@@ -351,6 +359,7 @@ impl ContainerClient {
             req.insert_header(headers::PREFER, constants::PREFER_MINIMAL);
         }
         req.insert_headers(&partition_key.into())?;
+        req.insert_headers(&ContentType::APPLICATION_JSON)?;
         req.set_json(&item)?;
         self.pipeline
             .send(options.method_options.context, &mut req, link)
@@ -440,6 +449,7 @@ impl ContainerClient {
         }
         req.insert_header(constants::IS_UPSERT, "true");
         req.insert_headers(&partition_key.into())?;
+        req.insert_headers(&ContentType::APPLICATION_JSON)?;
         req.set_json(&item)?;
         self.pipeline
             .send(
@@ -453,7 +463,7 @@ impl ContainerClient {
     /// Reads a specific item from the container.
     ///
     /// # Arguments
-    /// * `partition_key` - The partition key of the item to read.
+    /// * `partition_key` - The partition key of the item to read. See [`PartitionKey`] for more information on how to specify a partition key.
     /// * `item_id` - The id of the item to read.
     /// * `options` - Optional parameters for the request
     ///
@@ -606,6 +616,7 @@ impl ContainerClient {
             req.insert_header(headers::PREFER, constants::PREFER_MINIMAL);
         }
         req.insert_headers(&partition_key.into())?;
+        req.insert_headers(&ContentType::APPLICATION_JSON)?;
         req.set_json(&patch)?;
 
         self.pipeline
@@ -625,12 +636,18 @@ impl ContainerClient {
     /// # Arguments
     ///
     /// * `query` - The query to execute.
-    /// * `partition_key_strategy` - The partition key to scope the query on.
+    /// * `partition_key` - The partition key to scope the query on, or specify an empty key (`()`) to perform a cross-partition query.
     /// * `options` - Optional parameters for the request.
     ///
+    /// # Cross Partition Queries
+    ///
+    /// Cross-partition queries are significantly limited in the current version of the Cosmos DB SDK.
+    /// They are run on the gateway and limited to simple projections (`SELECT`) and filtering (`WHERE`).
+    /// For more details, see [the Cosmos DB documentation page on cross-partition queries](https://learn.microsoft.com/en-us/rest/api/cosmos-db/querying-cosmosdb-resources-using-the-rest-api#queries-that-cannot-be-served-by-gateway).
+    ///
     /// # Examples
     ///
-    /// The `query` and `partition_key_strategy` parameters accept anything that can be transformed [`Into`] their relevant types.
+    /// The `query` and `partition_key` parameters accept anything that can be transformed [`Into`] their relevant types.
     /// This allows simple queries without parameters to be expressed easily:
     ///
     /// ```rust,no_run
@@ -666,23 +683,38 @@ impl ContainerClient {
     /// ```
     ///
     /// See [`PartitionKey`](crate::PartitionKey) for more information on how to specify a partition key, and [`Query`] for more information on how to specify a query.
-    pub fn query_items<T: DeserializeOwned + Send>(
+    pub fn query_items<T: DeserializeOwned + 'static>(
         &self,
         query: impl Into<Query>,
-        partition_key: impl Into<QueryPartitionStrategy>,
+        partition_key: impl Into<PartitionKey>,
         options: Option<QueryOptions<'_>>,
     ) -> azure_core::Result<FeedPager<T>> {
-        let options = options.unwrap_or_default();
-        let url = self.pipeline.url(&self.items_link);
-        let mut base_request = Request::new(url, Method::Post);
-        let QueryPartitionStrategy::SinglePartition(partition_key) = partition_key.into();
-        base_request.insert_headers(&partition_key)?;
+        #[cfg_attr(not(feature = "preview_query_engine"), allow(unused_mut))]
+        let mut options = options.unwrap_or_default();
+        let partition_key = partition_key.into();
+        let query = query.into();
+
+        #[cfg(feature = "preview_query_engine")]
+        if partition_key.is_empty() {
+            if let Some(query_engine) = options.query_engine.take() {
+                return crate::query::executor::QueryExecutor::new(
+                    self.pipeline.clone(),
+                    self.link.clone(),
+                    query,
+                    options,
+                    query_engine,
+                )?
+                .into_stream();
+            }
+        }
 
+        let url = self.pipeline.url(&self.items_link);
         self.pipeline.send_query_request(
             options.method_options.context,
-            query.into(),
-            base_request,
+            query,
+            url,
             self.items_link.clone(),
+            |r| r.insert_headers(&partition_key),
         )
     }
 }

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

@@ -10,7 +10,11 @@ use crate::{
 };
 use azure_core::{
     credentials::TokenCredential,
-    http::{request::Request, response::Response, Method, Url},
+    http::{
+        request::{options::ContentType, Request},
+        response::Response,
+        Method, Url,
+    },
 };
 use serde::Serialize;
 use std::sync::Arc;
@@ -135,13 +139,13 @@ impl CosmosClient {
     ) -> azure_core::Result<FeedPager<DatabaseProperties>> {
         let options = options.unwrap_or_default();
         let url = self.pipeline.url(&self.databases_link);
-        let base_request = Request::new(url, Method::Post);
 
         self.pipeline.send_query_request(
             options.method_options.context,
             query.into(),
-            base_request,
+            url,
             self.databases_link.clone(),
+            |_| Ok(()),
         )
     }
 
@@ -167,6 +171,7 @@ impl CosmosClient {
         let url = self.pipeline.url(&self.databases_link);
         let mut req = Request::new(url, Method::Post);
         req.insert_headers(&options.throughput)?;
+        req.insert_headers(&ContentType::APPLICATION_JSON)?;
         req.set_json(&RequestBody { id })?;
 
         self.pipeline

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

@@ -11,7 +11,11 @@ use crate::{
     ThroughputOptions,
 };
 
-use azure_core::http::{request::Request, response::Response, Method};
+use azure_core::http::{
+    request::{options::ContentType, Request},
+    response::Response,
+    Method,
+};
 
 /// A client for working with a specific database in a Cosmos DB account.
 ///
@@ -110,13 +114,13 @@ impl DatabaseClient {
     ) -> azure_core::Result<FeedPager<ContainerProperties>> {
         let options = options.unwrap_or_default();
         let url = self.pipeline.url(&self.containers_link);
-        let base_request = Request::new(url, Method::Post);
 
         self.pipeline.send_query_request(
             options.method_options.context,
             query.into(),
-            base_request,
+            url,
             self.containers_link.clone(),
+            |_| Ok(()),
         )
     }
 
@@ -136,6 +140,7 @@ impl DatabaseClient {
         let url = self.pipeline.url(&self.containers_link);
         let mut req = Request::new(url, Method::Post);
         req.insert_headers(&options.throughput)?;
+        req.insert_headers(&ContentType::APPLICATION_JSON)?;
         req.set_json(&properties)?;
 
         self.pipeline

sdk/cosmos/azure_data_cosmos/src/constants.rs

@@ -13,6 +13,14 @@ use azure_core::http::{
 
 pub const QUERY: HeaderName = HeaderName::from_static("x-ms-documentdb-query");
 pub const PARTITION_KEY: HeaderName = HeaderName::from_static("x-ms-documentdb-partitionkey");
+pub const PARTITION_KEY_RANGE_ID: HeaderName =
+    HeaderName::from_static("x-ms-documentdb-partitionkeyrangeid");
+pub const QUERY_ENABLE_CROSS_PARTITION: HeaderName =
+    HeaderName::from_static("x-ms-documentdb-query-enablecrosspartition");
+pub const IS_QUERY_PLAN_REQUEST: HeaderName =
+    HeaderName::from_static("x-ms-cosmos-is-query-plan-request");
+pub const SUPPORTED_QUERY_FEATURES: HeaderName =
+    HeaderName::from_static("x-ms-cosmos-supported-query-features");
 pub const CONTINUATION: HeaderName = HeaderName::from_static("x-ms-continuation");
 pub const INDEX_METRICS: HeaderName = HeaderName::from_static("x-ms-cosmos-index-utilization");
 pub const QUERY_METRICS: HeaderName = HeaderName::from_static("x-ms-documentdb-query-metrics");

sdk/cosmos/azure_data_cosmos/src/feed.rs

@@ -11,6 +11,7 @@ use crate::constants;
 /// Cosmos DB queries can be executed using non-HTTP transports, depending on the circumstances.
 /// They may also produce results that don't directly correlate to specific HTTP responses (as in the case of cross-partition queries).
 /// Because of this, Cosmos DB query responses use `FeedPage` to represent the results, rather than a more generic type like [`Response`](azure_core::http::Response).
+#[derive(Debug)]
 pub struct FeedPage<T> {
     /// The items in the response.
     items: Vec<T>,
@@ -24,6 +25,16 @@ pub struct FeedPage<T> {
 }
 
 impl<T> FeedPage<T> {
+    /// Creates a new `FeedPage` instance.
+    #[cfg_attr(not(feature = "preview_query_engine"), allow(dead_code))]
+    pub(crate) fn new(items: Vec<T>, continuation: Option<String>, headers: Headers) -> Self {
+        Self {
+            items,
+            continuation,
+            headers,
+        }
+    }
+
     /// Gets the items in this page of results.
     pub fn items(&self) -> &[T] {
         &self.items

sdk/cosmos/azure_data_cosmos/src/lib.rs

@@ -16,7 +16,7 @@ mod feed;
 mod options;
 mod partition_key;
 pub(crate) mod pipeline;
-mod query;
+pub mod query;
 pub(crate) mod resource_context;
 pub(crate) mod utils;
 
@@ -27,6 +27,6 @@ pub use clients::CosmosClient;
 
 pub use options::*;
 pub use partition_key::*;
-pub use query::*;
+pub use query::Query;
 
 pub use feed::{FeedPage, FeedPager};

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

@@ -70,6 +70,25 @@ pub struct QueryDatabasesOptions<'a> {
 #[derive(Clone, Default)]
 pub struct QueryOptions<'a> {
     pub method_options: ClientMethodOptions<'a>,
+
+    /// An external query engine to use for executing the query.
+    ///
+    /// NOTE: This is an unstable feature and may change in the future.
+    /// Specifically, the query engine may be built-in to the SDK in the future, and this option may be removed entirely.
+    #[cfg(feature = "preview_query_engine")]
+    pub query_engine: Option<crate::query::QueryEngineRef>,
+}
+
+impl QueryOptions<'_> {
+    pub fn into_owned(self) -> QueryOptions<'static> {
+        QueryOptions {
+            method_options: ClientMethodOptions {
+                context: self.method_options.context.into_owned(),
+            },
+            #[cfg(feature = "preview_query_engine")]
+            query_engine: self.query_engine,
+        }
+    }
 }
 
 /// Options to be passed to [`ContainerClient::read()`](crate::clients::ContainerClient::read()).