viceroy_lib/
object_store.rs

1use {
2    crate::wiggle_abi::types::{FastlyStatus, KvError, KvInsertMode},
3    base64::prelude::*,
4    serde::Serialize,
5    std::{
6        collections::BTreeMap,
7        sync::{Arc, RwLock},
8        time::SystemTime,
9    },
10};
11
12#[derive(Debug, Clone)]
13pub struct ObjectValue {
14    pub body: Vec<u8>,
15    pub metadata: String,
16    pub metadata_len: usize,
17    pub generation: u64,
18    pub expiration: Option<SystemTime>,
19}
20
21#[derive(Clone, Debug, Default)]
22pub struct ObjectStores {
23    #[allow(clippy::type_complexity)]
24    stores: Arc<RwLock<BTreeMap<ObjectStoreKey, BTreeMap<ObjectKey, ObjectValue>>>>,
25}
26
27impl ObjectStores {
28    pub fn new() -> Self {
29        Self {
30            stores: Arc::new(RwLock::new(BTreeMap::new())),
31        }
32    }
33
34    pub(crate) fn store_exists(&self, obj_store_key: &str) -> Result<bool, ObjectStoreError> {
35        Ok(self
36            .stores
37            .read()
38            .map_err(|_| ObjectStoreError::PoisonedLock)?
39            .get(&ObjectStoreKey::new(obj_store_key))
40            .is_some())
41    }
42
43    pub fn lookup(
44        &self,
45        obj_store_key: ObjectStoreKey,
46        obj_key: ObjectKey,
47    ) -> Result<Option<ObjectValue>, KvStoreError> {
48        let mut res = Ok(None);
49
50        self.stores
51            .write()
52            .map_err(|_| KvStoreError::InternalError)?
53            .entry(obj_store_key)
54            .and_modify(|store| match store.get(&obj_key) {
55                Some(val) => {
56                    res = Ok(Some(val.clone()));
57                    // manages ttl
58                    if let Some(exp) = val.expiration {
59                        if SystemTime::now() >= exp {
60                            store.remove(&obj_key);
61                            res = Ok(None);
62                        }
63                    }
64                }
65                None => {
66                    res = Ok(None);
67                }
68            });
69
70        res
71    }
72
73    pub(crate) fn insert_empty_store(
74        &self,
75        obj_store_key: ObjectStoreKey,
76    ) -> Result<(), ObjectStoreError> {
77        self.stores
78            .write()
79            .map_err(|_| ObjectStoreError::PoisonedLock)?
80            .entry(obj_store_key)
81            .and_modify(|_| {})
82            .or_insert_with(BTreeMap::new);
83
84        Ok(())
85    }
86
87    pub fn insert(
88        &self,
89        obj_store_key: ObjectStoreKey,
90        obj_key: ObjectKey,
91        obj: Vec<u8>,
92        mode: KvInsertMode,
93        generation: Option<u64>,
94        metadata: Option<String>,
95        ttl: Option<std::time::Duration>,
96    ) -> Result<(), KvStoreError> {
97        // manages ttl
98        let existing = self
99            .lookup(obj_store_key.clone(), obj_key.clone())
100            .map_err(|_| KvStoreError::InternalError)?;
101
102        if let Some(g) = generation {
103            if let Some(val) = &existing {
104                if val.generation != g {
105                    return Err(KvStoreError::PreconditionFailed);
106                }
107            }
108        }
109
110        let out_obj = match mode {
111            KvInsertMode::Overwrite => obj,
112            KvInsertMode::Add => {
113                if existing.is_some() {
114                    // key exists, add fails
115                    return Err(KvStoreError::PreconditionFailed);
116                }
117                obj
118            }
119            KvInsertMode::Append => {
120                let mut out_obj;
121                match existing {
122                    None => {
123                        out_obj = obj;
124                    }
125                    Some(v) => {
126                        out_obj = v.body;
127                        out_obj.append(&mut obj.clone());
128                    }
129                }
130                out_obj
131            }
132            KvInsertMode::Prepend => {
133                let mut out_obj;
134                match existing {
135                    None => {
136                        out_obj = obj;
137                    }
138                    Some(mut v) => {
139                        out_obj = obj;
140                        out_obj.append(&mut v.body);
141                    }
142                }
143                out_obj
144            }
145        };
146
147        let exp = ttl.map(|t| SystemTime::now() + t);
148
149        let mut obj_val = ObjectValue {
150            body: out_obj,
151            metadata: String::new(),
152            metadata_len: 0,
153            generation: SystemTime::now()
154                .duration_since(SystemTime::UNIX_EPOCH)
155                .unwrap()
156                .as_nanos() as u64,
157            expiration: exp,
158        };
159
160        // magic number hack to ensure a case for integration tests
161        if obj_val.generation == 1337 {
162            obj_val.generation = 1338;
163        }
164
165        if let Some(m) = metadata {
166            obj_val.metadata_len = m.len();
167            obj_val.metadata = m;
168        }
169
170        self.stores
171            .write()
172            .map_err(|_| KvStoreError::InternalError)?
173            .entry(obj_store_key)
174            .and_modify(|store| {
175                store.insert(obj_key.clone(), obj_val.clone());
176            })
177            .or_insert_with(|| {
178                let mut store = BTreeMap::new();
179                store.insert(obj_key, obj_val);
180                store
181            });
182
183        Ok(())
184    }
185
186    pub fn delete(
187        &self,
188        obj_store_key: ObjectStoreKey,
189        obj_key: ObjectKey,
190    ) -> Result<bool, KvStoreError> {
191        let mut res = Ok(true);
192
193        self.stores
194            .write()
195            .map_err(|_| KvStoreError::InternalError)?
196            .entry(obj_store_key)
197            .and_modify(|store| match store.get(&obj_key) {
198                // 404 if the key doesn't exist, otherwise delete
199                Some(val) => {
200                    // manages ttl
201                    if let Some(exp) = val.expiration {
202                        if SystemTime::now() >= exp {
203                            res = Ok(false);
204                        }
205                    }
206                    store.remove(&obj_key);
207                }
208                None => res = Ok(false),
209            });
210
211        res
212    }
213
214    pub fn list(
215        &self,
216        obj_store_key: ObjectStoreKey,
217        cursor: Option<String>,
218        prefix: Option<String>,
219        limit: u32,
220    ) -> Result<Vec<u8>, KvStoreError> {
221        let mut res = Err(KvStoreError::InternalError);
222
223        let cursor = match cursor {
224            Some(c) => {
225                let cursor_bytes = BASE64_STANDARD
226                    .decode(c)
227                    .map_err(|_| KvStoreError::BadRequest)?;
228                let decoded =
229                    String::from_utf8(cursor_bytes).map_err(|_| KvStoreError::BadRequest)?;
230                Some(decoded)
231            }
232            None => None,
233        };
234
235        self.stores
236            .write()
237            .map_err(|_| KvStoreError::InternalError)?
238            .entry(obj_store_key.clone())
239            .and_modify(|store| {
240                // manages ttl
241                // a bit wasteful to run this loop twice, but we need mutable access to store,
242                // and it's already claimed in the filters below
243                let ttl_list = store.iter_mut().map(|(k, _)| k.clone()).collect::<Vec<_>>();
244                ttl_list.into_iter().for_each(|k| {
245                    let val = store.get(&k);
246                    if let Some(v) = val {
247                        if let Some(exp) = v.expiration {
248                            if SystemTime::now() >= exp {
249                                store.remove(&k);
250                            }
251                        }
252                    }
253                });
254
255                let mut list = store
256                    .iter_mut()
257                    .filter(|(k, _)| {
258                        if let Some(c) = &cursor {
259                            &k.0 > c
260                        } else {
261                            true
262                        }
263                    })
264                    .filter(|(k, _)| {
265                        if let Some(p) = &prefix {
266                            k.0.starts_with(p)
267                        } else {
268                            true
269                        }
270                    })
271                    .map(|(k, _)| String::from_utf8(k.0.as_bytes().to_vec()).unwrap())
272                    .collect::<Vec<_>>();
273
274                // limit
275                let old_len = list.len();
276                list.truncate(limit as usize);
277                let new_len = list.len();
278
279                let next_cursor = match old_len != new_len {
280                    true => Some(BASE64_STANDARD.encode(&list[new_len - 1])),
281                    false => None,
282                };
283
284                #[derive(Serialize)]
285                struct Metadata {
286                    limit: u32,
287                    #[serde(skip_serializing_if = "Option::is_none")]
288                    prefix: Option<String>,
289                    #[serde(skip_serializing_if = "Option::is_none")]
290                    next_cursor: Option<String>,
291                }
292                #[derive(Serialize)]
293                struct JsonOutput {
294                    data: Vec<String>,
295                    meta: Metadata,
296                }
297
298                let body = JsonOutput {
299                    data: list,
300                    meta: Metadata {
301                        limit,
302                        prefix,
303                        next_cursor,
304                    },
305                };
306
307                match serde_json::to_string(&body).map_err(|_| KvStoreError::InternalError) {
308                    Ok(s) => res = Ok(s.as_bytes().to_vec()),
309                    Err(e) => res = Err(e),
310                };
311            });
312        res
313    }
314}
315
316#[derive(Debug, PartialOrd, Ord, PartialEq, Eq, Clone, Default)]
317pub struct ObjectStoreKey(String);
318
319impl ObjectStoreKey {
320    pub fn new(key: impl ToString) -> Self {
321        Self(key.to_string())
322    }
323}
324
325#[derive(Debug, PartialOrd, Ord, PartialEq, Eq, Clone, Default)]
326pub struct ObjectKey(String);
327
328impl ObjectKey {
329    pub fn new(key: impl ToString) -> Result<Self, KeyValidationError> {
330        let key = key.to_string();
331        is_valid_key(&key)?;
332        Ok(Self(key))
333    }
334}
335
336#[derive(Debug, PartialOrd, Ord, PartialEq, Eq, thiserror::Error)]
337pub enum ObjectStoreError {
338    #[error("The object was not in the store")]
339    MissingObject,
340    #[error("Viceroy's ObjectStore lock was poisoned")]
341    PoisonedLock,
342    /// An Object Store with the given name was not found.
343    #[error("Unknown object-store: {0}")]
344    UnknownObjectStore(String),
345}
346
347impl From<&ObjectStoreError> for FastlyStatus {
348    fn from(e: &ObjectStoreError) -> Self {
349        use ObjectStoreError::*;
350        match e {
351            MissingObject => FastlyStatus::None,
352            PoisonedLock => panic!("{}", e),
353            UnknownObjectStore(_) => FastlyStatus::Inval,
354        }
355    }
356}
357
358#[derive(Debug, PartialOrd, Ord, PartialEq, Eq, thiserror::Error)]
359pub enum KvStoreError {
360    #[error("The error was not set")]
361    Uninitialized,
362    #[error("KV store cannot or will not process the request due to something that is perceived to be a client error")]
363    BadRequest,
364    #[error("KV store cannot fulfill the request, as definied by the client's prerequisites (ie. if-generation-match)")]
365    PreconditionFailed,
366    #[error("The size limit for a KV store key was exceeded")]
367    PayloadTooLarge,
368    #[error("The system encountered an unexpected internal error")]
369    InternalError,
370    #[error("Too many requests have been made to the KV store")]
371    TooManyRequests,
372}
373
374impl From<&KvError> for Result<Option<()>, KvStoreError> {
375    fn from(e: &KvError) -> Self {
376        Err(match e {
377            KvError::Ok => return Ok(Some(())),
378            KvError::NotFound => return Ok(None),
379            KvError::Uninitialized => KvStoreError::Uninitialized,
380            KvError::BadRequest => KvStoreError::BadRequest,
381            KvError::PreconditionFailed => KvStoreError::PreconditionFailed,
382            KvError::PayloadTooLarge => KvStoreError::PayloadTooLarge,
383            KvError::InternalError => KvStoreError::InternalError,
384            KvError::TooManyRequests => KvStoreError::TooManyRequests,
385        })
386    }
387}
388
389impl From<&KvStoreError> for KvError {
390    fn from(e: &KvStoreError) -> Self {
391        match e {
392            KvStoreError::Uninitialized => KvError::Uninitialized,
393            KvStoreError::BadRequest => KvError::BadRequest,
394            KvStoreError::PreconditionFailed => KvError::PreconditionFailed,
395            KvStoreError::PayloadTooLarge => KvError::PayloadTooLarge,
396            KvStoreError::InternalError => KvError::InternalError,
397            KvStoreError::TooManyRequests => KvError::TooManyRequests,
398        }
399    }
400}
401
402impl From<&KvStoreError> for FastlyStatus {
403    fn from(e: &KvStoreError) -> Self {
404        match e {
405            KvStoreError::Uninitialized => panic!("{}", e),
406            KvStoreError::BadRequest => FastlyStatus::Inval,
407            KvStoreError::PreconditionFailed => FastlyStatus::Inval,
408            KvStoreError::PayloadTooLarge => FastlyStatus::Inval,
409            KvStoreError::InternalError => FastlyStatus::Inval,
410            KvStoreError::TooManyRequests => FastlyStatus::Inval,
411        }
412    }
413}
414
415/// Keys in the Object Store must follow the following rules:
416///
417///   * Keys can contain any sequence of valid Unicode characters, of length 1-1024 bytes when
418///     UTF-8 encoded.
419///   * Keys cannot contain Carriage Return or Line Feed characters.
420///   * Keys cannot start with `.well-known/acme-challenge/`.
421///   * Keys cannot be named `.` or `..`.
422///   * Keys cannot use Unicode characters 0 through 32, 65534 and 65535 as
423///     single-character key names.  (0x0 through 0x20, 0xFFFE and 0xFFFF)
424fn is_valid_key(key: &str) -> Result<(), KeyValidationError> {
425    let len = key.as_bytes().len();
426    if len < 1 {
427        return Err(KeyValidationError::EmptyKey);
428    } else if len > 1024 {
429        return Err(KeyValidationError::Over1024Bytes);
430    }
431
432    if key.starts_with(".well-known/acme-challenge") {
433        return Err(KeyValidationError::StartsWithWellKnown);
434    }
435
436    if key.eq("..") || key.contains("../") || key.ends_with("/..") {
437        return Err(KeyValidationError::ContainsDotDot);
438    } else if key.eq(".") || key.contains("./") || key.ends_with("/.") {
439        return Err(KeyValidationError::ContainsDot);
440    } else if key.contains('\r') {
441        return Err(KeyValidationError::Contains("\r".to_owned()));
442    } else if key.contains('\n') {
443        return Err(KeyValidationError::Contains("\n".to_owned()));
444    } else if key.contains('#') {
445        return Err(KeyValidationError::Contains("#".to_owned()));
446    } else if key.contains(';') {
447        return Err(KeyValidationError::Contains(";".to_owned()));
448    } else if key.contains('?') {
449        return Err(KeyValidationError::Contains("?".to_owned()));
450    } else if key.contains('^') {
451        return Err(KeyValidationError::Contains("^".to_owned()));
452    } else if key.contains('|') {
453        return Err(KeyValidationError::Contains("|".to_owned()));
454    }
455
456    if key.len() == 1 {
457        let k = key.chars().next().unwrap();
458        match k {
459            '\u{0}'..='\u{20}' => {
460                return Err(KeyValidationError::Contains(k.escape_unicode().to_string()));
461            }
462            '\u{FFFE}'..='\u{FFFF}' => {
463                return Err(KeyValidationError::Contains(k.escape_unicode().to_string()));
464            }
465            _ => {}
466        }
467    }
468
469    Ok(())
470}
471
472#[derive(Debug, thiserror::Error)]
473pub enum KeyValidationError {
474    #[error("Keys for objects cannot be empty")]
475    EmptyKey,
476    #[error("Keys for objects cannot be over 1024 bytes in size")]
477    Over1024Bytes,
478    #[error("Keys for objects cannot start with `.well-known/acme-challenge`")]
479    StartsWithWellKnown,
480    #[error("Keys for objects cannot be named `.`")]
481    ContainsDot,
482    #[error("Keys for objects cannot be named `..`")]
483    ContainsDotDot,
484    #[error("Keys for objects cannot contain a `{0}`")]
485    Contains(String),
486}
487
488#[cfg(test)]
489mod tests {
490    use super::*;
491
492    const STORE_NAME: &'static str = "test_store";
493
494    #[test]
495    fn test_kv_store_exists() {
496        let stores = ObjectStores::default();
497        stores
498            .insert_empty_store(ObjectStoreKey(STORE_NAME.to_string()))
499            .unwrap();
500
501        let res = stores.store_exists(STORE_NAME);
502        match res {
503            Ok(true) => {}
504            Ok(false) => panic!("should have been Ok(true)"),
505            Err(e) => panic!("should not have been Err({:?})", e),
506        }
507    }
508
509    #[test]
510    fn test_kv_store_basics() {
511        let stores = ObjectStores::default();
512        stores
513            .insert_empty_store(ObjectStoreKey(STORE_NAME.to_string()))
514            .unwrap();
515
516        let key = "insert_key".to_string();
517        let val1 = "val1".to_string();
518
519        // insert
520        let res = stores.insert(
521            ObjectStoreKey(STORE_NAME.to_string()),
522            ObjectKey(key.clone()),
523            val1.clone().into(),
524            KvInsertMode::Overwrite,
525            None,
526            None,
527            None,
528        );
529        match res {
530            Err(e) => panic!("should not have been Err({:?})", e),
531            _ => {}
532        }
533
534        // lookup
535        let res = stores.lookup(
536            ObjectStoreKey(STORE_NAME.to_string()),
537            ObjectKey(key.clone()),
538        );
539        match res {
540            Ok(Some(ov)) => {
541                assert_eq!(ov.body, val1.as_bytes().to_vec())
542            }
543            Ok(None) => panic!("should have been Ok(Some(_))"),
544            Err(_) => panic!("should have been Ok(_)"),
545        }
546
547        // list
548        let limit = 1000;
549        let res = stores.list(ObjectStoreKey(STORE_NAME.to_string()), None, None, limit);
550        match res {
551            Ok(ov) => {
552                let val = format!(r#"{{"data":["{key}"],"meta":{{"limit":{limit}}}}}"#);
553                assert_eq!(std::str::from_utf8(&ov).unwrap(), val)
554            }
555            Err(e) => panic!("should not have been Err({:?})", e),
556        }
557
558        // delete
559        let res = stores.delete(
560            ObjectStoreKey(STORE_NAME.to_string()),
561            ObjectKey(key.clone()),
562        );
563        match res {
564            Ok(_) => {}
565            Err(e) => panic!("should not have been Err({:?})", e),
566        }
567    }
568
569    #[test]
570    fn test_kv_store_item_404s() {
571        let stores = ObjectStores::default();
572        stores
573            .insert_empty_store(ObjectStoreKey(STORE_NAME.to_string()))
574            .unwrap();
575
576        let res = stores.lookup(
577            ObjectStoreKey(STORE_NAME.to_string()),
578            ObjectKey("bad_key".to_string()),
579        );
580        match res {
581            Ok(Some(_)) => panic!("should have been Ok(None)"),
582            Ok(None) => {}
583            Err(e) => panic!("should not have been Err({:?})", e),
584        }
585
586        let res = stores.delete(
587            ObjectStoreKey(STORE_NAME.to_string()),
588            ObjectKey("bad_key".to_string()),
589        );
590        match res {
591            Ok(true) => panic!("should have been Ok(false)"),
592            Ok(false) => {}
593            Err(e) => panic!("should not have been Err({:?})", e),
594        }
595    }
596
597    #[test]
598    fn test_kv_store_item_insert_modes() {
599        let stores = ObjectStores::default();
600        stores
601            .insert_empty_store(ObjectStoreKey(STORE_NAME.to_string()))
602            .unwrap();
603
604        let key = "insert_key".to_string();
605        let val1 = "val1".to_string();
606        let val2 = "val2".to_string();
607        let val3 = "val3".to_string();
608
609        let res = stores.insert(
610            ObjectStoreKey(STORE_NAME.to_string()),
611            ObjectKey(key.clone()),
612            val1.clone().into(),
613            KvInsertMode::Add,
614            None,
615            None,
616            None,
617        );
618        assert!(res.is_ok());
619        // fail on Add, because key already exists
620        let res = stores.insert(
621            ObjectStoreKey(STORE_NAME.to_string()),
622            ObjectKey(key.clone()),
623            val1.clone().into(),
624            KvInsertMode::Add,
625            None,
626            None,
627            None,
628        );
629        match res {
630            Ok(_) => panic!("should not have been OK"),
631            Err(e) => assert_eq!(e, KvStoreError::PreconditionFailed),
632        }
633        // prepend val2
634        let res = stores.insert(
635            ObjectStoreKey(STORE_NAME.to_string()),
636            ObjectKey(key.clone()),
637            val2.clone().into(),
638            KvInsertMode::Prepend,
639            None,
640            None,
641            None,
642        );
643        match res {
644            Err(e) => panic!("should not have been Err({:?})", e),
645            _ => {}
646        }
647        // append val3
648        let res = stores.insert(
649            ObjectStoreKey(STORE_NAME.to_string()),
650            ObjectKey(key.clone()),
651            val3.clone().into(),
652            KvInsertMode::Append,
653            None,
654            None,
655            None,
656        );
657        match res {
658            Err(e) => panic!("should not have been Err({:?})", e),
659            _ => {}
660        }
661        let res = stores.lookup(
662            ObjectStoreKey(STORE_NAME.to_string()),
663            ObjectKey(key.clone()),
664        );
665        match res {
666            Ok(Some(ov)) => {
667                let val = format!("{val2}{val1}{val3}");
668                assert_eq!(ov.body, val.as_bytes().to_vec())
669            }
670            Ok(None) => panic!("should have been Ok(Some((_))"),
671            Err(e) => panic!("should not have been Err({:?})", e),
672        }
673
674        // overwrite val3
675        let res = stores.insert(
676            ObjectStoreKey(STORE_NAME.to_string()),
677            ObjectKey(key.clone()),
678            val3.clone().into(),
679            KvInsertMode::Overwrite,
680            None,
681            Some(val2.clone()),
682            None,
683        );
684        match res {
685            Err(e) => panic!("should not have been Err({:?})", e),
686            _ => {}
687        }
688
689        // test overwrite
690        let res = stores.lookup(
691            ObjectStoreKey(STORE_NAME.to_string()),
692            ObjectKey(key.clone()),
693        );
694        match res {
695            Ok(Some(ov)) => {
696                assert_eq!(ov.body, val3.as_bytes().to_vec());
697                assert_eq!(ov.metadata, val2);
698            }
699            Ok(None) => panic!("should have been Ok(Some(_))"),
700            Err(e) => panic!("should not have been Err({:?})", e),
701        }
702    }
703
704    #[test]
705    fn test_kv_store_item_insert_generation() {
706        let stores = ObjectStores::default();
707        stores
708            .insert_empty_store(ObjectStoreKey(STORE_NAME.to_string()))
709            .unwrap();
710
711        let key = "insert_key".to_string();
712        let val1 = "val1".to_string();
713
714        // insert val1
715        let res = stores.insert(
716            ObjectStoreKey(STORE_NAME.to_string()),
717            ObjectKey(key.clone()),
718            val1.clone().into(),
719            KvInsertMode::Overwrite,
720            None,
721            None,
722            None,
723        );
724        match res {
725            Err(e) => panic!("should not have been Err({:?})", e),
726            _ => {}
727        }
728
729        // test overwrite, get gen
730        let generation;
731        let res = stores.lookup(
732            ObjectStoreKey(STORE_NAME.to_string()),
733            ObjectKey(key.clone()),
734        );
735        match res {
736            Ok(Some(ov)) => {
737                assert_eq!(ov.body, val1.as_bytes().to_vec());
738                generation = ov.generation;
739            }
740            Ok(None) => panic!("should have been Ok(Some(_))"),
741            Err(e) => panic!("should not have been Err({:?})", e),
742        }
743
744        // test generation match failure
745        let res = stores.insert(
746            ObjectStoreKey(STORE_NAME.to_string()),
747            ObjectKey(key.clone()),
748            val1.clone().into(),
749            KvInsertMode::Overwrite,
750            Some(1337),
751            None,
752            None,
753        );
754        match res {
755            Err(KvStoreError::PreconditionFailed) => {}
756            _ => panic!("should have been Err(KvStoreError::PreconditionFailed)"),
757        }
758
759        // test generation match positive
760        let res = stores.insert(
761            ObjectStoreKey(STORE_NAME.to_string()),
762            ObjectKey(key.clone()),
763            val1.clone().into(),
764            KvInsertMode::Overwrite,
765            Some(generation),
766            None,
767            None,
768        );
769        match res {
770            Ok(_) => {}
771            Err(e) => panic!("should not have been Err({:?})", e),
772        }
773
774        // check result
775        let res = stores.lookup(
776            ObjectStoreKey(STORE_NAME.to_string()),
777            ObjectKey(key.clone()),
778        );
779        match res {
780            Ok(Some(ov)) => {
781                assert_eq!(ov.body, val1.as_bytes().to_vec());
782            }
783            Ok(None) => panic!("should have been Ok(Some(_))"),
784            Err(e) => panic!("should not have been Err({:?})", e),
785        }
786    }
787
788    #[test]
789    fn test_kv_store_item_list_advanced() {
790        let stores = ObjectStores::default();
791        stores
792            .insert_empty_store(ObjectStoreKey(STORE_NAME.to_string()))
793            .unwrap();
794
795        let key = "insert_key".to_string();
796        let prefix = "key".to_string();
797        let key1 = format!("{prefix}1").to_string();
798        let key2 = format!("{prefix}2").to_string();
799        let key3 = format!("{prefix}3").to_string();
800        let val1 = "val1".to_string();
801        let val2 = "val2".to_string();
802        let val3 = "val3".to_string();
803
804        // insert insert_key
805        let res = stores.insert(
806            ObjectStoreKey(STORE_NAME.to_string()),
807            ObjectKey(key.clone()),
808            val1.clone().into(),
809            KvInsertMode::Overwrite,
810            None,
811            None,
812            None,
813        );
814        match res {
815            Err(e) => panic!("should not have been Err({:?})", e),
816            _ => {}
817        }
818
819        // insert val1
820        let res = stores.insert(
821            ObjectStoreKey(STORE_NAME.to_string()),
822            ObjectKey(key1.clone()),
823            val1.clone().into(),
824            KvInsertMode::Overwrite,
825            None,
826            None,
827            None,
828        );
829        match res {
830            Err(e) => panic!("should not have been Err({:?})", e),
831            _ => {}
832        }
833        // insert val2
834        let res = stores.insert(
835            ObjectStoreKey(STORE_NAME.to_string()),
836            ObjectKey(key2.clone()),
837            val2.clone().into(),
838            KvInsertMode::Overwrite,
839            None,
840            None,
841            None,
842        );
843        match res {
844            Err(e) => panic!("should not have been Err({:?})", e),
845            _ => {}
846        }
847        // insert val3
848        let res = stores.insert(
849            ObjectStoreKey(STORE_NAME.to_string()),
850            ObjectKey(key3.clone()),
851            val3.clone().into(),
852            KvInsertMode::Overwrite,
853            None,
854            None,
855            None,
856        );
857        match res {
858            Err(e) => panic!("should not have been Err({:?})", e),
859            _ => {}
860        }
861
862        // list
863        let limit = 1000;
864        let res = stores.list(ObjectStoreKey(STORE_NAME.to_string()), None, None, limit);
865        match res {
866            Ok(ov) => {
867                let val = format!(
868                    r#"{{"data":["{key}","{key1}","{key2}","{key3}"],"meta":{{"limit":{limit}}}}}"#
869                );
870                assert_eq!(std::str::from_utf8(&ov).unwrap(), val)
871            }
872            Err(e) => panic!("should not have been Err({:?})", e),
873        }
874
875        // list w/prefix
876        let limit = 1000;
877        let res = stores.list(
878            ObjectStoreKey(STORE_NAME.to_string()),
879            None,
880            Some(prefix.clone()),
881            limit,
882        );
883        match res {
884            Ok(ov) => {
885                let val = format!(
886                    r#"{{"data":["{key1}","{key2}","{key3}"],"meta":{{"limit":{limit},"prefix":"{prefix}"}}}}"#
887                );
888                assert_eq!(std::str::from_utf8(&ov).unwrap(), val)
889            }
890            Err(e) => panic!("should not have been Err({:?})", e),
891        }
892
893        // list w/prefix&limit
894        let limit = 1;
895        let res = stores.list(
896            ObjectStoreKey(STORE_NAME.to_string()),
897            None,
898            Some(prefix.clone()),
899            limit,
900        );
901        match res {
902            Ok(ov) => {
903                let next_cursor = BASE64_STANDARD.encode(key1.clone());
904                let val = format!(
905                    r#"{{"data":["{key1}"],"meta":{{"limit":{limit},"prefix":"{prefix}","next_cursor":"{next_cursor}"}}}}"#
906                );
907                assert_eq!(std::str::from_utf8(&ov).unwrap(), val)
908            }
909            Err(e) => panic!("should not have been Err({:?})", e),
910        }
911
912        // list w/prefix&limit&cursor
913        let limit = 1;
914        let last_cursor = BASE64_STANDARD.encode(key1.clone());
915        let res = stores.list(
916            ObjectStoreKey(STORE_NAME.to_string()),
917            Some(last_cursor),
918            Some(prefix.clone()),
919            limit,
920        );
921        match res {
922            Ok(ov) => {
923                let next_cursor = BASE64_STANDARD.encode(key2.clone());
924                let val = format!(
925                    r#"{{"data":["{key2}"],"meta":{{"limit":{limit},"prefix":"{prefix}","next_cursor":"{next_cursor}"}}}}"#
926                );
927                assert_eq!(std::str::from_utf8(&ov).unwrap(), val)
928            }
929            Err(e) => panic!("should not have been Err({:?})", e),
930        }
931
932        // list w/prefix&limit&cursor
933        let limit = 1;
934        let last_cursor = BASE64_STANDARD.encode(key2.clone());
935        let res = stores.list(
936            ObjectStoreKey(STORE_NAME.to_string()),
937            Some(last_cursor),
938            Some(prefix.clone()),
939            limit,
940        );
941        match res {
942            Ok(ov) => {
943                let val = format!(
944                    r#"{{"data":["{key3}"],"meta":{{"limit":{limit},"prefix":"{prefix}"}}}}"#
945                );
946                assert_eq!(std::str::from_utf8(&ov).unwrap(), val)
947            }
948            Err(e) => panic!("should not have been Err({:?})", e),
949        }
950    }
951}
viceroy_lib/object_store.rs

viceroy_lib/
object_store.rs
