dbx_core/storage/
memory_wos.rs1use crate::error::DbxResult;
6use crate::storage::StorageBackend;
7use std::collections::{BTreeMap, HashMap};
8use std::ops::RangeBounds;
9use std::sync::RwLock;
10
11pub struct InMemoryWosBackend {
13 #[allow(clippy::type_complexity)]
14 tables: RwLock<HashMap<String, BTreeMap<Vec<u8>, Vec<u8>>>>,
15}
16
17impl InMemoryWosBackend {
18 pub fn new() -> Self {
20 Self {
21 tables: RwLock::new(HashMap::new()),
22 }
23 }
24}
25
26impl Default for InMemoryWosBackend {
27 fn default() -> Self {
28 Self::new()
29 }
30}
31
32impl StorageBackend for InMemoryWosBackend {
33 fn insert(&self, table: &str, key: &[u8], value: &[u8]) -> DbxResult<()> {
34 let mut tables = self.tables.write().unwrap();
35 tables
36 .entry(table.to_string())
37 .or_default()
38 .insert(key.to_vec(), value.to_vec());
39 Ok(())
40 }
41
42 fn get(&self, table: &str, key: &[u8]) -> DbxResult<Option<Vec<u8>>> {
43 let tables = self.tables.read().unwrap();
44 Ok(tables.get(table).and_then(|map| map.get(key).cloned()))
45 }
46
47 fn delete(&self, table: &str, key: &[u8]) -> DbxResult<bool> {
48 let mut tables = self.tables.write().unwrap();
49 if let Some(map) = tables.get_mut(table) {
50 Ok(map.remove(key).is_some())
51 } else {
52 Ok(false)
53 }
54 }
55
56 fn scan<R: RangeBounds<Vec<u8>> + Clone>(
57 &self,
58 table: &str,
59 range: R,
60 ) -> DbxResult<Vec<(Vec<u8>, Vec<u8>)>> {
61 let tables = self.tables.read().unwrap();
62 if let Some(map) = tables.get(table) {
63 Ok(map
64 .range(range)
65 .map(|(k, v)| (k.clone(), v.clone()))
66 .collect())
67 } else {
68 Ok(Vec::new())
69 }
70 }
71
72 fn scan_one<R: RangeBounds<Vec<u8>> + Clone>(
73 &self,
74 table: &str,
75 range: R,
76 ) -> DbxResult<Option<(Vec<u8>, Vec<u8>)>> {
77 let tables = self.tables.read().unwrap();
78 if let Some(map) = tables.get(table) {
79 Ok(map.range(range).map(|(k, v)| (k.clone(), v.clone())).next())
80 } else {
81 Ok(None)
82 }
83 }
84
85 fn flush(&self) -> DbxResult<()> {
86 Ok(())
88 }
89
90 fn count(&self, table: &str) -> DbxResult<usize> {
91 let tables = self.tables.read().unwrap();
92 Ok(tables.get(table).map(|m| m.len()).unwrap_or(0))
93 }
94
95 fn table_names(&self) -> DbxResult<Vec<String>> {
96 let tables = self.tables.read().unwrap();
97 Ok(tables.keys().cloned().collect())
98 }
99}
100
101#[cfg(test)]
102mod tests {
103 use super::*;
104
105 #[test]
106 fn test_insert_and_get() {
107 let backend = InMemoryWosBackend::new();
108 backend.insert("test", b"key1", b"value1").unwrap();
109
110 let result = backend.get("test", b"key1").unwrap();
111 assert_eq!(result, Some(b"value1".to_vec()));
112 }
113
114 #[test]
115 fn test_delete() {
116 let backend = InMemoryWosBackend::new();
117 backend.insert("test", b"key1", b"value1").unwrap();
118
119 assert!(backend.delete("test", b"key1").unwrap());
120 assert_eq!(backend.get("test", b"key1").unwrap(), None);
121 }
122
123 #[test]
124 fn test_scan() {
125 let backend = InMemoryWosBackend::new();
126 backend.insert("test", b"key1", b"value1").unwrap();
127 backend.insert("test", b"key2", b"value2").unwrap();
128 backend.insert("test", b"key3", b"value3").unwrap();
129
130 let results = backend
131 .scan("test", b"key1".to_vec()..b"key3".to_vec())
132 .unwrap();
133 assert_eq!(results.len(), 2);
134 }
135
136 #[test]
137 fn test_count() {
138 let backend = InMemoryWosBackend::new();
139 backend.insert("test", b"key1", b"value1").unwrap();
140 backend.insert("test", b"key2", b"value2").unwrap();
141
142 assert_eq!(backend.count("test").unwrap(), 2);
143 }
144}