use crate::error::DbxResult;
use crate::storage::StorageBackend;
use std::collections::{BTreeMap, HashMap};
use std::ops::RangeBounds;
use std::sync::RwLock;
pub struct InMemoryWosBackend {
#[allow(clippy::type_complexity)]
tables: RwLock<HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>>,
}
impl InMemoryWosBackend {
pub fn new() -> Self {
Self {
tables: RwLock::new(HashMap::new()),
}
}
}
impl Default for InMemoryWosBackend {
fn default() -> Self {
Self::new()
}
}
impl StorageBackend for InMemoryWosBackend {
fn insert(&self, table: &str, key: &[u8], value: &[u8]) -> DbxResult<()> {
let mut tables = self.tables.write().unwrap();
tables
.entry(table.to_string())
.or_default()
.insert(key.to_vec(), value.to_vec());
Ok(())
}
fn get(&self, table: &str, key: &[u8]) -> DbxResult<Option<Vec<u8>>> {
let tables = self.tables.read().unwrap();
Ok(tables.get(table).and_then(|map| map.get(key).cloned()))
}
fn delete(&self, table: &str, key: &[u8]) -> DbxResult<bool> {
let mut tables = self.tables.write().unwrap();
if let Some(map) = tables.get_mut(table) {
Ok(map.remove(key).is_some())
} else {
Ok(false)
}
}
fn scan<R: RangeBounds<Vec<u8>> + Clone>(
&self,
table: &str,
range: R,
) -> DbxResult<Vec<(Vec<u8>, Vec<u8>)>> {
let tables = self.tables.read().unwrap();
if let Some(map) = tables.get(table) {
Ok(map
.range(range)
.map(|(k, v)| (k.clone(), v.clone()))
.collect())
} else {
Ok(Vec::new())
}
}
fn scan_one<R: RangeBounds<Vec<u8>> + Clone>(
&self,
table: &str,
range: R,
) -> DbxResult<Option<(Vec<u8>, Vec<u8>)>> {
let tables = self.tables.read().unwrap();
if let Some(map) = tables.get(table) {
Ok(map.range(range).map(|(k, v)| (k.clone(), v.clone())).next())
} else {
Ok(None)
}
}
fn flush(&self) -> DbxResult<()> {
Ok(())
}
fn count(&self, table: &str) -> DbxResult<usize> {
let tables = self.tables.read().unwrap();
Ok(tables.get(table).map(|m| m.len()).unwrap_or(0))
}
fn table_names(&self) -> DbxResult<Vec<String>> {
let tables = self.tables.read().unwrap();
Ok(tables.keys().cloned().collect())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_insert_and_get() {
let backend = InMemoryWosBackend::new();
backend.insert("test", b"key1", b"value1").unwrap();
let result = backend.get("test", b"key1").unwrap();
assert_eq!(result, Some(b"value1".to_vec()));
}
#[test]
fn test_delete() {
let backend = InMemoryWosBackend::new();
backend.insert("test", b"key1", b"value1").unwrap();
assert!(backend.delete("test", b"key1").unwrap());
assert_eq!(backend.get("test", b"key1").unwrap(), None);
}
#[test]
fn test_scan() {
let backend = InMemoryWosBackend::new();
backend.insert("test", b"key1", b"value1").unwrap();
backend.insert("test", b"key2", b"value2").unwrap();
backend.insert("test", b"key3", b"value3").unwrap();
let results = backend
.scan("test", b"key1".to_vec()..b"key3".to_vec())
.unwrap();
assert_eq!(results.len(), 2);
}
#[test]
fn test_count() {
let backend = InMemoryWosBackend::new();
backend.insert("test", b"key1", b"value1").unwrap();
backend.insert("test", b"key2", b"value2").unwrap();
assert_eq!(backend.count("test").unwrap(), 2);
}
}