[Cosmos] Add APIs to perform single-partition queries against a container (#1814)

* add container_client and ContainerClient::read

* rudimentary initial query API

* re-do partition keys to hide serde_json

* switch to a `NullValue` marker shared by query and partition key, and update docs

* add tracing_subscriber with env filter to all cosmos examples

* fix doctests

* remove Send bound from WASM

* Add 'Sync' bound to from_response_body, but exclude bounds from wasm32

* copyright headers, copyright headers everywhere

* change Pageable back to an owning type

* pr feedback

* refactor partition_key parameter to partition_key_strategy

* refmt and move helper back to azure_data_cosmos

* fix docs

* remove unnecessary doctest from internal utility

* fix cosmos_query sample

* simplify doctests that don't run using panic!()

Author: analogrelay

.vscode/cspell.json

@@ -23,6 +23,7 @@
     "asyncoperation",
     "azsdk",
     "azurecli",
+    "Contoso",
     "cplusplus",
     "datalake",
     "datetime",

eng/dict/rust-custom.txt

@@ -3,3 +3,4 @@ newtype
 repr
 rustc
 rustls
+turbofish

sdk/cosmos/.dict.txt

@@ -1,4 +1,5 @@
 colls
+documentdb
 pkranges
 sprocs
 udfs

sdk/cosmos/azure_data_cosmos/Cargo.toml

@@ -22,13 +22,15 @@ typespec_client_core = { workspace = true, features = ["derive"] }
 tracing.workspace = true
 url.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 
 [dev-dependencies]
 tokio.workspace = true
-serde_json.workspace = true
 azure_identity.workspace = true
 clap.workspace = true
 time.workspace = true
+futures.workspace = true
+tracing-subscriber = { workspace = true, features = [ "env-filter", "fmt" ] }
 
 [lints]
 workspace = true

sdk/cosmos/azure_data_cosmos/examples/cosmos_connect.rs renamed to sdk/cosmos/azure_data_cosmos/examples/cosmos_metadata.rs

@@ -1,4 +1,10 @@
-use azure_data_cosmos::{clients::DatabaseClientMethods, CosmosClient, CosmosClientMethods};
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+use azure_data_cosmos::{
+    clients::{ContainerClientMethods, DatabaseClientMethods},
+    CosmosClient, CosmosClientMethods,
+};
 use clap::Parser;
 use std::sync::Arc;
 
@@ -11,6 +17,10 @@ pub struct Args {
     /// The database to fetch information for.
     database: String,
 
+    /// Optionally, the container to fetch information for.
+    #[clap(long, short)]
+    container: Option<String>,
+
     /// An authentication key to use when connecting to the Cosmos DB account. If omitted, the connection will use Entra ID.
     #[clap(long)]
     #[cfg(feature = "key_auth")]
@@ -19,13 +29,28 @@ pub struct Args {
 
 #[tokio::main]
 pub async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let _ = tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .init();
+
     let args = Args::parse();
 
     let client = create_client(&args);
 
     let db_client = client.database_client(&args.database);
-    let response = db_client.read(None).await?.deserialize_body().await?;
-    println!("{:?}", response);
+    if let Some(container_name) = args.container {
+        let container_client = db_client.container_client(container_name);
+        let response = container_client
+            .read(None)
+            .await?
+            .deserialize_body()
+            .await?;
+        println!("{:?}", response);
+        return Ok(());
+    } else {
+        let response = db_client.read(None).await?.deserialize_body().await?;
+        println!("{:?}", response);
+    }
     Ok(())
 }
 

sdk/cosmos/azure_data_cosmos/examples/cosmos_query.rs

@@ -0,0 +1,83 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+use azure_data_cosmos::{
+    clients::{ContainerClientMethods, DatabaseClientMethods},
+    CosmosClient, CosmosClientMethods, PartitionKey,
+};
+use azure_identity::DefaultAzureCredential;
+use clap::Parser;
+use futures::StreamExt;
+use std::sync::Arc;
+
+/// An example to show querying a Cosmos DB container.
+#[derive(Parser)]
+pub struct Args {
+    /// The Cosmos DB endpoint to connect to.
+    endpoint: String,
+
+    /// The database to query.
+    database: String,
+
+    /// The container to query.
+    container: String,
+
+    /// The query to execute.
+    #[clap(long, short)]
+    query: String,
+
+    /// The partition key to use when querying the container. Currently this only supports a single string partition key.
+    #[clap(long, short)]
+    partition_key: String,
+
+    /// An authentication key to use when connecting to the Cosmos DB account. If omitted, the connection will use Entra ID.
+    #[clap(long)]
+    #[cfg(feature = "key_auth")]
+    key: Option<String>,
+}
+
+#[tokio::main]
+pub async fn main() -> Result<(), Box<dyn std::error::Error>> {
+    let _ = tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .init();
+
+    let args = Args::parse();
+
+    let client = create_client(&args);
+
+    let db_client = client.database_client(&args.database);
+    let container_client = db_client.container_client(&args.container);
+
+    let pk = PartitionKey::from(args.partition_key);
+    let mut items_pager =
+        container_client.query_items::<serde_json::Value>(&args.query, pk, None)?;
+
+    while let Some(page) = items_pager.next().await {
+        let response = page?;
+        println!("Results Page");
+        println!("  Query Metrics: {:?}", response.query_metrics);
+        println!("  Index Metrics: {:?}", response.index_metrics);
+        println!("  Items:");
+        for item in response.items {
+            println!("    * {:?}", item);
+        }
+    }
+    Ok(())
+}
+
+#[cfg(feature = "key_auth")]
+fn create_client(args: &Args) -> CosmosClient {
+    if let Some(key) = args.key.as_ref() {
+        CosmosClient::with_key(&args.endpoint, key.clone(), None).unwrap()
+    } else {
+        let cred = DefaultAzureCredential::new().map(Arc::new).unwrap();
+        CosmosClient::new(&args.endpoint, cred, None).unwrap()
+    }
+}
+
+#[cfg(not(feature = "key_auth"))]
+fn create_client(args: &Args) -> CosmosClient {
+    let cred = DefaultAzureCredential::new().map(Arc::new).unwrap();
+    CosmosClient::new(&args.endpoint, cred, None).unwrap()
+}