nt_memory/agentdb/
storage.rs

1//! Persistent storage backend for vectors and metadata
2
3use serde::Serialize;
4use std::path::Path;
5
6/// Storage backend trait
7#[async_trait::async_trait]
8pub trait StorageBackend: Send + Sync {
9    /// Store key-value pair
10    async fn put(&self, key: &[u8], value: &[u8]) -> anyhow::Result<()>;
11
12    /// Get value by key
13    async fn get(&self, key: &[u8]) -> anyhow::Result<Option<Vec<u8>>>;
14
15    /// Delete key
16    async fn delete(&self, key: &[u8]) -> anyhow::Result<()>;
17
18    /// List all keys with prefix
19    async fn list_prefix(&self, prefix: &[u8]) -> anyhow::Result<Vec<Vec<u8>>>;
20
21    /// Batch operations
22    async fn batch_put(&self, items: Vec<(Vec<u8>, Vec<u8>)>) -> anyhow::Result<()>;
23}
24
25/// Sled-based persistent store
26pub struct PersistentStore {
27    db: sled::Db,
28}
29
30impl PersistentStore {
31    /// Open or create database
32    pub fn new<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
33        let db = sled::open(path)?;
34        Ok(Self { db })
35    }
36
37    /// Create in-memory database (for testing)
38    pub fn memory() -> anyhow::Result<Self> {
39        let db = sled::Config::new().temporary(true).open()?;
40        Ok(Self { db })
41    }
42
43    /// Flush to disk
44    pub async fn flush(&self) -> anyhow::Result<()> {
45        self.db.flush_async().await?;
46        Ok(())
47    }
48
49    /// Database size on disk
50    pub fn size_on_disk(&self) -> anyhow::Result<u64> {
51        Ok(self.db.size_on_disk()?)
52    }
53
54    /// Number of entries
55    pub fn len(&self) -> usize {
56        self.db.len()
57    }
58
59    /// Check if empty
60    pub fn is_empty(&self) -> bool {
61        self.len() == 0
62    }
63}
64
65#[async_trait::async_trait]
66impl StorageBackend for PersistentStore {
67    async fn put(&self, key: &[u8], value: &[u8]) -> anyhow::Result<()> {
68        self.db.insert(key, value)?;
69        Ok(())
70    }
71
72    async fn get(&self, key: &[u8]) -> anyhow::Result<Option<Vec<u8>>> {
73        Ok(self.db.get(key)?.map(|v| v.to_vec()))
74    }
75
76    async fn delete(&self, key: &[u8]) -> anyhow::Result<()> {
77        self.db.remove(key)?;
78        Ok(())
79    }
80
81    async fn list_prefix(&self, prefix: &[u8]) -> anyhow::Result<Vec<Vec<u8>>> {
82        let keys: Vec<Vec<u8>> = self
83            .db
84            .scan_prefix(prefix)
85            .keys()
86            .filter_map(|k| k.ok().map(|k| k.to_vec()))
87            .collect();
88
89        Ok(keys)
90    }
91
92    async fn batch_put(&self, items: Vec<(Vec<u8>, Vec<u8>)>) -> anyhow::Result<()> {
93        let mut batch = sled::Batch::default();
94
95        for (key, value) in items {
96            batch.insert(key, value);
97        }
98
99        self.db.apply_batch(batch)?;
100        Ok(())
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[tokio::test]
109    async fn test_persistent_store_operations() {
110        let store = PersistentStore::memory().unwrap();
111
112        // Put
113        store.put(b"key1", b"value1").await.unwrap();
114        assert_eq!(store.len(), 1);
115
116        // Get
117        let value = store.get(b"key1").await.unwrap();
118        assert_eq!(value, Some(b"value1".to_vec()));
119
120        // Delete
121        store.delete(b"key1").await.unwrap();
122        assert_eq!(store.len(), 0);
123    }
124
125    #[tokio::test]
126    async fn test_batch_operations() {
127        let store = PersistentStore::memory().unwrap();
128
129        let items = vec![
130            (b"key1".to_vec(), b"value1".to_vec()),
131            (b"key2".to_vec(), b"value2".to_vec()),
132            (b"key3".to_vec(), b"value3".to_vec()),
133        ];
134
135        store.batch_put(items).await.unwrap();
136        assert_eq!(store.len(), 3);
137    }
138
139    #[tokio::test]
140    async fn test_prefix_scan() {
141        let store = PersistentStore::memory().unwrap();
142
143        store.put(b"prefix:key1", b"value1").await.unwrap();
144        store.put(b"prefix:key2", b"value2").await.unwrap();
145        store.put(b"other:key3", b"value3").await.unwrap();
146
147        let keys = store.list_prefix(b"prefix:").await.unwrap();
148        assert_eq!(keys.len(), 2);
149    }
150}
nt_memory/agentdb/storage.rs

nt_memory/agentdb/
storage.rs
