Shared auth among cosmos examples

sdk/data_cosmos/examples/attachments_00.rs

@@ -3,6 +3,8 @@ use clap::Parser;
 use futures::StreamExt;
 use serde::{Deserialize, Serialize};
 
+mod util;
+
 // Now we create a sample struct.
 #[derive(Serialize, Deserialize, Clone, Debug)]
 struct MySampleStruct {
@@ -26,22 +28,18 @@ struct Args {
     database_name: String,
     /// Name of the collection in the database
     collection_name: String,
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
+    #[clap(flatten)]
+    auth: util::Auth,
 }
 
 // This example expects you to have created a collection
 // with partitionKey on "id".
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
     let args = Args::parse();
-    let authorization_token = AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    let client = CosmosClient::new(args.account, authorization_token, CosmosOptions::default())
+    let client = args
+        .auth
+        .into_client()?
         .database_client(args.database_name)
         .collection_client(args.collection_name);
 

sdk/data_cosmos/examples/cancellation.rs

@@ -1,30 +1,14 @@
-use azure_data_cosmos::prelude::*;
 use clap::Parser;
 use stop_token::prelude::*;
 use stop_token::StopSource;
 use tokio::time::{Duration, Instant};
 
-#[derive(Debug, Parser)]
-struct Args {
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
-}
+mod util;
 
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
     env_logger::init();
-    // First we retrieve the account name and access key from environment variables, and
-    // create an authorization token.
-    let args = Args::parse();
-    let authorization_token = AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    // Create a new Cosmos client.
-    let options = CosmosOptions::default();
-    let client = CosmosClient::new(args.account.clone(), authorization_token.clone(), options);
+    let client = util::Auth::parse().into_client()?;
 
     // Create a new database, and time out if it takes more than 1 second.
     let future = client.create_database("my_database").into_future();

sdk/data_cosmos/examples/collection.rs

@@ -1,43 +1,12 @@
-use azure_data_cosmos::prelude::*;
 use clap::Parser;
 use futures::stream::StreamExt;
-
-#[derive(Debug, Parser)]
-struct Args {
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
-}
+mod util;
 
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
-    // First we retrieve the account name and access key from environment variables.
-    // We expect access keys (ie, not resource constrained)
-    let args = Args::parse();
-
-    // This is how you construct an authorization token.
-    // Remember to pick the correct token type.
-    // Here we assume master.
-    // Most methods return a ```Result<_, azure_data_cosmos::Error>```.
-    // ```azure_data_cosmos::Error``` is an enum union of all the possible underlying
-    // errors, plus Azure specific ones. For example if a REST call returns the
-    // unexpected result (ie NotFound instead of Ok) we return an Err telling
-    // you that.
-    let authorization_token = AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    // Once we have an authorization token you can create a client instance. You can change the
-    // authorization token at later time if you need, for example, to escalate the privileges for a
-    // single operation.
-    // Here we are using reqwest but other clients are supported (check the documentation).
-    let client = CosmosClient::new(
-        args.account.clone(),
-        authorization_token,
-        CosmosOptions::default(),
-    );
-
+    let args = util::Auth::parse();
+    let account = args.account().clone();
+    let client = args.into_client()?;
     // The Cosmos' client exposes a lot of methods. This one lists the databases in the specified account.
     let databases = client
         .list_databases()
@@ -48,7 +17,7 @@ async fn main() -> azure_core::Result<()> {
 
     println!(
         "Account {} has {} database(s)",
-        args.account,
+        account,
         databases.databases.len()
     );
 

sdk/data_cosmos/examples/create_delete_database.rs

@@ -1,87 +1,60 @@
-use azure_data_cosmos::prelude::*;
 use clap::Parser;
 use futures::stream::StreamExt;
 
-#[derive(Debug, Parser)]
+mod util;
+
+#[derive(Debug, clap::Parser)]
 struct Args {
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
     /// The name of the database
     database_name: String,
+    #[clap(flatten)]
+    auth: util::Auth,
 }
 
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
-    // First we retrieve the account name and access key from environment variables.
-    // We expect access keys (ie, not resource constrained)
     let args = Args::parse();
+    let database_name = args.database_name;
+    let client = args.auth.into_client()?;
 
-    // This is how you construct an authorization token.
-    // Remember to pick the correct token type.
-    // Here we assume master.
-    // Most methods return a ```Result<_, azure_data_cosmos::Error>```.
-    // ```azure_data_cosmos::Error``` is an enum union of all the possible underlying
-    // errors, plus Azure specific ones. For example if a REST call returns the
-    // unexpected result (ie NotFound instead of Ok) we return an Err telling
-    // you that.
-    let authorization_token =
-        permission::AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    // Once we have an authorization token you can create a client instance. You can change the
-    // authorization token at later time if you need, for example, to escalate the privileges for a
-    // single operation.
-    let client = CosmosClient::new(args.account, authorization_token, CosmosOptions::default());
-
-    // The Cosmos' client exposes a lot of methods. This one lists the databases in the specified
-    // account. Database do not implement Display but deref to &str so you can pass it to methods
-    // both as struct or id.
-
+    // The Cosmos' client exposes a lot of methods. This one lists the databases in the specified account.
     let mut list_databases_stream = client.list_databases().into_stream();
     while let Some(list_databases_response) = list_databases_stream.next().await {
         println!("list_databases_response = {:#?}", list_databases_response?);
     }
     drop(list_databases_stream);
 
-    let db = client
-        .create_database(&args.database_name)
-        .into_future()
-        .await?;
+    let db = client.create_database(&database_name).into_future().await?;
     println!("created database = {:#?}", db);
 
     // create collection!
-    {
-        let database = client.database_client(args.database_name.clone());
-        let create_collection_response = database
-            .create_collection("panzadoro", "/id")
-            .into_future()
-            .await?;
-
-        println!(
-            "create_collection_response == {:#?}",
-            create_collection_response
-        );
+    let database = client.database_client(database_name.clone());
+    let create_collection_response = database
+        .create_collection("panzadoro", "/id")
+        .into_future()
+        .await?;
 
-        let db_collection = database.collection_client("panzadoro");
+    println!(
+        "create_collection_response == {:#?}",
+        create_collection_response
+    );
 
-        let get_collection_response = db_collection.get_collection().into_future().await?;
-        println!("get_collection_response == {:#?}", get_collection_response);
+    let db_collection = database.collection_client("panzadoro");
 
-        let mut stream = database.list_collections().into_stream();
-        while let Some(res) = stream.next().await {
-            let res = res?;
-            println!("res == {:#?}", res);
-        }
+    let get_collection_response = db_collection.get_collection().into_future().await?;
+    println!("get_collection_response == {:#?}", get_collection_response);
 
-        let delete_response = db_collection.delete_collection().into_future().await?;
-        println!("collection deleted: {:#?}", delete_response);
+    let mut stream = database.list_collections().into_stream();
+    while let Some(res) = stream.next().await {
+        let res = res?;
+        println!("res == {:#?}", res);
     }
 
+    let delete_response = db_collection.delete_collection().into_future().await?;
+    println!("collection deleted: {:#?}", delete_response);
+
     let resp = client
-        .database_client(args.database_name)
+        .database_client(database_name)
         .delete_database()
         .into_future()
         .await?;

sdk/data_cosmos/examples/database_00.rs

@@ -1,29 +1,12 @@
-use azure_data_cosmos::prelude::*;
 use clap::Parser;
 use futures::stream::StreamExt;
 use serde_json::Value;
 
-#[derive(Debug, Parser)]
-struct Args {
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
-}
+mod util;
 
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
-    // First we retrieve the account name and access key from environment variables.
-    // We expect access keys (ie, not resource constrained)
-
-    let args = Args::parse();
-
-    let authorization_token =
-        permission::AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    let client = CosmosClient::new(args.account, authorization_token, CosmosOptions::default());
+    let client = util::Auth::parse().into_client()?;
 
     let dbs = client
         .list_databases()

sdk/data_cosmos/examples/database_01.rs

@@ -1,25 +1,11 @@
-use azure_data_cosmos::prelude::*;
 use clap::Parser;
 use futures::stream::StreamExt;
 
-#[derive(Debug, Parser)]
-struct Args {
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
-}
+mod util;
 
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
-    // First we retrieve the account name and access key from environment variables.
-    // We expect access keys (ie, not resource constrained)
-    let args = Args::parse();
-    let authorization_token = AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    let client = CosmosClient::new(args.account, authorization_token, CosmosOptions::default());
+    let client = util::Auth::parse().into_client()?;
 
     let database = client.database_client("pollo");
     println!("database_name == {}", database.database_name());

sdk/data_cosmos/examples/document_00.rs

@@ -4,17 +4,9 @@ use serde::{Deserialize, Serialize};
 // Using the prelude module of the Cosmos crate makes easier to use the Rust Azure SDK for Cosmos DB.
 use azure_core::prelude::*;
 
-use azure_data_cosmos::prelude::*;
+mod util;
 
-#[derive(Debug, Parser)]
-struct Args {
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
-}
+use azure_data_cosmos::prelude::*;
 
 #[derive(Clone, Serialize, Deserialize, Debug)]
 struct MySampleStruct {
@@ -43,22 +35,7 @@ const COLLECTION: &str = "azuresdktc";
 // 4. Delete everything.
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
-    // Let's get Cosmos account and access key from env variables.
-    // This helps automated testing.
-    let args = Args::parse();
-
-    // First, we create an authorization token. There are two types of tokens, master and resource
-    // constrained. Please check the Azure documentation for details. You can change tokens
-    // at will and it's a good practice to raise your privileges only when needed.
-    let authorization_token = AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    // Next we will create a Cosmos client. You need an authorization_token but you can later
-    // change it if needed.
-    let client = CosmosClient::new(
-        args.account.clone(),
-        authorization_token.clone(),
-        CosmosOptions::default(),
-    );
+    let client = util::Auth::parse().into_client()?;
 
     // list_databases will give us the databases available in our account. If there is
     // an error (for example, the given key is not valid) you will receive a

sdk/data_cosmos/examples/document_entries_00.rs

@@ -5,18 +5,16 @@ use clap::Parser;
 use futures::stream::StreamExt;
 use serde::{Deserialize, Serialize};
 
-#[derive(Debug, Parser)]
+mod util;
+
+#[derive(Debug, clap::Parser)]
 struct Args {
-    /// Cosmos primary key name
-    #[clap(env = "COSMOS_PRIMARY_KEY")]
-    primary_key: String,
-    /// The cosmos account your're using
-    #[clap(env = "COSMOS_ACCOUNT")]
-    account: String,
     /// The name of the database
     database_name: String,
     /// The name of the collection
     collection_name: String,
+    #[clap(flatten)]
+    auth: util::Auth,
 }
 
 // Now we create a sample struct.
@@ -41,12 +39,14 @@ impl azure_data_cosmos::CosmosEntity for MySampleStruct {
 #[tokio::main]
 async fn main() -> azure_core::Result<()> {
     let args = Args::parse();
-    let authorization_token =
-        permission::AuthorizationToken::primary_from_base64(&args.primary_key)?;
-
-    let client = CosmosClient::new(args.account, authorization_token, CosmosOptions::default())
-        .database_client(args.database_name)
-        .collection_client(args.collection_name);
+    let database_name = args.database_name;
+    let collection_name = args.collection_name;
+
+    let client = args
+        .auth
+        .into_client()?
+        .database_client(database_name)
+        .collection_client(collection_name);
 
     let mut response = None;
     for i in 0u64..5 {