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    #[allow(clippy::missing_transmute_annotations)]
16    unsafe {
17        rusqlite::ffi::sqlite3_auto_extension(Some(std::mem::transmute(
18            sqlite3_vec_init as *const (),
19        )));
20    }
21}
22
23pub fn open_rw(path: &Path) -> Result<Connection, AppError> {
24    let conn = Connection::open(path)?;
25    apply_connection_pragmas(&conn)?;
26    apply_secure_permissions(path);
27    Ok(conn)
28}
29
30pub fn ensure_schema(conn: &mut Connection) -> Result<(), AppError> {
31    crate::migrations::runner()
32        .run(conn)
33        .map_err(|e| AppError::Internal(anyhow::anyhow!("migration failed: {e}")))?;
34    conn.execute_batch(&format!(
35        "PRAGMA user_version = {};",
36        crate::constants::SCHEMA_USER_VERSION
37    ))?;
38    Ok(())
39}
40
41/// Applies 600 permissions (owner read/write only) to the SQLite file and its WAL/SHM
42/// companion files on Unix to prevent leaking private memories in shared directories
43/// (e.g. multi-user /tmp, Dropbox, NFS). No-op on Windows. Failures are silent to avoid
44/// blocking the operation when the process does not own the file (e.g. read-only mount).
45#[allow(unused_variables)]
46fn apply_secure_permissions(path: &Path) {
47    #[cfg(unix)]
48    {
49        use std::os::unix::fs::PermissionsExt;
50        let candidates = [
51            path.to_path_buf(),
52            path.with_extension(format!(
53                "{}-wal",
54                path.extension()
55                    .and_then(|e| e.to_str())
56                    .unwrap_or("sqlite")
57            )),
58            path.with_extension(format!(
59                "{}-shm",
60                path.extension()
61                    .and_then(|e| e.to_str())
62                    .unwrap_or("sqlite")
63            )),
64        ];
65        for file in candidates.iter() {
66            if file.exists() {
67                if let Ok(meta) = std::fs::metadata(file) {
68                    let mut perms = meta.permissions();
69                    perms.set_mode(0o600);
70                    let _ = std::fs::set_permissions(file, perms);
71                }
72            }
73        }
74    }
75}
76
77pub fn open_ro(path: &Path) -> Result<Connection, AppError> {
78    let conn = Connection::open_with_flags(
79        path,
80        rusqlite::OpenFlags::SQLITE_OPEN_READ_ONLY | rusqlite::OpenFlags::SQLITE_OPEN_URI,
81    )?;
82    conn.execute_batch("PRAGMA foreign_keys = ON;")?;
83    Ok(conn)
84}