llm_kernel/store/
sqlite.rs

1use std::path::Path;
2
3use rusqlite::Connection;
4
5use crate::error::{KernelError, Result};
6
7/// Migration callback: given a connection and the current schema version,
8/// apply incremental migrations and return the new version.
9pub type MigrationFn = fn(&Connection, current_version: u32) -> Result<u32>;
10
11/// Initialize a SQLite database with WAL mode, foreign keys, and a custom schema.
12///
13/// Creates parent directories if needed. Applies standard PRAGMAs for safety
14/// and performance. Runs the provided DDL to create tables, then validates
15/// the schema version matches expectations.
16pub fn init_schema(path: &Path, schema_ddl: &str, expected_version: u32) -> Result<Connection> {
17    if let Some(parent) = path.parent() {
18        std::fs::create_dir_all(parent)?;
19    }
20
21    let conn = Connection::open(path).map_err(|e| KernelError::Store(e.to_string()))?;
22    apply_pragma(&conn)?;
23
24    // Check current schema version (stored as TEXT — parse to u32)
25    let current_version: u32 = conn
26        .query_row(
27            "SELECT value FROM _meta WHERE key = 'schema_version'",
28            [],
29            |row| row.get::<_, String>(0),
30        )
31        .ok()
32        .and_then(|s| s.parse().ok())
33        .unwrap_or(0);
34
35    if current_version != expected_version {
36        // Fresh or outdated — apply full DDL
37        conn.execute_batch(schema_ddl)
38            .map_err(|e| KernelError::Store(format!("Schema DDL failed: {}", e)))?;
39
40        conn.execute(
41            "INSERT OR REPLACE INTO _meta (key, value) VALUES ('schema_version', ?1)",
42            [expected_version.to_string()],
43        )
44        .map_err(|e| KernelError::Store(e.to_string()))?;
45    }
46
47    Ok(conn)
48}
49
50/// Initialize a SQLite database with optional incremental migration support.
51///
52/// Like [`init_schema`], but when the current version is behind `expected_version`,
53/// calls the `migrate` callback instead of replaying the full DDL. Falls back to
54/// full DDL when `migrate` is `None` or the DB is fresh (version 0).
55pub fn init_schema_with_migrations(
56    path: &Path,
57    schema_ddl: &str,
58    expected_version: u32,
59    migrate: Option<MigrationFn>,
60) -> Result<Connection> {
61    if let Some(parent) = path.parent() {
62        std::fs::create_dir_all(parent)?;
63    }
64
65    let conn = Connection::open(path).map_err(|e| KernelError::Store(e.to_string()))?;
66    apply_pragma(&conn)?;
67
68    // Query version — _meta may not exist yet (fresh DB), treated as version 0
69    let current_version: u32 = conn
70        .query_row(
71            "SELECT value FROM _meta WHERE key = 'schema_version'",
72            [],
73            |row| row.get::<_, String>(0),
74        )
75        .ok()
76        .and_then(|s| s.parse().ok())
77        .unwrap_or(0);
78
79    if current_version == expected_version {
80        return Ok(conn);
81    }
82
83    // Fresh DB or no migration callback → full DDL
84    if current_version == 0 || migrate.is_none() {
85        conn.execute_batch(schema_ddl)
86            .map_err(|e| KernelError::Store(format!("Schema DDL failed: {}", e)))?;
87
88        conn.execute(
89            "INSERT OR REPLACE INTO _meta (key, value) VALUES ('schema_version', ?1)",
90            [expected_version.to_string()],
91        )
92        .map_err(|e| KernelError::Store(e.to_string()))?;
93
94        return Ok(conn);
95    }
96
97    // Incremental migration — wrapped in a transaction for atomicity
98    let tx = conn
99        .unchecked_transaction()
100        .map_err(|e| KernelError::Store(format!("Transaction begin failed: {}", e)))?;
101    let new_version = migrate.unwrap()(&tx, current_version)?;
102    tx.execute(
103        "INSERT OR REPLACE INTO _meta (key, value) VALUES ('schema_version', ?1)",
104        [new_version.to_string()],
105    )
106    .map_err(|e| KernelError::Store(e.to_string()))?;
107    tx.commit()
108        .map_err(|e| KernelError::Store(format!("Transaction commit failed: {}", e)))?;
109
110    Ok(conn)
111}
112
113/// Create an in-memory SQLite connection with PRAGMAs applied.
114/// Useful for testing.
115pub fn init_in_memory(schema_ddl: &str) -> Result<Connection> {
116    let conn = Connection::open_in_memory().map_err(|e| KernelError::Store(e.to_string()))?;
117    apply_pragma(&conn)?;
118    conn.execute_batch(schema_ddl)
119        .map_err(|e| KernelError::Store(format!("Schema DDL failed: {}", e)))?;
120    Ok(conn)
121}
122
123fn apply_pragma(conn: &Connection) -> Result<()> {
124    conn.execute_batch(
125        "PRAGMA journal_mode = WAL;
126         PRAGMA foreign_keys = ON;
127         PRAGMA busy_timeout = 5000;
128         PRAGMA synchronous = NORMAL;",
129    )
130    .map_err(|e| KernelError::Store(format!("PRAGMA failed: {}", e)))?;
131    Ok(())
132}
133
134#[cfg(test)]
135mod tests {
136    use super::*;
137    use tempfile::TempDir;
138
139    #[test]
140    fn test_init_schema_creates_db() {
141        let dir = TempDir::new().unwrap();
142        let path = dir.path().join("test.db");
143        let ddl = "CREATE TABLE _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
144                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
145                   CREATE TABLE items (id TEXT PRIMARY KEY, name TEXT NOT NULL);";
146
147        let conn = init_schema(&path, ddl, 1).unwrap();
148        let version: String = conn
149            .query_row(
150                "SELECT value FROM _meta WHERE key = 'schema_version'",
151                [],
152                |row| row.get(0),
153            )
154            .unwrap();
155        assert_eq!(version, "1");
156    }
157
158    #[test]
159    fn test_init_in_memory() {
160        let ddl = "CREATE TABLE _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
161                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
162                   CREATE TABLE items (id TEXT PRIMARY KEY);";
163        let conn = init_in_memory(ddl).unwrap();
164        conn.execute("INSERT INTO items (id) VALUES ('a')", [])
165            .unwrap();
166    }
167
168    #[test]
169    fn test_init_schema_idempotent() {
170        let dir = TempDir::new().unwrap();
171        let path = dir.path().join("test.db");
172        let ddl = "CREATE TABLE IF NOT EXISTS _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
173                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
174                   CREATE TABLE IF NOT EXISTS items (id TEXT PRIMARY KEY);";
175
176        let conn1 = init_schema(&path, ddl, 1).unwrap();
177        drop(conn1);
178        let conn2 = init_schema(&path, ddl, 1).unwrap();
179
180        let count: u32 = conn2
181            .query_row(
182                "SELECT COUNT(*) FROM _meta WHERE key = 'schema_version'",
183                [],
184                |row| row.get(0),
185            )
186            .unwrap();
187        assert_eq!(count, 1);
188    }
189
190    fn noop_migrate(_conn: &Connection, _current: u32) -> Result<u32> {
191        Ok(2)
192    }
193
194    #[test]
195    fn migration_fresh_db_uses_full_ddl() {
196        let dir = TempDir::new().unwrap();
197        let path = dir.path().join("test.db");
198        let ddl = "CREATE TABLE _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
199                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
200                   CREATE TABLE items (id TEXT PRIMARY KEY);";
201
202        let conn = init_schema_with_migrations(&path, ddl, 2, Some(noop_migrate)).unwrap();
203        let version: String = conn
204            .query_row(
205                "SELECT value FROM _meta WHERE key = 'schema_version'",
206                [],
207                |row| row.get(0),
208            )
209            .unwrap();
210        // Fresh DB → full DDL sets version to 2 directly (migration not invoked)
211        assert_eq!(version, "2");
212    }
213
214    #[test]
215    fn migration_same_version_skips() {
216        let dir = TempDir::new().unwrap();
217        let path = dir.path().join("test.db");
218        let ddl = "CREATE TABLE IF NOT EXISTS _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
219                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
220                   CREATE TABLE IF NOT EXISTS items (id TEXT PRIMARY KEY);";
221
222        let conn1 = init_schema(&path, ddl, 2).unwrap();
223        drop(conn1);
224
225        let conn2 = init_schema_with_migrations(&path, ddl, 2, Some(noop_migrate)).unwrap();
226        let version: String = conn2
227            .query_row(
228                "SELECT value FROM _meta WHERE key = 'schema_version'",
229                [],
230                |row| row.get(0),
231            )
232            .unwrap();
233        assert_eq!(version, "2");
234    }
235
236    #[test]
237    fn migration_callback_invoked_on_version_mismatch() {
238        let dir = TempDir::new().unwrap();
239        let path = dir.path().join("test.db");
240        let ddl = "CREATE TABLE IF NOT EXISTS _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
241                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
242                   CREATE TABLE IF NOT EXISTS items (id TEXT PRIMARY KEY);";
243
244        // First init at version 1
245        let conn1 = init_schema(&path, ddl, 1).unwrap();
246        drop(conn1);
247
248        // Now "upgrade" to version 3 via migration
249        let migrate = |conn: &Connection, current: u32| -> Result<u32> {
250            assert_eq!(current, 1);
251            conn.execute_batch("ALTER TABLE items ADD COLUMN label TEXT DEFAULT ''")
252                .map_err(|e| KernelError::Store(e.to_string()))?;
253            Ok(3)
254        };
255
256        let conn2 =
257            init_schema_with_migrations(&path, ddl, 3, Some(migrate as MigrationFn)).unwrap();
258        let version: String = conn2
259            .query_row(
260                "SELECT value FROM _meta WHERE key = 'schema_version'",
261                [],
262                |row| row.get(0),
263            )
264            .unwrap();
265        assert_eq!(version, "3");
266
267        // Verify the migration actually added the column
268        let cols: Vec<String> = conn2
269            .prepare("PRAGMA table_info(items)")
270            .unwrap()
271            .query_map([], |r| r.get::<_, String>(1))
272            .unwrap()
273            .flatten()
274            .collect();
275        assert!(cols.contains(&"label".to_string()));
276    }
277
278    #[test]
279    fn migration_none_falls_back_to_full_ddl() {
280        let dir = TempDir::new().unwrap();
281        let path = dir.path().join("test.db");
282        let ddl = "CREATE TABLE IF NOT EXISTS _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
283                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
284                   CREATE TABLE IF NOT EXISTS items (id TEXT PRIMARY KEY);";
285
286        let conn1 = init_schema(&path, ddl, 1).unwrap();
287        drop(conn1);
288
289        let conn2 = init_schema_with_migrations(&path, ddl, 2, None).unwrap();
290        let version: String = conn2
291            .query_row(
292                "SELECT value FROM _meta WHERE key = 'schema_version'",
293                [],
294                |row| row.get(0),
295            )
296            .unwrap();
297        assert_eq!(version, "2");
298    }
299
300    #[test]
301    fn migration_failure_rolls_back() {
302        let dir = TempDir::new().unwrap();
303        let path = dir.path().join("test.db");
304        let ddl = "CREATE TABLE IF NOT EXISTS _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
305                   INSERT OR IGNORE INTO _meta (key, value) VALUES ('schema_version', '0');
306                   CREATE TABLE IF NOT EXISTS items (id TEXT PRIMARY KEY);";
307
308        // Init at version 1
309        let conn1 = init_schema(&path, ddl, 1).unwrap();
310        drop(conn1);
311
312        // Migration that fails — schema should remain at version 1
313        let fail_migrate = |_conn: &Connection, _current: u32| -> Result<u32> {
314            Err(KernelError::Store("migration failed".into()))
315        };
316
317        let result = init_schema_with_migrations(&path, ddl, 3, Some(fail_migrate as MigrationFn));
318        assert!(result.is_err());
319
320        // Verify version is still 1 (rolled back)
321        let conn2 = Connection::open(&path).unwrap();
322        let version: String = conn2
323            .query_row(
324                "SELECT value FROM _meta WHERE key = 'schema_version'",
325                [],
326                |row| row.get(0),
327            )
328            .unwrap();
329        assert_eq!(version, "1");
330    }
331}
llm_kernel/store/sqlite.rs

llm_kernel/store/
sqlite.rs
