prolly_store_cosmosdb/
lib.rs

1#![doc = include_str!("../README.md")]
2
3pub use prolly::{
4    RemoteBatchOp, RemoteManifestUpdate, RemoteNamedRoot, RemoteProllyStore, RemoteRootCondition,
5    RemoteRootWrite, RemoteStoreBackend, RemoteTransactionConflict, RemoteTransactionUpdate,
6};
7
8/// Cosmos DB adapter entry point.
9pub mod cosmosdb {
10    use std::collections::{HashMap, HashSet};
11    use std::error::Error as StdError;
12    use std::fmt;
13    use std::time::SystemTime;
14
15    use base64::engine::general_purpose::STANDARD as BASE64;
16    use base64::Engine as _;
17    use hmac::{Hmac, Mac};
18    use percent_encoding::{utf8_percent_encode, NON_ALPHANUMERIC};
19    use reqwest::{Method, StatusCode};
20    use serde::{Deserialize, Serialize};
21    use sha2::Sha256;
22
23    use crate::{
24        RemoteBatchOp, RemoteManifestUpdate, RemoteNamedRoot, RemoteRootCondition, RemoteRootWrite,
25        RemoteStoreBackend, RemoteTransactionConflict, RemoteTransactionUpdate,
26    };
27
28    /// Store adapter for Cosmos DB-backed prolly nodes and roots.
29    pub type CosmosDbStore = crate::RemoteProllyStore<CosmosDbBackend>;
30
31    /// Cosmos DB REST-backed backend.
32    ///
33    /// The container must use `/kind` as its partition key. The adapter stores
34    /// all documents for one backend instance under a single `kind` partition
35    /// value so Cosmos DB transactional batches can atomically commit nodes and
36    /// roots together. The logical document family lives in `family`.
37    #[derive(Clone, Debug)]
38    pub struct CosmosDbBackend {
39        http: reqwest::Client,
40        endpoint: String,
41        account_key: Vec<u8>,
42        database_id: String,
43        container_id: String,
44        container_link: String,
45        key_prefix: Vec<u8>,
46        partition_key: String,
47        read_parallelism: usize,
48    }
49
50    impl CosmosDbBackend {
51        /// Create a backend using Cosmos DB key authentication.
52        pub fn with_key(
53            endpoint: impl Into<String>,
54            account_key: &str,
55            database_id: impl Into<String>,
56            container_id: impl Into<String>,
57        ) -> Result<Self, CosmosDbBackendError> {
58            Self::with_http_client(
59                reqwest::Client::new(),
60                endpoint,
61                account_key,
62                database_id,
63                container_id,
64            )
65        }
66
67        /// Create a backend with a caller-provided HTTP client.
68        pub fn with_http_client(
69            http: reqwest::Client,
70            endpoint: impl Into<String>,
71            account_key: &str,
72            database_id: impl Into<String>,
73            container_id: impl Into<String>,
74        ) -> Result<Self, CosmosDbBackendError> {
75            let database_id = database_id.into();
76            let container_id = container_id.into();
77            let container_link = format!(
78                "dbs/{}/colls/{}",
79                encode_path_segment(&database_id),
80                encode_path_segment(&container_id)
81            );
82
83            Ok(Self {
84                http,
85                endpoint: endpoint.into().trim_end_matches('/').to_string(),
86                account_key: BASE64
87                    .decode(account_key)
88                    .map_err(CosmosDbBackendError::InvalidAccountKey)?,
89                database_id,
90                container_id,
91                container_link,
92                key_prefix: DEFAULT_KEY_PREFIX.to_vec(),
93                partition_key: DEFAULT_PARTITION_KEY.to_string(),
94                read_parallelism: DEFAULT_READ_PARALLELISM,
95            })
96        }
97
98        /// Return the Cosmos DB account endpoint.
99        pub fn endpoint(&self) -> &str {
100            &self.endpoint
101        }
102
103        /// Return the Cosmos DB database id.
104        pub fn database_id(&self) -> &str {
105            &self.database_id
106        }
107
108        /// Return the Cosmos DB container id.
109        pub fn container_id(&self) -> &str {
110            &self.container_id
111        }
112
113        /// Return the namespace prefix prepended to all logical keys.
114        pub fn key_prefix(&self) -> &[u8] {
115            &self.key_prefix
116        }
117
118        /// Return the `/kind` partition value used by this backend instance.
119        pub fn partition_key_value(&self) -> &str {
120            &self.partition_key
121        }
122
123        /// Set the namespace prefix prepended to all logical keys.
124        pub fn with_key_prefix(mut self, key_prefix: impl Into<Vec<u8>>) -> Self {
125            self.key_prefix = key_prefix.into();
126            self
127        }
128
129        /// Set the `/kind` partition value used by this backend instance.
130        ///
131        /// All nodes, roots, and hints for the backend must share this value for
132        /// Cosmos DB transactional batch support.
133        pub fn with_partition_key_value(mut self, partition_key: impl Into<String>) -> Self {
134            self.partition_key = partition_key.into();
135            self
136        }
137
138        /// Set the read parallelism advertised to async prolly traversals.
139        pub fn with_read_parallelism(mut self, read_parallelism: usize) -> Self {
140            self.read_parallelism = read_parallelism.max(1);
141            self
142        }
143
144        /// Delete every document under this backend's namespace prefix.
145        ///
146        /// This is primarily intended for isolated integration tests.
147        pub async fn clear_namespace(&self) -> Result<(), CosmosDbBackendError> {
148            if self.key_prefix.is_empty() {
149                return Err(CosmosDbBackendError::InvalidConfiguration(
150                    "refusing to clear an empty Cosmos DB key prefix".to_string(),
151                ));
152            }
153
154            for kind in [NODE_KIND, ROOT_KIND, HINT_KIND] {
155                let docs = self.query_kind(kind).await?;
156                for doc in docs {
157                    let logical_key = doc.logical_key()?;
158                    if logical_key.starts_with(&self.key_prefix) {
159                        self.delete_document(kind, &logical_key, None, true).await?;
160                    }
161                }
162            }
163
164            Ok(())
165        }
166
167        fn node_key(&self, key: &[u8]) -> Vec<u8> {
168            self.family_key(NODE_FAMILY, key)
169        }
170
171        fn root_key(&self, name: &[u8]) -> Vec<u8> {
172            self.family_key(ROOT_FAMILY, name)
173        }
174
175        fn hint_key(&self, namespace: &[u8], key: &[u8]) -> Vec<u8> {
176            let mut cosmos_key = self.family_key(HINT_FAMILY, &[]);
177            cosmos_key.extend_from_slice(&(namespace.len() as u64).to_be_bytes());
178            cosmos_key.extend_from_slice(namespace);
179            cosmos_key.extend_from_slice(key);
180            cosmos_key
181        }
182
183        fn family_key(&self, family: &[u8], suffix: &[u8]) -> Vec<u8> {
184            let mut key = Vec::with_capacity(self.key_prefix.len() + family.len() + suffix.len());
185            key.extend_from_slice(&self.key_prefix);
186            key.extend_from_slice(family);
187            key.extend_from_slice(suffix);
188            key
189        }
190
191        fn family_prefix(&self, family: &[u8]) -> Vec<u8> {
192            self.family_key(family, &[])
193        }
194
195        fn feed_link(&self) -> String {
196            format!("{}/docs", self.container_link)
197        }
198
199        fn document_link(&self, id: &str) -> String {
200            format!("{}/docs/{}", self.container_link, id)
201        }
202
203        fn resource_url(&self, link: &str) -> String {
204            format!("{}/{}", self.endpoint, link)
205        }
206
207        fn authorized_request(
208            &self,
209            method: Method,
210            resource_type: &'static str,
211            resource_link: &str,
212            url: String,
213        ) -> Result<reqwest::RequestBuilder, CosmosDbBackendError> {
214            let date = httpdate::fmt_http_date(SystemTime::now());
215            let auth =
216                self.authorization_header(method.as_str(), resource_type, resource_link, &date)?;
217            Ok(self
218                .http
219                .request(method, url)
220                .header("authorization", auth)
221                .header("x-ms-date", date)
222                .header("x-ms-version", COSMOS_API_VERSION))
223        }
224
225        fn authorization_header(
226            &self,
227            method: &str,
228            resource_type: &'static str,
229            resource_link: &str,
230            date: &str,
231        ) -> Result<String, CosmosDbBackendError> {
232            let payload = format!(
233                "{}\n{}\n{}\n{}\n\n",
234                method.to_ascii_lowercase(),
235                resource_type,
236                resource_link,
237                date.to_ascii_lowercase()
238            );
239            let mut mac = Hmac::<Sha256>::new_from_slice(&self.account_key)
240                .map_err(|err| CosmosDbBackendError::InvalidConfiguration(err.to_string()))?;
241            mac.update(payload.as_bytes());
242            let signature = BASE64.encode(mac.finalize().into_bytes());
243            let token = format!("type=master&ver=1.0&sig={signature}");
244            Ok(utf8_percent_encode(&token, NON_ALPHANUMERIC).to_string())
245        }
246
247        async fn read_document(
248            &self,
249            _kind: &'static str,
250            logical_key: &[u8],
251        ) -> Result<Option<CosmosReadDocument>, CosmosDbBackendError> {
252            let id = document_id(logical_key);
253            let link = self.document_link(&id);
254            let response = self
255                .authorized_request(Method::GET, DOCS_RESOURCE, &link, self.resource_url(&link))?
256                .header("x-ms-documentdb-partitionkey", self.partition_key_header())
257                .send()
258                .await
259                .map_err(CosmosDbBackendError::Http)?;
260
261            if response.status() == StatusCode::NOT_FOUND {
262                return Ok(None);
263            }
264            let response = ensure_status(response).await?;
265            let etag = response
266                .headers()
267                .get("etag")
268                .and_then(|value| value.to_str().ok())
269                .map(str::to_string)
270                .ok_or(CosmosDbBackendError::MissingEtag)?;
271            let document = response
272                .json::<CosmosProllyDocument>()
273                .await
274                .map_err(CosmosDbBackendError::Http)?;
275            Ok(Some(CosmosReadDocument { document, etag }))
276        }
277
278        async fn upsert_document(
279            &self,
280            kind: &'static str,
281            logical_key: &[u8],
282            value: &[u8],
283        ) -> Result<(), CosmosDbBackendError> {
284            let doc = self.document(kind, logical_key, value);
285            let link = self.feed_link();
286            let response = self
287                .authorized_request(
288                    Method::POST,
289                    DOCS_RESOURCE,
290                    &self.container_link,
291                    self.resource_url(&link),
292                )?
293                .header("content-type", "application/json")
294                .header("x-ms-documentdb-partitionkey", self.partition_key_header())
295                .header("x-ms-documentdb-is-upsert", "True")
296                .json(&doc)
297                .send()
298                .await
299                .map_err(CosmosDbBackendError::Http)?;
300            ensure_status(response).await?;
301            Ok(())
302        }
303
304        async fn create_document_if_absent(
305            &self,
306            kind: &'static str,
307            logical_key: &[u8],
308            value: &[u8],
309        ) -> Result<bool, CosmosDbBackendError> {
310            let doc = self.document(kind, logical_key, value);
311            let link = self.feed_link();
312            let response = self
313                .authorized_request(
314                    Method::POST,
315                    DOCS_RESOURCE,
316                    &self.container_link,
317                    self.resource_url(&link),
318                )?
319                .header("content-type", "application/json")
320                .header("if-none-match", "*")
321                .header("x-ms-documentdb-partitionkey", self.partition_key_header())
322                .json(&doc)
323                .send()
324                .await
325                .map_err(CosmosDbBackendError::Http)?;
326            if is_conflict_status(response.status()) {
327                return Ok(false);
328            }
329            ensure_status(response).await?;
330            Ok(true)
331        }
332
333        async fn replace_document_if_match(
334            &self,
335            kind: &'static str,
336            logical_key: &[u8],
337            value: &[u8],
338            etag: &str,
339        ) -> Result<bool, CosmosDbBackendError> {
340            let id = document_id(logical_key);
341            let doc = self.document(kind, logical_key, value);
342            let link = self.document_link(&id);
343            let response = self
344                .authorized_request(Method::PUT, DOCS_RESOURCE, &link, self.resource_url(&link))?
345                .header("content-type", "application/json")
346                .header("if-match", etag)
347                .header("x-ms-documentdb-partitionkey", self.partition_key_header())
348                .json(&doc)
349                .send()
350                .await
351                .map_err(CosmosDbBackendError::Http)?;
352            if is_conflict_status(response.status()) {
353                return Ok(false);
354            }
355            ensure_status(response).await?;
356            Ok(true)
357        }
358
359        async fn delete_document(
360            &self,
361            _kind: &'static str,
362            logical_key: &[u8],
363            etag: Option<&str>,
364            ignore_missing: bool,
365        ) -> Result<bool, CosmosDbBackendError> {
366            let id = document_id(logical_key);
367            let link = self.document_link(&id);
368            let mut request = self
369                .authorized_request(
370                    Method::DELETE,
371                    DOCS_RESOURCE,
372                    &link,
373                    self.resource_url(&link),
374                )?
375                .header("x-ms-documentdb-partitionkey", self.partition_key_header());
376            if let Some(etag) = etag {
377                request = request.header("if-match", etag);
378            }
379
380            let response = request.send().await.map_err(CosmosDbBackendError::Http)?;
381            if response.status() == StatusCode::NOT_FOUND && ignore_missing {
382                return Ok(true);
383            }
384            if is_conflict_status(response.status()) {
385                return Ok(false);
386            }
387            ensure_status(response).await?;
388            Ok(true)
389        }
390
391        async fn query_kind(
392            &self,
393            kind: &'static str,
394        ) -> Result<Vec<CosmosProllyDocument>, CosmosDbBackendError> {
395            let mut documents = Vec::new();
396            let mut continuation = None;
397
398            loop {
399                let link = self.feed_link();
400                let body = serde_json::json!({
401                    "query": "SELECT * FROM c WHERE c.kind = @kind AND c.family = @family",
402                    "parameters": [
403                        { "name": "@kind", "value": self.partition_key },
404                        { "name": "@family", "value": kind }
405                    ]
406                });
407                let mut request = self
408                    .authorized_request(
409                        Method::POST,
410                        DOCS_RESOURCE,
411                        &self.container_link,
412                        self.resource_url(&link),
413                    )?
414                    .header("content-type", "application/query+json")
415                    .header("x-ms-documentdb-isquery", "True")
416                    .header("x-ms-documentdb-partitionkey", self.partition_key_header())
417                    .header("x-ms-max-item-count", "100")
418                    .json(&body);
419                if let Some(token) = continuation.as_deref() {
420                    request = request.header("x-ms-continuation", token);
421                }
422
423                let response =
424                    ensure_status(request.send().await.map_err(CosmosDbBackendError::Http)?)
425                        .await?;
426                continuation = response
427                    .headers()
428                    .get("x-ms-continuation")
429                    .and_then(|value| value.to_str().ok())
430                    .map(str::to_string);
431                let page = response
432                    .json::<CosmosFeed>()
433                    .await
434                    .map_err(CosmosDbBackendError::Http)?;
435                documents.extend(page.documents);
436
437                if continuation.is_none() {
438                    break;
439                }
440            }
441
442            Ok(documents)
443        }
444
445        fn document(
446            &self,
447            kind: &'static str,
448            logical_key: &[u8],
449            value: &[u8],
450        ) -> CosmosProllyDocument {
451            CosmosProllyDocument::new(&self.partition_key, kind, logical_key, value)
452        }
453
454        fn partition_key_header(&self) -> String {
455            partition_key(&self.partition_key)
456        }
457
458        fn batch_partition_key(&self) -> String {
459            self.partition_key_header()
460        }
461
462        async fn execute_transaction_batch(
463            &self,
464            operations: &[CosmosBatchOperation],
465        ) -> Result<Vec<CosmosBatchOperationResponse>, CosmosDbBackendError> {
466            let link = self.feed_link();
467            let response = self
468                .authorized_request(
469                    Method::POST,
470                    DOCS_RESOURCE,
471                    &self.container_link,
472                    self.resource_url(&link),
473                )?
474                .header("content-type", "application/json")
475                .header("x-ms-documentdb-partitionkey", self.partition_key_header())
476                .header("x-ms-cosmos-is-batch-request", "True")
477                .header("x-ms-cosmos-batch-atomic", "True")
478                .json(operations)
479                .send()
480                .await
481                .map_err(CosmosDbBackendError::Http)?;
482
483            let response = ensure_status(response).await?;
484            response
485                .json::<Vec<CosmosBatchOperationResponse>>()
486                .await
487                .map_err(CosmosDbBackendError::Http)
488        }
489
490        async fn push_root_condition_operation(
491            &self,
492            operations: &mut Vec<CosmosBatchOperation>,
493            operation_conditions: &mut Vec<Option<RemoteRootCondition>>,
494            condition: &RemoteRootCondition,
495        ) -> Result<(), CosmosDbBackendError> {
496            let logical_key = self.root_key(&condition.name);
497            match condition.expected.as_deref() {
498                Some(expected) => {
499                    let Some(current) = self.read_document(ROOT_KIND, &logical_key).await? else {
500                        return Err(CosmosDbBackendError::RootConditionConflict(
501                            RemoteTransactionConflict::new(
502                                condition.name.clone(),
503                                condition.expected.clone(),
504                                None,
505                            ),
506                        ));
507                    };
508                    let current_value = current.document.value_bytes()?;
509                    if current_value.as_slice() != expected {
510                        return Err(CosmosDbBackendError::RootConditionConflict(
511                            RemoteTransactionConflict::new(
512                                condition.name.clone(),
513                                condition.expected.clone(),
514                                Some(current_value),
515                            ),
516                        ));
517                    }
518
519                    operations.push(CosmosBatchOperation::read(
520                        document_id(&logical_key),
521                        self.batch_partition_key(),
522                        current.etag,
523                    ));
524                    operation_conditions.push(Some(condition.clone()));
525                }
526                None => {
527                    let doc = self.document(ROOT_KIND, &logical_key, &[]);
528                    operations.push(CosmosBatchOperation::create_if_absent(
529                        self.batch_partition_key(),
530                        doc,
531                    ));
532                    operation_conditions.push(Some(condition.clone()));
533                    operations.push(CosmosBatchOperation::delete(
534                        document_id(&logical_key),
535                        self.batch_partition_key(),
536                        None,
537                    ));
538                    operation_conditions.push(Some(condition.clone()));
539                }
540            }
541            Ok(())
542        }
543
544        async fn push_root_write_operation(
545            &self,
546            operations: &mut Vec<CosmosBatchOperation>,
547            operation_conditions: &mut Vec<Option<RemoteRootCondition>>,
548            write: &RemoteRootWrite,
549            condition: Option<&RemoteRootCondition>,
550        ) -> Result<(), CosmosDbBackendError> {
551            let name = root_write_name(write);
552            let logical_key = self.root_key(name);
553            match (
554                condition.and_then(|condition| condition.expected.as_deref()),
555                write,
556            ) {
557                (Some(expected), RemoteRootWrite::Put { manifest, .. }) => {
558                    let Some(current) = self.read_document(ROOT_KIND, &logical_key).await? else {
559                        return Err(CosmosDbBackendError::RootConditionConflict(
560                            RemoteTransactionConflict::new(
561                                name.to_vec(),
562                                Some(expected.to_vec()),
563                                None,
564                            ),
565                        ));
566                    };
567                    let current_value = current.document.value_bytes()?;
568                    if current_value.as_slice() != expected {
569                        return Err(CosmosDbBackendError::RootConditionConflict(
570                            RemoteTransactionConflict::new(
571                                name.to_vec(),
572                                Some(expected.to_vec()),
573                                Some(current_value),
574                            ),
575                        ));
576                    }
577
578                    operations.push(CosmosBatchOperation::replace(
579                        document_id(&logical_key),
580                        self.batch_partition_key(),
581                        current.etag,
582                        self.document(ROOT_KIND, &logical_key, manifest),
583                    ));
584                    operation_conditions.push(condition.cloned());
585                }
586                (Some(expected), RemoteRootWrite::Delete { .. }) => {
587                    let Some(current) = self.read_document(ROOT_KIND, &logical_key).await? else {
588                        return Err(CosmosDbBackendError::RootConditionConflict(
589                            RemoteTransactionConflict::new(
590                                name.to_vec(),
591                                Some(expected.to_vec()),
592                                None,
593                            ),
594                        ));
595                    };
596                    let current_value = current.document.value_bytes()?;
597                    if current_value.as_slice() != expected {
598                        return Err(CosmosDbBackendError::RootConditionConflict(
599                            RemoteTransactionConflict::new(
600                                name.to_vec(),
601                                Some(expected.to_vec()),
602                                Some(current_value),
603                            ),
604                        ));
605                    }
606
607                    operations.push(CosmosBatchOperation::delete(
608                        document_id(&logical_key),
609                        self.batch_partition_key(),
610                        Some(current.etag),
611                    ));
612                    operation_conditions.push(condition.cloned());
613                }
614                (None, RemoteRootWrite::Put { manifest, .. }) if condition.is_some() => {
615                    operations.push(CosmosBatchOperation::create_if_absent(
616                        self.batch_partition_key(),
617                        self.document(ROOT_KIND, &logical_key, manifest),
618                    ));
619                    operation_conditions.push(condition.cloned());
620                }
621                (None, RemoteRootWrite::Delete { .. }) if condition.is_some() => {
622                    self.push_root_condition_operation(
623                        operations,
624                        operation_conditions,
625                        condition.expect("condition checked"),
626                    )
627                    .await?;
628                }
629                (None, RemoteRootWrite::Put { manifest, .. }) => {
630                    operations.push(CosmosBatchOperation::upsert(
631                        self.batch_partition_key(),
632                        self.document(ROOT_KIND, &logical_key, manifest),
633                    ));
634                    operation_conditions.push(None);
635                }
636                (None, RemoteRootWrite::Delete { .. }) => {
637                    if let Some(current) = self.read_document(ROOT_KIND, &logical_key).await? {
638                        operations.push(CosmosBatchOperation::delete(
639                            document_id(&logical_key),
640                            self.batch_partition_key(),
641                            Some(current.etag),
642                        ));
643                        operation_conditions.push(None);
644                    }
645                }
646            }
647            Ok(())
648        }
649
650        async fn conflict_from_batch_response(
651            &self,
652            responses: &[CosmosBatchOperationResponse],
653            operation_conditions: &[Option<RemoteRootCondition>],
654            root_conditions: &[RemoteRootCondition],
655        ) -> Result<Option<RemoteTransactionConflict>, CosmosDbBackendError> {
656            for (response, condition) in responses.iter().zip(operation_conditions) {
657                if response.is_success() {
658                    continue;
659                }
660                let Some(condition) = condition else {
661                    continue;
662                };
663                let current = self.get_root_manifest(&condition.name).await?;
664                return Ok(Some(RemoteTransactionConflict::new(
665                    condition.name.clone(),
666                    condition.expected.clone(),
667                    current,
668                )));
669            }
670
671            for condition in root_conditions {
672                let current = self.get_root_manifest(&condition.name).await?;
673                if current != condition.expected {
674                    return Ok(Some(RemoteTransactionConflict::new(
675                        condition.name.clone(),
676                        condition.expected.clone(),
677                        current,
678                    )));
679                }
680            }
681
682            Ok(None)
683        }
684    }
685
686    impl RemoteStoreBackend for CosmosDbBackend {
687        type Error = CosmosDbBackendError;
688
689        async fn get_node(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
690            self.read_document(NODE_KIND, &self.node_key(key))
691                .await?
692                .map(|doc| doc.document.value_bytes())
693                .transpose()
694        }
695
696        async fn put_node(&self, key: &[u8], value: &[u8]) -> Result<(), Self::Error> {
697            self.upsert_document(NODE_KIND, &self.node_key(key), value)
698                .await
699        }
700
701        async fn delete_node(&self, key: &[u8]) -> Result<(), Self::Error> {
702            self.delete_document(NODE_KIND, &self.node_key(key), None, true)
703                .await?;
704            Ok(())
705        }
706
707        async fn batch_nodes(&self, ops: &[RemoteBatchOp<'_>]) -> Result<(), Self::Error> {
708            for op in ops {
709                match op {
710                    RemoteBatchOp::Upsert { key, value } => self.put_node(key, value).await?,
711                    RemoteBatchOp::Delete { key } => self.delete_node(key).await?,
712                }
713            }
714            Ok(())
715        }
716
717        async fn batch_get_nodes_ordered(
718            &self,
719            keys: &[&[u8]],
720        ) -> Result<Vec<Option<Vec<u8>>>, Self::Error> {
721            let mut values = Vec::with_capacity(keys.len());
722            for key in keys {
723                values.push(self.get_node(key).await?);
724            }
725            Ok(values)
726        }
727
728        async fn batch_put_nodes(&self, entries: &[(&[u8], &[u8])]) -> Result<(), Self::Error> {
729            for (key, value) in entries {
730                self.put_node(key, value).await?;
731            }
732            Ok(())
733        }
734
735        async fn list_node_cids(&self) -> Result<Vec<Vec<u8>>, Self::Error> {
736            let prefix = self.family_prefix(NODE_FAMILY);
737            let mut cids = self
738                .query_kind(NODE_KIND)
739                .await?
740                .into_iter()
741                .map(|doc| doc.logical_key())
742                .collect::<Result<Vec<_>, _>>()?
743                .into_iter()
744                .filter_map(|key| {
745                    key.strip_prefix(prefix.as_slice())
746                        .filter(|cid| cid.len() == 32)
747                        .map(<[u8]>::to_vec)
748                })
749                .collect::<Vec<_>>();
750            cids.sort();
751            Ok(cids)
752        }
753
754        fn read_parallelism(&self) -> usize {
755            self.read_parallelism
756        }
757
758        fn supports_hints(&self) -> bool {
759            true
760        }
761
762        async fn get_hint(
763            &self,
764            namespace: &[u8],
765            key: &[u8],
766        ) -> Result<Option<Vec<u8>>, Self::Error> {
767            self.read_document(HINT_KIND, &self.hint_key(namespace, key))
768                .await?
769                .map(|doc| doc.document.value_bytes())
770                .transpose()
771        }
772
773        async fn put_hint(
774            &self,
775            namespace: &[u8],
776            key: &[u8],
777            value: &[u8],
778        ) -> Result<(), Self::Error> {
779            self.upsert_document(HINT_KIND, &self.hint_key(namespace, key), value)
780                .await
781        }
782
783        async fn batch_put_nodes_with_hint(
784            &self,
785            entries: &[(&[u8], &[u8])],
786            namespace: &[u8],
787            key: &[u8],
788            value: &[u8],
789        ) -> Result<(), Self::Error> {
790            self.batch_put_nodes(entries).await?;
791            self.put_hint(namespace, key, value).await
792        }
793
794        async fn get_root_manifest(&self, name: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
795            self.read_document(ROOT_KIND, &self.root_key(name))
796                .await?
797                .map(|doc| doc.document.value_bytes())
798                .transpose()
799        }
800
801        async fn put_root_manifest(&self, name: &[u8], manifest: &[u8]) -> Result<(), Self::Error> {
802            self.upsert_document(ROOT_KIND, &self.root_key(name), manifest)
803                .await
804        }
805
806        async fn delete_root_manifest(&self, name: &[u8]) -> Result<(), Self::Error> {
807            self.delete_document(ROOT_KIND, &self.root_key(name), None, true)
808                .await?;
809            Ok(())
810        }
811
812        async fn compare_and_swap_root_manifest(
813            &self,
814            name: &[u8],
815            expected: Option<&[u8]>,
816            new: Option<&[u8]>,
817        ) -> Result<RemoteManifestUpdate, Self::Error> {
818            let logical_key = self.root_key(name);
819            match (expected, new) {
820                (None, Some(manifest)) => {
821                    if self
822                        .create_document_if_absent(ROOT_KIND, &logical_key, manifest)
823                        .await?
824                    {
825                        Ok(RemoteManifestUpdate::Applied)
826                    } else {
827                        Ok(RemoteManifestUpdate::Conflict {
828                            current: self.get_root_manifest(name).await?,
829                        })
830                    }
831                }
832                (Some(expected), Some(manifest)) => {
833                    let Some(current) = self.read_document(ROOT_KIND, &logical_key).await? else {
834                        return Ok(RemoteManifestUpdate::Conflict { current: None });
835                    };
836                    let current_value = current.document.value_bytes()?;
837                    if current_value.as_slice() != expected {
838                        return Ok(RemoteManifestUpdate::Conflict {
839                            current: Some(current_value),
840                        });
841                    }
842                    if self
843                        .replace_document_if_match(ROOT_KIND, &logical_key, manifest, &current.etag)
844                        .await?
845                    {
846                        Ok(RemoteManifestUpdate::Applied)
847                    } else {
848                        Ok(RemoteManifestUpdate::Conflict {
849                            current: self.get_root_manifest(name).await?,
850                        })
851                    }
852                }
853                (Some(expected), None) => {
854                    let Some(current) = self.read_document(ROOT_KIND, &logical_key).await? else {
855                        return Ok(RemoteManifestUpdate::Conflict { current: None });
856                    };
857                    let current_value = current.document.value_bytes()?;
858                    if current_value.as_slice() != expected {
859                        return Ok(RemoteManifestUpdate::Conflict {
860                            current: Some(current_value),
861                        });
862                    }
863                    if self
864                        .delete_document(ROOT_KIND, &logical_key, Some(&current.etag), false)
865                        .await?
866                    {
867                        Ok(RemoteManifestUpdate::Applied)
868                    } else {
869                        Ok(RemoteManifestUpdate::Conflict {
870                            current: self.get_root_manifest(name).await?,
871                        })
872                    }
873                }
874                (None, None) => {
875                    let current = self.get_root_manifest(name).await?;
876                    if current.is_none() {
877                        Ok(RemoteManifestUpdate::Applied)
878                    } else {
879                        Ok(RemoteManifestUpdate::Conflict { current })
880                    }
881                }
882            }
883        }
884
885        async fn list_root_manifests(&self) -> Result<Vec<RemoteNamedRoot>, Self::Error> {
886            let prefix = self.family_prefix(ROOT_FAMILY);
887            let mut roots = self
888                .query_kind(ROOT_KIND)
889                .await?
890                .into_iter()
891                .filter_map(|doc| {
892                    let logical_key = match doc.logical_key() {
893                        Ok(key) => key,
894                        Err(err) => return Some(Err(err)),
895                    };
896                    let Some(name) = logical_key.strip_prefix(prefix.as_slice()) else {
897                        return None;
898                    };
899                    let manifest = match doc.value_bytes() {
900                        Ok(value) => value,
901                        Err(err) => return Some(Err(err)),
902                    };
903                    Some(Ok(RemoteNamedRoot::new(name.to_vec(), manifest)))
904                })
905                .collect::<Result<Vec<_>, CosmosDbBackendError>>()?;
906            roots.sort_by(|left, right| left.name.cmp(&right.name));
907            Ok(roots)
908        }
909
910        fn supports_transactions(&self) -> bool {
911            true
912        }
913
914        async fn commit_transaction(
915            &self,
916            node_writes: &[RemoteBatchOp<'_>],
917            root_conditions: &[RemoteRootCondition],
918            root_writes: &[RemoteRootWrite],
919        ) -> Result<RemoteTransactionUpdate, Self::Error> {
920            let conditions_by_name = root_conditions
921                .iter()
922                .map(|condition| (condition.name.as_slice(), condition))
923                .collect::<HashMap<_, _>>();
924            let written_roots = root_writes
925                .iter()
926                .map(|write| root_write_name(write))
927                .collect::<HashSet<_>>();
928
929            let mut operations = Vec::new();
930            let mut operation_conditions = Vec::new();
931
932            for condition in root_conditions {
933                if !written_roots.contains(condition.name.as_slice()) {
934                    if let Err(err) = self
935                        .push_root_condition_operation(
936                            &mut operations,
937                            &mut operation_conditions,
938                            condition,
939                        )
940                        .await
941                    {
942                        return match err {
943                            CosmosDbBackendError::RootConditionConflict(conflict) => {
944                                Ok(RemoteTransactionUpdate::Conflict(conflict))
945                            }
946                            err => Err(err),
947                        };
948                    }
949                }
950            }
951
952            for write in root_writes {
953                if let Err(err) = self
954                    .push_root_write_operation(
955                        &mut operations,
956                        &mut operation_conditions,
957                        write,
958                        conditions_by_name.get(root_write_name(write)).copied(),
959                    )
960                    .await
961                {
962                    return match err {
963                        CosmosDbBackendError::RootConditionConflict(conflict) => {
964                            Ok(RemoteTransactionUpdate::Conflict(conflict))
965                        }
966                        err => Err(err),
967                    };
968                }
969            }
970
971            for write in node_writes {
972                match write {
973                    RemoteBatchOp::Upsert { key, value } => {
974                        let logical_key = self.node_key(key);
975                        operations.push(CosmosBatchOperation::upsert(
976                            self.batch_partition_key(),
977                            self.document(NODE_KIND, &logical_key, value),
978                        ));
979                        operation_conditions.push(None);
980                    }
981                    RemoteBatchOp::Delete { key } => {
982                        let logical_key = self.node_key(key);
983                        if let Some(current) = self.read_document(NODE_KIND, &logical_key).await? {
984                            operations.push(CosmosBatchOperation::delete(
985                                document_id(&logical_key),
986                                self.batch_partition_key(),
987                                Some(current.etag),
988                            ));
989                            operation_conditions.push(None);
990                        }
991                    }
992                }
993            }
994
995            if operations.len() > COSMOS_BATCH_OPERATION_LIMIT {
996                return Err(CosmosDbBackendError::TransactionTooLarge {
997                    operations: operations.len(),
998                    limit: COSMOS_BATCH_OPERATION_LIMIT,
999                });
1000            }
1001            if operations.is_empty() {
1002                return Ok(RemoteTransactionUpdate::Applied);
1003            }
1004
1005            let responses = self.execute_transaction_batch(&operations).await?;
1006            if responses
1007                .iter()
1008                .all(CosmosBatchOperationResponse::is_success)
1009            {
1010                return Ok(RemoteTransactionUpdate::Applied);
1011            }
1012
1013            if let Some(conflict) = self
1014                .conflict_from_batch_response(&responses, &operation_conditions, root_conditions)
1015                .await?
1016            {
1017                return Ok(RemoteTransactionUpdate::Conflict(conflict));
1018            }
1019
1020            Err(batch_response_error(&responses))
1021        }
1022    }
1023
1024    /// Error returned by the Cosmos DB backend.
1025    #[derive(Debug)]
1026    pub enum CosmosDbBackendError {
1027        /// HTTP request failed.
1028        Http(reqwest::Error),
1029        /// Cosmos DB returned an unexpected status code.
1030        UnexpectedStatus { status: StatusCode, body: String },
1031        /// Account key was not valid base64.
1032        InvalidAccountKey(base64::DecodeError),
1033        /// Stored document key was not valid hex.
1034        InvalidKeyHex(hex::FromHexError),
1035        /// Stored document value was not valid base64.
1036        InvalidValueBase64(base64::DecodeError),
1037        /// A point read response did not include an ETag.
1038        MissingEtag,
1039        /// Backend configuration is unsafe or invalid.
1040        InvalidConfiguration(String),
1041        /// The staged transaction exceeds Cosmos DB transactional batch limits.
1042        TransactionTooLarge { operations: usize, limit: usize },
1043        /// A root condition failed while building a transactional batch.
1044        RootConditionConflict(RemoteTransactionConflict),
1045    }
1046
1047    impl fmt::Display for CosmosDbBackendError {
1048        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1049            match self {
1050                Self::Http(err) => write!(f, "Cosmos DB HTTP error: {err}"),
1051                Self::UnexpectedStatus { status, body } => {
1052                    write!(f, "Cosmos DB returned {status}: {body}")
1053                }
1054                Self::InvalidAccountKey(err) => write!(f, "invalid Cosmos DB account key: {err}"),
1055                Self::InvalidKeyHex(err) => write!(f, "invalid Cosmos DB document key: {err}"),
1056                Self::InvalidValueBase64(err) => {
1057                    write!(f, "invalid Cosmos DB document value: {err}")
1058                }
1059                Self::MissingEtag => f.write_str("Cosmos DB response missing ETag"),
1060                Self::InvalidConfiguration(message) => f.write_str(message),
1061                Self::TransactionTooLarge { operations, limit } => write!(
1062                    f,
1063                    "Cosmos DB transaction has {operations} operations, exceeding the limit of {limit}"
1064                ),
1065                Self::RootConditionConflict(conflict) => write!(
1066                    f,
1067                    "Cosmos DB root condition conflict for {:?}",
1068                    String::from_utf8_lossy(&conflict.name)
1069                ),
1070            }
1071        }
1072    }
1073
1074    impl StdError for CosmosDbBackendError {
1075        fn source(&self) -> Option<&(dyn StdError + 'static)> {
1076            match self {
1077                Self::Http(err) => Some(err),
1078                Self::InvalidAccountKey(err) => Some(err),
1079                Self::InvalidKeyHex(err) => Some(err),
1080                Self::InvalidValueBase64(err) => Some(err),
1081                _ => None,
1082            }
1083        }
1084    }
1085
1086    #[derive(Debug, Clone, Serialize, Deserialize)]
1087    struct CosmosProllyDocument {
1088        id: String,
1089        kind: String,
1090        #[serde(default, skip_serializing_if = "Option::is_none")]
1091        family: Option<String>,
1092        key: String,
1093        value: String,
1094    }
1095
1096    impl CosmosProllyDocument {
1097        fn new(
1098            partition_key: &str,
1099            family: &'static str,
1100            logical_key: &[u8],
1101            value: &[u8],
1102        ) -> Self {
1103            Self {
1104                id: document_id(logical_key),
1105                kind: partition_key.to_string(),
1106                family: Some(family.to_string()),
1107                key: hex::encode(logical_key),
1108                value: BASE64.encode(value),
1109            }
1110        }
1111
1112        fn logical_key(&self) -> Result<Vec<u8>, CosmosDbBackendError> {
1113            hex::decode(&self.key).map_err(CosmosDbBackendError::InvalidKeyHex)
1114        }
1115
1116        fn value_bytes(&self) -> Result<Vec<u8>, CosmosDbBackendError> {
1117            BASE64
1118                .decode(&self.value)
1119                .map_err(CosmosDbBackendError::InvalidValueBase64)
1120        }
1121    }
1122
1123    #[derive(Debug, Serialize)]
1124    struct CosmosBatchOperation {
1125        #[serde(rename = "operationType")]
1126        operation_type: &'static str,
1127        #[serde(skip_serializing_if = "Option::is_none")]
1128        id: Option<String>,
1129        #[serde(rename = "partitionKey")]
1130        partition_key: String,
1131        #[serde(rename = "ifMatch")]
1132        if_match: String,
1133        #[serde(rename = "ifNoneMatch")]
1134        if_none_match: String,
1135        #[serde(rename = "resourceBody", skip_serializing_if = "Option::is_none")]
1136        resource_body: Option<CosmosProllyDocument>,
1137    }
1138
1139    impl CosmosBatchOperation {
1140        fn create_if_absent(partition_key: String, document: CosmosProllyDocument) -> Self {
1141            Self {
1142                operation_type: "Create",
1143                id: None,
1144                partition_key,
1145                if_match: String::new(),
1146                if_none_match: "*".to_string(),
1147                resource_body: Some(document),
1148            }
1149        }
1150
1151        fn upsert(partition_key: String, document: CosmosProllyDocument) -> Self {
1152            Self {
1153                operation_type: "Upsert",
1154                id: None,
1155                partition_key,
1156                if_match: String::new(),
1157                if_none_match: String::new(),
1158                resource_body: Some(document),
1159            }
1160        }
1161
1162        fn replace(
1163            id: String,
1164            partition_key: String,
1165            etag: String,
1166            document: CosmosProllyDocument,
1167        ) -> Self {
1168            Self {
1169                operation_type: "Replace",
1170                id: Some(id),
1171                partition_key,
1172                if_match: etag,
1173                if_none_match: String::new(),
1174                resource_body: Some(document),
1175            }
1176        }
1177
1178        fn read(id: String, partition_key: String, etag: String) -> Self {
1179            Self {
1180                operation_type: "Read",
1181                id: Some(id),
1182                partition_key,
1183                if_match: etag,
1184                if_none_match: String::new(),
1185                resource_body: None,
1186            }
1187        }
1188
1189        fn delete(id: String, partition_key: String, etag: Option<String>) -> Self {
1190            Self {
1191                operation_type: "Delete",
1192                id: Some(id),
1193                partition_key,
1194                if_match: etag.unwrap_or_default(),
1195                if_none_match: String::new(),
1196                resource_body: None,
1197            }
1198        }
1199    }
1200
1201    #[derive(Debug, Deserialize)]
1202    struct CosmosBatchOperationResponse {
1203        #[serde(rename = "statusCode")]
1204        status_code: u16,
1205    }
1206
1207    impl CosmosBatchOperationResponse {
1208        fn is_success(&self) -> bool {
1209            StatusCode::from_u16(self.status_code).is_ok_and(|status| status.is_success())
1210        }
1211    }
1212
1213    struct CosmosReadDocument {
1214        document: CosmosProllyDocument,
1215        etag: String,
1216    }
1217
1218    #[derive(Debug, Deserialize)]
1219    struct CosmosFeed {
1220        #[serde(rename = "Documents", alias = "documents")]
1221        documents: Vec<CosmosProllyDocument>,
1222    }
1223
1224    async fn ensure_status(
1225        response: reqwest::Response,
1226    ) -> Result<reqwest::Response, CosmosDbBackendError> {
1227        if response.status().is_success() {
1228            return Ok(response);
1229        }
1230
1231        let status = response.status();
1232        let body = response.text().await.unwrap_or_default();
1233        Err(CosmosDbBackendError::UnexpectedStatus { status, body })
1234    }
1235
1236    fn is_conflict_status(status: StatusCode) -> bool {
1237        matches!(
1238            status,
1239            StatusCode::CONFLICT | StatusCode::PRECONDITION_FAILED | StatusCode::NOT_FOUND
1240        )
1241    }
1242
1243    fn partition_key(value: &str) -> String {
1244        serde_json::to_string(&[value]).expect("serialize Cosmos DB partition key")
1245    }
1246
1247    fn document_id(logical_key: &[u8]) -> String {
1248        format!("k{}", hex::encode(logical_key))
1249    }
1250
1251    fn encode_path_segment(segment: &str) -> String {
1252        utf8_percent_encode(segment, NON_ALPHANUMERIC).to_string()
1253    }
1254
1255    fn root_write_name(write: &RemoteRootWrite) -> &[u8] {
1256        match write {
1257            RemoteRootWrite::Put { name, .. } | RemoteRootWrite::Delete { name } => name,
1258        }
1259    }
1260
1261    fn batch_response_error(responses: &[CosmosBatchOperationResponse]) -> CosmosDbBackendError {
1262        let (index, response) = responses
1263            .iter()
1264            .enumerate()
1265            .find(|(_, response)| !response.is_success())
1266            .unwrap_or_else(|| {
1267                (
1268                    0,
1269                    responses
1270                        .first()
1271                        .expect("transactional batch response is not empty"),
1272                )
1273            });
1274        let status =
1275            StatusCode::from_u16(response.status_code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
1276        CosmosDbBackendError::UnexpectedStatus {
1277            status,
1278            body: format!(
1279                "Cosmos DB transactional batch operation {index} returned status {}; response={response:?}",
1280                response.status_code
1281            ),
1282        }
1283    }
1284
1285    const COSMOS_API_VERSION: &str = "2018-12-31";
1286    const DOCS_RESOURCE: &str = "docs";
1287
1288    const DEFAULT_KEY_PREFIX: &[u8] = b"prolly:";
1289    const DEFAULT_PARTITION_KEY: &str = "prolly";
1290    const DEFAULT_READ_PARALLELISM: usize = 16;
1291    const COSMOS_BATCH_OPERATION_LIMIT: usize = 100;
1292
1293    const NODE_KIND: &str = "node";
1294    const ROOT_KIND: &str = "root";
1295    const HINT_KIND: &str = "hint";
1296
1297    const NODE_FAMILY: &[u8] = b"node:";
1298    const ROOT_FAMILY: &[u8] = b"root:";
1299    const HINT_FAMILY: &[u8] = b"hint:";
1300
1301    /// Default `/kind` partition value used by the adapter.
1302    pub const DEFAULT_PARTITION: &str = DEFAULT_PARTITION_KEY;
1303    /// Default logical partition for immutable nodes.
1304    pub const NODE_PARTITION: &str = DEFAULT_PARTITION_KEY;
1305    /// Default logical partition for named root manifests.
1306    pub const ROOT_PARTITION: &str = DEFAULT_PARTITION_KEY;
1307    /// Default logical partition for hints.
1308    pub const HINT_PARTITION: &str = DEFAULT_PARTITION_KEY;
1309
1310    #[cfg(test)]
1311    mod tests {
1312        use serde_json::json;
1313
1314        use super::*;
1315
1316        #[test]
1317        fn document_layout_uses_shared_partition_and_family() {
1318            let document = CosmosProllyDocument::new("tenant-a", NODE_KIND, b"node:abc", b"value");
1319
1320            assert_eq!(document.kind, "tenant-a");
1321            assert_eq!(document.family.as_deref(), Some(NODE_KIND));
1322            assert_eq!(document.key, hex::encode(b"node:abc"));
1323            assert_eq!(document.value_bytes().unwrap(), b"value");
1324        }
1325
1326        #[test]
1327        fn transactional_batch_operation_uses_cosmos_rest_shape() {
1328            let document = CosmosProllyDocument::new("prolly", ROOT_KIND, b"root:main", b"root");
1329            let operation =
1330                CosmosBatchOperation::create_if_absent(partition_key("prolly"), document);
1331            let value = serde_json::to_value(operation).unwrap();
1332
1333            assert_eq!(
1334                value,
1335                json!({
1336                    "operationType": "Create",
1337                    "partitionKey": "[\"prolly\"]",
1338                    "ifMatch": "",
1339                    "ifNoneMatch": "*",
1340                    "resourceBody": {
1341                        "id": document_id(b"root:main"),
1342                        "kind": "prolly",
1343                        "family": "root",
1344                        "key": hex::encode(b"root:main"),
1345                        "value": BASE64.encode(b"root")
1346                    }
1347                })
1348            );
1349        }
1350    }
1351}
1352
1353pub use cosmosdb::*;
prolly_store_cosmosdb/lib.rs

prolly_store_cosmosdb/
lib.rs
