[Cosmos] Control-Plane Over Data-Plane APIs (Database/Container CRUD) (Azure#1853)

Author: analogrelay

sdk/cosmos/azure_data_cosmos/Cargo.toml

@@ -13,8 +13,6 @@ documentation = "https://docs.rs/azure_data_cosmos"
 keywords = ["sdk", "azure", "rest", "cloud", "cosmos", "database"]
 categories = ["api-bindings"]
 
-# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
-
 [dependencies]
 async-trait.workspace = true
 azure_core.workspace = true

sdk/cosmos/azure_data_cosmos/docs/control-plane-warning.md

@@ -0,0 +1,5 @@
+<div class="warning">
+
+This is a control-plane API and requires that you authenticate using a key. To use Entra ID to perform this operation, you must use the [Azure Resource Manager APIs](https://learn.microsoft.com/en-us/azure/templates/microsoft.documentdb/databaseaccounts).
+
+</div>

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

@@ -2,44 +2,125 @@ use std::error::Error;
 
 use azure_data_cosmos::{
     clients::{ContainerClientMethods, DatabaseClientMethods},
+    models::{ContainerProperties, PartitionKeyDefinition},
     CosmosClient, CosmosClientMethods, PartitionKey,
 };
-use clap::Args;
+use clap::{Args, Subcommand};
 
-/// Creates a new item.
+/// Creates a new item, database, or container.
 #[derive(Clone, Args)]
 pub struct CreateCommand {
-    /// The database in which to create the item.
-    database: String,
+    #[command(subcommand)]
+    subcommand: Subcommands,
+}
+
+#[derive(Clone, Subcommand)]
+pub enum Subcommands {
+    /// Create an item in a container.
+    Item {
+        /// The database in which to create the item.
+        database: String,
+
+        /// The container in which to create the item.
+        container: String,
+
+        /// The partition key of the new item.
+        #[clap(long, short)]
+        partition_key: String,
+
+        /// The JSON of the new item.
+        #[clap(long, short)]
+        json: String,
+    },
+
+    /// Create a database (does not support Entra ID).
+    Database {
+        /// The ID of the new database to create.
+        id: String,
+    },
 
-    /// The container in which to create the item.
-    container: String,
+    /// Create a container (does not support Entra ID).
+    Container {
+        /// The ID of the database to create the container in.
+        database: String,
 
-    /// The partition key of the new item.
-    #[clap(long, short)]
-    partition_key: String,
+        /// The ID of the new container to create.
+        #[clap(long, short)]
+        id: Option<String>,
 
-    /// The JSON of the new item.
-    #[clap(long, short)]
-    json: String,
+        /// The path to the partition key properties (supports up to 3).
+        #[clap(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)]
+        json: Option<String>,
+    },
 }
 
 impl CreateCommand {
     pub async fn run(self, client: CosmosClient) -> Result<(), Box<dyn Error>> {
-        let db_client = client.database_client(&self.database);
-        let container_client = db_client.container_client(&self.container);
-
-        let pk = PartitionKey::from(&self.partition_key);
-        let item: serde_json::Value = serde_json::from_str(&self.json)?;
-
-        let created = container_client
-            .create_item(pk, item, None)
-            .await?
-            .deserialize_body()
-            .await?
-            .unwrap();
-        println!("Created item:");
-        println!("{:#?}", created);
-        Ok(())
+        match self.subcommand {
+            Subcommands::Item {
+                database,
+                container,
+                partition_key,
+                json,
+            } => {
+                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?
+                    .deserialize_body()
+                    .await?
+                    .unwrap();
+                println!("Created item:");
+                println!("{:#?}", created);
+                Ok(())
+            }
+
+            Subcommands::Database { id } => {
+                let db = client
+                    .create_database(id, None)
+                    .await?
+                    .deserialize_body()
+                    .await?
+                    .unwrap();
+                println!("Created database:");
+                println!("{:#?}", db);
+                Ok(())
+            }
+
+            Subcommands::Container {
+                database,
+                id,
+                partition_key,
+                json,
+            } => {
+                let properties = match json {
+                    Some(j) => serde_json::from_str(&j).unwrap(),
+                    None => ContainerProperties {
+                        id: id.expect("the ID is required when not using '--json'"),
+                        partition_key: PartitionKeyDefinition::new(partition_key),
+                        ..Default::default()
+                    },
+                };
+                let container = client
+                    .database_client(database)
+                    .create_container(properties, None)
+                    .await?
+                    .deserialize_body()
+                    .await?
+                    .unwrap();
+                println!("Created container:");
+                println!("{:#?}", container);
+                Ok(())
+            }
+        }
     }
 }

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

@@ -5,39 +5,87 @@ use azure_data_cosmos::{
     clients::{ContainerClientMethods, DatabaseClientMethods},
     CosmosClient, CosmosClientMethods,
 };
-use clap::Args;
+use clap::{Args, Subcommand};
 
-/// Deletes an item.
+/// Deletes an item, database, or container.
 #[derive(Clone, Args)]
 pub struct DeleteCommand {
-    /// The database containing the item.
-    database: String,
+    #[command(subcommand)]
+    subcommand: Subcommands,
+}
+
+#[derive(Clone, Subcommand)]
+pub enum Subcommands {
+    /// Delete an item in a container.
+    Item {
+        /// The database containing the item.
+        database: String,
+
+        /// The container containing the item.
+        container: String,
+
+        /// The ID of the item.
+        #[clap(long, short)]
+        item_id: String,
 
-    /// The container containing the item.
-    container: String,
+        /// The partition key of the item.
+        #[clap(long, short)]
+        partition_key: String,
+    },
 
-    /// The ID of the item.
-    #[clap(long, short)]
-    item_id: String,
+    /// Create a database (does not support Entra ID).
+    Database {
+        /// The ID of the database to delete.
+        id: String,
+    },
 
-    /// The partition key of the item.
-    #[clap(long, short)]
-    partition_key: String,
+    /// Create a container (does not support Entra ID).
+    Container {
+        /// The ID of the database the container is in.
+        database: String,
+
+        /// The ID of the container to delete
+        id: String,
+    },
 }
 
 impl DeleteCommand {
     pub async fn run(self, client: CosmosClient) -> Result<(), Box<dyn Error>> {
-        let db_client = client.database_client(&self.database);
-        let container_client = db_client.container_client(&self.container);
-
-        let response = container_client
-            .delete_item(&self.partition_key, &self.item_id, None)
-            .await;
-        match response {
-            Err(e) if e.http_status() == Some(StatusCode::NotFound) => println!("Item not found!"),
-            Ok(_) => println!("Item deleted"),
-            Err(e) => return Err(e.into()),
-        };
-        Ok(())
+        match self.subcommand {
+            Subcommands::Item {
+                database,
+                container,
+                item_id,
+                partition_key,
+            } => {
+                let db_client = client.database_client(database);
+                let container_client = db_client.container_client(container);
+
+                let response = container_client
+                    .delete_item(partition_key, item_id, None)
+                    .await;
+                match response {
+                    Err(e) if e.http_status() == Some(StatusCode::NotFound) => {
+                        println!("Item not found!")
+                    }
+                    Ok(_) => println!("Item deleted"),
+                    Err(e) => return Err(e.into()),
+                };
+                Ok(())
+            }
+
+            Subcommands::Database { id } => {
+                let db_client = client.database_client(id);
+                db_client.delete(None).await?;
+                Ok(())
+            }
+
+            Subcommands::Container { database, id } => {
+                let db_client = client.database_client(database);
+                let container_client = db_client.container_client(id);
+                container_client.delete(None).await?;
+                Ok(())
+            }
+        }
     }
 }

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

@@ -34,7 +34,6 @@ struct SharedArgs {
 
     /// 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>,
 }
 
@@ -62,7 +61,7 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
         return Ok(());
     };
 
-    let client = create_client(&args.shared_args);
+    let client = create_client(&args.shared_args)?;
 
     match cmd {
         Subcommands::Create(cmd) => cmd.run(client).await,
@@ -75,18 +74,19 @@ pub async fn main() -> Result<(), Box<dyn Error>> {
     }
 }
 
-#[cfg(feature = "key_auth")]
-fn create_client(args: &SharedArgs) -> CosmosClient {
+fn create_client(args: &SharedArgs) -> Result<CosmosClient, Box<dyn Error>> {
     if let Some(key) = args.key.as_ref() {
-        CosmosClient::with_key(&args.endpoint, key.clone(), None).unwrap()
+        #[cfg(feature = "key_auth")]
+        {
+            Ok(CosmosClient::with_key(&args.endpoint, key.clone(), None)?)
+        }
+        #[cfg(not(feature = "key_auth"))]
+        {
+            let _ = key; // Mark 'key' as used to make the compiler happy.
+            Err("cannot authenticate with a key unless the 'key_auth' feature is enabled".into())
+        }
     } else {
         let cred = DefaultAzureCredential::new().unwrap();
-        CosmosClient::new(&args.endpoint, cred, None).unwrap()
+        Ok(CosmosClient::new(&args.endpoint, cred, None)?)
     }
 }
-
-#[cfg(not(feature = "key_auth"))]
-fn create_client(args: &SharedArgs) -> CosmosClient {
-    let cred = DefaultAzureCredential::new().unwrap();
-    CosmosClient::new(&args.endpoint, cred, None).unwrap()
-}