orchai/src-tauri/src/db.rs

use rusqlite::{Connection, Result};
use std::path::Path;

const MIGRATION_001: &str = include_str!("../migrations/001_init.sql");
const MIGRATION_002: &str = include_str!("../migrations/002_add_last_polled.sql");
const MIGRATION_003: &str = include_str!("../migrations/003_add_agents.sql");
const MIGRATION_004: &str = include_str!("../migrations/004_default_agents.sql");
const MIGRATION_005: &str = include_str!("../migrations/005_orchestration_modules_chat_tasks.sql");
const MIGRATION_006: &str = include_str!("../migrations/006_processed_tickets_unique_index.sql");
const MIGRATION_007: &str = include_str!("../migrations/007_normalize_timestamps_rfc3339.sql");

pub fn init(db_path: &Path) -> Result<Connection> {
    let conn = Connection::open(db_path)?;
    configure(&conn)?;
    migrate(&conn)?;
    Ok(conn)
}

#[cfg(test)]
pub fn init_in_memory() -> Result<Connection> {
    let conn = Connection::open_in_memory()?;
    configure(&conn)?;
    migrate(&conn)?;
    Ok(conn)
}

fn configure(conn: &Connection) -> Result<()> {
    conn.pragma_update(None, "journal_mode", "wal")?;
    conn.pragma_update(None, "foreign_keys", "ON")?;
    Ok(())
}

fn migrate(conn: &Connection) -> Result<()> {
    let version: i32 = conn.pragma_query_value(None, "user_version", |row| row.get(0))?;

    if version < 1 {
        conn.execute_batch(MIGRATION_001)?;
        conn.pragma_update(None, "user_version", 1)?;
    }
    if version < 2 {
        conn.execute_batch(MIGRATION_002)?;
        conn.pragma_update(None, "user_version", 2)?;
    }
    if version < 3 {
        conn.execute_batch(MIGRATION_003)?;
        conn.pragma_update(None, "user_version", 3)?;
    }
    if version < 4 {
        conn.execute_batch(MIGRATION_004)?;
        conn.pragma_update(None, "user_version", 4)?;
    }
    if version < 5 {
        conn.execute_batch(MIGRATION_005)?;
        conn.pragma_update(None, "user_version", 5)?;
    }
    if version < 6 {
        conn.execute_batch(MIGRATION_006)?;
        conn.pragma_update(None, "user_version", 6)?;
    }
    if version < 7 {
        conn.execute_batch(MIGRATION_007)?;
        conn.pragma_update(None, "user_version", 7)?;
    }

    Ok(())
}

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

    #[test]
    fn test_init_in_memory_creates_tables() {
        let conn = init_in_memory().expect("should initialize");

        // Verify all application tables exist
        let tables: Vec<String> = conn
            .prepare("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name")
            .unwrap()
            .query_map([], |row| row.get(0))
            .unwrap()
            .collect::<Result<Vec<String>, _>>()
            .unwrap();

        assert_eq!(
            tables,
            vec![
                "agents",
                "notifications",
                "processed_tickets",
                "project_agent_tasks",
                "project_live_messages",
                "project_live_sessions",
                "project_modules",
                "projects",
                "tuleap_credentials",
                "watched_trackers",
                "worktrees",
            ]
        );
    }

    #[test]
    fn test_init_in_memory_enables_foreign_keys() {
        let conn = init_in_memory().expect("should initialize");
        let fk_enabled: i32 = conn
            .query_row("PRAGMA foreign_keys", [], |row| row.get(0))
            .unwrap();
        assert_eq!(fk_enabled, 1);
    }

    #[test]
    fn test_migration_is_idempotent() {
        let conn = init_in_memory().expect("should initialize");
        // Running init again on same connection should not fail
        let version: i32 = conn
            .pragma_query_value(None, "user_version", |row| row.get(0))
            .unwrap();
        assert_eq!(version, 7);
    }

    #[test]
    fn test_default_agents_are_seeded() {
        let conn = init_in_memory().expect("should initialize");

        let analyst_defaults: i32 = conn
            .query_row(
                "SELECT COUNT(*) FROM agents WHERE role = 'analyst' AND is_default = 1",
                [],
                |row| row.get(0),
            )
            .unwrap();
        let developer_defaults: i32 = conn
            .query_row(
                "SELECT COUNT(*) FROM agents WHERE role = 'developer' AND is_default = 1",
                [],
                |row| row.get(0),
            )
            .unwrap();

        assert_eq!(analyst_defaults, 1);
        assert_eq!(developer_defaults, 1);
    }

    #[test]
    fn test_processed_tickets_unique_index_exists() {
        let conn = init_in_memory().expect("should initialize");

        let unique_idx_count: i32 = conn
            .query_row(
                "SELECT COUNT(*) FROM pragma_index_list('processed_tickets') \
                 WHERE name = 'idx_processed_tickets_tracker_artifact_unique' AND \"unique\" = 1",
                [],
                |row| row.get(0),
            )
            .unwrap();

        assert_eq!(unique_idx_count, 1);
    }
}