spin_key_value_azure/
store.rs

1use anyhow::{Context, Result};
2use async_trait::async_trait;
3use azure_data_cosmos::{
4    CosmosEntity,
5    prelude::{
6        AuthorizationToken, CollectionClient, CosmosClient, CosmosClientBuilder, Operation, Param,
7        Query,
8    },
9};
10use futures::StreamExt;
11use serde::{Deserialize, Serialize};
12use spin_factor_key_value::{
13    Cas, Error, Store, StoreManager, SwapError, log_cas_error, log_error, log_error_v3, v3,
14};
15use std::sync::{Arc, Mutex};
16
17pub struct KeyValueAzureCosmos {
18    client: CollectionClient,
19    /// An optional app id
20    ///
21    /// If provided, the store will handle multiple stores per container using a
22    /// partition key of `/$app_id/$store_name`, otherwise there will be one container
23    /// per store, and the partition key will be `/id`.
24    app_id: Option<String>,
25}
26
27/// Azure Cosmos Key / Value runtime config literal options for authentication
28#[derive(Clone, Debug)]
29pub struct KeyValueAzureCosmosRuntimeConfigOptions {
30    key: String,
31}
32
33impl KeyValueAzureCosmosRuntimeConfigOptions {
34    pub fn new(key: String) -> Self {
35        Self { key }
36    }
37}
38
39/// Azure Cosmos Key / Value enumeration for the possible authentication options
40#[derive(Clone, Debug)]
41pub enum KeyValueAzureCosmosAuthOptions {
42    /// Runtime Config values indicates the account and key have been specified directly
43    RuntimeConfigValues(KeyValueAzureCosmosRuntimeConfigOptions),
44    /// Environmental indicates that the environment variables of the process should be used to
45    /// create the TokenCredential for the Cosmos client. This will use the Azure Rust SDK's
46    /// DefaultCredentialChain to derive the TokenCredential based on what environment variables
47    /// have been set.
48    ///
49    /// Service Principal with client secret:
50    /// - `AZURE_TENANT_ID`: ID of the service principal's Azure tenant.
51    /// - `AZURE_CLIENT_ID`: the service principal's client ID.
52    /// - `AZURE_CLIENT_SECRET`: one of the service principal's secrets.
53    ///
54    /// Service Principal with certificate:
55    /// - `AZURE_TENANT_ID`: ID of the service principal's Azure tenant.
56    /// - `AZURE_CLIENT_ID`: the service principal's client ID.
57    /// - `AZURE_CLIENT_CERTIFICATE_PATH`: path to a PEM or PKCS12 certificate file including the private key.
58    /// - `AZURE_CLIENT_CERTIFICATE_PASSWORD`: (optional) password for the certificate file.
59    ///
60    /// Workload Identity (Kubernetes, injected by the Workload Identity mutating webhook):
61    /// - `AZURE_TENANT_ID`: ID of the service principal's Azure tenant.
62    /// - `AZURE_CLIENT_ID`: the service principal's client ID.
63    /// - `AZURE_FEDERATED_TOKEN_FILE`: TokenFilePath is the path of a file containing a Kubernetes service account token.
64    ///
65    /// Managed Identity (User Assigned or System Assigned identities):
66    /// - `AZURE_CLIENT_ID`: (optional) if using a user assigned identity, this will be the client ID of the identity.
67    ///
68    /// Azure CLI:
69    /// - `AZURE_TENANT_ID`: (optional) use a specific tenant via the Azure CLI.
70    ///
71    /// Common across each:
72    /// - `AZURE_AUTHORITY_HOST`: (optional) the host for the identity provider. For example, for Azure public cloud the host defaults to "https://login.microsoftonline.com".
73    ///   See also: https://github.com/Azure/azure-sdk-for-rust/blob/main/sdk/identity/README.md
74    Environmental,
75}
76
77impl KeyValueAzureCosmos {
78    pub fn new(
79        account: String,
80        database: String,
81        container: String,
82        auth_options: KeyValueAzureCosmosAuthOptions,
83        app_id: Option<String>,
84    ) -> Result<Self> {
85        let token = match auth_options {
86            KeyValueAzureCosmosAuthOptions::RuntimeConfigValues(config) => {
87                AuthorizationToken::primary_key(config.key).map_err(log_error)?
88            }
89            KeyValueAzureCosmosAuthOptions::Environmental => {
90                AuthorizationToken::from_token_credential(
91                    azure_identity::create_default_credential()?,
92                )
93            }
94        };
95        let cosmos_client = cosmos_client(account, token)?;
96        let database_client = cosmos_client.database_client(database);
97        let client = database_client.collection_client(container);
98
99        Ok(Self { client, app_id })
100    }
101}
102
103fn cosmos_client(account: impl Into<String>, token: AuthorizationToken) -> Result<CosmosClient> {
104    if cfg!(feature = "connection-pooling") {
105        let client = reqwest::ClientBuilder::new()
106            .build()
107            .context("failed to build reqwest client")?;
108        let transport_options = azure_core::TransportOptions::new(std::sync::Arc::new(client));
109        Ok(CosmosClientBuilder::new(account, token)
110            .transport(transport_options)
111            .build())
112    } else {
113        Ok(CosmosClient::new(account, token))
114    }
115}
116
117#[async_trait]
118impl StoreManager for KeyValueAzureCosmos {
119    async fn get(&self, name: &str) -> Result<Arc<dyn Store>, Error> {
120        Ok(Arc::new(AzureCosmosStore {
121            client: self.client.clone(),
122            store_id: self.app_id.as_ref().map(|i| format!("{i}/{name}")),
123        }))
124    }
125
126    fn is_defined(&self, _store_name: &str) -> bool {
127        true
128    }
129
130    fn summary(&self, _store_name: &str) -> Option<String> {
131        let database = self.client.database_client().database_name();
132        let collection = self.client.collection_name();
133        Some(format!(
134            "Azure CosmosDB database: {database}, collection: {collection}"
135        ))
136    }
137}
138
139#[derive(Clone)]
140struct AzureCosmosStore {
141    client: CollectionClient,
142    /// An optional store id to use as a partition key for all operations.
143    ///
144    /// If the store ID is not set, the store will use `/id` (the row key) as
145    /// the partition key. For example, if `store.set("my_key", "my_value")` is
146    /// called, the partition key will be `my_key` if the store ID is set to
147    /// `None`. If the store ID is set to `Some("myappid/default"), the
148    /// partition key will be `myappid/default`.
149    store_id: Option<String>,
150}
151
152#[async_trait]
153impl Store for AzureCosmosStore {
154    async fn get(&self, key: &str, max_result_bytes: usize) -> Result<Option<Vec<u8>>, Error> {
155        let pair = self.get_entity::<Pair>(key).await?;
156        let value = pair.map(|p| p.value);
157
158        // Currently there's no way to stream a single query result using the
159        // `azure_data_cosmos` crate without buffering, so the damage (in terms
160        // of host memory usage) is already done, but we can still enforce the
161        // limit:
162        if std::mem::size_of::<Option<Vec<u8>>>() + value.as_ref().map(|v| v.len()).unwrap_or(0)
163            > max_result_bytes
164        {
165            Err(Error::Other(format!(
166                "query result exceeds limit of {max_result_bytes} bytes"
167            )))
168        } else {
169            Ok(value)
170        }
171    }
172
173    async fn set(&self, key: &str, value: &[u8]) -> Result<(), Error> {
174        let illegal_chars = ['/', '\\', '?', '#'];
175
176        if key.contains(|c| illegal_chars.contains(&c)) {
177            return Err(Error::Other(format!(
178                "Key contains an illegal character. Keys must not include any of: {}",
179                illegal_chars.iter().collect::<String>()
180            )));
181        }
182
183        let pair = Pair {
184            id: key.to_string(),
185            value: value.to_vec(),
186            store_id: self.store_id.clone(),
187        };
188        self.client
189            .create_document(pair)
190            .is_upsert(true)
191            .await
192            .map_err(log_error)?;
193        Ok(())
194    }
195
196    async fn delete(&self, key: &str) -> Result<(), Error> {
197        let document_client = self
198            .client
199            .document_client(key, &self.store_id.clone().unwrap_or(key.to_string()))
200            .map_err(log_error)?;
201        if let Err(e) = document_client.delete_document().await
202            && e.as_http_error().map(|e| e.status() != 404).unwrap_or(true)
203        {
204            return Err(log_error(e));
205        }
206        Ok(())
207    }
208
209    async fn exists(&self, key: &str) -> Result<bool, Error> {
210        let mut stream = self
211            .client
212            .query_documents(get_id_query(key, self.store_id.as_deref()))
213            .query_cross_partition(true)
214            .max_item_count(1)
215            .into_stream::<Key>();
216
217        match stream.next().await {
218            Some(Ok(res)) => Ok(!res.results.is_empty()),
219            Some(Err(e)) => Err(log_error(e)),
220            None => Ok(false),
221        }
222    }
223
224    async fn get_keys(&self, max_result_bytes: usize) -> Result<Vec<String>, Error> {
225        self.get_keys(max_result_bytes).await
226    }
227
228    async fn get_keys_async(
229        &self,
230        max_result_bytes: usize,
231    ) -> (
232        tokio::sync::mpsc::Receiver<String>,
233        tokio::sync::oneshot::Receiver<Result<(), v3::Error>>,
234    ) {
235        let (keys_tx, keys_rx) = tokio::sync::mpsc::channel(4);
236        let (err_tx, err_rx) = tokio::sync::oneshot::channel();
237
238        let query = self
239            .client
240            .query_documents(get_keys_query(self.store_id.as_deref()))
241            .query_cross_partition(true);
242
243        let the_work = async move {
244            let mut stream = query.into_stream::<Key>();
245            while let Some(resp) = stream.next().await {
246                let resp = resp.map_err(log_error_v3)?;
247
248                if resp.results.iter().map(|(k, _)| k.id.len()).sum::<usize>() > max_result_bytes {
249                    return Err(v3::Error::Other(format!(
250                        "query exceeds limit of {max_result_bytes} bytes"
251                    )));
252                }
253
254                for (key, _) in resp.results {
255                    keys_tx.send(key.id).await.map_err(log_error_v3)?;
256                }
257            }
258            Ok(())
259        };
260        tokio::spawn(async move {
261            let res = the_work.await;
262            _ = err_tx.send(res);
263        });
264
265        (keys_rx, err_rx)
266    }
267
268    async fn get_many(
269        &self,
270        keys: Vec<String>,
271        max_result_bytes: usize,
272    ) -> Result<Vec<(String, Option<Vec<u8>>)>, Error> {
273        let stmt = get_in_query(keys, self.store_id.as_deref());
274        let query = self
275            .client
276            .query_documents(stmt)
277            .query_cross_partition(true);
278
279        let mut res = Vec::new();
280        let mut stream = query.into_stream::<Pair>();
281        let mut byte_count = std::mem::size_of::<Vec<(String, Option<Vec<u8>>)>>();
282        while let Some(resp) = stream.next().await {
283            let resp = resp.map_err(log_error)?.results;
284            byte_count += resp
285                .iter()
286                .map(|(pair, _)| {
287                    std::mem::size_of::<(String, Option<Vec<u8>>)>()
288                        + pair.id.len()
289                        + pair.value.len()
290                })
291                .sum::<usize>();
292            if byte_count > max_result_bytes {
293                return Err(Error::Other(format!(
294                    "query result exceeds limit of {max_result_bytes} bytes"
295                )));
296            }
297            res.extend(
298                resp.into_iter()
299                    .map(|(pair, _)| (pair.id, Some(pair.value))),
300            );
301        }
302        Ok(res)
303    }
304
305    async fn set_many(&self, key_values: Vec<(String, Vec<u8>)>) -> Result<(), Error> {
306        for (key, value) in key_values {
307            self.set(key.as_ref(), &value).await?
308        }
309        Ok(())
310    }
311
312    async fn delete_many(&self, keys: Vec<String>) -> Result<(), Error> {
313        for key in keys {
314            self.delete(key.as_ref()).await?
315        }
316        Ok(())
317    }
318
319    /// Increments a numerical value.
320    ///
321    /// The initial value for the item must be set through this interface, as this sets the
322    /// number value if it does not exist. If the value was previously set using
323    /// the `set` interface, this will fail due to a type mismatch.
324    // TODO: The function should parse the new value from the return response
325    // rather than sending an additional new request. However, the current SDK
326    // version does not support this.
327    async fn increment(&self, key: String, delta: i64) -> Result<i64, Error> {
328        let operations = vec![Operation::incr("/value", delta).map_err(log_error)?];
329        match self
330            .client
331            .document_client(&key, &self.store_id.clone().unwrap_or(key.to_string()))
332            .map_err(log_error)?
333            .patch_document(operations)
334            .await
335        {
336            Err(e) => {
337                if e.as_http_error()
338                    .map(|e| e.status() == 404)
339                    .unwrap_or(false)
340                {
341                    let counter = Counter {
342                        id: key.clone(),
343                        value: delta,
344                        store_id: self.store_id.clone(),
345                    };
346                    if let Err(e) = self.client.create_document(counter).is_upsert(false).await {
347                        if e.as_http_error()
348                            .map(|e| e.status())
349                            .unwrap_or(azure_core::StatusCode::Continue)
350                            == 409
351                        {
352                            // Conflict trying to create counter, retry increment
353                            self.increment(key, delta).await?;
354                        } else {
355                            return Err(log_error(e));
356                        }
357                    }
358                    Ok(delta)
359                } else {
360                    Err(log_error(e))
361                }
362            }
363            Ok(_) => self
364                .get_entity::<Counter>(key.as_ref())
365                .await?
366                .map(|c| c.value)
367                .ok_or(Error::Other(
368                    "increment returned an empty value after patching, which indicates a bug"
369                        .to_string(),
370                )),
371        }
372    }
373
374    async fn new_compare_and_swap(
375        &self,
376        bucket_rep: u32,
377        key: &str,
378    ) -> Result<Arc<dyn spin_factor_key_value::Cas>, Error> {
379        Ok(Arc::new(CompareAndSwap {
380            key: key.to_string(),
381            client: self.client.clone(),
382            etag: Mutex::new(None),
383            bucket_rep,
384            store_id: self.store_id.clone(),
385        }))
386    }
387}
388
389struct CompareAndSwap {
390    key: String,
391    client: CollectionClient,
392    bucket_rep: u32,
393    etag: Mutex<Option<String>>,
394    store_id: Option<String>,
395}
396
397#[async_trait]
398impl Cas for CompareAndSwap {
399    /// `current` will fetch the current value for the key and store the etag for the record. The
400    /// etag will be used to perform and optimistic concurrency update using the `if-match` header.
401    async fn current(&self, max_result_bytes: usize) -> Result<Option<Vec<u8>>, Error> {
402        let mut stream = self
403            .client
404            .query_documents(get_query(&self.key, self.store_id.as_deref()))
405            .query_cross_partition(true)
406            .max_item_count(1)
407            .into_stream::<Pair>();
408
409        let current_value: Option<(Vec<u8>, Option<String>)> = match stream.next().await {
410            Some(r) => {
411                let r = r.map_err(log_error)?;
412                match r.results.first() {
413                    Some((item, Some(attr))) => {
414                        Some((item.clone().value, Some(attr.etag().to_string())))
415                    }
416                    Some((item, None)) => Some((item.clone().value, None)),
417                    _ => None,
418                }
419            }
420            None => None,
421        };
422
423        let value = match current_value {
424            Some((value, etag)) => {
425                self.etag.lock().unwrap().clone_from(&etag);
426                Some(value)
427            }
428            None => None,
429        };
430
431        // Currently there's no way to stream a single query result using the
432        // `azure_data_cosmos` crate without buffering, so the damage (in terms
433        // of host memory usage) is already done, but we can still enforce the
434        // limit:
435        if std::mem::size_of::<Option<Vec<u8>>>() + value.as_ref().map(|v| v.len()).unwrap_or(0)
436            > max_result_bytes
437        {
438            Err(Error::Other(format!(
439                "query result exceeds limit of {max_result_bytes} bytes"
440            )))
441        } else {
442            Ok(value)
443        }
444    }
445
446    /// `swap` updates the value for the key using the etag saved in the `current` function for
447    /// optimistic concurrency.
448    async fn swap(&self, value: Vec<u8>) -> Result<(), SwapError> {
449        let pair = Pair {
450            id: self.key.clone(),
451            value,
452            store_id: self.store_id.clone(),
453        };
454
455        let doc_client = self
456            .client
457            .document_client(&self.key, &pair.partition_key())
458            .map_err(log_cas_error)?;
459
460        let etag_value = self.etag.lock().unwrap().clone();
461        match etag_value {
462            Some(etag) => {
463                // attempt to replace the document if the etag matches
464                doc_client
465                    .replace_document(pair)
466                    .if_match_condition(azure_core::request_options::IfMatchCondition::Match(etag))
467                    .await
468                    .map_err(|e| SwapError::CasFailed(format!("{e:?}")))
469                    .map(drop)
470            }
471            None => {
472                // if we have no etag, then we assume the document does not yet exist and must insert; no upserts.
473                self.client
474                    .create_document(pair)
475                    .await
476                    .map_err(|e| SwapError::CasFailed(format!("{e:?}")))
477                    .map(drop)
478            }
479        }
480    }
481
482    async fn bucket_rep(&self) -> u32 {
483        self.bucket_rep
484    }
485
486    async fn key(&self) -> String {
487        self.key.clone()
488    }
489}
490
491impl AzureCosmosStore {
492    async fn get_entity<F>(&self, key: &str) -> Result<Option<F>, Error>
493    where
494        F: CosmosEntity + Send + Sync + serde::de::DeserializeOwned + Clone,
495    {
496        let query = self
497            .client
498            .query_documents(get_query(key, self.store_id.as_deref()))
499            .query_cross_partition(true)
500            .max_item_count(1);
501
502        // There can be no duplicated keys, so we create the stream and only take the first result.
503        let mut stream = query.into_stream::<F>();
504        let Some(res) = stream.next().await else {
505            return Ok(None);
506        };
507        Ok(res
508            .map_err(log_error)?
509            .results
510            .first()
511            .map(|(p, _)| p.clone()))
512    }
513
514    async fn get_keys(&self, max_result_bytes: usize) -> Result<Vec<String>, Error> {
515        let query = self
516            .client
517            .query_documents(get_keys_query(self.store_id.as_deref()))
518            .query_cross_partition(true);
519        let mut res = Vec::new();
520
521        let mut stream = query.into_stream::<Key>();
522        let mut byte_count = std::mem::size_of::<Vec<String>>();
523        while let Some(resp) = stream.next().await {
524            let resp = resp.map_err(log_error)?.results;
525            byte_count += resp
526                .iter()
527                .map(|(key, _)| std::mem::size_of::<String>() + key.id.len())
528                .sum::<usize>();
529            if byte_count > max_result_bytes {
530                return Err(Error::Other(format!(
531                    "query result exceeds limit of {max_result_bytes} bytes"
532                )));
533            }
534            res.extend(resp.into_iter().map(|(key, _)| key.id));
535        }
536
537        Ok(res)
538    }
539}
540
541fn get_query(key: &str, store_id: Option<&str>) -> Query {
542    build_query(
543        "SELECT * FROM c",
544        vec!["c.id=@id"],
545        vec![Param::new("@id".to_owned(), key)],
546        store_id,
547    )
548}
549
550fn get_id_query(key: &str, store_id: Option<&str>) -> Query {
551    build_query(
552        "SELECT c.id, c.store_id FROM c",
553        vec!["c.id=@id"],
554        vec![Param::new("@id".to_owned(), key)],
555        store_id,
556    )
557}
558
559fn get_keys_query(store_id: Option<&str>) -> Query {
560    build_query("SELECT c.id, c.store_id FROM c", vec![], vec![], store_id)
561}
562
563fn get_in_query(keys: Vec<String>, store_id: Option<&str>) -> Query {
564    let params: Vec<Param> = keys
565        .into_iter()
566        .enumerate()
567        .map(|(i, key)| Param::new(format!("@id{i}"), key))
568        .collect();
569    let names = params.iter().map(Param::name).collect::<Vec<&str>>();
570    let in_condition = format!("c.id IN ({})", names.join(", "));
571    build_query(
572        "SELECT * FROM c",
573        vec![in_condition.as_str()],
574        params,
575        store_id,
576    )
577}
578
579/// Builds a parameterized query from a select clause and `AND`-joined conditions, adding a
580/// `c.store_id=@store_id` condition when a store id is set.
581fn build_query(
582    select: &str,
583    mut conditions: Vec<&str>,
584    mut params: Vec<Param>,
585    store_id: Option<&str>,
586) -> Query {
587    if let Some(store_id) = store_id {
588        conditions.push("c.store_id=@store_id");
589        params.push(Param::new("@store_id".to_owned(), store_id));
590    }
591    let mut sql = select.to_owned();
592    if !conditions.is_empty() {
593        sql.push_str(" WHERE ");
594        sql.push_str(&conditions.join(" AND "));
595    }
596    Query::with_params(sql, params)
597}
598
599// Pair structure for key value operations
600#[derive(Serialize, Deserialize, Clone, Debug)]
601pub struct Pair {
602    pub id: String,
603    pub value: Vec<u8>,
604    #[serde(skip_serializing_if = "Option::is_none")]
605    pub store_id: Option<String>,
606}
607
608impl CosmosEntity for Pair {
609    type Entity = String;
610
611    fn partition_key(&self) -> Self::Entity {
612        self.store_id.clone().unwrap_or_else(|| self.id.clone())
613    }
614}
615
616// Counter structure for increment operations
617#[derive(Serialize, Deserialize, Clone, Debug)]
618pub struct Counter {
619    pub id: String,
620    pub value: i64,
621    #[serde(skip_serializing_if = "Option::is_none")]
622    pub store_id: Option<String>,
623}
624
625impl CosmosEntity for Counter {
626    type Entity = String;
627
628    fn partition_key(&self) -> Self::Entity {
629        self.store_id.clone().unwrap_or_else(|| self.id.clone())
630    }
631}
632
633// Key structure for operations with generic value types
634#[derive(Serialize, Deserialize, Clone, Debug)]
635pub struct Key {
636    pub id: String,
637    #[serde(skip_serializing_if = "Option::is_none")]
638    pub store_id: Option<String>,
639}
640
641impl CosmosEntity for Key {
642    type Entity = String;
643
644    fn partition_key(&self) -> Self::Entity {
645        self.store_id.clone().unwrap_or_else(|| self.id.clone())
646    }
647}
spin_key_value_azure/store.rs

spin_key_value_azure/
store.rs
