wraith/src-tauri/src/db/mod.rs

use rusqlite::Connection;
use std::path::Path;
use std::sync::{Arc, Mutex};

/// Cheap-to-clone handle to a single SQLite connection protected by a mutex.
///
/// Using a single shared connection (rather than a pool) is correct for
/// desktop use: SQLite's WAL mode allows concurrent reads from the OS level,
/// and the mutex ensures we never issue overlapping writes from Tauri's
/// async command threads.
#[derive(Clone)]
pub struct Database {
    conn: Arc<Mutex<Connection>>,
}

impl Database {
    /// Open (or create) the SQLite database at `path` and apply PRAGMAs.
    pub fn open(path: &Path) -> Result<Self, Box<dyn std::error::Error>> {
        let conn = Connection::open(path)?;

        conn.execute_batch(
            "PRAGMA journal_mode=WAL;
             PRAGMA busy_timeout=5000;
             PRAGMA foreign_keys=ON;",
        )?;

        Ok(Self {
            conn: Arc::new(Mutex::new(conn)),
        })
    }

    /// Acquire a lock on the underlying connection.
    ///
    /// Panics if the mutex was poisoned (which only happens if a thread
    /// panicked while holding the lock — a non-recoverable situation anyway).
    pub fn conn(&self) -> std::sync::MutexGuard<'_, Connection> {
        self.conn.lock().unwrap()
    }

    /// Run all embedded SQL migrations.
    ///
    /// The migration file is compiled into the binary via `include_str!`,
    /// so there is no runtime file-system dependency on it.
    pub fn migrate(&self) -> Result<(), Box<dyn std::error::Error>> {
        let conn = self.conn();
        conn.execute_batch(include_str!("migrations/001_initial.sql"))?;
        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn open_in_memory() {
        let db = Database::open(Path::new(":memory:"));
        assert!(db.is_ok());
    }

    #[test]
    fn migrate_creates_tables() {
        let db = Database::open(Path::new(":memory:")).unwrap();
        db.migrate().unwrap();

        let conn = db.conn();
        let tables: Vec<String> = conn
            .prepare("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
            .unwrap()
            .query_map([], |row| row.get(0))
            .unwrap()
            .collect::<Result<Vec<_>, _>>()
            .unwrap();

        assert!(tables.contains(&"connections".to_string()));
        assert!(tables.contains(&"credentials".to_string()));
        assert!(tables.contains(&"groups".to_string()));
        assert!(tables.contains(&"host_keys".to_string()));
        assert!(tables.contains(&"settings".to_string()));
        assert!(tables.contains(&"themes".to_string()));
        assert!(tables.contains(&"ssh_keys".to_string()));
        assert!(tables.contains(&"connection_history".to_string()));
    }

    #[test]
    fn migrate_is_idempotent() {
        let db = Database::open(Path::new(":memory:")).unwrap();
        db.migrate().unwrap();
        db.migrate().unwrap(); // second run must not error
    }

    #[test]
    fn foreign_keys_enabled() {
        let db = Database::open(Path::new(":memory:")).unwrap();
        let conn = db.conn();
        let fk_enabled: i64 = conn
            .query_row("PRAGMA foreign_keys", [], |row| row.get(0))
            .unwrap();
        assert_eq!(fk_enabled, 1);
    }

    #[test]
    fn wal_mode_enabled() {
        let db = Database::open(Path::new(":memory:")).unwrap();
        // In-memory databases use "memory" journal mode, but WAL is set
        // for file-backed DBs. Just verify the pragma doesn't error.
        let conn = db.conn();
        let mode: String = conn
            .query_row("PRAGMA journal_mode", [], |row| row.get(0))
            .unwrap();
        // In-memory always reports "memory"; file-backed would report "wal".
        assert!(!mode.is_empty());
    }

    #[test]
    fn clone_shares_connection() {
        let db = Database::open(Path::new(":memory:")).unwrap();
        db.migrate().unwrap();

        let db2 = db.clone();
        // Write through one handle, read through the other.
        db.conn()
            .execute(
                "INSERT INTO settings (key, value) VALUES ('test', 'yes')",
                [],
            )
            .unwrap();

        let val: String = db2
            .conn()
            .query_row("SELECT value FROM settings WHERE key = 'test'", [], |r| {
                r.get(0)
            })
            .unwrap();
        assert_eq!(val, "yes");
    }
}