llm_kernel/store/
kv.rs

1//! Generic key-value store trait and a SQLite implementation.
2//!
3//! [`KvStore`] is a small, sync, byte-oriented abstraction so that callers
4//! (LLM response cache, embedding cache, session state, rate-limit counters)
5//! can depend on `Arc<dyn KvStore>` without binding to SQLite. [`SqliteKvStore`]
6//! is the bundled implementation over a single guarded connection.
7
8use std::sync::Mutex;
9
10use rusqlite::Connection;
11
12use crate::error::{KernelError, Result};
13
14/// Generic byte-oriented key-value store.
15///
16/// Sync by design — SQLite is synchronous, and local-disk calls inside an async
17/// context are acceptable. The trait is object-safe, so a cache or session
18/// store can hold a `Box<dyn KvStore>` / `Arc<dyn KvStore>`.
19pub trait KvStore: Send + Sync {
20    /// Fetch the value for `key`, or `None` if it is absent.
21    fn get(&self, key: &str) -> Result<Option<Vec<u8>>>;
22    /// Store `value` under `key`, replacing any existing value.
23    fn put(&self, key: &str, value: &[u8]) -> Result<()>;
24    /// Remove `key`. Returns `true` if a value was present.
25    fn delete(&self, key: &str) -> Result<bool>;
26}
27
28/// DDL for the single `kv` table. Idempotent.
29const KV_DDL: &str = "CREATE TABLE IF NOT EXISTS kv (
30    key   TEXT PRIMARY KEY,
31    value BLOB NOT NULL
32);";
33
34/// SQLite-backed [`KvStore`] over one connection guarded by a mutex.
35///
36/// The connection is opened with WAL journaling and a 5 s busy timeout for
37/// safe concurrent access from multiple threads in the same process.
38pub struct SqliteKvStore {
39    conn: Mutex<Connection>,
40}
41
42impl SqliteKvStore {
43    /// Open (or create) a KV store at `path`, applying the schema if needed.
44    pub fn open(path: &std::path::Path) -> Result<Self> {
45        let conn = Connection::open(path).map_err(|e| KernelError::Store(e.to_string()))?;
46        Self::init(&conn)?;
47        Ok(Self {
48            conn: Mutex::new(conn),
49        })
50    }
51
52    /// Create an in-memory KV store — useful for tests and ephemeral caches.
53    pub fn open_in_memory() -> Result<Self> {
54        let conn = Connection::open_in_memory().map_err(|e| KernelError::Store(e.to_string()))?;
55        Self::init(&conn)?;
56        Ok(Self {
57            conn: Mutex::new(conn),
58        })
59    }
60
61    fn init(conn: &Connection) -> Result<()> {
62        conn.execute_batch("PRAGMA journal_mode = WAL; PRAGMA busy_timeout = 5000;")
63            .map_err(|e| KernelError::Store(e.to_string()))?;
64        conn.execute_batch(KV_DDL)
65            .map_err(|e| KernelError::Store(e.to_string()))?;
66        Ok(())
67    }
68
69    /// Recover the connection guard even if a previous holder panicked, so a
70    /// poisoned mutex never permanently wedges the store.
71    fn lock(&self) -> std::sync::MutexGuard<'_, Connection> {
72        self.conn.lock().unwrap_or_else(|e| e.into_inner())
73    }
74}
75
76impl KvStore for SqliteKvStore {
77    fn get(&self, key: &str) -> Result<Option<Vec<u8>>> {
78        let conn = self.lock();
79        match conn.query_row(
80            "SELECT value FROM kv WHERE key = ?1",
81            rusqlite::params![key],
82            |r| r.get::<_, Vec<u8>>(0),
83        ) {
84            Ok(v) => Ok(Some(v)),
85            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
86            Err(e) => Err(KernelError::Store(e.to_string())),
87        }
88    }
89
90    fn put(&self, key: &str, value: &[u8]) -> Result<()> {
91        let conn = self.lock();
92        conn.execute(
93            "INSERT OR REPLACE INTO kv (key, value) VALUES (?1, ?2)",
94            rusqlite::params![key, value],
95        )
96        .map_err(|e| KernelError::Store(e.to_string()))?;
97        Ok(())
98    }
99
100    fn delete(&self, key: &str) -> Result<bool> {
101        let conn = self.lock();
102        let n = conn
103            .execute("DELETE FROM kv WHERE key = ?1", rusqlite::params![key])
104            .map_err(|e| KernelError::Store(e.to_string()))?;
105        Ok(n > 0)
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn round_trip_in_memory() {
115        let kv = SqliteKvStore::open_in_memory().unwrap();
116        assert!(kv.get("missing").unwrap().is_none());
117        kv.put("a", b"value-a").unwrap();
118        assert_eq!(kv.get("a").unwrap(), Some(b"value-a".to_vec()));
119        // Overwrite.
120        kv.put("a", b"value-a2").unwrap();
121        assert_eq!(kv.get("a").unwrap(), Some(b"value-a2".to_vec()));
122        // Delete.
123        assert!(kv.delete("a").unwrap());
124        assert!(kv.get("a").unwrap().is_none());
125        // Delete missing returns false.
126        assert!(!kv.delete("a").unwrap());
127    }
128
129    #[test]
130    fn open_on_file_persists() {
131        let dir = tempfile::TempDir::new().unwrap();
132        let path = dir.path().join("kv.db");
133        {
134            let kv = SqliteKvStore::open(&path).unwrap();
135            kv.put("k", b"v").unwrap();
136        }
137        let kv = SqliteKvStore::open(&path).unwrap();
138        assert_eq!(kv.get("k").unwrap(), Some(b"v".to_vec()));
139    }
140
141    /// A blanket `dyn KvStore` works (object-safety).
142    #[test]
143    fn trait_object_round_trip() {
144        let kv: Box<dyn KvStore> = Box::new(SqliteKvStore::open_in_memory().unwrap());
145        kv.put("x", b"y").unwrap();
146        assert_eq!(kv.get("x").unwrap(), Some(b"y".to_vec()));
147    }
148}
llm_kernel/store/kv.rs

llm_kernel/store/
kv.rs
