use thiserror::Error;
#[non_exhaustive]
#[derive(Debug, Error)]
pub enum StorageError {
#[error("Failed to store data: {0}")]
StoreFailed(String),
#[error("Failed to load data: {0}")]
LoadFailed(String),
#[error("Failed to delete data: {0}")]
DeleteFailed(String),
#[error("Key not found: {0}")]
KeyNotFound(String),
#[error("Corrupted data: {0}")]
CorruptedData(String),
#[error("Storage unavailable: {0}")]
Unavailable(String),
}
pub type StorageResult<T> = std::result::Result<T, StorageError>;
pub trait MlsStorage: Send + Sync {
fn store(&self, key_type: &str, key_id: &str, data: &[u8]) -> StorageResult<()>;
fn load(&self, key_type: &str, key_id: &str) -> StorageResult<Option<Vec<u8>>>;
fn delete(&self, key_type: &str, key_id: &str) -> StorageResult<()>;
fn list_keys(&self, key_type: &str) -> StorageResult<Vec<String>>;
fn exists(&self, key_type: &str, key_id: &str) -> StorageResult<bool> {
Ok(self.load(key_type, key_id)?.is_some())
}
fn clear_type(&self, key_type: &str) -> StorageResult<()> {
for key_id in self.list_keys(key_type)? {
self.delete(key_type, &key_id)?;
}
Ok(())
}
}
#[derive(Debug, Default)]
pub struct InMemoryStorage {
data: std::sync::RwLock<std::collections::HashMap<String, Vec<u8>>>,
}
impl InMemoryStorage {
pub fn new() -> Self {
Self::default()
}
fn make_key(key_type: &str, key_id: &str) -> String {
format!("{}:{}", key_type, key_id)
}
}
impl MlsStorage for InMemoryStorage {
fn store(&self, key_type: &str, key_id: &str, data: &[u8]) -> StorageResult<()> {
let mut storage = self
.data
.write()
.map_err(|e| StorageError::StoreFailed(e.to_string()))?;
storage.insert(Self::make_key(key_type, key_id), data.to_vec());
Ok(())
}
fn load(&self, key_type: &str, key_id: &str) -> StorageResult<Option<Vec<u8>>> {
let storage = self
.data
.read()
.map_err(|e| StorageError::LoadFailed(e.to_string()))?;
Ok(storage.get(&Self::make_key(key_type, key_id)).cloned())
}
fn delete(&self, key_type: &str, key_id: &str) -> StorageResult<()> {
let mut storage = self
.data
.write()
.map_err(|e| StorageError::DeleteFailed(e.to_string()))?;
storage.remove(&Self::make_key(key_type, key_id));
Ok(())
}
fn list_keys(&self, key_type: &str) -> StorageResult<Vec<String>> {
let storage = self
.data
.read()
.map_err(|e| StorageError::LoadFailed(e.to_string()))?;
let prefix = format!("{}:", key_type);
let keys: Vec<String> = storage
.keys()
.filter_map(|k| {
if k.starts_with(&prefix) {
Some(k[prefix.len()..].to_string())
} else {
None
}
})
.collect();
Ok(keys)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_in_memory_storage_basic_operations() {
let storage = InMemoryStorage::new();
storage.store("test", "key1", b"value1").unwrap();
let loaded = storage.load("test", "key1").unwrap();
assert_eq!(loaded, Some(b"value1".to_vec()));
let loaded = storage.load("test", "key2").unwrap();
assert_eq!(loaded, None);
storage.delete("test", "key1").unwrap();
let loaded = storage.load("test", "key1").unwrap();
assert_eq!(loaded, None);
}
#[test]
fn test_in_memory_storage_list_keys() {
let storage = InMemoryStorage::new();
storage.store("type_a", "key1", b"v1").unwrap();
storage.store("type_a", "key2", b"v2").unwrap();
storage.store("type_b", "key3", b"v3").unwrap();
let mut keys_a = storage.list_keys("type_a").unwrap();
keys_a.sort();
assert_eq!(keys_a, vec!["key1", "key2"]);
let keys_b = storage.list_keys("type_b").unwrap();
assert_eq!(keys_b, vec!["key3"]);
let keys_c = storage.list_keys("type_c").unwrap();
assert!(keys_c.is_empty());
}
#[test]
fn test_in_memory_storage_exists() {
let storage = InMemoryStorage::new();
assert!(!storage.exists("test", "key1").unwrap());
storage.store("test", "key1", b"value").unwrap();
assert!(storage.exists("test", "key1").unwrap());
storage.delete("test", "key1").unwrap();
assert!(!storage.exists("test", "key1").unwrap());
}
#[test]
fn test_in_memory_storage_clear_type() {
let storage = InMemoryStorage::new();
storage.store("type_a", "key1", b"v1").unwrap();
storage.store("type_a", "key2", b"v2").unwrap();
storage.store("type_b", "key3", b"v3").unwrap();
storage.clear_type("type_a").unwrap();
assert!(storage.list_keys("type_a").unwrap().is_empty());
assert_eq!(storage.list_keys("type_b").unwrap().len(), 1);
}
}