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 = 3;
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    // v2 -> v3: composite indexes for relation-filtered directed edge lookups
44    // (powers edges_for_node_dir / neighbors_weighted with a relation filter).
45    if v < 3 {
46        tx.execute_batch(
47            "CREATE INDEX IF NOT EXISTS idx_edges_src_rel ON edges(source, relation);
48             CREATE INDEX IF NOT EXISTS idx_edges_tgt_rel ON edges(target, relation);",
49        )
50        .map_err(|e| KernelError::Store(format!("migration v2->v3 failed: {e}")))?;
51        v = 3;
52    }
53    tx.execute(
54        "UPDATE _meta SET value = ?1 WHERE key = 'graph_schema_version'",
55        params![v.to_string()],
56    )
57    .map_err(|e| KernelError::Store(e.to_string()))?;
58    tx.commit()
59        .map_err(|e| KernelError::Store(format!("migration commit failed: {e}")))?;
60    Ok(v)
61}
62
63/// Apply the full knowledge graph schema (tables, indexes, FTS5 triggers) to a connection.
64///
65/// Idempotent — uses `IF NOT EXISTS` for all DDL. Safe to call on every startup.
66pub fn init_graph_schema(conn: &Connection) -> Result<()> {
67    // WAL auto-checkpoint for better concurrency
68    let _ = conn.execute_batch("PRAGMA wal_autocheckpoint=100;");
69
70    conn.execute_batch(
71        "CREATE TABLE IF NOT EXISTS nodes (
72            id           TEXT PRIMARY KEY,
73            type         TEXT NOT NULL,
74            title        TEXT NOT NULL,
75            tags         TEXT NOT NULL DEFAULT '',
76            projects     TEXT NOT NULL DEFAULT '',
77            agents       TEXT NOT NULL DEFAULT '',
78            created      TEXT NOT NULL,
79            updated      TEXT NOT NULL,
80            body         TEXT NOT NULL DEFAULT '',
81            importance   REAL NOT NULL DEFAULT 0.5,
82            access_count INTEGER NOT NULL DEFAULT 0,
83            accessed_at  TEXT NOT NULL DEFAULT ''
84        );
85
86        -- FTS5 full-text search with trigram tokenizer
87        CREATE VIRTUAL TABLE IF NOT EXISTS nodes_fts
88            USING fts5(title, body, tags, content=nodes, content_rowid=rowid, tokenize='trigram');
89
90        -- Keep FTS in sync with node changes
91        CREATE TRIGGER IF NOT EXISTS nodes_ai AFTER INSERT ON nodes BEGIN
92            INSERT INTO nodes_fts(rowid, title, body, tags)
93            VALUES (new.rowid, new.title, new.body, new.tags);
94        END;
95        CREATE TRIGGER IF NOT EXISTS nodes_ad AFTER DELETE ON nodes BEGIN
96            INSERT INTO nodes_fts(nodes_fts, rowid, title, body, tags)
97            VALUES('delete', old.rowid, old.title, old.body, old.tags);
98        END;
99        CREATE TRIGGER IF NOT EXISTS nodes_au AFTER UPDATE ON nodes BEGIN
100            INSERT INTO nodes_fts(nodes_fts, rowid, title, body, tags)
101            VALUES('delete', old.rowid, old.title, old.body, old.tags);
102            INSERT INTO nodes_fts(rowid, title, body, tags)
103            VALUES (new.rowid, new.title, new.body, new.tags);
104        END;
105
106        CREATE TABLE IF NOT EXISTS edges (
107            id       TEXT PRIMARY KEY,
108            source   TEXT NOT NULL,
109            target   TEXT NOT NULL,
110            relation TEXT NOT NULL DEFAULT 'related',
111            weight   REAL NOT NULL DEFAULT 1.0,
112            ts       TEXT NOT NULL
113        );
114
115        CREATE INDEX IF NOT EXISTS idx_edges_source  ON edges(source);
116        CREATE INDEX IF NOT EXISTS idx_edges_target  ON edges(target);
117        CREATE UNIQUE INDEX IF NOT EXISTS idx_edges_src_tgt_rel ON edges(source, target, relation);
118        CREATE INDEX IF NOT EXISTS idx_edges_src_rel ON edges(source, relation);
119        CREATE INDEX IF NOT EXISTS idx_edges_tgt_rel ON edges(target, relation);
120        CREATE INDEX IF NOT EXISTS idx_nodes_type    ON nodes(type);
121        CREATE INDEX IF NOT EXISTS idx_nodes_updated ON nodes(updated DESC);
122        CREATE INDEX IF NOT EXISTS idx_nodes_title_updated ON nodes(title, updated DESC);
123        CREATE INDEX IF NOT EXISTS idx_nodes_importance ON nodes(importance DESC);
124        CREATE INDEX IF NOT EXISTS idx_nodes_accessed ON nodes(accessed_at DESC);
125        CREATE INDEX IF NOT EXISTS idx_nodes_created ON nodes(created);
126
127        -- Schema version tracking
128        CREATE TABLE IF NOT EXISTS _meta (key TEXT PRIMARY KEY, value TEXT NOT NULL);
129        INSERT OR IGNORE INTO _meta (key, value) VALUES ('graph_schema_version', '3');
130        ",
131    )
132    .map_err(|e| KernelError::Store(format!("Graph schema init failed: {}", e)))?;
133
134    Ok(())
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140
141    fn mem_db() -> Connection {
142        let conn = Connection::open_in_memory().unwrap();
143        init_graph_schema(&conn).unwrap();
144        conn
145    }
146
147    #[test]
148    fn schema_creates_tables() {
149        let conn = mem_db();
150        let tables: Vec<String> = conn
151            .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
152            .unwrap()
153            .query_map([], |r| r.get(0))
154            .unwrap()
155            .flatten()
156            .collect();
157        assert!(tables.contains(&"nodes".to_string()));
158        assert!(tables.contains(&"edges".to_string()));
159        assert!(tables.contains(&"_meta".to_string()));
160    }
161
162    #[test]
163    fn schema_is_idempotent() {
164        let conn = Connection::open_in_memory().unwrap();
165        init_graph_schema(&conn).unwrap();
166        init_graph_schema(&conn).unwrap();
167        let count: i64 = conn
168            .query_row(
169                "SELECT COUNT(*) FROM _meta WHERE key = 'graph_schema_version'",
170                [],
171                |r| r.get(0),
172            )
173            .unwrap();
174        assert_eq!(count, 1);
175    }
176
177    #[test]
178    fn fts_table_exists() {
179        let conn = mem_db();
180        let name: String = conn
181            .query_row(
182                "SELECT name FROM sqlite_master WHERE type='table' AND name='nodes_fts'",
183                [],
184                |r| r.get(0),
185            )
186            .unwrap();
187        assert_eq!(name, "nodes_fts");
188    }
189
190    /// Helper: force the recorded version down to `v` (simulates an older DB).
191    fn set_version(conn: &Connection, v: u32) {
192        conn.execute(
193            "UPDATE _meta SET value = ?1 WHERE key = 'graph_schema_version'",
194            params![v.to_string()],
195        )
196        .unwrap();
197    }
198
199    fn has_index(conn: &Connection, name: &str) -> bool {
200        let count: i64 = conn
201            .query_row(
202                "SELECT COUNT(*) FROM sqlite_master WHERE type='index' AND name = ?1",
203                params![name],
204                |r| r.get(0),
205            )
206            .unwrap();
207        count > 0
208    }
209
210    #[test]
211    fn schema_version_reads_meta() {
212        let conn = mem_db();
213        assert_eq!(schema_version(&conn).unwrap(), GRAPH_SCHEMA_VERSION);
214        set_version(&conn, 1);
215        assert_eq!(schema_version(&conn).unwrap(), 1);
216    }
217
218    /// AC5: a v1 database migrates up to the current version, applying the
219    /// v1→v2 step (the `idx_nodes_created` index becomes observable).
220    #[test]
221    fn migrate_advances_v1_to_current() {
222        let conn = mem_db();
223        // Drop the v2 index, then rewind the version to simulate a v1 DB.
224        conn.execute_batch("DROP INDEX IF EXISTS idx_nodes_created;")
225            .unwrap();
226        set_version(&conn, 1);
227        assert!(!has_index(&conn, "idx_nodes_created"));
228
229        let new_version = migrate_graph(&conn, 1).unwrap();
230        assert_eq!(new_version, GRAPH_SCHEMA_VERSION);
231        assert_eq!(schema_version(&conn).unwrap(), GRAPH_SCHEMA_VERSION);
232        assert!(has_index(&conn, "idx_nodes_created"));
233    }
234
235    /// AC5: a v2 database migrates up to v3, adding the relation-filter
236    /// composite indexes used by directed edge lookups.
237    #[test]
238    fn migrate_advances_v2_to_v3() {
239        let conn = mem_db();
240        // init_graph_schema already created the v3 indexes; drop them and
241        // rewind the version to simulate a v2 DB.
242        conn.execute_batch(
243            "DROP INDEX IF EXISTS idx_edges_src_rel; DROP INDEX IF EXISTS idx_edges_tgt_rel;",
244        )
245        .unwrap();
246        set_version(&conn, 2);
247        assert!(!has_index(&conn, "idx_edges_src_rel"));
248        assert!(!has_index(&conn, "idx_edges_tgt_rel"));
249
250        let new_version = migrate_graph(&conn, 2).unwrap();
251        assert_eq!(new_version, GRAPH_SCHEMA_VERSION);
252        assert!(has_index(&conn, "idx_edges_src_rel"));
253        assert!(has_index(&conn, "idx_edges_tgt_rel"));
254    }
255
256    /// AC5: migrating an already-current DB is a no-op.
257    #[test]
258    fn migrate_is_noop_when_current() {
259        let conn = mem_db();
260        let v = migrate_graph(&conn, GRAPH_SCHEMA_VERSION).unwrap();
261        assert_eq!(v, GRAPH_SCHEMA_VERSION);
262    }
263
264    /// AC5: a migration whose step fails rolls back, leaving the version
265    /// unchanged. We force failure by dropping the `nodes` table so the
266    /// `CREATE INDEX … ON nodes` step cannot succeed.
267    #[test]
268    fn migrate_rolls_back_on_failure() {
269        let conn = mem_db();
270        set_version(&conn, 1);
271        // Sabotage the migration target so the v1→v2 CREATE INDEX fails.
272        conn.execute_batch("DROP TABLE nodes;").unwrap();
273
274        let result = migrate_graph(&conn, 1);
275        assert!(result.is_err(), "expected migration to fail");
276        // Version unchanged because the transaction rolled back.
277        assert_eq!(schema_version(&conn).unwrap(), 1);
278    }
279}
llm_kernel/graph/schema.rs

llm_kernel/graph/
schema.rs
