eshanized_polaris_core/
storage.rs1use crate::errors::PolarisResult;
4use async_trait::async_trait;
5use bytes::Bytes;
6use std::sync::Arc;
7
8#[async_trait]
10pub trait Storage: Send + Sync {
11 async fn get(&self, key: &str) -> PolarisResult<Option<Bytes>>;
13
14 async fn set(&self, key: &str, value: Bytes) -> PolarisResult<()>;
16
17 async fn delete(&self, key: &str) -> PolarisResult<()>;
19
20 async fn exists(&self, key: &str) -> PolarisResult<bool>;
22
23 async fn list(&self, prefix: &str) -> PolarisResult<Vec<String>>;
25}
26
27#[derive(Debug, Clone)]
29pub struct InMemoryStorage {
30 data: Arc<dashmap::DashMap<String, Bytes>>,
31}
32
33impl InMemoryStorage {
34 pub fn new() -> Self {
36 Self {
37 data: Arc::new(dashmap::DashMap::new()),
38 }
39 }
40}
41
42impl Default for InMemoryStorage {
43 fn default() -> Self {
44 Self::new()
45 }
46}
47
48#[async_trait]
49impl Storage for InMemoryStorage {
50 async fn get(&self, key: &str) -> PolarisResult<Option<Bytes>> {
51 Ok(self.data.get(key).map(|v| v.clone()))
52 }
53
54 async fn set(&self, key: &str, value: Bytes) -> PolarisResult<()> {
55 self.data.insert(key.to_string(), value);
56 Ok(())
57 }
58
59 async fn delete(&self, key: &str) -> PolarisResult<()> {
60 self.data.remove(key);
61 Ok(())
62 }
63
64 async fn exists(&self, key: &str) -> PolarisResult<bool> {
65 Ok(self.data.contains_key(key))
66 }
67
68 async fn list(&self, prefix: &str) -> PolarisResult<Vec<String>> {
69 Ok(self
70 .data
71 .iter()
72 .filter(|entry| entry.key().starts_with(prefix))
73 .map(|entry| entry.key().clone())
74 .collect())
75 }
76}
77
78#[cfg(test)]
87mod tests {
88 use super::*;
89
90 #[tokio::test]
91 async fn test_in_memory_storage() {
92 let storage = InMemoryStorage::new();
93
94 storage
96 .set("key1", Bytes::from("value1"))
97 .await
98 .unwrap();
99 let value = storage.get("key1").await.unwrap();
100 assert_eq!(value, Some(Bytes::from("value1")));
101
102 assert!(storage.exists("key1").await.unwrap());
104 assert!(!storage.exists("key2").await.unwrap());
105
106 storage.delete("key1").await.unwrap();
108 assert!(!storage.exists("key1").await.unwrap());
109 }
110
111 #[tokio::test]
112 async fn test_storage_list() {
113 let storage = InMemoryStorage::new();
114
115 storage.set("prefix:key1", Bytes::from("v1")).await.unwrap();
116 storage.set("prefix:key2", Bytes::from("v2")).await.unwrap();
117 storage.set("other:key3", Bytes::from("v3")).await.unwrap();
118
119 let keys = storage.list("prefix:").await.unwrap();
120 assert_eq!(keys.len(), 2);
121 assert!(keys.contains(&"prefix:key1".to_string()));
122 assert!(keys.contains(&"prefix:key2".to_string()));
123 }
124}