Skip to main content

sqlite_graphrag/storage/
connection.rs

1//! SQLite connection setup with PRAGMAs and 0600 permissions.
2//!
3//! Opens (or creates) the database file, loads the `sqlite-vec` extension,
4//! applies WAL/journal PRAGMAs, and enforces 0600 file permissions on Unix.
5
6use crate::errors::AppError;
7use crate::pragmas::apply_connection_pragmas;
8use rusqlite::Connection;
9use sqlite_vec::sqlite3_vec_init;
10use std::path::Path;
11
12/// Register sqlite-vec GLOBALLY before any connection is opened.
13/// Must be called once at program start.
14pub fn register_vec_extension() {
15    // SAFETY: sqlite3_auto_extension is a C FFI function that registers a callback
16    // invoked when SQLite opens any new connection. Soundness assumptions:
17    // 1. `sqlite3_vec_init` has the exact ABI signature `extern "C" fn(...) -> i32`
18    //    expected by SQLite's auto-extension API (verified by sqlite-vec crate).
19    // 2. The transmute from `*const ()` to the expected fn pointer is valid because
20    //    both have identical layout on supported platforms (Linux, macOS, Windows).
21    // 3. This function is only called once at program start (asserted by callers
22    //    in main.rs:80 before any connection is opened).
23    #[allow(clippy::missing_transmute_annotations)]
24    unsafe {
25        rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
26            sqlite3_vec_init as *const (),
27        )));
28    }
29}
30
31pub fn open_rw(path: &Path) -> Result<Connection, AppError> {
32    let conn = Connection::open(path)?;
33    apply_connection_pragmas(&conn)?;
34    apply_secure_permissions(path);
35    Ok(conn)
36}
37
38pub fn ensure_schema(conn: &mut Connection) -> Result<(), AppError> {
39    crate::migrations::runner()
40        .run(conn)
41        .map_err(|e| AppError::Internal(anyhow::anyhow!("migration failed: {e}")))?;
42    conn.execute_batch(&format!(
43        "PRAGMA user_version = {};",
44        crate::constants::SCHEMA_USER_VERSION
45    ))?;
46    Ok(())
47}
48
49/// Applies 600 permissions (owner read/write only) to the SQLite file and its WAL/SHM
50/// companion files on Unix to prevent leaking private memories in shared directories
51/// (e.g. multi-user /tmp, Dropbox, NFS). No-op on Windows. Failures are silent to avoid
52/// blocking the operation when the process does not own the file (e.g. read-only mount).
53#[allow(unused_variables)]
54fn apply_secure_permissions(path: &Path) {
55    #[cfg(unix)]
56    {
57        use std::os::unix::fs::PermissionsExt;
58        let candidates = [
59            path.to_path_buf(),
60            path.with_extension(format!(
61                "{}-wal",
62                path.extension()
63                    .and_then(|e| e.to_str())
64                    .unwrap_or("sqlite")
65            )),
66            path.with_extension(format!(
67                "{}-shm",
68                path.extension()
69                    .and_then(|e| e.to_str())
70                    .unwrap_or("sqlite")
71            )),
72        ];
73        for file in candidates.iter() {
74            if file.exists() {
75                if let Ok(meta) = std::fs::metadata(file) {
76                    let mut perms = meta.permissions();
77                    perms.set_mode(0o600);
78                    let _ = std::fs::set_permissions(file, perms);
79                }
80            }
81        }
82    }
83}
84
85pub fn open_ro(path: &Path) -> Result<Connection, AppError> {
86    let conn = Connection::open_with_flags(
87        path,
88        rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_URI,
89    )?;
90    conn.execute_batch("PRAGMA foreign_keys = ON;")?;
91    Ok(conn)
92}