Support CosmosClient Creation via Connection String (#2641)

* add method to init client from connection string

* consolidate connection string logic

* add unit tests for connection string

* cleanup + CHANGELOG

* pr comments

* pr comments

* update comment

* suppress bare url lint warning

* suppress again

* fix cspell errors

* remove unused import

* pr comments

Author: nehrao1

.vscode/cspell.json

@@ -24,6 +24,8 @@
   ],
   "words": [
     "aarch",
+    "accountendpoint",
+    "accountkey",
     "amqp",
     "asyncoperation",
     "azsdk",

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+* Added a function `CosmosClient::with_connection_string` to enable `CosmosClient` creation via connection string. ([#2641](https://github.com/Azure/azure-sdk-for-rust/pull/2641))
+
 ### Breaking Changes
 
 ### Bugs Fixed

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

@@ -96,6 +96,36 @@ impl CosmosClient {
         })
     }
 
+    /// Creates a new CosmosClient, using a connection string.
+    ///
+    /// # Arguments
+    ///
+    /// * `connection_string` - the connection string to use for the client, e.g. `AccountEndpoint=https://accountname.documents.azure.com:443/‌​;AccountKey=accountk‌​ey`
+    /// * `options` - Optional configuration for the client.
+    ///
+    /// # Examples
+    ///
+    /// ```rust,no_run
+    /// use azure_data_cosmos::CosmosClient;
+    /// use azure_core::credentials::Secret;
+    ///
+    /// let client = CosmosClient::with_connection_string(
+    ///     "AccountEndpoint=https://accountname.documents.azure.com:443/‌​;AccountKey=accountk‌​ey",
+    ///     None)
+    ///     .unwrap();
+    /// ```
+    #[cfg(feature = "key_auth")]
+    pub fn with_connection_string(
+        connection_string: Secret,
+        options: Option<CosmosClientOptions>,
+    ) -> Result<Self, azure_core::Error> {
+        let connection_str = crate::ConnectionString::try_from(&connection_string)?;
+        let endpoint = connection_str.account_endpoint;
+        let key = connection_str.account_key;
+
+        Self::with_key(endpoint.as_str(), key, options)
+    }
+
     /// Gets a [`DatabaseClient`] that can be used to access the database with the specified ID.
     ///
     /// # Arguments

sdk/cosmos/azure_data_cosmos/src/connection_string.rs

@@ -0,0 +1,156 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+use std::str::FromStr;
+
+use azure_core::{credentials::Secret, fmt::SafeDebug, Error};
+
+/// Represents a Cosmos DB connection string.
+#[derive(Clone, PartialEq, Eq, SafeDebug)]
+pub struct ConnectionString {
+    pub account_endpoint: String,
+    pub account_key: Secret,
+}
+
+impl TryFrom<&Secret> for ConnectionString {
+    type Error = azure_core::Error;
+    fn try_from(secret: &Secret) -> Result<Self, Self::Error> {
+        secret.secret().parse()
+    }
+}
+
+impl FromStr for ConnectionString {
+    type Err = azure_core::Error;
+    fn from_str(connection_string: &str) -> Result<Self, Self::Err> {
+        if connection_string.is_empty() {
+            return Err(Error::new(
+                azure_core::error::ErrorKind::Other,
+                "connection string cannot be empty",
+            ));
+        }
+
+        let splat = connection_string.split(';');
+        let mut account_endpoint = None;
+        let mut account_key = None;
+        for part in splat {
+            let part = part.trim();
+            if part.is_empty() {
+                continue;
+            }
+
+            let (key, value) = part.split_once('=').ok_or(Error::new(
+                azure_core::error::ErrorKind::Other,
+                "invalid connection string",
+            ))?;
+
+            if key.eq_ignore_ascii_case("AccountEndpoint") {
+                account_endpoint = Some(value.to_string())
+            }
+
+            if key.eq_ignore_ascii_case("AccountKey") {
+                account_key = Some(Secret::new(value.to_string()))
+            }
+        }
+
+        let Some(endpoint) = account_endpoint else {
+            return Err(Error::new(
+                azure_core::error::ErrorKind::Other,
+                "invalid connection string, missing 'AccountEndpoint'",
+            ));
+        };
+
+        let Some(key) = account_key else {
+            return Err(Error::new(
+                azure_core::error::ErrorKind::Other,
+                "invalid connection string, missing 'AccountKey'",
+            ));
+        };
+
+        Ok(Self {
+            account_endpoint: endpoint,
+            account_key: key,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::ConnectionString;
+    use azure_core::credentials::Secret;
+
+    #[test]
+    pub fn test_valid_connection_string() {
+        let connection_string =
+            "AccountEndpoint=https://accountname.documents.azure.com:443/;AccountKey=key";
+        let secret = Secret::new(connection_string);
+        let connection_str = ConnectionString::try_from(&secret).unwrap();
+
+        assert_eq!(
+            "https://accountname.documents.azure.com:443/",
+            connection_str.account_endpoint
+        );
+
+        assert_eq!("key", connection_str.account_key.secret());
+    }
+
+    #[test]
+    pub fn test_valid_connection_string_mismatched_case() {
+        let connection_string =
+            "accountendpoint=https://accountname.documents.azure.com:443/;accountkey=key";
+        let secret = Secret::new(connection_string);
+        let connection_str = ConnectionString::try_from(&secret).unwrap();
+
+        // should pass, we don't expect connection string keys to be case sensitive
+        assert_eq!(
+            "https://accountname.documents.azure.com:443/",
+            connection_str.account_endpoint
+        );
+
+        assert_eq!("key", connection_str.account_key.secret());
+    }
+
+    #[test]
+    pub fn test_empty_connection_string() {
+        test_bad_connection_string("", "connection string cannot be empty")
+    }
+
+    #[test]
+    pub fn test_malformed_connection_string() {
+        test_bad_connection_string(
+            "AccountEndpointhttps://accountname.documents.azure.com:443AccountKeyaccountkey",
+            "invalid connection string",
+        );
+    }
+
+    #[test]
+    pub fn test_partially_malformed_connection_string() {
+        test_bad_connection_string(
+            "AccountEndpointhttps://accountname.documents.azure.com:443/AccountKey=accountkey",
+            "invalid connection string, missing 'AccountEndpoint'",
+        );
+    }
+
+    #[test]
+    pub fn test_connection_string_missing_account_endpoint() {
+        test_bad_connection_string(
+            "AccountKey=key",
+            "invalid connection string, missing 'AccountEndpoint'",
+        );
+    }
+
+    #[test]
+    pub fn test_connection_string_missing_account_key() {
+        test_bad_connection_string(
+            "AccountEndpoint=https://accountname.documents.azure.com:443/;",
+            "invalid connection string, missing 'AccountKey'",
+        );
+    }
+
+    fn test_bad_connection_string(connection_string: &str, expected_error_message: &str) {
+        let secret = Secret::new(connection_string.to_owned());
+        let connection_str = ConnectionString::try_from(&secret);
+        let err = connection_str.unwrap_err();
+        let actual_error_message = format!("{}", err);
+        assert_eq!(expected_error_message, actual_error_message.as_str())
+    }
+}

sdk/cosmos/azure_data_cosmos/src/lib.rs

@@ -11,6 +11,7 @@
 #![cfg_attr(docsrs, feature(doc_cfg_hide))]
 
 pub mod clients;
+mod connection_string;
 pub mod constants;
 mod feed;
 mod options;
@@ -25,6 +26,7 @@ pub mod models;
 #[doc(inline)]
 pub use clients::CosmosClient;
 
+pub use connection_string::*;
 pub use options::*;
 pub use partition_key::*;
 pub use query::Query;

sdk/cosmos/azure_data_cosmos/tests/framework/test_account.rs

@@ -2,11 +2,11 @@
 
 #![cfg_attr(not(feature = "key_auth"), allow(dead_code))]
 
-use std::{borrow::Cow, sync::Arc};
+use std::{borrow::Cow, str::FromStr, sync::Arc};
 
 use azure_core::{credentials::Secret, http::TransportOptions, test::TestMode};
 use azure_core_test::TestContext;
-use azure_data_cosmos::{CosmosClientOptions, Query};
+use azure_data_cosmos::{ConnectionString, CosmosClientOptions, Query};
 use reqwest::ClientBuilder;
 
 /// Represents a Cosmos DB account for testing purposes.
@@ -74,39 +74,16 @@ impl TestAccount {
         options: Option<TestAccountOptions>,
     ) -> Result<Self, Box<dyn std::error::Error>> {
         let options = options.unwrap_or_default();
-        let splat = connection_string.split(';');
-        let mut account_endpoint = None;
-        let mut account_key = None;
-        for part in splat {
-            let part = part.trim();
-            if part.is_empty() {
-                continue;
-            }
-
-            let (key, value) = part.split_once('=').ok_or("invalid connection string")?;
-            match key {
-                "AccountEndpoint" => account_endpoint = Some(value.to_string()),
-                "AccountKey" => account_key = Some(Secret::new(value.to_string())),
-                _ => {}
-            }
-        }
-
-        let Some(endpoint) = account_endpoint else {
-            return Err("invalid connection string, missing 'AccountEndpoint'".into());
-        };
-
-        let Some(key) = account_key else {
-            return Err("invalid connection string, missing 'AccountKey'".into());
-        };
+        let connection_str = ConnectionString::from_str(connection_string)?;
 
         // We need the context_id to be constant, so that record/replay work.
         let context_id = context.name().to_string();
 
         Ok(TestAccount {
             context,
             context_id,
-            endpoint,
-            key,
+            endpoint: connection_str.account_endpoint.to_string(),
+            key: connection_str.account_key,
             options,
         })
     }