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 visit_value(
45        &self,
46        key: &[u8],
47        control: &crate::read_control::StorageReadControl,
48        visit: &mut crate::read_control::ValueReadVisitor<'_>,
49    ) -> StorageBackendResult<()> {
50        control.check()?;
51        let state = self.inner.lock();
52        control.check()?;
53        visit(state.map.get(key).map(Vec::as_slice))?;
54        control.check()
55    }
56
57    fn visit_prefix_after(
58        &self,
59        prefix: &[u8],
60        after: Option<&[u8]>,
61        limit: usize,
62        control: &crate::read_control::StorageReadControl,
63        visit: &mut crate::read_control::KeyValueReadVisitor<'_>,
64    ) -> StorageBackendResult<()> {
65        use std::ops::Bound::{Excluded, Included, Unbounded};
66        control.check()?;
67        if limit == 0 {
68            return Ok(());
69        }
70        let state = self.inner.lock();
71        let lower = match after {
72            Some(after) if after >= prefix => Excluded(after),
73            _ => Included(prefix),
74        };
75        for (key, value) in state.map.range::<[u8], _>((lower, Unbounded)).take(limit) {
76            control.check()?;
77            if !key.starts_with(prefix) {
78                break;
79            }
80            visit(key, value)?;
81        }
82        control.check()
83    }
84
85    fn get(&self, key: &[u8]) -> StorageBackendResult<Option<Vec<u8>>> {
86        Ok(self.inner.lock().map.get(key).cloned())
87    }
88
89    fn contains_key(&self, key: &[u8]) -> StorageBackendResult<bool> {
90        Ok(self.inner.lock().map.contains_key(key))
91    }
92
93    fn put(&self, key: &[u8], value: &[u8]) -> StorageBackendResult<()> {
94        let mut inner = self.inner.lock();
95        prepare_write(&mut inner)?;
96        inner.map.insert(key.to_vec(), value.to_vec());
97        finish_autocommit_write(&mut inner);
98        Ok(())
99    }
100
101    fn delete(&self, key: &[u8]) -> StorageBackendResult<()> {
102        let mut inner = self.inner.lock();
103        prepare_write(&mut inner)?;
104        inner.map.remove(key);
105        finish_autocommit_write(&mut inner);
106        Ok(())
107    }
108
109    fn scan_prefix(&self, prefix: &[u8]) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
110        Ok(self
111            .inner
112            .lock()
113            .map
114            .range(prefix.to_vec()..)
115            .take_while(|(key, _)| key.starts_with(prefix))
116            .map(|(key, value)| (key.clone(), value.clone()))
117            .collect())
118    }
119
120    fn scan_prefix_after(
121        &self,
122        prefix: &[u8],
123        after: Option<&[u8]>,
124        limit: usize,
125    ) -> StorageBackendResult<Vec<(Vec<u8>, Vec<u8>)>> {
126        use std::ops::Bound::{Excluded, Included, Unbounded};
127
128        if limit == 0 {
129            return Ok(Vec::new());
130        }
131        let inner = self.inner.lock();
132        let lower = match after {
133            Some(after) if after >= prefix => Excluded(after.to_vec()),
134            Some(_) | None => Included(prefix.to_vec()),
135        };
136        Ok(inner
137            .map
138            .range((lower, Unbounded))
139            .take_while(|(key, _)| key.starts_with(prefix))
140            .take(limit)
141            .map(|(key, value)| (key.clone(), value.clone()))
142            .collect())
143    }
144
145    fn scan_prefix_keys_after(
146        &self,
147        prefix: &[u8],
148        after: Option<&[u8]>,
149        limit: usize,
150    ) -> StorageBackendResult<Vec<Vec<u8>>> {
151        use std::ops::Bound::{Excluded, Included, Unbounded};
152
153        if limit == 0 {
154            return Ok(Vec::new());
155        }
156        let inner = self.inner.lock();
157        let lower = match after {
158            Some(after) if after >= prefix => Excluded(after.to_vec()),
159            Some(_) | None => Included(prefix.to_vec()),
160        };
161        Ok(inner
162            .map
163            .range((lower, Unbounded))
164            .take_while(|(key, _)| key.starts_with(prefix))
165            .take(limit)
166            .map(|(key, _)| key.clone())
167            .collect())
168    }
169
170    fn first_prefix_after(
171        &self,
172        prefix: &[u8],
173        after: Option<&[u8]>,
174    ) -> StorageBackendResult<Option<(Vec<u8>, Vec<u8>)>> {
175        use std::ops::Bound::{Excluded, Included, Unbounded};
176
177        let inner = self.inner.lock();
178        let lower = match after {
179            Some(after) if after >= prefix => Excluded(after.to_vec()),
180            Some(_) | None => Included(prefix.to_vec()),
181        };
182        Ok(inner
183            .map
184            .range((lower, Unbounded))
185            .next()
186            .filter(|(key, _)| key.starts_with(prefix))
187            .map(|(key, value)| (key.clone(), value.clone())))
188    }
189
190    fn delete_prefix(&self, prefix: &[u8]) -> StorageBackendResult<usize> {
191        let mut inner = self.inner.lock();
192        prepare_write(&mut inner)?;
193        let keys = inner
194            .map
195            .range(prefix.to_vec()..)
196            .take_while(|(key, _)| key.starts_with(prefix))
197            .map(|(key, _)| key.clone())
198            .collect::<Vec<_>>();
199        for key in &keys {
200            inner.map.remove(key);
201        }
202        finish_autocommit_write(&mut inner);
203        Ok(keys.len())
204    }
205
206    fn batch(&self) -> Box<dyn KeyValueBatch + '_> {
207        Box::new(MemoryKeyValueBatch {
208            store: self,
209            operations: Vec::new(),
210        })
211    }
212
213    fn begin_transaction(&self) -> StorageBackendResult<()> {
214        let mut inner = self.inner.lock();
215        if !inner.transactions.is_empty() {
216            return Err(StorageBackendError::Other(
217                "a KeyValue transaction is already open".into(),
218            ));
219        }
220        let snapshot = inner.map.clone();
221        inner.transactions.push(snapshot);
222        inner.transaction_read_only = false;
223        inner.transaction_written = false;
224        Ok(())
225    }
226
227    fn begin_read_transaction(&self) -> StorageBackendResult<()> {
228        let mut inner = self.inner.lock();
229        if !inner.transactions.is_empty() {
230            return Err(StorageBackendError::Other(
231                "a KeyValue transaction is already open".into(),
232            ));
233        }
234        let snapshot = inner.map.clone();
235        inner.transactions.push(snapshot);
236        inner.transaction_read_only = true;
237        inner.transaction_written = false;
238        Ok(())
239    }
240
241    fn in_transaction(&self) -> bool {
242        !self.inner.lock().transactions.is_empty()
243    }
244
245    fn transaction_has_written(&self) -> StorageBackendResult<bool> {
246        Ok(self.inner.lock().transaction_written)
247    }
248
249    fn change_version(&self) -> StorageBackendResult<Option<u64>> {
250        Ok(Some(self.inner.lock().change_version))
251    }
252
253    fn commit_transaction(&self) -> StorageBackendResult<()> {
254        let mut inner = self.inner.lock();
255        inner.transactions.pop().ok_or_else(|| {
256            StorageBackendError::Other("no open KeyValue transaction to commit".into())
257        })?;
258        if inner.transaction_written {
259            inner.change_version = inner.change_version.wrapping_add(1);
260        }
261        inner.transaction_read_only = false;
262        inner.transaction_written = false;
263        inner.savepoints.clear();
264        Ok(())
265    }
266
267    fn rollback_transaction(&self) -> StorageBackendResult<()> {
268        let mut inner = self.inner.lock();
269        let snapshot = inner.transactions.pop().ok_or_else(|| {
270            StorageBackendError::Other("no open KeyValue transaction to roll back".into())
271        })?;
272        inner.map = snapshot;
273        inner.transaction_read_only = false;
274        inner.transaction_written = false;
275        inner.savepoints.clear();
276        Ok(())
277    }
278
279    fn savepoint(&self, name: &str) -> StorageBackendResult<()> {
280        let mut inner = self.inner.lock();
281        if inner.transactions.is_empty() {
282            return Err(StorageBackendError::Other(
283                "cannot create a savepoint outside a KeyValue transaction".into(),
284            ));
285        }
286        let snapshot = inner.map.clone();
287        inner.savepoints.push(MemorySavepoint {
288            name: name.to_string(),
289            snapshot,
290        });
291        Ok(())
292    }
293
294    fn release_savepoint(&self, name: &str) -> StorageBackendResult<()> {
295        let mut inner = self.inner.lock();
296        let position = inner
297            .savepoints
298            .iter()
299            .rposition(|savepoint| savepoint.name == name)
300            .ok_or_else(|| StorageBackendError::Other(format!("unknown savepoint `{name}`")))?;
301        inner.savepoints.truncate(position);
302        Ok(())
303    }
304
305    fn rollback_to_savepoint(&self, name: &str) -> StorageBackendResult<()> {
306        let mut inner = self.inner.lock();
307        let position = inner
308            .savepoints
309            .iter()
310            .rposition(|savepoint| savepoint.name == name)
311            .ok_or_else(|| StorageBackendError::Other(format!("unknown savepoint `{name}`")))?;
312        inner.map = inner.savepoints[position].snapshot.clone();
313        inner.savepoints.truncate(position + 1);
314        Ok(())
315    }
316}
317
318struct MemoryKeyValueBatch<'a> {
319    store: &'a MemoryKeyValueStore,
320    operations: Vec<KeyValueBatchOperation>,
321}
322
323impl KeyValueBatch for MemoryKeyValueBatch<'_> {
324    fn put(&mut self, key: &[u8], value: &[u8]) -> StorageBackendResult<()> {
325        self.operations
326            .push(KeyValueBatchOperation::Put(key.to_vec(), value.to_vec()));
327        Ok(())
328    }
329
330    fn delete(&mut self, key: &[u8]) -> StorageBackendResult<()> {
331        self.operations
332            .push(KeyValueBatchOperation::Delete(key.to_vec()));
333        Ok(())
334    }
335
336    fn delete_prefix(&mut self, prefix: &[u8]) -> StorageBackendResult<()> {
337        self.operations
338            .push(KeyValueBatchOperation::DeletePrefix(prefix.to_vec()));
339        Ok(())
340    }
341
342    fn commit(self: Box<Self>) -> StorageBackendResult<()> {
343        let mut inner = self.store.inner.lock();
344        prepare_write(&mut inner)?;
345        for operation in self.operations {
346            match operation {
347                KeyValueBatchOperation::Put(key, value) => {
348                    inner.map.insert(key, value);
349                }
350                KeyValueBatchOperation::Delete(key) => {
351                    inner.map.remove(&key);
352                }
353                KeyValueBatchOperation::DeletePrefix(prefix) => {
354                    let keys = inner
355                        .map
356                        .range(prefix.clone()..)
357                        .take_while(|(key, _)| key.starts_with(&prefix))
358                        .map(|(key, _)| key.clone())
359                        .collect::<Vec<_>>();
360                    for key in keys {
361                        inner.map.remove(&key);
362                    }
363                }
364            }
365        }
366        finish_autocommit_write(&mut inner);
367        Ok(())
368    }
369}
370
371fn prepare_write(inner: &mut MemoryKeyValueState) -> StorageBackendResult<()> {
372    if !inner.transactions.is_empty() && inner.transaction_read_only {
373        return Err(StorageBackendError::Other(
374            "cannot write in a read-only KeyValue transaction".into(),
375        ));
376    }
377    if !inner.transactions.is_empty() {
378        inner.transaction_written = true;
379    }
380    Ok(())
381}
382
383fn finish_autocommit_write(inner: &mut MemoryKeyValueState) {
384    if inner.transactions.is_empty() {
385        inner.change_version = inner.change_version.wrapping_add(1);
386    }
387}
uqa_storage/key_value/memory_store.rs

uqa_storage/key_value/
memory_store.rs
