Skip to main content

code_kb_core/
db.rs

1pub use rusqlite::Connection;
2use rusqlite::OpenFlags;
3use std::path::Path;
4use thiserror::Error;
5
6#[derive(Debug, Error)]
7pub enum DbError {
8    #[error("Failed to open SQLite database at '{0}': {1}")]
9    OpenFailed(String, rusqlite::Error),
10    #[error("Failed to configure connection pragmas: {0}")]
11    PragmaFailed(rusqlite::Error),
12    #[error("Database file does not exist: {0}")]
13    NotFound(String),
14    #[error("Database error: {0}")]
15    Sqlite(#[from] rusqlite::Error),
16}
17
18/// Opens a read-only SQLite connection configured for low-overhead WAL reads.
19pub fn open_read_only(path: &Path) -> Result<Connection, DbError> {
20    if !path.exists() {
21        return Err(DbError::NotFound(path.display().to_string()));
22    }
23
24    let flags = OpenFlags::SQLITE_OPEN_READ_ONLY | OpenFlags::SQLITE_OPEN_NO_MUTEX;
25    let conn = Connection::open_with_flags(path, flags)
26        .map_err(|e| DbError::OpenFailed(path.display().to_string(), e))?;
27
28    #[cfg(windows)]
29    conn.execute_batch(
30        "PRAGMA busy_timeout = 5000;
31         PRAGMA query_only = ON;
32         PRAGMA cache_size = -4000;
33         PRAGMA mmap_size = 0;",
34    )
35    .map_err(DbError::PragmaFailed)?;
36
37    #[cfg(not(windows))]
38    conn.execute_batch(
39        "PRAGMA busy_timeout = 5000;
40         PRAGMA query_only = ON;
41         PRAGMA cache_size = -4000;
42         PRAGMA mmap_size = 268435456;",
43    )
44    .map_err(DbError::PragmaFailed)?;
45
46    Ok(conn)
47}
48
49/// Safely flushes all committed transactions from the WAL file into the main database file
50/// and truncates the WAL to zero bytes. Returns an error if the database is busy and unable to truncate.
51pub fn checkpoint_truncate(conn: &Connection) -> Result<(), DbError> {
52    let busy: i32 = conn
53        .query_row("PRAGMA wal_checkpoint(TRUNCATE);", [], |r| r.get(0))
54        .map_err(DbError::PragmaFailed)?;
55    if busy != 0 {
56        return Err(DbError::PragmaFailed(rusqlite::Error::SqliteFailure(
57            rusqlite::ffi::Error::new(rusqlite::ffi::SQLITE_BUSY),
58            Some("wal_checkpoint(TRUNCATE) failed: database busy".to_string()),
59        )));
60    }
61    Ok(())
62}
63
64/// Opens a read-write SQLite connection (used when creating fresh or test databases).
65pub fn open_read_write(path: &Path) -> Result<Connection, DbError> {
66    let conn =
67        Connection::open(path).map_err(|e| DbError::OpenFailed(path.display().to_string(), e))?;
68
69    conn.execute_batch(
70        "PRAGMA journal_mode = WAL;
71         PRAGMA busy_timeout = 5000;
72         PRAGMA synchronous = NORMAL;
73         PRAGMA foreign_keys = ON;",
74    )
75    .map_err(DbError::PragmaFailed)?;
76
77    Ok(conn)
78}
79
80/// SQL predicate that is true for a symbols row that is a local variable or a parameter:
81/// a `variable` declared inside a function, method, or constructor, directly or through
82/// enclosing variables such as closures.
83pub fn local_variable_predicate(alias: &str) -> String {
84    format!(
85        "({alias}.kind = 'variable' AND EXISTS (
86            WITH RECURSIVE ancestor(symbol_id, kind, parent_symbol_id) AS (
87                SELECT p.symbol_id, p.kind, p.parent_symbol_id FROM symbols p
88                WHERE p.symbol_id = {alias}.parent_symbol_id
89                UNION
90                SELECT p.symbol_id, p.kind, p.parent_symbol_id FROM symbols p
91                JOIN ancestor a ON p.symbol_id = a.parent_symbol_id
92                WHERE a.kind = 'variable'
93            )
94            SELECT 1 FROM ancestor WHERE kind IN ('function', 'method', 'constructor')))"
95    )
96}
97
98/// Identifies the rule the FTS content was built under. A stored marker that differs from this
99/// value means the index predates the rule and must be repopulated once.
100const FTS_RULE: &str = "exclude-locals-v1";
101
102fn stored_fts_rule(conn: &Connection) -> Option<String> {
103    conn.query_row(
104        "SELECT value FROM artifact_metadata WHERE key = 'fts_rule'",
105        [],
106        |r| r.get(0),
107    )
108    .ok()
109}
110
111fn fts_content_is_missing(conn: &Connection) -> bool {
112    let has_symbols = conn
113        .query_row("SELECT 1 FROM symbols LIMIT 1", [], |_| Ok(true))
114        .unwrap_or(false);
115    if !has_symbols {
116        return false;
117    }
118    !conn
119        .query_row("SELECT 1 FROM symbols_fts_docsize LIMIT 1", [], |_| {
120            Ok(true)
121        })
122        .unwrap_or(false)
123}
124
125/// Ensures the `symbols_fts` FTS5 virtual table and synchronization triggers exist in the SQLite
126/// database, and that it holds every symbol except locals and parameters. An index built under an
127/// earlier rule, or one left without content, is repopulated once.
128/// julie may write a local before its enclosing function, so the insert trigger also drops any
129/// same-file variable that became local when its parent arrived.
130pub fn ensure_fts_index(conn: &Connection) -> Result<(), rusqlite::Error> {
131    let symbols_table_exists: bool = conn
132        .query_row(
133            "SELECT 1 FROM sqlite_master WHERE type='table' AND name='symbols'",
134            [],
135            |_| Ok(true),
136        )
137        .unwrap_or(false);
138
139    if !symbols_table_exists {
140        return Ok(());
141    }
142
143    if stored_fts_rule(conn).as_deref() == Some(FTS_RULE) && !fts_content_is_missing(conn) {
144        return Ok(());
145    }
146
147    let is_local = local_variable_predicate("s");
148    let new_is_local = local_variable_predicate("new");
149    let child_is_local = local_variable_predicate("c");
150    let already_indexed = "EXISTS (SELECT 1 FROM symbols_fts_docsize d WHERE d.id = old.rowid)";
151
152    conn.execute_batch(&format!(
153        "CREATE VIRTUAL TABLE IF NOT EXISTS symbols_fts USING fts5(
154            name,
155            signature,
156            doc_comment,
157            content='symbols',
158            content_rowid='rowid',
159            tokenize='porter unicode61'
160        );
161
162        DROP TRIGGER IF EXISTS symbols_ai;
163        DROP TRIGGER IF EXISTS symbols_ad;
164        DROP TRIGGER IF EXISTS symbols_au;
165
166        CREATE TRIGGER symbols_ai AFTER INSERT ON symbols BEGIN
167            INSERT INTO symbols_fts(rowid, name, signature, doc_comment)
168            SELECT new.rowid, new.name, new.signature, new.doc_comment
169            WHERE NOT {new_is_local};
170            INSERT INTO symbols_fts(symbols_fts, rowid, name, signature, doc_comment)
171            SELECT 'delete', c.rowid, c.name, c.signature, c.doc_comment
172            FROM symbols c
173            WHERE new.kind IN ('function', 'method', 'constructor')
174              AND c.path = new.path
175              AND c.kind = 'variable'
176              AND {child_is_local}
177              AND EXISTS (SELECT 1 FROM symbols_fts_docsize d WHERE d.id = c.rowid);
178        END;
179
180        CREATE TRIGGER symbols_ad AFTER DELETE ON symbols BEGIN
181            INSERT INTO symbols_fts(symbols_fts, rowid, name, signature, doc_comment)
182            SELECT 'delete', old.rowid, old.name, old.signature, old.doc_comment
183            WHERE {already_indexed};
184        END;
185
186        CREATE TRIGGER symbols_au AFTER UPDATE ON symbols BEGIN
187            INSERT INTO symbols_fts(symbols_fts, rowid, name, signature, doc_comment)
188            SELECT 'delete', old.rowid, old.name, old.signature, old.doc_comment
189            WHERE {already_indexed};
190            INSERT INTO symbols_fts(rowid, name, signature, doc_comment)
191            SELECT new.rowid, new.name, new.signature, new.doc_comment
192            WHERE NOT {new_is_local};
193        END;"
194    ))?;
195
196    conn.execute(
197        "INSERT INTO symbols_fts(symbols_fts) VALUES('delete-all')",
198        [],
199    )?;
200    conn.execute(
201        &format!(
202            "INSERT INTO symbols_fts(rowid, name, signature, doc_comment)
203             SELECT s.rowid, s.name, s.signature, s.doc_comment
204             FROM symbols s WHERE NOT {is_local}"
205        ),
206        [],
207    )?;
208
209    conn.execute(
210        "CREATE TABLE IF NOT EXISTS artifact_metadata (key TEXT PRIMARY KEY, value TEXT)",
211        [],
212    )?;
213    conn.execute(
214        "INSERT INTO artifact_metadata (key, value) VALUES ('fts_rule', ?1)
215         ON CONFLICT(key) DO UPDATE SET value = excluded.value",
216        [FTS_RULE],
217    )?;
218
219    Ok(())
220}
221
222/// Ensures the FTS5 index on `symbols` exists at the specified database file path.
223pub fn ensure_fts_index_path(path: &Path) -> Result<(), DbError> {
224    if !path.exists() {
225        return Err(DbError::NotFound(path.display().to_string()));
226    }
227    let conn = open_read_write(path)?;
228    ensure_fts_index(&conn).map_err(DbError::PragmaFailed)?;
229    Ok(())
230}
231
232/// Retargets the `root_path` key in `artifact_metadata` to a new workspace canonical root.
233/// This is essential when cloning or copying an artifact database (e.g. into a git worktree),
234/// ensuring `julie-extract update`, `delete`, and `scan` recognize the new root without root mismatch errors.
235pub fn retarget_artifact_root(db_path: &Path, new_root: &Path) -> Result<(), DbError> {
236    if !db_path.exists() {
237        return Err(DbError::NotFound(db_path.display().to_string()));
238    }
239    let conn = open_read_write(db_path)?;
240    conn.execute(
241        "CREATE TABLE IF NOT EXISTS artifact_metadata (key TEXT PRIMARY KEY, value TEXT)",
242        [],
243    )?;
244
245    let existing_root: Option<String> = conn
246        .query_row(
247            "SELECT value FROM artifact_metadata WHERE key = 'root_path'",
248            [],
249            |r| r.get(0),
250        )
251        .ok();
252
253    let root_str = if existing_root
254        .as_deref()
255        .is_some_and(|ex| ex.starts_with(r"\\?\") || ex.starts_with(r"\\.\"))
256        || (existing_root.is_none() && cfg!(windows))
257    {
258        std::fs::canonicalize(new_root)
259            .map(|p| p.to_string_lossy().to_string())
260            .unwrap_or_else(|_| {
261                let s = new_root.to_string_lossy();
262                format!(r"\\?\{s}")
263            })
264    } else {
265        new_root.to_string_lossy().to_string()
266    };
267
268    conn.execute(
269        "INSERT INTO artifact_metadata (key, value) VALUES ('root_path', ?1)
270         ON CONFLICT(key) DO UPDATE SET value = excluded.value",
271        rusqlite::params![root_str],
272    )?;
273
274    Ok(())
275}
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280
281    #[test]
282    fn test_open_read_write_and_read_only() {
283        let dir = crate::safe_tempdir();
284        let db_path = dir.path().join("test.db");
285        let conn_rw = open_read_write(&db_path).unwrap();
286        conn_rw
287            .execute("CREATE TABLE test (id INTEGER PRIMARY KEY, name TEXT);", [])
288            .unwrap();
289        conn_rw
290            .execute("INSERT INTO test (name) VALUES ('alpha');", [])
291            .unwrap();
292        drop(conn_rw);
293
294        let conn_ro = open_read_only(&db_path).unwrap();
295        let name: String = conn_ro
296            .query_row("SELECT name FROM test WHERE id = 1", [], |r| r.get(0))
297            .unwrap();
298        assert_eq!(name, "alpha");
299
300        // Verifying query_only prevents writes
301        let write_res = conn_ro.execute("INSERT INTO test (name) VALUES ('beta');", []);
302        assert!(write_res.is_err());
303    }
304
305    #[test]
306    fn test_fts5_support() {
307        let conn = rusqlite::Connection::open_in_memory().unwrap();
308        conn.execute("CREATE VIRTUAL TABLE test_fts USING fts5(content);", [])
309            .unwrap();
310        conn.execute(
311            "INSERT INTO test_fts (content) VALUES ('hello world token search');",
312            [],
313        )
314        .unwrap();
315        let count: i64 = conn
316            .query_row(
317                "SELECT count(*) FROM test_fts WHERE test_fts MATCH 'token'",
318                [],
319                |r| r.get(0),
320            )
321            .unwrap();
322        assert_eq!(count, 1);
323    }
324
325    #[test]
326    fn test_ensure_fts_index_lifecycle() {
327        let dir = crate::safe_tempdir();
328        let db_path = dir.path().join("fts_lifecycle.db");
329        let conn = open_read_write(&db_path).unwrap();
330
331        conn.execute_batch(
332            "CREATE TABLE symbols (
333                symbol_id TEXT PRIMARY KEY,
334                path TEXT,
335                name TEXT,
336                kind TEXT,
337                parent_symbol_id TEXT,
338                signature TEXT,
339                doc_comment TEXT
340            );
341            INSERT INTO symbols VALUES ('1', 'src/pay.rs', 'PaymentGateway', 'trait', NULL, 'pub trait PaymentGateway', 'Core payment provider interface');
342            INSERT INTO symbols VALUES ('2', 'src/pay.rs', 'StripeClient', 'struct', NULL, 'pub struct StripeClient', 'Handles HTTP requests to stripe API');",
343        )
344        .unwrap();
345
346        ensure_fts_index(&conn).unwrap();
347
348        let count: i64 = conn
349            .query_row(
350                "SELECT count(*) FROM symbols_fts WHERE symbols_fts MATCH 'payment'",
351                [],
352                |r| r.get(0),
353            )
354            .unwrap();
355        assert_eq!(count, 1);
356
357        conn.execute(
358            "INSERT INTO symbols VALUES ('3', 'src/pay.rs', 'RefundHandler', 'function', NULL, 'pub fn handle_refund()', 'Processes transaction refunds');",
359            [],
360        )
361        .unwrap();
362
363        let count: i64 = conn
364            .query_row(
365                "SELECT count(*) FROM symbols_fts WHERE symbols_fts MATCH 'refund'",
366                [],
367                |r| r.get(0),
368            )
369            .unwrap();
370        assert_eq!(count, 1);
371
372        conn.execute("DELETE FROM symbols WHERE symbol_id = '3';", [])
373            .unwrap();
374        let count: i64 = conn
375            .query_row(
376                "SELECT count(*) FROM symbols_fts WHERE symbols_fts MATCH 'refund'",
377                [],
378                |r| r.get(0),
379            )
380            .unwrap();
381        assert_eq!(count, 0);
382    }
383
384    #[test]
385    fn ensure_fts_index_drops_a_local_indexed_before_its_parent() {
386        let dir = crate::safe_tempdir();
387        let db_path = dir.path().join("fts_rule.db");
388        let conn = open_read_write(&db_path).unwrap();
389        conn.execute_batch(
390            "CREATE TABLE symbols (
391                symbol_id TEXT PRIMARY KEY,
392                path TEXT,
393                name TEXT,
394                kind TEXT,
395                parent_symbol_id TEXT,
396                signature TEXT,
397                doc_comment TEXT
398            );
399            INSERT INTO symbols VALUES ('1', 'src/db.rs', 'open_conn', 'function', NULL, 'fn open_conn()', '');",
400        )
401        .unwrap();
402
403        ensure_fts_index(&conn).unwrap();
404
405        let rule: String = conn
406            .query_row(
407                "SELECT value FROM artifact_metadata WHERE key = 'fts_rule'",
408                [],
409                |r| r.get(0),
410            )
411            .unwrap();
412        assert_eq!(rule, FTS_RULE);
413
414        conn.execute(
415            "INSERT INTO symbols VALUES ('2', 'src/db.rs', 'drifted', 'variable', '3', 'let drifted = 1', '')",
416            [],
417        )
418        .unwrap();
419        conn.execute(
420            "INSERT INTO symbols VALUES ('3', 'src/db.rs', 'later', 'function', NULL, 'fn later()', '')",
421            [],
422        )
423        .unwrap();
424
425        let count: i64 = conn
426            .query_row(
427                "SELECT count(*) FROM symbols_fts WHERE symbols_fts MATCH 'drifted'",
428                [],
429                |r| r.get(0),
430            )
431            .unwrap();
432        assert_eq!(count, 0);
433    }
434
435    #[test]
436    fn ensure_fts_index_repopulates_an_emptied_index() {
437        let dir = crate::safe_tempdir();
438        let db_path = dir.path().join("fts_empty.db");
439        let conn = open_read_write(&db_path).unwrap();
440        conn.execute_batch(
441            "CREATE TABLE symbols (
442                symbol_id TEXT PRIMARY KEY,
443                path TEXT,
444                name TEXT,
445                kind TEXT,
446                parent_symbol_id TEXT,
447                signature TEXT,
448                doc_comment TEXT
449            );
450            INSERT INTO symbols VALUES ('1', 'src/pay.rs', 'RefundHandler', 'function', NULL, 'fn handle_refund()', '');",
451        )
452        .unwrap();
453
454        ensure_fts_index(&conn).unwrap();
455        conn.execute(
456            "INSERT INTO symbols_fts(symbols_fts) VALUES('delete-all')",
457            [],
458        )
459        .unwrap();
460
461        ensure_fts_index(&conn).unwrap();
462
463        let count: i64 = conn
464            .query_row(
465                "SELECT count(*) FROM symbols_fts WHERE symbols_fts MATCH 'refund'",
466                [],
467                |r| r.get(0),
468            )
469            .unwrap();
470        assert_eq!(count, 1);
471    }
472
473    #[test]
474    fn test_mmap_size_configuration() {
475        let dir = crate::safe_tempdir();
476        let db_path = dir.path().join("mmap_test.db");
477        let conn_rw = open_read_write(&db_path).unwrap();
478        conn_rw.execute("CREATE TABLE t (x INT);", []).unwrap();
479        drop(conn_rw);
480
481        let conn_ro = open_read_only(&db_path).unwrap();
482        let mmap_size: i64 = conn_ro
483            .query_row("PRAGMA mmap_size;", [], |r| r.get(0))
484            .unwrap();
485        #[cfg(windows)]
486        assert_eq!(
487            mmap_size, 0,
488            "mmap_size must be 0 on Windows to prevent file locks"
489        );
490        #[cfg(not(windows))]
491        assert_eq!(
492            mmap_size, 268435456,
493            "mmap_size should be 256MB on non-Windows"
494        );
495    }
496
497    #[test]
498    fn test_checkpoint_truncate() {
499        let dir = crate::safe_tempdir();
500        let db_path = dir.path().join("wal_checkpoint.db");
501        let conn_rw = open_read_write(&db_path).unwrap();
502        conn_rw
503            .execute("CREATE TABLE items (id INTEGER PRIMARY KEY, val TEXT);", [])
504            .unwrap();
505        conn_rw
506            .execute("INSERT INTO items (val) VALUES ('persisted_val');", [])
507            .unwrap();
508        checkpoint_truncate(&conn_rw).expect("checkpoint_truncate should succeed");
509        drop(conn_rw);
510
511        let conn_ro = open_read_only(&db_path).unwrap();
512        let val: String = conn_ro
513            .query_row("SELECT val FROM items WHERE id = 1", [], |r| r.get(0))
514            .unwrap();
515        assert_eq!(val, "persisted_val");
516    }
517
518    #[test]
519    fn test_retarget_artifact_root() {
520        let dir = crate::safe_tempdir();
521        let db_path = dir.path().join("retarget.db");
522
523        // 1. Missing db file -> NotFound error
524        let missing_path = dir.path().join("missing.db");
525        assert!(matches!(
526            retarget_artifact_root(&missing_path, dir.path()),
527            Err(DbError::NotFound(_))
528        ));
529
530        // 2. Db without existing artifact_metadata table -> creates table and sets root_path
531        {
532            let conn = open_read_write(&db_path).unwrap();
533            conn.execute("CREATE TABLE other (x INT);", []).unwrap();
534        }
535        let fresh_root = dir.path().join("fresh_root");
536        std::fs::create_dir_all(&fresh_root).unwrap();
537        retarget_artifact_root(&db_path, &fresh_root).expect("Must create table and succeed");
538        {
539            let conn = open_read_only(&db_path).unwrap();
540            let val: String = conn
541                .query_row(
542                    "SELECT value FROM artifact_metadata WHERE key = 'root_path'",
543                    [],
544                    |r| r.get(0),
545                )
546                .unwrap();
547            assert!(
548                crate::workspace::paths_equal(Path::new(&val), &fresh_root),
549                "Paths must be equal: got {val}, expected {}",
550                fresh_root.display()
551            );
552        }
553
554        // 3. Db with existing artifact_metadata -> updates root_path
555        {
556            let conn = open_read_write(&db_path).unwrap();
557            conn.execute(
558                "UPDATE artifact_metadata SET value = '/old/root' WHERE key = 'root_path';",
559                [],
560            )
561            .unwrap();
562        }
563
564        let new_root = dir.path().join("new_root");
565        std::fs::create_dir_all(&new_root).unwrap();
566        retarget_artifact_root(&db_path, &new_root).expect("retargeting must succeed");
567
568        {
569            let conn = open_read_only(&db_path).unwrap();
570            let val: String = conn
571                .query_row(
572                    "SELECT value FROM artifact_metadata WHERE key = 'root_path'",
573                    [],
574                    |r| r.get(0),
575                )
576                .unwrap();
577            assert!(
578                crate::workspace::paths_equal(Path::new(&val), &new_root),
579                "Paths must be equal: got {val}, expected {}",
580                new_root.display()
581            );
582        }
583
584        // 4. Verbatim prefix preservation on Windows when existing root started with \\?\
585        #[cfg(windows)]
586        {
587            {
588                let conn = open_read_write(&db_path).unwrap();
589                conn.execute(
590                    "UPDATE artifact_metadata SET value = '\\\\?\\C:\\old\\root' WHERE key = 'root_path';",
591                    [],
592                )
593                .unwrap();
594            }
595
596            retarget_artifact_root(&db_path, &new_root).expect("retargeting verbatim must succeed");
597
598            let conn = open_read_only(&db_path).unwrap();
599            let val: String = conn
600                .query_row(
601                    "SELECT value FROM artifact_metadata WHERE key = 'root_path'",
602                    [],
603                    |r| r.get(0),
604                )
605                .unwrap();
606            assert!(
607                val.starts_with(r"\\?\"),
608                "Must preserve \\\\?\\ prefix when existing root had it: got {val}"
609            );
610        }
611    }
612}