Azure Cosmos SDK upgrade for kv store - #3559
Conversation
c54b8f4 to
0ef22f4
Compare
0ef22f4 to
047b4d5
Compare
There was a problem hiding this comment.
Pull request overview
Note
Copilot was unable to run its full agentic suite in this review.
Modernizes the Azure Cosmos-backed key-value store to use newer Azure SDK APIs and adds explicit Azure AD authentication options plus region-aware routing.
Changes:
- Migrates Cosmos DB interactions from older
CollectionClient/document APIs toContainerClientitem APIs with lazy client initialization. - Adds explicit Azure AD credential selection (
auth_type+ optionalclient_id) and optional region configuration for proximity-based routing. - Updates Azure SDK dependency versions and removes the old connection-pooling feature wiring.
Reviewed changes
Copilot reviewed 3 out of 4 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
| crates/key-value-azure/src/store.rs | Refactors Cosmos client initialization + CRUD/CAS operations to new SDK APIs; adds auth parsing/utilities and unit tests. |
| crates/key-value-azure/src/lib.rs | Extends runtime config with region, auth_type, and client_id; wires them into store construction and exports AzureCredentialKind. |
| crates/key-value-azure/Cargo.toml | Updates Azure SDK dependencies to newer major versions; removes old connection-pooling feature block. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
35982d4 to
198e1a1
Compare
64546c4 to
b5f6d83
Compare
b5f6d83 to
4e6e1ba
Compare
cc3f18e to
964720e
Compare
ChihweiLHBird
changed the title
a8d1b60 to
13d2cea
Compare
Signed-off-by: Zhiwei Liang <zhiwei.liang@zliang.me>
13d2cea to
7da00eb
Compare
|
@ChihweiLHBird I rebased this to address merge conflicts, and squashed the commits (in case I ever have to rebase it again). Sorry for letting it linger! |
|
No problem, and thank you for the help, @itowlson! |
|
@ChihweiLHBird Would you be open to extending this PR to cover our other Azure dependency, the KeyVault stuff in the |
|
Hi @itowlson, if variables-azure crate doesn't depend on this PR, I think we can work on it in a follow up PR? It may be easier to review that way. |
|
I'm looking at the runtime config stuff and I'm not clear if this is a breaking change. It sounds like ye olde key-style configuration is the same but what used to be "environmental" configuration may be different - you need to provide the It's not a problem if it is breaking, it just means we'll need to hold it for the next major release. (And it makes sense, because from what I've seen of the Azure identity client, it is much stricter about being specific about your auth type than before.) |
|
@itowlson Unfortunately yes, it's a breaking change by the Azure SDK redesign. And only users with Developer tools: Azure CLI ( |
|
@ChihweiLHBird no worries. I've added a label so that we're aware and don't merge it until we're okay with breakage. It would be super useful for docs purposes if you could jot down a quick note on what the runtime config TOML will look like once this is merged - your code is clearer than what we had before, but I'm still gonna need some help! Thanks! |
itowlson
left a comment
itowlson
left a comment
There was a problem hiding this comment.
I'm no expert on this stuff but it looks good to me - thanks for slogging through it all. It looks rather nicer than what we had before (although I guess some of that is down to good work on the SDK), and I feel much more comfortable with the auth stuff. I had a couple of comments / questions but nothing blocking!
| /// 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] |
There was a problem hiding this comment.
Do you think we should force an explicit choice rather than defaulting? (This is not a suggestion: it's me not knowing what Azure tooling normally does.)
| 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\", \ |
There was a problem hiding this comment.
Might be more readable to do this with a r#"..."# string
| app_id: Option<String>, | ||
| } | ||
|
|
||
| /// Azure Cosmos Key / Value runtime config literal options for authentication |
There was a problem hiding this comment.
Good choice on breaking all this out into the auth module!
| && 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); |
There was a problem hiding this comment.
consider self.partition_key(key) since we seem to be writing this out a lot!
|
|
||
| async fn get_keys(&self, max_result_bytes: usize) -> Result<Vec<String>, Error> { | ||
| self.get_keys(max_result_bytes).await | ||
| let mut stream = self |
There was a problem hiding this comment.
This looks like a lot of it is very similar to get_keys_async - could we have it sit on top of that? Or pull out the common code?
Closes #3021.
Summary
Migrates
spin-key-value-azurefrom the legacy Azure Cosmos SDK to the currentazure_data_cosmos 0.37(driver0.6.0) /azure_core 1.0/azure_identity 1.0stack. Reworks client construction and the item/query APIs, keeps account-key auth, and replaces the old ambient-auth path with explicit Azure AD credential selection.What Changed
AccountReference, async container-client init, and SDK-managed routing (RoutingStrategy::ProximityTo); moved CRUD/batch/increment/CAS to the new APIs (getis now aread_itempoint read) and dropped the old custom connection pooling andreqwesttransport setup.default-features = false+key_auth/native_tls/hmac_rustso the Azure stack uses OpenSSL instead of pulling reqwest's aws-lc-rs rustls provider alongside Spin's workspace-pinnedring.auth_type, including user-assigned managed identity byclient_id; account-key auth is unchanged.regionruntime config (defaults toEast US) used as the routing proximity anchor, and account validation atspin up.currentandswap) — to retryable CAS failures, and everything else to hard errors.get-manynow returns(key, none)for missing keys per thewasi:keyvaluebatch contract instead of dropping them.Runtime Config
Existing account-key configuration works unchanged. When
keyis omitted,auth_typeselects the AAD credential:developer_tools(default),managed_identity(optionally withclient_id),workload_identity, orservice_principal(readsAZURE_TENANT_ID/AZURE_CLIENT_ID/AZURE_CLIENT_SECRET). There is no fallback between types. Optionally setregionto bias routing toward the closest replica.