llm_kernel/graph/
schema.rs

1//! Schema initialization for the knowledge graph SQLite database.
2
3use rusqlite::{Connection, params};
4
5use crate::error::{KernelError, Result};
6
7/// Current graph schema version. Increment when adding migrations.
8pub const GRAPH_SCHEMA_VERSION: u32 = 2;
9
10/// Read the recorded graph schema version from `_meta`, or `0` if unset.
11pub fn schema_version(conn: &Connection) -> Result<u32> {
12    Ok(conn
13        .query_row(
14            "SELECT value FROM _meta WHERE key = 'graph_schema_version'",
15            [],
16            |row| row.get::<_, String>(0),
17        )
18        .ok()
19        .and_then(|s| s.parse().ok())
20        .unwrap_or(0))
21}
22
23/// Incremental migration for the knowledge graph schema.
24///
25/// Applies version-to-version migrations from `current` up to
26/// [`GRAPH_SCHEMA_VERSION`] inside a single transaction. On failure the
27/// transaction rolls back and the recorded version is unchanged. Returns the
28/// new version (equal to `current` when already up to date).
29pub fn migrate_graph(conn: &Connection, current: u32) -> Result<u32> {
30    if current >= GRAPH_SCHEMA_VERSION {
31        return Ok(current);
32    }
33    let tx = conn
34        .unchecked_transaction()
35        .map_err(|e| KernelError::Store(format!("migration begin failed: {e}")))?;
36    let mut v = current;
37    // v1 -> v2: index nodes by creation timestamp (used by recency ordering).
38    if v < 2 {
39        tx.execute_batch("CREATE INDEX IF NOT EXISTS idx_nodes_created ON nodes(created);")
40            .map_err(|e| KernelError::Store(format!("migration v1->v2 failed: {e}")))?;
41        v = 2;
42    }
43    tx.execute(
44        "UPDATE _meta SET value = ?1 WHERE key = 'graph_schema_version'",
45        params![v.to_string()],
46    )
47    .map_err(|e| KernelError::Store(e.to_string()))?;
48    tx.commit()
49        .map_err(|e| KernelError::Store(format!("migration commit failed: {e}")))?;
50    Ok(v)
51}
52
53/// Apply the full knowledge graph schema (tables, indexes, FTS5 triggers) to a connection.
54///
55/// Idempotent — uses `IF NOT EXISTS` for all DDL. Safe to call on every startup.
56pub fn init_graph_schema(conn: &Connection) -> Result<()> {
57    // WAL auto-checkpoint for better concurrency
58    let _ = conn.execute_batch("PRAGMA wal_autocheckpoint=100;");
59
60    conn.execute_batch(
61        "CREATE TABLE IF NOT EXISTS nodes (
62            id           TEXT PRIMARY KEY,
63            type         TEXT NOT NULL,
64            title        TEXT NOT NULL,
65            tags         TEXT NOT NULL DEFAULT '',
66            projects     TEXT NOT NULL DEFAULT '',
67            agents       TEXT NOT NULL DEFAULT '',
68            created      TEXT NOT NULL,
69            updated      TEXT NOT NULL,
70            body         TEXT NOT NULL DEFAULT '',
71            importance   REAL NOT NULL DEFAULT 0.5,
72            access_count INTEGER NOT NULL DEFAULT 0,
73            accessed_at  TEXT NOT NULL DEFAULT ''
74        );
75
76        -- FTS5 full-text search with trigram tokenizer
77        CREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts
78            USING fts5(title, body, tags, content=nodes, content_rowid=rowid, tokenize='trigram');
79
80        -- Keep FTS in sync with node changes
81        CREATE TRIGGER IF NOT EXISTS nodes_ai AFTER INSERT ON nodes BEGIN
82            INSERT INTO nodes_fts(rowid, title, body, tags)
83            VALUES (new.rowid, new.title, new.body, new.tags);
84        END;
85        CREATE TRIGGER IF NOT EXISTS nodes_ad AFTER DELETE ON nodes BEGIN
86            INSERT INTO nodes_fts(nodes_fts, rowid, title, body, tags)
87            VALUES('delete', old.rowid, old.title, old.body, old.tags);
88        END;
89        CREATE TRIGGER IF NOT EXISTS nodes_au AFTER UPDATE ON nodes BEGIN
90            INSERT INTO nodes_fts(nodes_fts, rowid, title, body, tags)
91            VALUES('delete', old.rowid, old.title, old.body, old.tags);
92            INSERT INTO nodes_fts(rowid, title, body, tags)
93            VALUES (new.rowid, new.title, new.body, new.tags);
94        END;
95
96        CREATE TABLE IF NOT EXISTS edges (
97            id       TEXT PRIMARY KEY,
98            source   TEXT NOT NULL,
99            target   TEXT NOT NULL,
100            relation TEXT NOT NULL DEFAULT 'related',
101            weight   REAL NOT NULL DEFAULT 1.0,
102            ts       TEXT NOT NULL
103        );
104
105        CREATE INDEX IF NOT EXISTS idx_edges_source  ON edges(source);
106        CREATE INDEX IF NOT EXISTS idx_edges_target  ON edges(target);
107        CREATE UNIQUE INDEX IF NOT EXISTS idx_edges_src_tgt_rel ON edges(source, target, relation);
108        CREATE INDEX IF NOT EXISTS idx_nodes_type    ON nodes(type);
109        CREATE INDEX IF NOT EXISTS idx_nodes_updated ON nodes(updated DESC);
110        CREATE INDEX IF NOT EXISTS idx_nodes_title_updated ON nodes(title, updated DESC);
111        CREATE INDEX IF NOT EXISTS idx_nodes_importance ON nodes(importance DESC);
112        CREATE INDEX IF NOT EXISTS idx_nodes_accessed ON nodes(accessed_at DESC);
113        CREATE INDEX IF NOT EXISTS idx_nodes_created ON nodes(created);
114
115        -- Schema version tracking
116        CREATE TABLE IF NOT EXISTS _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
117        INSERT OR IGNORE INTO _meta (key, value) VALUES ('graph_schema_version', '2');
118        ",
119    )
120    .map_err(|e| KernelError::Store(format!("Graph schema init failed: {}", e)))?;
121
122    Ok(())
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    fn mem_db() -> Connection {
130        let conn = Connection::open_in_memory().unwrap();
131        init_graph_schema(&conn).unwrap();
132        conn
133    }
134
135    #[test]
136    fn schema_creates_tables() {
137        let conn = mem_db();
138        let tables: Vec<String> = conn
139            .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
140            .unwrap()
141            .query_map([], |r| r.get(0))
142            .unwrap()
143            .flatten()
144            .collect();
145        assert!(tables.contains(&"nodes".to_string()));
146        assert!(tables.contains(&"edges".to_string()));
147        assert!(tables.contains(&"_meta".to_string()));
148    }
149
150    #[test]
151    fn schema_is_idempotent() {
152        let conn = Connection::open_in_memory().unwrap();
153        init_graph_schema(&conn).unwrap();
154        init_graph_schema(&conn).unwrap();
155        let count: i64 = conn
156            .query_row(
157                "SELECT COUNT(*) FROM _meta WHERE key = 'graph_schema_version'",
158                [],
159                |r| r.get(0),
160            )
161            .unwrap();
162        assert_eq!(count, 1);
163    }
164
165    #[test]
166    fn fts_table_exists() {
167        let conn = mem_db();
168        let name: String = conn
169            .query_row(
170                "SELECT name FROM sqlite_master WHERE type='table' AND name='nodes_fts'",
171                [],
172                |r| r.get(0),
173            )
174            .unwrap();
175        assert_eq!(name, "nodes_fts");
176    }
177
178    /// Helper: force the recorded version down to `v` (simulates an older DB).
179    fn set_version(conn: &Connection, v: u32) {
180        conn.execute(
181            "UPDATE _meta SET value = ?1 WHERE key = 'graph_schema_version'",
182            params![v.to_string()],
183        )
184        .unwrap();
185    }
186
187    fn has_index(conn: &Connection, name: &str) -> bool {
188        let count: i64 = conn
189            .query_row(
190                "SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name = ?1",
191                params![name],
192                |r| r.get(0),
193            )
194            .unwrap();
195        count > 0
196    }
197
198    #[test]
199    fn schema_version_reads_meta() {
200        let conn = mem_db();
201        assert_eq!(schema_version(&conn).unwrap(), GRAPH_SCHEMA_VERSION);
202        set_version(&conn, 1);
203        assert_eq!(schema_version(&conn).unwrap(), 1);
204    }
205
206    /// AC5: a v1 database migrates up to the current version, applying the
207    /// v1→v2 step (the `idx_nodes_created` index becomes observable).
208    #[test]
209    fn migrate_advances_v1_to_current() {
210        let conn = mem_db();
211        // Drop the v2 index, then rewind the version to simulate a v1 DB.
212        conn.execute_batch("DROP INDEX IF EXISTS idx_nodes_created;")
213            .unwrap();
214        set_version(&conn, 1);
215        assert!(!has_index(&conn, "idx_nodes_created"));
216
217        let new_version = migrate_graph(&conn, 1).unwrap();
218        assert_eq!(new_version, GRAPH_SCHEMA_VERSION);
219        assert_eq!(schema_version(&conn).unwrap(), GRAPH_SCHEMA_VERSION);
220        assert!(has_index(&conn, "idx_nodes_created"));
221    }
222
223    /// AC5: migrating an already-current DB is a no-op.
224    #[test]
225    fn migrate_is_noop_when_current() {
226        let conn = mem_db();
227        let v = migrate_graph(&conn, GRAPH_SCHEMA_VERSION).unwrap();
228        assert_eq!(v, GRAPH_SCHEMA_VERSION);
229    }
230
231    /// AC5: a migration whose step fails rolls back, leaving the version
232    /// unchanged. We force failure by dropping the `nodes` table so the
233    /// `CREATE INDEX … ON nodes` step cannot succeed.
234    #[test]
235    fn migrate_rolls_back_on_failure() {
236        let conn = mem_db();
237        set_version(&conn, 1);
238        // Sabotage the migration target so the v1→v2 CREATE INDEX fails.
239        conn.execute_batch("DROP TABLE nodes;").unwrap();
240
241        let result = migrate_graph(&conn, 1);
242        assert!(result.is_err(), "expected migration to fail");
243        // Version unchanged because the transaction rolled back.
244        assert_eq!(schema_version(&conn).unwrap(), 1);
245    }
246}
llm_kernel/graph/schema.rs

llm_kernel/graph/
schema.rs
