[Cosmos] "Enable Content Response on Write" option and integration tests for items and query (#1963)

* format-specific deserializers for Response<T>

* add enable_content_response_on_write option

* stash not content on write

* remove changes outside cosmos for now

* don't try to start test proxy on live tests

* [Cosmos] add integration tests for items

* [Cosmos] add integration tests for query

* spelling fix

* show 'use' statements in examples

* switch clap attributes to arg/command

Author: analogrelay

sdk/core/azure_core_test/src/recorded.rs

@@ -19,7 +19,14 @@ static TEST_PROXY: OnceCell<Result<Arc<Proxy>>> = OnceCell::const_new();
 ///
 /// The [Test Proxy](https://github.com/Azure/azure-sdk-tools/blob/main/tools/test-proxy/Azure.Sdk.Tools.TestProxy/README.md) service will be started as needed.
 /// Every `#[recorded::test]` will call this automatically, but it can also be called manually by any other test e.g., those attributed with `#[tokio::test]`.
-pub async fn start(ctx: &TestContext, options: Option<ProxyOptions>) -> Result<Recording> {
+///
+/// This function will return `Ok(None)` if the test is running in live mode and should not use the test proxy at all.
+pub async fn start(ctx: &TestContext, options: Option<ProxyOptions>) -> Result<Option<Recording>> {
+    // Live tests don't use test-proxy.
+    if ctx.test_mode() == azure_core::test::TestMode::Live {
+        return Ok(None);
+    }
+
     let proxy = TEST_PROXY
         .get_or_init(|| async move {
             proxy::start(ctx.test_data_dir(), options)
@@ -31,9 +38,9 @@ pub async fn start(ctx: &TestContext, options: Option<ProxyOptions>) -> Result<R
         .map_err(|err| azure_core::Error::new(err.kind().clone(), err))?;
 
     let span = debug_span!(target: crate::SPAN_TARGET, "recording", mode = ?ctx.test_mode(), test = ?ctx.test_name());
-    Ok(Recording {
+    Ok(Some(Recording {
         proxy: proxy.clone(),
         span,
         mode: ctx.test_mode(),
-    })
+    }))
 }

sdk/cosmos/azure_data_cosmos/examples/cosmos/create.rs

@@ -2,7 +2,7 @@ use std::error::Error;
 
 use azure_data_cosmos::{
     models::{ContainerProperties, PartitionKeyDefinition, ThroughputProperties},
-    CosmosClient, CreateContainerOptions, CreateDatabaseOptions, PartitionKey,
+    CosmosClient, CreateContainerOptions, CreateDatabaseOptions, ItemOptions, PartitionKey,
 };
 use clap::{Args, Subcommand};
 
@@ -26,20 +26,24 @@ pub enum Subcommands {
         container: String,
 
         /// The partition key of the new item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         partition_key: String,
 
         /// The JSON of the new item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         json: String,
+
+        /// If set, the updated item will be included in the response.
+        #[arg(long)]
+        show_updated: bool,
     },
 
     /// Create a database (does not support Entra ID).
     Database {
         /// The ID of the new database to create.
         id: String,
 
-        #[clap(flatten)]
+        #[command(flatten)]
         throughput_options: ThroughputOptions,
     },
 
@@ -48,19 +52,19 @@ pub enum Subcommands {
         /// The ID of the database to create the container in.
         database: String,
 
-        #[clap(flatten)]
+        #[command(flatten)]
         throughput_options: ThroughputOptions,
 
         /// The ID of the new container to create.
-        #[clap(long, short)]
+        #[arg(long, short)]
         id: Option<String>,
 
         /// The path to the partition key properties (supports up to 3).
-        #[clap(long, short)]
+        #[arg(long, short)]
         partition_key: Vec<String>,
 
         /// The JSON for a ContainerProperties value. The 'id' and 'partition key' options are ignored if this is set.
-        #[clap(long)]
+        #[arg(long)]
         json: Option<String>,
     },
 }
@@ -73,21 +77,30 @@ impl CreateCommand {
                 container,
                 partition_key,
                 json,
+                show_updated,
             } => {
                 let db_client = client.database_client(&database);
                 let container_client = db_client.container_client(&container);
 
                 let pk = PartitionKey::from(&partition_key);
                 let item: serde_json::Value = serde_json::from_str(&json)?;
 
-                let created = container_client
-                    .create_item(pk, item, None)
-                    .await?
-                    .into_body()
-                    .await?
-                    .unwrap();
-                println!("Created item:");
-                println!("{:#?}", created);
+                let options = ItemOptions {
+                    enable_content_response_on_write: show_updated,
+                    ..Default::default()
+                };
+
+                let response = container_client
+                    .create_item(pk, item, Some(options))
+                    .await?;
+
+                println!("Created item successfully");
+
+                if show_updated {
+                    let created = response.into_json_body::<serde_json::Value>().await?;
+                    println!("Newly created item:");
+                    println!("{:#?}", created);
+                }
                 Ok(())
             }
 
@@ -106,8 +119,7 @@ impl CreateCommand {
                     .create_database(&id, options)
                     .await?
                     .into_body()
-                    .await?
-                    .unwrap();
+                    .await?;
                 println!("Created database:");
                 println!("{:#?}", db);
                 Ok(())
@@ -152,8 +164,7 @@ impl CreateCommand {
                     .create_container(properties, options)
                     .await?
                     .into_body()
-                    .await?
-                    .unwrap();
+                    .await?;
                 println!("Created container:");
                 println!("{:#?}", container);
                 Ok(())

sdk/cosmos/azure_data_cosmos/examples/cosmos/delete.rs

@@ -22,11 +22,11 @@ pub enum Subcommands {
         container: String,
 
         /// The ID of the item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         item_id: String,
 
         /// The partition key of the item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         partition_key: String,
     },
 

sdk/cosmos/azure_data_cosmos/examples/cosmos/main.rs

@@ -22,7 +22,7 @@ mod utils;
 /// It exists for illustrative purposes and to simplify ad-hoc end-to-end testing.
 #[derive(Clone, Parser)]
 struct ProgramArgs {
-    #[clap(flatten)]
+    #[command(flatten)]
     shared_args: SharedArgs,
 
     #[command(subcommand)]
@@ -35,7 +35,7 @@ struct SharedArgs {
     endpoint: String,
 
     /// An authentication key to use when connecting to the Cosmos DB account. If omitted, the connection will use Entra ID.
-    #[clap(long)]
+    #[arg(long)]
     key: Option<String>,
 }
 

sdk/cosmos/azure_data_cosmos/examples/cosmos/metadata.rs

@@ -13,7 +13,7 @@ pub struct MetadataCommand {
     database: String,
 
     /// Optionally, the container to fetch information for.
-    #[clap(long, short)]
+    #[arg(long, short)]
     container: Option<String>,
 }
 

sdk/cosmos/azure_data_cosmos/examples/cosmos/patch.rs

@@ -17,15 +17,15 @@ pub struct PatchCommand {
     container: String,
 
     /// The ID of the item.
-    #[clap(long, short)]
+    #[arg(long, short)]
     item_id: String,
 
     /// The partition key of the new item.
-    #[clap(long, short)]
+    #[arg(long, short)]
     partition_key: String,
 
     /// A JSON patch operation to apply to the item, can be specified multiple times. See https://learn.microsoft.com/en-us/azure/cosmos-db/partial-document-update
-    #[clap(long, short)]
+    #[arg(long, short)]
     operation: Vec<String>,
 }
 

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

@@ -24,7 +24,7 @@ enum Subcommands {
         query: String,
 
         /// The partition key to use when querying the container. Currently this only supports a single string partition key.
-        #[clap(long, short)]
+        #[arg(long, short)]
         partition_key: String,
     },
     Databases {

sdk/cosmos/azure_data_cosmos/examples/cosmos/read.rs

@@ -32,11 +32,11 @@ enum Subcommands {
         container: String,
 
         /// The ID of the item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         item_id: String,
 
         /// The partition key of the item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         partition_key: String,
     },
 }
@@ -61,7 +61,7 @@ impl ReadCommand {
                         println!("Item not found!")
                     }
                     Ok(r) => {
-                        let item: serde_json::Value = r.into_body().await?.unwrap();
+                        let item: serde_json::Value = r.into_json_body().await?;
                         println!("Found item:");
                         println!("{:#?}", item);
                     }

sdk/cosmos/azure_data_cosmos/examples/cosmos/replace.rs

@@ -1,7 +1,7 @@
 use std::error::Error;
 
 use azure_core::StatusCode;
-use azure_data_cosmos::{CosmosClient, PartitionKey};
+use azure_data_cosmos::{CosmosClient, ItemOptions, PartitionKey};
 use clap::{Args, Subcommand};
 
 use crate::utils::ThroughputOptions;
@@ -23,22 +23,26 @@ pub enum Subcommands {
         container: String,
 
         /// The ID of the item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         item_id: String,
 
         /// The partition key of the new item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         partition_key: String,
 
         /// The JSON of the new item.
-        #[clap(long, short)]
+        #[arg(long, short)]
         json: String,
+
+        /// If set, the updated item will be included in the response.
+        #[arg(long)]
+        show_updated: bool,
     },
     DatabaseThroughput {
         /// The database to update throughput for.
         database: String,
 
-        #[clap(flatten)]
+        #[command(flatten)]
         throughput_options: ThroughputOptions,
     },
     ContainerThroughput {
@@ -48,7 +52,7 @@ pub enum Subcommands {
         /// The container to update throughput for.
         container: String,
 
-        #[clap(flatten)]
+        #[command(flatten)]
         throughput_options: ThroughputOptions,
     },
 }
@@ -62,24 +66,34 @@ impl ReplaceCommand {
                 item_id,
                 partition_key,
                 json,
+                show_updated,
             } => {
                 let db_client = client.database_client(&database);
                 let container_client = db_client.container_client(&container);
 
                 let pk = PartitionKey::from(&partition_key);
                 let item: serde_json::Value = serde_json::from_str(&json)?;
 
+                let options = ItemOptions {
+                    enable_content_response_on_write: show_updated,
+                    ..Default::default()
+                };
+
                 let response = container_client
-                    .replace_item(pk, &item_id, item, None)
+                    .replace_item(pk, &item_id, item, Some(options))
                     .await;
                 match response {
                     Err(e) if e.http_status() == Some(StatusCode::NotFound) => {
                         println!("Item not found!")
                     }
                     Ok(r) => {
-                        let item: serde_json::Value = r.into_body().await?.unwrap();
-                        println!("Replaced item:");
-                        println!("{:#?}", item);
+                        println!("Replaced item successfully");
+
+                        if show_updated {
+                            let created: serde_json::Value = r.into_json_body().await?;
+                            println!("Newly replaced item:");
+                            println!("{:#?}", created);
+                        }
                     }
                     Err(e) => return Err(e.into()),
                 };

sdk/cosmos/azure_data_cosmos/examples/cosmos/upsert.rs

@@ -1,6 +1,6 @@
 use std::error::Error;
 
-use azure_data_cosmos::{CosmosClient, PartitionKey};
+use azure_data_cosmos::{CosmosClient, ItemOptions, PartitionKey};
 use clap::Args;
 
 /// Creates a new item or replaces an existing item, if a matching item already exists.
@@ -13,12 +13,16 @@ pub struct UpsertCommand {
     container: String,
 
     /// The partition key of the new item.
-    #[clap(long, short)]
+    #[arg(long, short)]
     partition_key: String,
 
     /// The JSON of the new item.
-    #[clap(long, short)]
+    #[arg(long, short)]
     json: String,
+
+    /// If set, the updated item will be included in the response.
+    #[arg(long)]
+    show_updated: bool,
 }
 
 impl UpsertCommand {
@@ -29,14 +33,21 @@ impl UpsertCommand {
         let pk = PartitionKey::from(&self.partition_key);
         let item: serde_json::Value = serde_json::from_str(&self.json)?;
 
-        let created = container_client
-            .upsert_item(pk, item, None)
-            .await?
-            .into_body()
-            .await?
-            .unwrap();
-        println!("Created item:");
-        println!("{:#?}", created);
+        let options = ItemOptions {
+            enable_content_response_on_write: self.show_updated,
+            ..Default::default()
+        };
+
+        let response = container_client
+            .upsert_item(pk, item, Some(options))
+            .await?;
+        println!("Item updated successfully");
+
+        if self.show_updated {
+            let created: serde_json::Value = response.into_json_body().await?;
+            println!("Updated item:");
+            println!("{:#?}", created);
+        }
         Ok(())
     }
 }