uqa_storage/key_value/
memory_store.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! In-memory implementation of the backend-neutral key/value traits.
8
9use super::{
10    BTreeMap, KeyValueBatch, KeyValueBatchOperation, KeyValueStore, Mutex, StorageBackendError,
11    StorageBackendResult,
12};
13
14/// In-memory Key/Value store used by trait-level tests and future non-SQL
15/// fixtures.
16#[derive(Debug, Default)]
17pub struct MemoryKeyValueStore {
18    inner: Mutex<MemoryKeyValueState>,
19}
20
21#[derive(Debug, Default, Clone)]
22struct MemoryKeyValueState {
23    map: BTreeMap<Vec<u8>, Vec<u8>>,
24    transactions: Vec<BTreeMap<Vec<u8>, Vec<u8>>>,
25    savepoints: Vec<MemorySavepoint>,
26    transaction_read_only: bool,
27    transaction_written: bool,
28    change_version: u64,
29}
30
31#[derive(Debug, Clone)]
32struct MemorySavepoint {
33    name: String,
34    snapshot: BTreeMap<Vec<u8>, Vec<u8>>,
35}
36
37impl MemoryKeyValueStore {
38    pub fn new() -> Self {
39        Self::default()
40    }
41}
42
43impl KeyValueStore for MemoryKeyValueStore {
44    fn get(&self, key: &[u8]) -> StorageBackendResult<Option<Vec<u8>>> {
45        Ok(self.inner.lock().map.get(key).cloned())
46    }
47
48    fn contains_key(&self, key: &[u8]) -> StorageBackendResult<bool> {
49        Ok(self.inner.lock().map.contains_key(key))
50    }
51
52    fn put(&self, key: &[u8], value: &[u8]) -> StorageBackendResult<()> {
53        let mut inner = self.inner.lock();
54        prepare_write(&mut inner)?;
55        inner.map.insert(key.to_vec(), value.to_vec());
56        finish_autocommit_write(&mut inner);
57        Ok(())
58    }
59
60    fn delete(&self, key: &[u8]) -> StorageBackendResult<()> {
61        let mut inner = self.inner.lock();
62        prepare_write(&mut inner)?;
63        inner.map.remove(key);
64        finish_autocommit_write(&mut inner);
65        Ok(())
66    }
67
68    fn scan_prefix(&self, prefix: &[u8]) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
69        Ok(self
70            .inner
71            .lock()
72            .map
73            .range(prefix.to_vec()..)
74            .take_while(|(key, _)| key.starts_with(prefix))
75            .map(|(key, value)| (key.clone(), value.clone()))
76            .collect())
77    }
78
79    fn scan_prefix_after(
80        &self,
81        prefix: &[u8],
82        after: Option<&[u8]>,
83        limit: usize,
84    ) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
85        use std::ops::Bound::{Excluded, Included, Unbounded};
86
87        if limit == 0 {
88            return Ok(Vec::new());
89        }
90        let inner = self.inner.lock();
91        let lower = match after {
92            Some(after) if after >= prefix => Excluded(after.to_vec()),
93            Some(_) | None => Included(prefix.to_vec()),
94        };
95        Ok(inner
96            .map
97            .range((lower, Unbounded))
98            .take_while(|(key, _)| key.starts_with(prefix))
99            .take(limit)
100            .map(|(key, value)| (key.clone(), value.clone()))
101            .collect())
102    }
103
104    fn scan_prefix_keys_after(
105        &self,
106        prefix: &[u8],
107        after: Option<&[u8]>,
108        limit: usize,
109    ) -> StorageBackendResult<Vec<Vec<u8>>> {
110        use std::ops::Bound::{Excluded, Included, Unbounded};
111
112        if limit == 0 {
113            return Ok(Vec::new());
114        }
115        let inner = self.inner.lock();
116        let lower = match after {
117            Some(after) if after >= prefix => Excluded(after.to_vec()),
118            Some(_) | None => Included(prefix.to_vec()),
119        };
120        Ok(inner
121            .map
122            .range((lower, Unbounded))
123            .take_while(|(key, _)| key.starts_with(prefix))
124            .take(limit)
125            .map(|(key, _)| key.clone())
126            .collect())
127    }
128
129    fn first_prefix_after(
130        &self,
131        prefix: &[u8],
132        after: Option<&[u8]>,
133    ) -> StorageBackendResult<Option<(Vec<u8>, Vec<u8>)>> {
134        use std::ops::Bound::{Excluded, Included, Unbounded};
135
136        let inner = self.inner.lock();
137        let lower = match after {
138            Some(after) if after >= prefix => Excluded(after.to_vec()),
139            Some(_) | None => Included(prefix.to_vec()),
140        };
141        Ok(inner
142            .map
143            .range((lower, Unbounded))
144            .next()
145            .filter(|(key, _)| key.starts_with(prefix))
146            .map(|(key, value)| (key.clone(), value.clone())))
147    }
148
149    fn delete_prefix(&self, prefix: &[u8]) -> StorageBackendResult<usize> {
150        let mut inner = self.inner.lock();
151        prepare_write(&mut inner)?;
152        let keys = inner
153            .map
154            .range(prefix.to_vec()..)
155            .take_while(|(key, _)| key.starts_with(prefix))
156            .map(|(key, _)| key.clone())
157            .collect::<Vec<_>>();
158        for key in &keys {
159            inner.map.remove(key);
160        }
161        finish_autocommit_write(&mut inner);
162        Ok(keys.len())
163    }
164
165    fn batch(&self) -> Box<dyn KeyValueBatch + '_> {
166        Box::new(MemoryKeyValueBatch {
167            store: self,
168            operations: Vec::new(),
169        })
170    }
171
172    fn begin_transaction(&self) -> StorageBackendResult<()> {
173        let mut inner = self.inner.lock();
174        if !inner.transactions.is_empty() {
175            return Err(StorageBackendError::Other(
176                "a KeyValue transaction is already open".into(),
177            ));
178        }
179        let snapshot = inner.map.clone();
180        inner.transactions.push(snapshot);
181        inner.transaction_read_only = false;
182        inner.transaction_written = false;
183        Ok(())
184    }
185
186    fn begin_read_transaction(&self) -> StorageBackendResult<()> {
187        let mut inner = self.inner.lock();
188        if !inner.transactions.is_empty() {
189            return Err(StorageBackendError::Other(
190                "a KeyValue transaction is already open".into(),
191            ));
192        }
193        let snapshot = inner.map.clone();
194        inner.transactions.push(snapshot);
195        inner.transaction_read_only = true;
196        inner.transaction_written = false;
197        Ok(())
198    }
199
200    fn in_transaction(&self) -> bool {
201        !self.inner.lock().transactions.is_empty()
202    }
203
204    fn transaction_has_written(&self) -> StorageBackendResult<bool> {
205        Ok(self.inner.lock().transaction_written)
206    }
207
208    fn change_version(&self) -> StorageBackendResult<Option<u64>> {
209        Ok(Some(self.inner.lock().change_version))
210    }
211
212    fn commit_transaction(&self) -> StorageBackendResult<()> {
213        let mut inner = self.inner.lock();
214        inner.transactions.pop().ok_or_else(|| {
215            StorageBackendError::Other("no open KeyValue transaction to commit".into())
216        })?;
217        if inner.transaction_written {
218            inner.change_version = inner.change_version.wrapping_add(1);
219        }
220        inner.transaction_read_only = false;
221        inner.transaction_written = false;
222        inner.savepoints.clear();
223        Ok(())
224    }
225
226    fn rollback_transaction(&self) -> StorageBackendResult<()> {
227        let mut inner = self.inner.lock();
228        let snapshot = inner.transactions.pop().ok_or_else(|| {
229            StorageBackendError::Other("no open KeyValue transaction to roll back".into())
230        })?;
231        inner.map = snapshot;
232        inner.transaction_read_only = false;
233        inner.transaction_written = false;
234        inner.savepoints.clear();
235        Ok(())
236    }
237
238    fn savepoint(&self, name: &str) -> StorageBackendResult<()> {
239        let mut inner = self.inner.lock();
240        if inner.transactions.is_empty() {
241            return Err(StorageBackendError::Other(
242                "cannot create a savepoint outside a KeyValue transaction".into(),
243            ));
244        }
245        let snapshot = inner.map.clone();
246        inner.savepoints.push(MemorySavepoint {
247            name: name.to_string(),
248            snapshot,
249        });
250        Ok(())
251    }
252
253    fn release_savepoint(&self, name: &str) -> StorageBackendResult<()> {
254        let mut inner = self.inner.lock();
255        let position = inner
256            .savepoints
257            .iter()
258            .rposition(|savepoint| savepoint.name == name)
259            .ok_or_else(|| StorageBackendError::Other(format!("unknown savepoint `{name}`")))?;
260        inner.savepoints.truncate(position);
261        Ok(())
262    }
263
264    fn rollback_to_savepoint(&self, name: &str) -> StorageBackendResult<()> {
265        let mut inner = self.inner.lock();
266        let position = inner
267            .savepoints
268            .iter()
269            .rposition(|savepoint| savepoint.name == name)
270            .ok_or_else(|| StorageBackendError::Other(format!("unknown savepoint `{name}`")))?;
271        inner.map = inner.savepoints[position].snapshot.clone();
272        inner.savepoints.truncate(position + 1);
273        Ok(())
274    }
275}
276
277struct MemoryKeyValueBatch<'a> {
278    store: &'a MemoryKeyValueStore,
279    operations: Vec<KeyValueBatchOperation>,
280}
281
282impl KeyValueBatch for MemoryKeyValueBatch<'_> {
283    fn put(&mut self, key: &[u8], value: &[u8]) -> StorageBackendResult<()> {
284        self.operations
285            .push(KeyValueBatchOperation::Put(key.to_vec(), value.to_vec()));
286        Ok(())
287    }
288
289    fn delete(&mut self, key: &[u8]) -> StorageBackendResult<()> {
290        self.operations
291            .push(KeyValueBatchOperation::Delete(key.to_vec()));
292        Ok(())
293    }
294
295    fn delete_prefix(&mut self, prefix: &[u8]) -> StorageBackendResult<()> {
296        self.operations
297            .push(KeyValueBatchOperation::DeletePrefix(prefix.to_vec()));
298        Ok(())
299    }
300
301    fn commit(self: Box<Self>) -> StorageBackendResult<()> {
302        let mut inner = self.store.inner.lock();
303        prepare_write(&mut inner)?;
304        for operation in self.operations {
305            match operation {
306                KeyValueBatchOperation::Put(key, value) => {
307                    inner.map.insert(key, value);
308                }
309                KeyValueBatchOperation::Delete(key) => {
310                    inner.map.remove(&key);
311                }
312                KeyValueBatchOperation::DeletePrefix(prefix) => {
313                    let keys = inner
314                        .map
315                        .range(prefix.clone()..)
316                        .take_while(|(key, _)| key.starts_with(&prefix))
317                        .map(|(key, _)| key.clone())
318                        .collect::<Vec<_>>();
319                    for key in keys {
320                        inner.map.remove(&key);
321                    }
322                }
323            }
324        }
325        finish_autocommit_write(&mut inner);
326        Ok(())
327    }
328}
329
330fn prepare_write(inner: &mut MemoryKeyValueState) -> StorageBackendResult<()> {
331    if !inner.transactions.is_empty() && inner.transaction_read_only {
332        return Err(StorageBackendError::Other(
333            "cannot write in a read-only KeyValue transaction".into(),
334        ));
335    }
336    if !inner.transactions.is_empty() {
337        inner.transaction_written = true;
338    }
339    Ok(())
340}
341
342fn finish_autocommit_write(inner: &mut MemoryKeyValueState) {
343    if inner.transactions.is_empty() {
344        inner.change_version = inner.change_version.wrapping_add(1);
345    }
346}
uqa_storage/key_value/memory_store.rs

uqa_storage/key_value/
memory_store.rs
