1use std::collections::HashMap;
2use std::sync::Arc;
3
4use async_trait::async_trait;
5use serde::de::DeserializeOwned;
6use serde::Serialize;
7use tokio::sync::RwLock;
8
9pub mod namespaced;
10pub use namespaced::NamespacedStore;
11
12#[derive(Debug, thiserror::Error)]
13pub enum StoreError {
14 #[error("Serialization error: {0}")]
15 Serialization(String),
16
17 #[error("Database error: {0}")]
18 Database(String),
19
20 #[error("Not found: {0}")]
21 NotFound(String),
22}
23
24impl From<sled::Error> for StoreError {
25 fn from(e: sled::Error) -> Self {
26 Self::Database(e.to_string())
27 }
28}
29
30#[async_trait]
32pub trait Store: Send + Sync {
33 async fn get_raw(&self, key: &str) -> Result<Option<Vec<u8>>, StoreError>;
34 async fn set_raw(&self, key: &str, value: Vec<u8>) -> Result<(), StoreError>;
35 async fn delete(&self, key: &str) -> Result<(), StoreError>;
36 async fn exists(&self, key: &str) -> Result<bool, StoreError>;
37 async fn scan_prefix_raw(&self, prefix: &str) -> Result<Vec<(String, Vec<u8>)>, StoreError>;
38 async fn list_keys(&self, prefix: &str) -> Result<Vec<String>, StoreError>;
39 async fn clear(&self) -> Result<(), StoreError>;
40}
41
42#[async_trait]
44pub trait StoreExt: Store {
45 async fn get<T: DeserializeOwned + Send>(&self, key: &str) -> Result<Option<T>, StoreError> {
46 match self.get_raw(key).await? {
47 Some(bytes) => {
48 let value =
49 serde_json::from_slice(&bytes).map_err(|e| StoreError::Serialization(e.to_string()))?;
50 Ok(Some(value))
51 }
52 None => Ok(None),
53 }
54 }
55
56 async fn set<T: Serialize + Send + Sync>(&self, key: &str, value: &T) -> Result<(), StoreError> {
57 let bytes =
58 serde_json::to_vec(value).map_err(|e| StoreError::Serialization(e.to_string()))?;
59 self.set_raw(key, bytes).await
60 }
61
62 async fn scan_prefix<T: DeserializeOwned + Send>(
63 &self,
64 prefix: &str,
65 ) -> Result<Vec<(String, T)>, StoreError> {
66 let mut results = self.scan_prefix_unsorted(prefix).await?;
67 results.sort_by(|a, b| a.0.cmp(&b.0));
68 Ok(results)
69 }
70
71 async fn scan_prefix_unsorted<T: DeserializeOwned + Send>(
74 &self,
75 prefix: &str,
76 ) -> Result<Vec<(String, T)>, StoreError> {
77 let raw = self.scan_prefix_raw(prefix).await?;
78 let mut results = Vec::new();
79 for (key, bytes) in raw {
80 let value = serde_json::from_slice(&bytes)
81 .map_err(|e| StoreError::Serialization(e.to_string()))?;
82 results.push((key, value));
83 }
84 Ok(results)
85 }
86}
87
88impl<T: Store + ?Sized> StoreExt for T {}
89
90#[derive(Clone)]
93pub struct InMemoryStore {
94 data: Arc<RwLock<HashMap<String, Vec<u8>>>>,
95}
96
97impl InMemoryStore {
98 #[must_use]
99 pub fn new() -> Self {
100 Self {
101 data: Arc::new(RwLock::new(HashMap::new())),
102 }
103 }
104}
105
106impl Default for InMemoryStore {
107 fn default() -> Self {
108 Self::new()
109 }
110}
111
112#[async_trait]
113impl Store for InMemoryStore {
114 async fn get_raw(&self, key: &str) -> Result<Option<Vec<u8>>, StoreError> {
115 let data = self.data.read().await;
116 Ok(data.get(key).cloned())
117 }
118
119 async fn set_raw(&self, key: &str, value: Vec<u8>) -> Result<(), StoreError> {
120 let mut data = self.data.write().await;
121 data.insert(key.to_string(), value);
122 Ok(())
123 }
124
125 async fn delete(&self, key: &str) -> Result<(), StoreError> {
126 let mut data = self.data.write().await;
127 data.remove(key);
128 Ok(())
129 }
130
131 async fn exists(&self, key: &str) -> Result<bool, StoreError> {
132 let data = self.data.read().await;
133 Ok(data.contains_key(key))
134 }
135
136 async fn scan_prefix_raw(&self, prefix: &str) -> Result<Vec<(String, Vec<u8>)>, StoreError> {
137 let data = self.data.read().await;
138 let mut results = Vec::new();
139 for (key, bytes) in data.iter() {
140 if key.starts_with(prefix) {
141 results.push((key.clone(), bytes.clone()));
142 }
143 }
144 results.sort_by(|a, b| a.0.cmp(&b.0));
145 Ok(results)
146 }
147
148 async fn list_keys(&self, prefix: &str) -> Result<Vec<String>, StoreError> {
149 let data = self.data.read().await;
150 let mut keys: Vec<String> = data
151 .keys()
152 .filter(|k| k.starts_with(prefix))
153 .cloned()
154 .collect();
155 keys.sort();
156 Ok(keys)
157 }
158
159 async fn clear(&self) -> Result<(), StoreError> {
160 let mut data = self.data.write().await;
161 data.clear();
162 Ok(())
163 }
164}
165
166#[derive(Clone)]
169pub struct SledStore {
170 db: sled::Db,
171}
172
173impl SledStore {
174 pub fn new(path: impl AsRef<std::path::Path>) -> Result<Self, StoreError> {
177 let db = sled::open(path).map_err(|e| StoreError::Database(e.to_string()))?;
178 Ok(Self { db })
179 }
180
181 pub fn new_temporary() -> Result<Self, StoreError> {
184 let db = sled::Config::new()
185 .temporary(true)
186 .open()
187 .map_err(|e| StoreError::Database(e.to_string()))?;
188 Ok(Self { db })
189 }
190}
191
192#[async_trait]
193impl Store for SledStore {
194 async fn get_raw(&self, key: &str) -> Result<Option<Vec<u8>>, StoreError> {
195 match self.db.get(key.as_bytes())? {
196 Some(ivec) => Ok(Some(ivec.to_vec())),
197 None => Ok(None),
198 }
199 }
200
201 async fn set_raw(&self, key: &str, value: Vec<u8>) -> Result<(), StoreError> {
202 self.db.insert(key.as_bytes(), value)?;
203 Ok(())
204 }
205
206 async fn delete(&self, key: &str) -> Result<(), StoreError> {
207 self.db.remove(key.as_bytes())?;
208 Ok(())
209 }
210
211 async fn exists(&self, key: &str) -> Result<bool, StoreError> {
212 Ok(self.db.contains_key(key.as_bytes())?)
213 }
214
215 async fn scan_prefix_raw(&self, prefix: &str) -> Result<Vec<(String, Vec<u8>)>, StoreError> {
216 let mut results = Vec::new();
217 for result in self.db.scan_prefix(prefix.as_bytes()) {
218 let (key_bytes, value_bytes) = result?;
219 let key = String::from_utf8_lossy(&key_bytes).to_string();
220 let value = value_bytes.to_vec();
221 results.push((key, value));
222 }
223 results.sort_by(|a, b| a.0.cmp(&b.0));
224 Ok(results)
225 }
226
227 async fn list_keys(&self, prefix: &str) -> Result<Vec<String>, StoreError> {
228 let mut keys = Vec::new();
229 for result in self.db.scan_prefix(prefix.as_bytes()) {
230 let (key_bytes, _) = result?;
231 let key = String::from_utf8_lossy(&key_bytes).to_string();
232 keys.push(key);
233 }
234 keys.sort();
235 Ok(keys)
236 }
237
238 async fn clear(&self) -> Result<(), StoreError> {
239 self.db.clear()?;
240 Ok(())
241 }
242}
243
244#[cfg(test)]
247#[path = "lib_test.rs"]
248mod tests;
249