spin up time account validation

Signed-off-by: Zhiwei Liang <zhiwei.liang@zliang.me>

Author: ChihweiLHBird

crates/key-value-azure/src/auth.rs

@@ -54,7 +54,8 @@ pub enum AzureCredentialKind {
     /// Service principal authenticated with a client secret. Reads
     /// `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, and `AZURE_CLIENT_SECRET` from the
     /// environment — the same variables the legacy SDK's `EnvironmentCredential`
-    /// used, so existing deployments keep working without config changes.
+    /// used, so the environment carries over once `auth_type = "service_principal"`
+    /// is set.
     ServicePrincipal,
 }
 
@@ -93,7 +94,7 @@ impl AzureCredentialKind {
     ///
     /// This runs the credential's own setup (which may fail — e.g. if the
     /// environment for workload identity or service principal is absent), so it
-    /// is called lazily when the Cosmos client is first built.
+    /// is called at store construction, failing `spin up` on misconfiguration.
     pub(crate) fn credential(&self) -> azure_core::Result<Arc<dyn TokenCredential>> {
         match self {
             Self::DeveloperTools => Ok(azure_identity::DeveloperToolsCredential::new(None)?),
@@ -116,7 +117,7 @@ impl AzureCredentialKind {
                 // azure_identity 1.0 removed the env-driven `EnvironmentCredential`,
                 // so read the same variables it used and pass them to
                 // `ClientSecretCredential` explicitly. A missing variable surfaces
-                // here (lazily, when the client is first built) as a clear error.
+                // here (at store construction, during `spin up`) as a clear error.
                 let tenant_id = service_principal_env("AZURE_TENANT_ID")?;
                 let client_id = service_principal_env("AZURE_CLIENT_ID")?;
                 let secret = service_principal_env("AZURE_CLIENT_SECRET")?;

crates/key-value-azure/src/store.rs

@@ -18,10 +18,10 @@ use crate::auth::KeyValueAzureCosmosAuthOptions;
 
 pub struct KeyValueAzureCosmos {
     /// Parameters for initializing the Cosmos DB client
-    account: String,
+    account_ref: AccountReference,
+
     database: String,
     container: String,
-    auth_options: KeyValueAzureCosmosAuthOptions,
     region: Region,
 
     /// The Cosmos DB client
@@ -43,54 +43,48 @@ impl KeyValueAzureCosmos {
         region: Region,
         app_id: Option<String>,
     ) -> Result<Self> {
+        let endpoint: azure_data_cosmos::AccountEndpoint =
+            format!("https://{account}.documents.azure.com/")
+                .parse()
+                .map_err(log_error)?;
+
+        let account_ref = match auth_options {
+            KeyValueAzureCosmosAuthOptions::RuntimeConfigValues(config) => {
+                AccountReference::with_authentication_key(
+                    endpoint,
+                    Secret::from(config.key.clone()),
+                )
+            }
+            KeyValueAzureCosmosAuthOptions::AadCredential(kind) => {
+                let credential = kind.credential().map_err(log_error)?;
+                AccountReference::with_credential(endpoint, credential)
+            }
+        };
+
         Ok(Self {
-            account,
+            account_ref,
             database,
             container,
-            auth_options,
             region,
             client: tokio::sync::OnceCell::new(),
             app_id,
         })
     }
 }
 
-async fn build_cosmos_client(
-    account: &str,
-    auth_options: &KeyValueAzureCosmosAuthOptions,
-    region: &Region,
-) -> Result<CosmosClient, Error> {
-    let endpoint: azure_data_cosmos::AccountEndpoint =
-        format!("https://{account}.documents.azure.com/")
-            .parse()
-            .map_err(log_error)?;
-
-    let account_ref = match auth_options {
-        KeyValueAzureCosmosAuthOptions::RuntimeConfigValues(config) => {
-            AccountReference::with_authentication_key(endpoint, Secret::from(config.key.clone()))
-        }
-        KeyValueAzureCosmosAuthOptions::AadCredential(kind) => {
-            let credential = kind.credential().map_err(log_error)?;
-            AccountReference::with_credential(endpoint, credential)
-        }
-    };
-
-    let routing_strategy = RoutingStrategy::ProximityTo(region.clone());
-
-    CosmosClient::builder()
-        .build(account_ref, routing_strategy)
-        .await
-        .map_err(log_error)
-}
-
 #[async_trait]
 impl StoreManager for KeyValueAzureCosmos {
     async fn get(&self, name: &str) -> Result<Arc<dyn Store>, Error> {
         let client = self
             .client
             .get_or_try_init(|| async {
-                return build_cosmos_client(&self.account, &self.auth_options, &self.region)
-                    .await?
+                return CosmosClient::builder()
+                    .build(
+                        self.account_ref.clone(),
+                        RoutingStrategy::ProximityTo(self.region.clone()),
+                    )
+                    .await
+                    .map_err(log_error)?
                     .database_client(&self.database)
                     .container_client(&self.container)
                     .await