zesdex/crates/zesdex-libs/src/database.rs

#![allow(
    clippy::cast_possible_truncation,
    clippy::cast_sign_loss,
    clippy::cast_precision_loss,
    clippy::cast_possible_wrap
)]
//! SQLite database connection pool initialisation and schema migrations.
//!
//! Uses `r2d2` + `r2d2_sqlite` for connection pooling with the same
//! `rusqlite` backend the rest of the project uses, avoiding native
//! library conflicts between `rusqlite` and `sqlx`.

use anyhow::{Context, Result};
use std::sync::Arc;
use std::sync::Mutex;

/// A shared SQLite connection wrapped for thread-safe access.
/// Uses a simple Mutex-guarded connection rather than a full pool,
/// since the daemon is single-threaded for database operations.
#[derive(Clone)]
pub struct DbConn {
    conn: Arc<Mutex<rusqlite::Connection>>,
}

impl DbConn {
    /// Execute a closure with a reference to the underlying connection.
    pub fn with<F, T>(&self, f: F) -> Result<T>
    where
        F: FnOnce(&rusqlite::Connection) -> Result<T>,
    {
        let conn = self.conn.lock().map_err(|e| anyhow::anyhow!("db lock poisoned: {e}"))?;
        f(&conn)
    }
}

/// Embedded SQL schema for all zesdex tables.
///
/// Uses `CREATE TABLE IF NOT EXISTS` so repeated runs are idempotent.
const SCHEMA_SQL: &str = r#"
CREATE TABLE IF NOT EXISTS sessions (
    id          TEXT PRIMARY KEY,
    created_at  INTEGER NOT NULL,
    updated_at  INTEGER NOT NULL,
    title       TEXT NOT NULL DEFAULT '',
    model       TEXT NOT NULL DEFAULT '',
    workspace_roots TEXT NOT NULL DEFAULT '[]',
    message_count   INTEGER NOT NULL DEFAULT 0,
    token_count     INTEGER NOT NULL DEFAULT 0,
    archived    INTEGER NOT NULL DEFAULT 0,
    summary     TEXT
);

CREATE TABLE IF NOT EXISTS settings (
    id          INTEGER PRIMARY KEY CHECK (id = 1),
    data        TEXT NOT NULL DEFAULT '{}',
    updated_at  INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS conversations (
    session_id  TEXT PRIMARY KEY,
    data        TEXT NOT NULL DEFAULT '{}',
    updated_at  INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS memories (
    name        TEXT PRIMARY KEY,
    data        TEXT NOT NULL DEFAULT '{}',
    updated_at  INTEGER NOT NULL
);

CREATE TABLE IF NOT EXISTS edit_logs (
    id          INTEGER PRIMARY KEY AUTOINCREMENT,
    session_id  TEXT NOT NULL,
    entry       TEXT NOT NULL,
    created_at  INTEGER NOT NULL
);

CREATE INDEX IF NOT EXISTS idx_edit_logs_session
    ON edit_logs (session_id);
"#;

/// Initialise a shared SQLite connection at the given path.
///
/// Opens (or creates) the database, enables WAL mode, and returns a
/// thread-safe `DbConn` handle.
///
/// # Errors
///
/// Returns an error if the database cannot be opened or created.
pub fn init_db(db_path: &str) -> Result<DbConn> {
    let conn = rusqlite::Connection::open(db_path)
        .with_context(|| format!("failed to open SQLite database at '{db_path}'"))?;

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

    tracing::info!("connected to SQLite database at '{db_path}'");
    Ok(DbConn {
        conn: Arc::new(Mutex::new(conn)),
    })
}

/// Run embedded SQL schema migrations.
///
/// Executes the [`SCHEMA_SQL`] string which creates all tables using
/// `CREATE TABLE IF NOT EXISTS`, making it safe to call on every startup.
///
/// # Errors
///
/// Returns an error if any SQL statement fails.
pub fn run_migrations(db: &DbConn) -> Result<()> {
    db.with(|conn| {
        conn.execute_batch(SCHEMA_SQL)
            .context("failed to execute database schema migrations")
    })?;
    tracing::info!("database schema migrations applied");
    Ok(())
}

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

    #[test]
    fn test_init_db_and_migrate() {
        let tmp = std::env::temp_dir().join(format!("zesdex-test-db-{}", uuid::Uuid::new_v4()));
        let db_path = tmp.to_str().unwrap().to_string();

        let db = init_db(&db_path).unwrap();
        run_migrations(&db).unwrap();

        // Verify sessions table exists
        db.with(|conn| {
            let count: i64 = conn
                .query_row("SELECT COUNT(*) FROM sessions", [], |row| row.get(0))
                .unwrap();
            assert_eq!(count, 0);
            Ok(())
        })
        .unwrap();

        let _ = std::fs::remove_file(&db_path);
    }
}