Skip to main content

origin_storage/
memory.rs

1use crate::{Record, Storage, StorageKey};
2use async_trait::async_trait;
3use origin_domain::Result;
4use std::collections::HashMap;
5use tokio::sync::RwLock;
6
7/// Non-persistent storage for tests and headless runs.
8#[derive(Debug, Default)]
9pub struct MemoryStorage {
10    records: RwLock<HashMap<StorageKey, Record>>,
11}
12
13impl MemoryStorage {
14    pub fn new() -> Self {
15        Self::default()
16    }
17}
18
19#[async_trait]
20impl Storage for MemoryStorage {
21    async fn get(&self, key: &StorageKey) -> Result<Option<Record>> {
22        Ok(self.records.read().await.get(key).cloned())
23    }
24
25    async fn put(&self, key: &StorageKey, record: Record) -> Result<()> {
26        self.records.write().await.insert(key.clone(), record);
27        Ok(())
28    }
29
30    async fn delete(&self, key: &StorageKey) -> Result<()> {
31        self.records.write().await.remove(key);
32        Ok(())
33    }
34
35    async fn keys(&self, namespace: &str) -> Result<Vec<StorageKey>> {
36        Ok(self
37            .records
38            .read()
39            .await
40            .keys()
41            .filter(|key| key.namespace() == namespace)
42            .cloned()
43            .collect())
44    }
45
46    async fn clear(&self, namespace: &str) -> Result<()> {
47        self.records
48            .write()
49            .await
50            .retain(|key, _| key.namespace() != namespace);
51        Ok(())
52    }
53
54    async fn clear_prefix(&self, prefix: &str) -> Result<usize> {
55        let mut records = self.records.write().await;
56        let before = records.len();
57        records.retain(|key, _| !key.namespace().starts_with(prefix));
58        Ok(before - records.len())
59    }
60}