diff --git a/Cargo.lock b/Cargo.lock index f38da1bdf0..5e99354388 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,13 +2,22 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "addr2line" +version = "0.25.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b5d307320b3181d6d7954e663bd7c774a838b8220fe0593c86d9fb09f498b4b" +dependencies = [ + "gimli 0.32.3", +] + [[package]] name = "addr2line" version = "0.26.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "59317f77929f0e679d39364702289274de2f0f0b22cbf50b2b8cff2169a0b27a" dependencies = [ - "gimli", + "gimli 0.33.0", ] [[package]] @@ -872,50 +881,98 @@ dependencies = [ [[package]] name = "azure_core" -version = "0.21.0" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b552ad43a45a746461ec3d3a51dfb6466b4759209414b439c165eb6a6b7729e" +checksum = "4e41cbd819986ba41904c207d8ffc4106f8f8352a548d773e9554906379bb2fb" dependencies = [ + "async-lock", "async-trait", - "base64 0.22.1", + "azure_core_macros", "bytes", - "dyn-clone", "futures", - "getrandom 0.2.17", "hmac 0.12.1", - "http-types", - "once_cell", - "paste", "pin-project", - "rand 0.8.6", - "reqwest 0.12.28", "rustc_version", "serde", "serde_json", "sha2 0.10.9", + "tokio", + "tracing", + "typespec", + "typespec_client_core", +] + +[[package]] +name = "azure_core_macros" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9b52dba6a345f3ad2d42ff8d0d63df9d0994cfa29657bf18ffdbf149f78a4f5" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", + "tracing", +] + +[[package]] +name = "azure_data_cosmos" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f9d01906458ceafb6cf34b2217b79bfa2c757f507f1942706e134916ba94b0c" +dependencies = [ + "async-lock", + "async-trait", + "azure_core 1.1.0", + "azure_data_cosmos_driver", + "base64 0.22.1", + "futures", + "pin-project", + "reqwest 0.13.4", + "serde", + "serde_json", "time", "tracing", "url", - "uuid", ] [[package]] -name = "azure_data_cosmos" -version = "0.21.0" +name = "azure_data_cosmos_driver" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0aa5603f2de38c21165a1b5dfed94d64b1ab265526b0686e8557c907a53a0ee2" +checksum = "68790fbe29d23e2c28d2e80820701e5bb7646aff33473b3a975e9646c5726611" dependencies = [ + "arc-swap", + "async-lock", "async-trait", - "azure_core 0.21.0", + "azure_core 1.1.0", + "azure_data_cosmos_macros", + "backtrace", + "base64 0.22.1", "bytes", + "crossbeam-epoch", "futures", + "h2 0.4.15", + "rand 0.10.2", + "reqwest 0.13.4", "serde", "serde_json", - "thiserror 1.0.69", "time", + "tokio", "tracing", "url", "uuid", + "windows", +] + +[[package]] +name = "azure_data_cosmos_macros" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7171e936ab196e5889af4cbee33f775a8a7548c03be8076207e560776d7b87ba" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.118", ] [[package]] @@ -940,23 +997,21 @@ dependencies = [ [[package]] name = "azure_identity" -version = "0.21.0" +version = "1.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "88ddd80344317c40c04b603807b63a5cefa532f1b43522e72f480a988141f744" +checksum = "32edf96b356ca7c51d7590c4925cc36efc3947a5da4468e8e0b25c56ecbb3de5" dependencies = [ "async-lock", - "async-process", "async-trait", - "azure_core 0.21.0", + "azure_core 1.1.0", "futures", - "oauth2", "pin-project", "serde", + "serde_json", "time", + "tokio", "tracing", - "tz-rs", "url", - "uuid", ] [[package]] @@ -981,6 +1036,21 @@ dependencies = [ "fastrand 2.4.1", ] +[[package]] +name = "backtrace" +version = "0.3.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb531853791a215d7c62a30daf0dde835f381ab5de4589cfe7c649d2cbe92bd6" +dependencies = [ + "addr2line 0.25.1", + "cfg-if", + "libc", + "miniz_oxide", + "object 0.37.3", + "rustc-demangle", + "windows-link", +] + [[package]] name = "base16ct" version = "0.2.0" @@ -1995,7 +2065,7 @@ dependencies = [ "cranelift-control", "cranelift-entity", "cranelift-isle", - "gimli", + "gimli 0.33.0", "hashbrown 0.17.1", "libm", "log", @@ -2504,7 +2574,6 @@ version = "0.5.8" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" dependencies = [ - "powerfmt", "serde_core", ] @@ -3783,6 +3852,12 @@ dependencies = [ "syn 2.0.118", ] +[[package]] +name = "gimli" +version = "0.32.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e629b9b98ef3dd8afe6ca2bd0f89306cec16d43d907889945bc5d6687f2f13c7" + [[package]] name = "gimli" version = "0.33.0" @@ -5983,9 +6058,9 @@ dependencies = [ [[package]] name = "num-conv" -version = "0.1.0" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51d515d32fb182ee37cda2ccdcb92950d6a3c2893aa280e540671c2cd0f3b1d9" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" [[package]] name = "num-integer" @@ -6145,6 +6220,15 @@ dependencies = [ "cc", ] +[[package]] +name = "object" +version = "0.37.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff76201f031d8863c38aa7f905eca4f53abbfa15f609db4277d44cd8938f33fe" +dependencies = [ + "memchr", +] + [[package]] name = "object" version = "0.39.1" @@ -7776,10 +7860,12 @@ dependencies = [ "bytes", "futures-core", "futures-util", + "h2 0.4.15", "http 1.4.2", "http-body 1.0.1", "http-body-util", "hyper 1.10.1", + "hyper-rustls 0.27.9", "hyper-tls", "hyper-util", "js-sys", @@ -9595,11 +9681,11 @@ version = "4.1.0-pre0" dependencies = [ "anyhow", "async-trait", - "azure_core 0.21.0", + "azure_core 1.1.0", "azure_data_cosmos", - "azure_identity 0.21.0", + "azure_identity 1.0.0", "futures", - "reqwest 0.12.28", + "itertools 0.15.0", "serde", "spin-factor-key-value", "tokio", @@ -10581,12 +10667,11 @@ dependencies = [ [[package]] name = "time" -version = "0.3.45" +version = "0.3.54" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f9e442fc33d7fdb45aa9bfeb312c095964abdf596f7567261062b2a7107aaabd" +checksum = "3e1d5e639ff6bab73cb6885cc7e7b1de96c3f32c68ec55f3952614bec1092244" dependencies = [ "deranged", - "itoa", "js-sys", "libc", "num-conv", @@ -10599,15 +10684,15 @@ dependencies = [ [[package]] name = "time-core" -version = "0.1.7" +version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8b36ee98fd31ec7426d599183e8fe26932a8dc1fb76ddb6214d05493377d34ca" +checksum = "9e1c906769ad99c88eaa54e728060edef082f8e358ff32030cb7c7d315e81109" [[package]] name = "time-macros" -version = "0.2.25" +version = "0.2.32" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71e552d1249bf61ac2a52db88179fd0673def1e1ad8243a00d9ec9ed71fee3dd" +checksum = "7e689342a48d2ea927c87ea50cabf8594854bf940e9310208848d680d668ed85" dependencies = [ "num-conv", "time-core", @@ -11138,6 +11223,57 @@ version = "1.20.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" +[[package]] +name = "typespec" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "753a2fe021e407d4fc9ee6f4f0a33403cc306d5c54c4e4ebe1b8cbde0ca052b9" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures", + "serde", + "serde_json", + "url", +] + +[[package]] +name = "typespec_client_core" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0373af0f9d4f580b3a1a9d9639cedaabe015ed262b35bfbe13941bfb14fe1ea6" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "dyn-clone", + "futures", + "pin-project", + "rand 0.10.2", + "reqwest 0.13.4", + "serde", + "serde_json", + "time", + "tokio", + "tracing", + "typespec", + "typespec_macros", + "url", + "uuid", +] + +[[package]] +name = "typespec_macros" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c608f4427943f8adb211abc95c87672b1b98847152783507d54e3246e502f60" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "syn 2.0.118", +] + [[package]] name = "tz-rs" version = "0.6.14" @@ -11361,6 +11497,7 @@ checksum = "bf80a72845275afea99e7f2b434723d3bc7e38470fcd1c7ed39a599c73319a53" dependencies = [ "getrandom 0.4.3", "js-sys", + "rand 0.10.2", "wasm-bindgen", ] @@ -12080,7 +12217,7 @@ version = "47.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06d9f263ccb212017dae1d109b9997f218a7675db1e9a7bedf07ae1643322bf3" dependencies = [ - "addr2line", + "addr2line 0.26.1", "async-trait", "bitflags 2.13.0", "bumpalo", @@ -12090,13 +12227,13 @@ dependencies = [ "encoding_rs", "futures", "fxprof-processed-profile", - "gimli", + "gimli 0.33.0", "ittapi", "libc", "log", "mach2", "memfd", - "object", + "object 0.39.1", "once_cell", "postcard", "pulley-interpreter", @@ -12139,11 +12276,11 @@ dependencies = [ "cranelift-bforest", "cranelift-bitset", "cranelift-entity", - "gimli", + "gimli 0.33.0", "hashbrown 0.17.1", "indexmap 2.14.0", "log", - "object", + "object 0.39.1", "postcard", "rustc-demangle", "semver", @@ -12224,10 +12361,10 @@ dependencies = [ "cranelift-entity", "cranelift-frontend", "cranelift-native", - "gimli", + "gimli 0.33.0", "itertools 0.14.0", "log", - "object", + "object 0.39.1", "pulley-interpreter", "smallvec", "target-lexicon", @@ -12261,7 +12398,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f1f60ecba3497cfc26284903c2e79023fd5ea991b95b0e8dafe994fd9fca4c4e" dependencies = [ "cc", - "object", + "object 0.39.1", "rustix", "wasmtime-internal-versioned-export-macros", ] @@ -12287,7 +12424,7 @@ dependencies = [ "cfg-if", "cranelift-codegen", "log", - "object", + "object 0.39.1", "wasmtime-environ", ] diff --git a/crates/key-value-azure/Cargo.toml b/crates/key-value-azure/Cargo.toml index 75c1c691c1..e06119e88b 100644 --- a/crates/key-value-azure/Cargo.toml +++ b/crates/key-value-azure/Cargo.toml @@ -11,19 +11,17 @@ rust-version.workspace = true [dependencies] anyhow = { workspace = true } async-trait = { workspace = true } -azure_core = "0.21.0" -azure_data_cosmos = "0.21.0" -azure_identity = "0.21.0" +azure_core = { version = "1.0", default-features = false } +# Use native-tls (OpenSSL) rather than the SDK-default rustls. The default +# `rustls` feature drags in reqwest's aws-lc-rs provider, which collides with +# Spin's workspace-pinned `ring` provider and panics at runtime. +azure_data_cosmos = { version = "0.37", default-features = false, features = ["key_auth", "native_tls", "hmac_rust"] } +azure_identity = { version = "1.0", default-features = false, features = ["tokio"] } futures = { workspace = true } -reqwest = { version = "0.12", default-features = false } +itertools = { workspace = true } serde = { workspace = true } spin-factor-key-value = { path = "../factor-key-value" } tokio = { workspace = true, features = ["rt", "sync"] } [lints] workspace = true - -[features] -# Enables reusing connections to the Azure Cosmos DB service. -connection-pooling = [] -default = ["connection-pooling"] diff --git a/crates/key-value-azure/src/auth.rs b/crates/key-value-azure/src/auth.rs new file mode 100644 index 0000000000..07640c5f7b --- /dev/null +++ b/crates/key-value-azure/src/auth.rs @@ -0,0 +1,147 @@ +use anyhow::Result; +use azure_core::credentials::TokenCredential; +use std::sync::Arc; + +/// Azure Cosmos Key / Value runtime config literal options for authentication +#[derive(Clone, Debug)] +pub struct KeyValueAzureCosmosRuntimeConfigOptions { + pub(crate) key: String, +} + +impl KeyValueAzureCosmosRuntimeConfigOptions { + pub fn new(key: String) -> Self { + Self { key } + } +} + +/// Azure Cosmos Key / Value enumeration for the possible authentication options +#[derive(Clone, Debug)] +pub enum KeyValueAzureCosmosAuthOptions { + /// Runtime Config values indicates the account and key have been specified directly + RuntimeConfigValues(KeyValueAzureCosmosRuntimeConfigOptions), + /// An Azure AD token credential, used when the runtime config omits `key`. + /// + /// The specific credential is chosen by the operator via the `auth_type` + /// runtime-config field (defaulting to developer tools for local + /// development). There is deliberately no fallback *between* credential + /// types: `azure_identity` 1.0 removed `EnvironmentCredential` / + /// `DefaultAzureCredential` because silently trying a different identity + /// after one fails is a security footgun (Azure/azure-sdk-for-rust#2283). + /// This mirrors their recommended "specific credential" pattern. + AadCredential(AzureCredentialKind), +} + +/// The specific Azure AD credential to use when authenticating to Cosmos +/// without an account key. +/// +/// Each variant maps to exactly one `azure_identity` credential; the operator +/// names the one matching their deployment via the `auth_type` runtime-config +/// field. Modeled on `azure_identity`'s `specific_credential.rs` example. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub enum AzureCredentialKind { + /// Developer tools: Azure CLI (`az login`), then Azure Developer CLI + /// (`azd auth login`). Intended for local development; the default when + /// `auth_type` is omitted. + #[default] + DeveloperTools, + /// Managed identity (Azure VM, App Service, or AKS with managed identity). + /// + /// `client_id` optionally selects a *user-assigned* managed identity by its + /// client ID; when `None`, the SDK uses the *system-assigned* identity. + ManagedIdentity { client_id: Option }, + /// Workload identity (AKS federated token). + WorkloadIdentity, + /// 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 the environment carries over once `auth_type = "service_principal"` + /// is set. + ServicePrincipal, +} + +impl AzureCredentialKind { + /// Parses the `auth_type` runtime-config value into a credential kind. + /// + /// `None` (the field omitted) defaults to [`AzureCredentialKind::DeveloperTools`], + /// intended for local development. An unrecognized value is an error rather + /// than a silent fallback. + /// + /// `client_id` selects a user-assigned managed identity by its client ID. It + /// only applies to `managed_identity`; with any other `auth_type` it is + /// ignored. + pub fn from_auth_type(auth_type: Option<&str>, client_id: Option) -> Result { + // Case-insensitive, but the value must be one of the canonical + // snake_case names: a non-canonical form (e.g. a space or hyphen + // separator) is rejected rather than silently normalized, so the + // accepted set matches exactly what the docs and the error below list. + // `client_id` is consumed only by the `managed_identity` arm; for any + // other auth type it is simply dropped. + match auth_type.map(|s| s.to_lowercase()).as_deref() { + None => Ok(Self::default()), + Some("developer_tools") => Ok(Self::DeveloperTools), + Some("managed_identity") => Ok(Self::ManagedIdentity { client_id }), + Some("workload_identity") => Ok(Self::WorkloadIdentity), + Some("service_principal") => Ok(Self::ServicePrincipal), + Some(other) => anyhow::bail!( + "unknown Azure Cosmos `auth_type` {other:?}; expected one of \ + \"managed_identity\", \"workload_identity\", \"service_principal\", \ + or \"developer_tools\" (or set `key` for account-key auth)" + ), + } + } + + /// Constructs the corresponding `azure_identity` token credential. + /// + /// 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 at store construction, failing `spin up` on misconfiguration. + pub(crate) fn credential(&self) -> azure_core::Result> { + match self { + Self::DeveloperTools => Ok(azure_identity::DeveloperToolsCredential::new(None)?), + Self::ManagedIdentity { client_id } => { + // Pass options only when a user-assigned client ID was given; + // `None` keeps the SDK default of the system-assigned identity. + let options = + client_id + .as_ref() + .map(|id| azure_identity::ManagedIdentityCredentialOptions { + user_assigned_id: Some(azure_identity::UserAssignedId::ClientId( + id.clone(), + )), + ..Default::default() + }); + Ok(azure_identity::ManagedIdentityCredential::new(options)?) + } + Self::WorkloadIdentity => Ok(azure_identity::WorkloadIdentityCredential::new(None)?), + Self::ServicePrincipal => { + // 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 (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")?; + Ok(azure_identity::ClientSecretCredential::new( + &tenant_id, + client_id, + secret.into(), + None, + )?) + } + } + } +} + +/// Reads a required Service Principal environment variable, mapping a missing +/// value to a clear credential error. +fn service_principal_env(name: &str) -> azure_core::Result { + std::env::var(name).map_err(|_| { + azure_core::Error::with_message( + azure_core::error::ErrorKind::Credential, + format!( + "Azure Cosmos `service_principal` auth requires the `{name}` \ + environment variable to be set" + ), + ) + }) +} diff --git a/crates/key-value-azure/src/lib.rs b/crates/key-value-azure/src/lib.rs index 82195fc5ab..f59321c9e2 100644 --- a/crates/key-value-azure/src/lib.rs +++ b/crates/key-value-azure/src/lib.rs @@ -1,11 +1,14 @@ +mod auth; mod store; +use azure_data_cosmos::options::Region; use serde::Deserialize; use spin_factor_key_value::runtime_config::spin::MakeKeyValueStore; -pub use store::{ - KeyValueAzureCosmos, KeyValueAzureCosmosAuthOptions, KeyValueAzureCosmosRuntimeConfigOptions, +pub use auth::{ + AzureCredentialKind, KeyValueAzureCosmosAuthOptions, KeyValueAzureCosmosRuntimeConfigOptions, }; +pub use store::KeyValueAzureCosmos; /// A key-value store that uses Azure Cosmos as the backend. pub struct AzureKeyValueStore { @@ -15,7 +18,7 @@ pub struct AzureKeyValueStore { impl AzureKeyValueStore { /// Creates a new `AzureKeyValueStore`. /// - /// When `app_id` is provided, the store will a partition key of `$app_id/$store_name`, + /// When `app_id` is provided, the store will use a partition key of `$app_id/$store_name`, /// otherwise the partition key will be `id`. pub fn new(app_id: Option) -> Self { Self { app_id } @@ -32,8 +35,32 @@ pub struct AzureCosmosKeyValueRuntimeConfig { /// The Azure Cosmos DB database. database: String, /// The Azure Cosmos DB container where data is stored. - /// The CosmosDB container must be created with the default partition key, /id + /// The container's partition key path must be `/id` (the default) — or + /// `/store_id` if the store is constructed with an `app_id`. container: String, + + /// Optional. The Azure region the spin application is running in (or the + /// closest Azure region to it), used as the proximity-sorting anchor + /// for the Azure SDK's region selection. When omitted, defaults to + /// East US. + region: Option, + + /// Optional. When `key` is omitted, selects which Azure AD credential to + /// use: "managed_identity", "workload_identity", "service_principal", or + /// "developer_tools". When omitted, defaults to developer tools (Azure CLI + /// / azd), intended for local development. Ignored when `key` is set. + /// + /// "service_principal" reads `AZURE_TENANT_ID`, `AZURE_CLIENT_ID`, and + /// `AZURE_CLIENT_SECRET` from the environment. + /// + /// There is intentionally no automatic fallback between credential types; + /// name the one matching your deployment. + auth_type: Option, + + /// Optional. Only used with `auth_type = "managed_identity"`: the client ID + /// of a user-assigned managed identity to authenticate. When omitted, the + /// system-assigned identity is used. Ignored with any other `auth_type`. + client_id: Option, } impl MakeKeyValueStore for AzureKeyValueStore { @@ -51,13 +78,23 @@ impl MakeKeyValueStore for AzureKeyValueStore { Some(key) => KeyValueAzureCosmosAuthOptions::RuntimeConfigValues( KeyValueAzureCosmosRuntimeConfigOptions::new(key), ), - None => KeyValueAzureCosmosAuthOptions::Environmental, + None => { + KeyValueAzureCosmosAuthOptions::AadCredential(AzureCredentialKind::from_auth_type( + runtime_config.auth_type.as_deref(), + runtime_config.client_id, + )?) + } }; + let region = runtime_config + .region + .map(Region::from) + .unwrap_or(Region::EAST_US); KeyValueAzureCosmos::new( runtime_config.account, runtime_config.database, runtime_config.container, auth_options, + region, self.app_id.clone(), ) } diff --git a/crates/key-value-azure/src/store.rs b/crates/key-value-azure/src/store.rs index 51ff7eff9c..457705265d 100644 --- a/crates/key-value-azure/src/store.rs +++ b/crates/key-value-azure/src/store.rs @@ -1,20 +1,33 @@ -use anyhow::{Context, Result}; +use anyhow::Result; use async_trait::async_trait; +use azure_core::credentials::Secret; +use azure_core::http::Etag; +use azure_data_cosmos::models::{PatchInstructions, PatchOperation}; +use azure_data_cosmos::options::{ItemWriteOptions, Precondition, Region}; use azure_data_cosmos::{ - CosmosEntity, - prelude::{ - AuthorizationToken, CollectionClient, CosmosClient, CosmosClientBuilder, Operation, Query, - }, + AccountReference, ContainerClient, CosmosClient, FeedScope, Query, RoutingStrategy, }; -use futures::StreamExt; +use futures::TryStreamExt; +use itertools::Itertools; use serde::{Deserialize, Serialize}; use spin_factor_key_value::{ - Cas, Error, Store, StoreManager, SwapError, log_cas_error, log_error, log_error_v3, v3, + Cas, Error, Store, StoreManager, SwapError, log_error, log_error_v3, v3, }; +use std::collections::HashMap; use std::sync::{Arc, Mutex}; +use crate::auth::KeyValueAzureCosmosAuthOptions; + pub struct KeyValueAzureCosmos { - client: CollectionClient, + /// Parameters for initializing the Cosmos DB client + account_ref: AccountReference, + + database: String, + container: String, + region: Region, + + /// The Cosmos DB client + client: tokio::sync::OnceCell, /// An optional app id /// /// If provided, the store will handle multiple stores per container using a @@ -23,101 +36,63 @@ pub struct KeyValueAzureCosmos { app_id: Option, } -/// Azure Cosmos Key / Value runtime config literal options for authentication -#[derive(Clone, Debug)] -pub struct KeyValueAzureCosmosRuntimeConfigOptions { - key: String, -} - -impl KeyValueAzureCosmosRuntimeConfigOptions { - pub fn new(key: String) -> Self { - Self { key } - } -} - -/// Azure Cosmos Key / Value enumeration for the possible authentication options -#[derive(Clone, Debug)] -pub enum KeyValueAzureCosmosAuthOptions { - /// Runtime Config values indicates the account and key have been specified directly - RuntimeConfigValues(KeyValueAzureCosmosRuntimeConfigOptions), - /// Environmental indicates that the environment variables of the process should be used to - /// create the TokenCredential for the Cosmos client. This will use the Azure Rust SDK's - /// DefaultCredentialChain to derive the TokenCredential based on what environment variables - /// have been set. - /// - /// Service Principal with client secret: - /// - `AZURE_TENANT_ID`: ID of the service principal's Azure tenant. - /// - `AZURE_CLIENT_ID`: the service principal's client ID. - /// - `AZURE_CLIENT_SECRET`: one of the service principal's secrets. - /// - /// Service Principal with certificate: - /// - `AZURE_TENANT_ID`: ID of the service principal's Azure tenant. - /// - `AZURE_CLIENT_ID`: the service principal's client ID. - /// - `AZURE_CLIENT_CERTIFICATE_PATH`: path to a PEM or PKCS12 certificate file including the private key. - /// - `AZURE_CLIENT_CERTIFICATE_PASSWORD`: (optional) password for the certificate file. - /// - /// Workload Identity (Kubernetes, injected by the Workload Identity mutating webhook): - /// - `AZURE_TENANT_ID`: ID of the service principal's Azure tenant. - /// - `AZURE_CLIENT_ID`: the service principal's client ID. - /// - `AZURE_FEDERATED_TOKEN_FILE`: TokenFilePath is the path of a file containing a Kubernetes service account token. - /// - /// Managed Identity (User Assigned or System Assigned identities): - /// - `AZURE_CLIENT_ID`: (optional) if using a user assigned identity, this will be the client ID of the identity. - /// - /// Azure CLI: - /// - `AZURE_TENANT_ID`: (optional) use a specific tenant via the Azure CLI. - /// - /// Common across each: - /// - `AZURE_AUTHORITY_HOST`: (optional) the host for the identity provider. For example, for Azure public cloud the host defaults to "https://login.microsoftonline.com". - /// See also: https://github.com/Azure/azure-sdk-for-rust/blob/main/sdk/identity/README.md - Environmental, -} - impl KeyValueAzureCosmos { pub fn new( account: String, database: String, container: String, auth_options: KeyValueAzureCosmosAuthOptions, + region: Region, app_id: Option, ) -> Result { - let token = match auth_options { + 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) => { - AuthorizationToken::primary_key(config.key).map_err(log_error)? + AccountReference::with_authentication_key(endpoint, Secret::from(config.key)) } - KeyValueAzureCosmosAuthOptions::Environmental => { - AuthorizationToken::from_token_credential( - azure_identity::create_default_credential()?, - ) + KeyValueAzureCosmosAuthOptions::AadCredential(kind) => { + let credential = kind.credential().map_err(log_error)?; + AccountReference::with_credential(endpoint, credential) } }; - let cosmos_client = cosmos_client(account, token)?; - let database_client = cosmos_client.database_client(database); - let client = database_client.collection_client(container); - - Ok(Self { client, app_id }) - } -} -fn cosmos_client(account: impl Into, token: AuthorizationToken) -> Result { - if cfg!(feature = "connection-pooling") { - let client = reqwest::ClientBuilder::new() - .build() - .context("failed to build reqwest client")?; - let transport_options = azure_core::TransportOptions::new(std::sync::Arc::new(client)); - Ok(CosmosClientBuilder::new(account, token) - .transport(transport_options) - .build()) - } else { - Ok(CosmosClient::new(account, token)) + Ok(Self { + account_ref, + database, + container, + region, + client: tokio::sync::OnceCell::new(), + app_id, + }) } } #[async_trait] impl StoreManager for KeyValueAzureCosmos { async fn get(&self, name: &str) -> Result, Error> { + let client = self + .client + .get_or_try_init(|| async { + 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 + .map_err(log_error); + }) + .await? + .clone(); Ok(Arc::new(AzureCosmosStore { - client: self.client.clone(), + client, store_id: self.app_id.as_ref().map(|i| format!("{i}/{name}")), })) } @@ -127,17 +102,16 @@ impl StoreManager for KeyValueAzureCosmos { } fn summary(&self, _store_name: &str) -> Option { - let database = self.client.database_client().database_name(); - let collection = self.client.collection_name(); Some(format!( - "Azure CosmosDB database: {database}, collection: {collection}" + "Azure CosmosDB database: {}, container: {}", + self.database, self.container )) } } #[derive(Clone)] struct AzureCosmosStore { - client: CollectionClient, + client: ContainerClient, /// An optional store id to use as a partition key for all operations. /// /// If the store ID is not set, the store will use `/id` (the row key) as @@ -151,18 +125,18 @@ struct AzureCosmosStore { #[async_trait] impl Store for AzureCosmosStore { async fn get(&self, key: &str, max_result_bytes: usize) -> Result>, Error> { - let pair = self.get_entity::(key).await?; - let value = pair.map(|p| p.value); - - // Currently there's no way to stream a single query result using the - // `azure_data_cosmos` crate without buffering, so the damage (in terms - // of host memory usage) is already done, but we can still enforce the - // limit: - if std::mem::size_of::>>() + value.as_ref().map(|v| v.len()).unwrap_or(0) + let partition_key = partition_key(self.store_id.as_deref(), key); + let value = match self.client.read_item(partition_key, key, None).await { + Ok(response) => Some(response.into_model::().map_err(log_error)?.value), + Err(e) if e.status().is_not_found() => None, + Err(e) => return Err(log_error(e)), + }; + + if std::mem::size_of::>>() + value.as_ref().map(Vec::len).unwrap_or(0) > max_result_bytes { Err(Error::Other(format!( - "query result exceeds limit of {max_result_bytes} bytes" + "read result exceeds limit of {max_result_bytes} bytes" ))) } else { Ok(value) @@ -184,44 +158,63 @@ impl Store for AzureCosmosStore { value: value.to_vec(), store_id: self.store_id.clone(), }; + + let partition_key = partition_key(self.store_id.as_deref(), key); self.client - .create_document(pair) - .is_upsert(true) + .upsert_item(partition_key, key, pair, None) .await .map_err(log_error)?; Ok(()) } async fn delete(&self, key: &str) -> Result<(), Error> { - let document_client = self - .client - .document_client(key, &self.store_id.clone().unwrap_or(key.to_string())) - .map_err(log_error)?; - if let Err(e) = document_client.delete_document().await - && e.as_http_error().map(|e| e.status() != 404).unwrap_or(true) - { - return Err(log_error(e)); + let partition_key = partition_key(self.store_id.as_deref(), key); + + match self.client.delete_item(partition_key, key, None).await { + Ok(_) => Ok(()), + Err(e) if e.status().is_not_found() => Ok(()), + Err(e) => Err(log_error(e)), } - Ok(()) } async fn exists(&self, key: &str) -> Result { let mut stream = self .client - .query_documents(Query::new(self.get_id_query(key))) - .query_cross_partition(true) - .max_item_count(1) - .into_stream::(); - - match stream.next().await { - Some(Ok(res)) => Ok(!res.results.is_empty()), - Some(Err(e)) => Err(log_error(e)), - None => Ok(false), - } + .query_items::( + Query::from(self.get_id_query(key)), + FeedScope::partition(partition_key(self.store_id.as_deref(), key)), + None, + ) + .await + .map_err(log_error)?; + Ok(stream.try_next().await.map_err(log_error)?.is_some()) } async fn get_keys(&self, max_result_bytes: usize) -> Result, Error> { - self.get_keys(max_result_bytes).await + let mut stream = self + .client + .query_items::( + Query::from(self.get_keys_query()), + FeedScope::full_container(), + None, + ) + .await + .map_err(log_error)?; + + let mut result = Vec::new(); + let mut byte_count = std::mem::size_of::>(); + + while let Some(key) = stream.try_next().await.map_err(log_error)? { + byte_count += std::mem::size_of::() + key.id.len(); + if byte_count > max_result_bytes { + return Err(Error::Other(format!( + "query result exceeds limit of {max_result_bytes} bytes" + ))); + } + result.push(key.id); + } + + Ok(result) } async fn get_keys_async( @@ -234,25 +227,25 @@ impl Store for AzureCosmosStore { let (keys_tx, keys_rx) = tokio::sync::mpsc::channel(4); let (err_tx, err_rx) = tokio::sync::oneshot::channel(); - let query = self - .client - .query_documents(Query::new(self.get_keys_query())) - .query_cross_partition(true); + let client = self.client.clone(); + let query = self.get_keys_query(); let the_work = async move { - let mut stream = query.into_stream::(); - while let Some(resp) = stream.next().await { - let resp = resp.map_err(log_error_v3)?; - - if resp.results.iter().map(|(k, _)| k.id.len()).sum::() > max_result_bytes { + let mut stream = client + .query_items::(Query::from(query), FeedScope::full_container(), None) + .await + .map_err(log_error_v3)?; + + let mut byte_count = std::mem::size_of::>(); + while let Some(key) = stream.try_next().await.map_err(log_error_v3)? { + byte_count += std::mem::size_of::() + key.id.len(); + if byte_count > max_result_bytes { return Err(v3::Error::Other(format!( "query exceeds limit of {max_result_bytes} bytes" ))); } - for (key, _) in resp.results { - keys_tx.send(key.id).await.map_err(log_error_v3)?; - } + keys_tx.send(key.id).await.map_err(log_error_v3)?; } Ok(()) }; @@ -269,36 +262,41 @@ impl Store for AzureCosmosStore { keys: Vec, max_result_bytes: usize, ) -> Result>)>, Error> { - let stmt = Query::new(self.get_in_query(keys)); - let query = self + // Deduplicate first: a repeated key would otherwise be counted once + // against `max_result_bytes` but returned once per occurrence. + let keys: Vec = keys.into_iter().unique().collect(); + + let mut stream = self .client - .query_documents(stmt) - .query_cross_partition(true); + .query_items::( + Query::from(self.get_in_query(&keys)), + FeedScope::full_container(), + None, + ) + .await + .map_err(log_error)?; - let mut res = Vec::new(); - let mut stream = query.into_stream::(); + let mut found = HashMap::new(); let mut byte_count = std::mem::size_of::>)>>(); - while let Some(resp) = stream.next().await { - let resp = resp.map_err(log_error)?.results; - byte_count += resp - .iter() - .map(|(pair, _)| { - std::mem::size_of::<(String, Option>)>() - + pair.id.len() - + pair.value.len() - }) - .sum::(); + while let Some(pair) = stream.try_next().await.map_err(log_error)? { + byte_count += + std::mem::size_of::<(String, Option>)>() + pair.id.len() + pair.value.len(); + if byte_count > max_result_bytes { return Err(Error::Other(format!( "query result exceeds limit of {max_result_bytes} bytes" ))); } - res.extend( - resp.into_iter() - .map(|(pair, _)| (pair.id, Some(pair.value))), - ); + found.insert(pair.id, pair.value); } - Ok(res) + + Ok(keys + .into_iter() + .map(|key| { + let value = found.remove(&key); + (key, value) + }) + .collect()) } async fn set_many(&self, key_values: Vec<(String, Vec)>) -> Result<(), Error> { @@ -320,53 +318,34 @@ impl Store for AzureCosmosStore { /// The initial value for the item must be set through this interface, as this sets the /// number value if it does not exist. If the value was previously set using /// the `set` interface, this will fail due to a type mismatch. - // TODO: The function should parse the new value from the return response - // rather than sending an additional new request. However, the current SDK - // version does not support this. async fn increment(&self, key: String, delta: i64) -> Result { - let operations = vec![Operation::incr("/value", delta).map_err(log_error)?]; + let patch = + PatchInstructions::default().with_operation(PatchOperation::increment("/value", delta)); + let partition_key = partition_key(self.store_id.as_deref(), &key); + match self .client - .document_client(&key, &self.store_id.clone().unwrap_or(key.to_string())) - .map_err(log_error)? - .patch_document(operations) + .patch_item(partition_key.clone(), &key, patch, None) .await { - Err(e) => { - if e.as_http_error() - .map(|e| e.status() == 404) - .unwrap_or(false) + Err(e) if e.status().is_not_found() => { + let counter = Counter { + id: key.clone(), + value: delta, + store_id: self.store_id.clone(), + }; + match self + .client + .create_item(partition_key, &key, counter, None) + .await { - let counter = Counter { - id: key.clone(), - value: delta, - store_id: self.store_id.clone(), - }; - if let Err(e) = self.client.create_document(counter).is_upsert(false).await { - if e.as_http_error() - .map(|e| e.status()) - .unwrap_or(azure_core::StatusCode::Continue) - == 409 - { - // Conflict trying to create counter, retry increment - self.increment(key, delta).await?; - } else { - return Err(log_error(e)); - } - } - Ok(delta) - } else { - Err(log_error(e)) + Ok(_) => Ok(delta), + Err(e) if e.status().is_conflict() => self.increment(key, delta).await, + Err(e) => Err(log_error(e)), } } - Ok(_) => self - .get_entity::(key.as_ref()) - .await? - .map(|c| c.value) - .ok_or(Error::Other( - "increment returned an empty value after patching, which indicates a bug" - .to_string(), - )), + Err(e) => Err(log_error(e)), + Ok(response) => Ok(response.into_model::().map_err(log_error)?.value), } } @@ -387,63 +366,33 @@ impl Store for AzureCosmosStore { struct CompareAndSwap { key: String, - client: CollectionClient, + client: ContainerClient, bucket_rep: u32, - etag: Mutex>, + etag: Mutex>, store_id: Option, } -impl CompareAndSwap { - fn get_query(&self) -> String { - let mut query = format!("SELECT * FROM c WHERE c.id='{}'", self.key); - self.append_store_id(&mut query, true); - query - } - - fn append_store_id(&self, query: &mut String, condition_already_exists: bool) { - append_store_id_condition(query, self.store_id.as_deref(), condition_already_exists); - } -} - #[async_trait] impl Cas for CompareAndSwap { /// `current` will fetch the current value for the key and store the etag for the record. The /// etag will be used to perform and optimistic concurrency update using the `if-match` header. async fn current(&self, max_result_bytes: usize) -> Result>, Error> { - let mut stream = self - .client - .query_documents(Query::new(self.get_query())) - .query_cross_partition(true) - .max_item_count(1) - .into_stream::(); - - let current_value: Option<(Vec, Option)> = match stream.next().await { - Some(r) => { - let r = r.map_err(log_error)?; - match r.results.first() { - Some((item, Some(attr))) => { - Some((item.clone().value, Some(attr.etag().to_string()))) - } - Some((item, None)) => Some((item.clone().value, None)), - _ => None, - } - } - None => None, - }; + let partition_key = partition_key(self.store_id.as_deref(), &self.key); + let result = self.client.read_item(partition_key, &self.key, None).await; - let value = match current_value { - Some((value, etag)) => { - self.etag.lock().unwrap().clone_from(&etag); - Some(value) + let value = match result { + Ok(response) => { + *self.etag.lock().unwrap() = response.headers().etag().cloned(); + Some(response.into_model::().map_err(log_error)?.value) } - None => None, + Err(e) if e.status().is_not_found() => { + *self.etag.lock().unwrap() = None; + None + } + Err(e) => return Err(log_error(e)), }; - // Currently there's no way to stream a single query result using the - // `azure_data_cosmos` crate without buffering, so the damage (in terms - // of host memory usage) is already done, but we can still enforce the - // limit: - if std::mem::size_of::>>() + value.as_ref().map(|v| v.len()).unwrap_or(0) + if std::mem::size_of::>>() + value.as_ref().map(Vec::len).unwrap_or(0) > max_result_bytes { Err(Error::Other(format!( @@ -463,31 +412,37 @@ impl Cas for CompareAndSwap { store_id: self.store_id.clone(), }; - let doc_client = self - .client - .document_client(&self.key, &pair.partition_key()) - .map_err(log_cas_error)?; + let partition_key = partition_key(self.store_id.as_deref(), &self.key); + let etag = self.etag.lock().unwrap().clone(); + let had_etag = etag.is_some(); - let etag_value = self.etag.lock().unwrap().clone(); - match etag_value { + let response = match etag { Some(etag) => { - // attempt to replace the document if the etag matches - doc_client - .replace_document(pair) - .if_match_condition(azure_core::request_options::IfMatchCondition::Match(etag)) + let opts = + ItemWriteOptions::default().with_precondition(Precondition::IfMatch(etag)); + self.client + .replace_item(partition_key, &self.key, &pair, Some(opts)) .await - .map_err(|e| SwapError::CasFailed(format!("{e:?}"))) - .map(drop) } None => { - // if we have no etag, then we assume the document does not yet exist and must insert; no upserts. self.client - .create_document(pair) + .create_item(partition_key, &self.key, &pair, None) .await - .map_err(|e| SwapError::CasFailed(format!("{e:?}"))) - .map(drop) } - } + }; + + response.map(drop).map_err(|e| { + let msg = format!("{e:?}"); + let status = e.status(); + if status.is_precondition_failed() + || status.is_conflict() + || (had_etag && status.is_not_found()) + { + SwapError::CasFailed(msg) + } else { + SwapError::Other(msg) + } + }) } async fn bucket_rep(&self) -> u32 { @@ -500,87 +455,31 @@ impl Cas for CompareAndSwap { } impl AzureCosmosStore { - async fn get_entity(&self, key: &str) -> Result, Error> - where - F: CosmosEntity + Send + Sync + serde::de::DeserializeOwned + Clone, - { - let query = self - .client - .query_documents(Query::new(self.get_query(key))) - .query_cross_partition(true) - .max_item_count(1); - - // There can be no duplicated keys, so we create the stream and only take the first result. - let mut stream = query.into_stream::(); - let Some(res) = stream.next().await else { - return Ok(None); - }; - Ok(res - .map_err(log_error)? - .results - .first() - .map(|(p, _)| p.clone())) - } - - async fn get_keys(&self, max_result_bytes: usize) -> Result, Error> { - let query = self - .client - .query_documents(Query::new(self.get_keys_query())) - .query_cross_partition(true); - let mut res = Vec::new(); - - let mut stream = query.into_stream::(); - let mut byte_count = std::mem::size_of::>(); - while let Some(resp) = stream.next().await { - let resp = resp.map_err(log_error)?.results; - byte_count += resp - .iter() - .map(|(key, _)| std::mem::size_of::() + key.id.len()) - .sum::(); - if byte_count > max_result_bytes { - return Err(Error::Other(format!( - "query result exceeds limit of {max_result_bytes} bytes" - ))); - } - res.extend(resp.into_iter().map(|(key, _)| key.id)); - } - - Ok(res) - } - - fn get_query(&self, key: &str) -> String { - let mut query = format!("SELECT * FROM c WHERE c.id='{key}'"); - self.append_store_id(&mut query, true); - query - } - fn get_id_query(&self, key: &str) -> String { let mut query = format!("SELECT c.id, c.store_id FROM c WHERE c.id='{key}'"); - self.append_store_id(&mut query, true); + append_store_id_condition(&mut query, self.store_id.as_deref(), true); query } fn get_keys_query(&self) -> String { let mut query = "SELECT c.id, c.store_id FROM c".to_owned(); - self.append_store_id(&mut query, false); + append_store_id_condition(&mut query, self.store_id.as_deref(), false); query } - fn get_in_query(&self, keys: Vec) -> String { - let in_clause: String = keys - .into_iter() - .map(|k| format!("'{k}'")) - .collect::>() - .join(", "); + fn get_in_query(&self, keys: &[String]) -> String { + let in_clause = keys.iter().map(|k| format!("'{k}'")).join(", "); let mut query = format!("SELECT * FROM c WHERE c.id IN ({in_clause})"); - self.append_store_id(&mut query, true); + append_store_id_condition(&mut query, self.store_id.as_deref(), true); query } +} - fn append_store_id(&self, query: &mut String, condition_already_exists: bool) { - append_store_id_condition(query, self.store_id.as_deref(), condition_already_exists); - } +fn partition_key(store_id: Option<&str>, key: &str) -> String { + store_id + .map(str::to_string) + .unwrap_or_else(|| key.to_string()) } /// Appends an option store id condition to the query. @@ -601,7 +500,7 @@ fn append_store_id_condition( } } -// Pair structure for key value operations +/// Pair structure for key value operations #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Pair { pub id: String, @@ -609,16 +508,7 @@ pub struct Pair { #[serde(skip_serializing_if = "Option::is_none")] pub store_id: Option, } - -impl CosmosEntity for Pair { - type Entity = String; - - fn partition_key(&self) -> Self::Entity { - self.store_id.clone().unwrap_or_else(|| self.id.clone()) - } -} - -// Counter structure for increment operations +/// Counter structure for increment operations #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Counter { pub id: String, @@ -627,26 +517,10 @@ pub struct Counter { pub store_id: Option, } -impl CosmosEntity for Counter { - type Entity = String; - - fn partition_key(&self) -> Self::Entity { - self.store_id.clone().unwrap_or_else(|| self.id.clone()) - } -} - -// Key structure for operations with generic value types +/// Key structure for operations with generic value types #[derive(Serialize, Deserialize, Clone, Debug)] pub struct Key { pub id: String, #[serde(skip_serializing_if = "Option::is_none")] pub store_id: Option, } - -impl CosmosEntity for Key { - type Entity = String; - - fn partition_key(&self) -> Self::Entity { - self.store_id.clone().unwrap_or_else(|| self.id.clone()) - } -}