aztec_pxe/stores/
sled_store.rs

1//! Persistent KV store backed by sled — survives process restarts.
2//!
3//! Provides a `SledKvStore` that implements the `KvStore` trait using
4//! the `sled` embedded database, matching Phase 3 requirement for
5//! persistent storage.
6
7use std::path::Path;
8
9use async_trait::async_trait;
10use aztec_core::error::Error;
11
12use super::kv::KvStore;
13
14/// Persistent key-value store backed by sled.
15///
16/// Stores data on disk in a directory specified at creation time.
17/// Suitable for production use where PXE state must survive restarts.
18pub struct SledKvStore {
19    db: sled::Db,
20}
21
22impl SledKvStore {
23    /// Open or create a sled database at the given path.
24    pub fn open(path: impl AsRef<Path>) -> Result<Self, Error> {
25        let db = sled::open(path.as_ref())
26            .map_err(|e| Error::InvalidData(format!("failed to open sled database: {e}")))?;
27        Ok(Self { db })
28    }
29
30    /// Open a temporary sled database (for testing).
31    pub fn open_temporary() -> Result<Self, Error> {
32        let config = sled::Config::new().temporary(true);
33        let db = config.open().map_err(|e| {
34            Error::InvalidData(format!("failed to open temporary sled database: {e}"))
35        })?;
36        Ok(Self { db })
37    }
38
39    /// Flush all pending writes to disk.
40    pub fn flush(&self) -> Result<(), Error> {
41        self.db
42            .flush()
43            .map_err(|e| Error::InvalidData(format!("sled flush failed: {e}")))?;
44        Ok(())
45    }
46
47    /// Get the on-disk size estimate in bytes.
48    pub fn size_on_disk(&self) -> Result<u64, Error> {
49        self.db
50            .size_on_disk()
51            .map_err(|e| Error::InvalidData(format!("sled size_on_disk failed: {e}")))
52    }
53}
54
55#[async_trait]
56impl KvStore for SledKvStore {
57    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Error> {
58        self.db
59            .get(key)
60            .map(|opt| opt.map(|ivec| ivec.to_vec()))
61            .map_err(|e| Error::InvalidData(format!("sled get failed: {e}")))
62    }
63
64    async fn put(&self, key: &[u8], value: &[u8]) -> Result<(), Error> {
65        self.db
66            .insert(key, value)
67            .map_err(|e| Error::InvalidData(format!("sled put failed: {e}")))?;
68        Ok(())
69    }
70
71    async fn delete(&self, key: &[u8]) -> Result<(), Error> {
72        self.db
73            .remove(key)
74            .map_err(|e| Error::InvalidData(format!("sled delete failed: {e}")))?;
75        Ok(())
76    }
77
78    async fn list_prefix(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>, Error> {
79        let results: Result<Vec<_>, _> = self
80            .db
81            .scan_prefix(prefix)
82            .map(|result| {
83                result
84                    .map(|(k, v)| (k.to_vec(), v.to_vec()))
85                    .map_err(|e| Error::InvalidData(format!("sled scan failed: {e}")))
86            })
87            .collect();
88        results
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    fn make_store() -> SledKvStore {
97        SledKvStore::open_temporary().unwrap()
98    }
99
100    #[tokio::test]
101    async fn put_get_delete() {
102        let store = make_store();
103        assert!(store.get(b"key1").await.unwrap().is_none());
104
105        store.put(b"key1", b"value1").await.unwrap();
106        assert_eq!(store.get(b"key1").await.unwrap().unwrap(), b"value1");
107
108        store.delete(b"key1").await.unwrap();
109        assert!(store.get(b"key1").await.unwrap().is_none());
110    }
111
112    #[tokio::test]
113    async fn list_prefix_filtering() {
114        let store = make_store();
115        store.put(b"contract:a", b"1").await.unwrap();
116        store.put(b"contract:b", b"2").await.unwrap();
117        store.put(b"key:x", b"3").await.unwrap();
118
119        let results = store.list_prefix(b"contract:").await.unwrap();
120        assert_eq!(results.len(), 2);
121        assert_eq!(results[0].0, b"contract:a");
122        assert_eq!(results[1].0, b"contract:b");
123    }
124
125    #[tokio::test]
126    async fn overwrite_value() {
127        let store = make_store();
128        store.put(b"k", b"v1").await.unwrap();
129        store.put(b"k", b"v2").await.unwrap();
130        assert_eq!(store.get(b"k").await.unwrap().unwrap(), b"v2");
131    }
132
133    #[tokio::test]
134    async fn flush_succeeds() {
135        let store = make_store();
136        store.put(b"k", b"v").await.unwrap();
137        store.flush().unwrap();
138    }
139
140    #[tokio::test]
141    async fn size_on_disk_returns_value() {
142        let store = make_store();
143        let size = store.size_on_disk().unwrap();
144        // Just verify it returns a value
145        let _ = size;
146    }
147}