track/db/

mod.rs

//! Database initialization and management for the track CLI.
//!
//! This module handles SQLite database creation, schema initialization, migrations,
//! and application state management. The database stores all task, TODO, link, scrap,
//! and Git repository information.

use crate::models::{TaskStatus, TodoStatus, VcsMode};
use crate::utils::Result;
use directories::ProjectDirs;
use rusqlite::{params, Connection, OptionalExtension};
use std::path::PathBuf;
use std::time::Duration;

pub mod migrate;
pub mod row_mapping;

/// Revision numbers for each section, used for change detection.
///
/// Each section has a revision number that is incremented whenever
/// data in that section is modified. This allows efficient change
/// detection without complex queries.
#[derive(Clone, Debug, PartialEq, Default)]
pub struct SectionRevs {
    /// Task metadata (description, ticket, alias) revision
    pub task: i64,
    /// TODOs section revision
    pub todos: i64,
    /// Scraps section revision
    pub scraps: i64,
    /// Links section revision
    pub links: i64,
    /// Repositories section revision
    pub repos: i64,
    /// Worktrees section revision
    pub worktrees: i64,
}

/// Database connection and management.
///
/// Handles SQLite database operations including schema initialization,
/// migrations, and application state persistence.
pub struct Database {
    conn: Connection,
}

impl Database {
    /// Creates a new database instance with the default file location.
    ///
    /// The database file is stored in the platform-specific data directory.
    /// The schema is automatically initialized if the database is new.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The data directory cannot be determined
    /// - The database file cannot be created or opened
    /// - Schema initialization fails
    pub fn new() -> Result<Self> {
        let db_path = Self::get_db_path()?;

        // Create parent directory if it doesn't exist
        if let Some(parent) = db_path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        let conn = Connection::open(&db_path)?;
        Self::configure_connection(&conn)?;
        let db = Database { conn };
        db.initialize_schema()?;
        Ok(db)
    }

    /// Creates a new in-memory database (primarily for testing).
    ///
    /// # Errors
    ///
    /// Returns an error if schema initialization fails.
    #[allow(dead_code)]
    pub fn new_in_memory() -> Result<Self> {
        let conn = Connection::open_in_memory()?;
        Self::configure_connection(&conn)?;
        let db = Database { conn };
        db.initialize_schema()?;
        Ok(db)
    }

    /// Configures the SQLite connection for optimal concurrent access.
    ///
    /// Enables WAL (Write-Ahead Logging) mode for better read/write concurrency
    /// and sets a busy timeout to automatically retry on lock contention.
    fn configure_connection(conn: &Connection) -> Result<()> {
        // Enable WAL mode for better concurrent access
        // WAL allows readers to proceed while a writer is active
        conn.pragma_update(None, "journal_mode", "WAL")?;

        // Set busy timeout to 5 seconds
        // SQLite will automatically retry if the database is locked
        conn.busy_timeout(Duration::from_secs(5))?;

        Ok(())
    }

    fn get_db_path() -> Result<PathBuf> {
        let proj_dirs = ProjectDirs::from("", "", "track")
            .ok_or(crate::utils::TrackError::DataDirectoryUnavailable)?;

        Ok(proj_dirs.data_dir().join("track.db"))
    }

    fn initialize_schema(&self) -> Result<()> {
        let schema = format!(
            r#"
            CREATE TABLE IF NOT EXISTS app_state (
                key TEXT PRIMARY KEY,
                value TEXT
            );

            CREATE TABLE IF NOT EXISTS tasks (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                status TEXT NOT NULL DEFAULT '{task_active}' CHECK (status IN ('{task_active}', '{task_archived}')),
                ticket_id TEXT UNIQUE,
                ticket_url TEXT,
                created_at TEXT NOT NULL
            );

            CREATE TABLE IF NOT EXISTS todos (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                content TEXT NOT NULL,
                status TEXT NOT NULL DEFAULT '{todo_pending}' CHECK (status IN ('{todo_pending}', '{todo_done}', '{todo_cancelled}')),
                worktree_requested INTEGER NOT NULL DEFAULT 0,
                requires_workspace INTEGER NOT NULL DEFAULT 1,
                created_at TEXT NOT NULL,
                completed_at TEXT,
                FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS links (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                url TEXT NOT NULL,
                title TEXT NOT NULL,
                created_at TEXT NOT NULL,
                FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS scraps (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                content TEXT NOT NULL,
                created_at TEXT NOT NULL,
                active_todo_id INTEGER,
                FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS worktrees (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                path TEXT NOT NULL,
                branch TEXT NOT NULL,
                base_repo TEXT,
                status TEXT NOT NULL DEFAULT 'active',
                created_at TEXT NOT NULL,
                todo_id INTEGER,
                is_base INTEGER DEFAULT 0,
                FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE,
                FOREIGN KEY (todo_id) REFERENCES todos(id) ON DELETE SET NULL
            );

            CREATE TABLE IF NOT EXISTS repo_links (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                worktree_id INTEGER NOT NULL,
                url TEXT NOT NULL,
                kind TEXT NOT NULL,
                created_at TEXT NOT NULL,
                FOREIGN KEY (worktree_id) REFERENCES worktrees(id) ON DELETE CASCADE
            );

            CREATE TABLE IF NOT EXISTS task_repos (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                task_id INTEGER NOT NULL,
                repo_path TEXT NOT NULL,
                created_at TEXT NOT NULL,
                FOREIGN KEY (task_id) REFERENCES tasks(id) ON DELETE CASCADE,
                UNIQUE(task_id, repo_path)
            );

            CREATE INDEX IF NOT EXISTS idx_todos_task_id ON todos(task_id);
            CREATE INDEX IF NOT EXISTS idx_links_task_id ON links(task_id);
            CREATE INDEX IF NOT EXISTS idx_scraps_task_id ON scraps(task_id);
            CREATE INDEX IF NOT EXISTS idx_worktrees_task_id ON worktrees(task_id);
            CREATE INDEX IF NOT EXISTS idx_task_repos_task_id ON task_repos(task_id);
            "#,
            task_active = TaskStatus::ACTIVE,
            task_archived = TaskStatus::ARCHIVED,
            todo_pending = TodoStatus::PENDING,
            todo_done = TodoStatus::DONE,
            todo_cancelled = TodoStatus::CANCELLED,
        );

        self.conn.execute_batch(&schema)?;

        migrate::migrate_schema(&self.conn)?;

        Ok(())
    }

    /// Returns a reference to the underlying SQLite connection.
    pub fn get_connection(&self) -> &Connection {
        &self.conn
    }

    /// Execute operations within an IMMEDIATE transaction.
    ///
    /// This method wraps the provided closure in a SQLite transaction that:
    /// - Uses `BEGIN IMMEDIATE` to acquire the write lock upfront
    /// - Automatically commits on success
    /// - Automatically rolls back on error
    ///
    /// Using IMMEDIATE mode prevents race conditions in read-modify-write
    /// sequences by acquiring the write lock before reading, ensuring no
    /// other process can modify the data between the read and write steps.
    ///
    /// # Arguments
    ///
    /// * `f` - A closure that performs database operations and returns a Result
    ///
    /// # Example
    ///
    /// ```ignore
    /// db.with_transaction(|| {
    ///     let next_index = get_next_index()?;
    ///     insert_with_index(next_index)?;
    ///     Ok(())
    /// })?;
    /// ```
    pub fn with_transaction<T, F>(&self, f: F) -> Result<T>
    where
        F: FnOnce() -> Result<T>,
    {
        self.conn.execute("BEGIN IMMEDIATE", [])?;
        match f() {
            Ok(value) => {
                self.conn.execute("COMMIT", [])?;
                Ok(value)
            }
            Err(e) => {
                // Attempt rollback, but ignore errors (connection may be in bad state)
                let _ = self.conn.execute("ROLLBACK", []);
                Err(e)
            }
        }
    }

    pub fn get_app_state(&self, key: &str) -> Result<Option<String>> {
        let mut stmt = self
            .conn
            .prepare("SELECT value FROM app_state WHERE key = ?1")?;
        let result = stmt.query_row(params![key], |row| row.get(0)).optional()?;
        Ok(result)
    }

    pub fn set_app_state(&self, key: &str, value: &str) -> Result<()> {
        self.conn.execute(
            "INSERT OR REPLACE INTO app_state (key, value) VALUES (?1, ?2)",
            params![key, value],
        )?;
        Ok(())
    }

    /// Returns the configured VCS backend (`jj` by default).
    pub fn get_vcs_mode(&self) -> Result<VcsMode> {
        match self.get_app_state(VcsMode::KEY)? {
            Some(value) => value
                .parse()
                .map_err(crate::utils::TrackError::InvalidVcsMode),
            None => Ok(VcsMode::default()),
        }
    }

    /// Persists the VCS backend preference.
    pub fn set_vcs_mode(&self, mode: VcsMode) -> Result<()> {
        self.set_app_state(VcsMode::KEY, mode.as_str())
    }

    /// Gets the ID of the current active task.
    ///
    /// # Returns
    ///
    /// `Some(task_id)` if a task is currently active, `None` otherwise.
    pub fn get_current_task_id(&self) -> Result<Option<i64>> {
        match self.get_app_state("current_task_id")? {
            Some(id_str) => Ok(Some(id_str.parse().map_err(|_| {
                crate::utils::TrackError::InvalidAppStateValue {
                    key: "current_task_id".to_string(),
                    detail: "expected integer".to_string(),
                }
            })?)),
            None => Ok(None),
        }
    }

    /// Sets the current active task.
    ///
    /// # Arguments
    ///
    /// * `task_id` - The ID of the task to set as current
    pub fn set_current_task_id(&self, task_id: i64) -> Result<()> {
        self.set_app_state("current_task_id", &task_id.to_string())
    }

    /// Clears the current active task.
    pub fn clear_current_task_id(&self) -> Result<()> {
        self.conn
            .execute("DELETE FROM app_state WHERE key = 'current_task_id'", [])?;
        Ok(())
    }

    /// Increments the revision number for a section and returns the new value.
    ///
    /// # Arguments
    ///
    /// * `section` - The section name (e.g., "todos", "scraps", "links", "repos", "worktrees", "task")
    pub fn increment_rev(&self, section: &str) -> Result<i64> {
        let key = format!("rev:{}", section);
        let current: i64 = self
            .get_app_state(&key)?
            .and_then(|s| s.parse().ok())
            .unwrap_or(0);
        let new_rev = current + 1;
        self.set_app_state(&key, &new_rev.to_string())?;
        Ok(new_rev)
    }

    /// Gets the current revision number for a section.
    ///
    /// # Arguments
    ///
    /// * `section` - The section name
    pub fn get_rev(&self, section: &str) -> Result<i64> {
        let key = format!("rev:{}", section);
        Ok(self
            .get_app_state(&key)?
            .and_then(|s| s.parse().ok())
            .unwrap_or(0))
    }

    /// Gets all section revision numbers at once.
    pub fn get_all_revs(&self) -> Result<SectionRevs> {
        Ok(SectionRevs {
            task: self.get_rev("task")?,
            todos: self.get_rev("todos")?,
            scraps: self.get_rev("scraps")?,
            links: self.get_rev("links")?,
            repos: self.get_rev("repos")?,
            worktrees: self.get_rev("worktrees")?,
        })
    }
}

impl crate::ports::AppStateStore for Database {
    fn get_current_task_id(&self) -> Result<Option<i64>> {
        <Database>::get_current_task_id(self)
    }

    fn set_current_task_id(&self, task_id: i64) -> Result<()> {
        <Database>::set_current_task_id(self, task_id)
    }

    fn clear_current_task_id(&self) -> Result<()> {
        <Database>::clear_current_task_id(self)
    }
}

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

    #[test]
    fn test_new_in_memory() {
        let db = Database::new_in_memory().unwrap();
        // Verify connection is valid by querying
        let result: i64 = db
            .get_connection()
            .query_row("SELECT 1", [], |row| row.get(0))
            .unwrap();
        assert_eq!(result, 1);
    }

    #[test]
    fn test_wal_mode_enabled() {
        let db = Database::new_in_memory().unwrap();
        // Note: In-memory databases use "memory" journal mode, not "wal"
        // This test verifies the pragma is set without error
        // For file-based databases, WAL mode will be properly enabled
        let mode: String = db
            .get_connection()
            .query_row("PRAGMA journal_mode", [], |row| row.get(0))
            .unwrap();
        // In-memory databases return "memory" as journal mode
        assert!(mode == "memory" || mode == "wal");
    }

    #[test]
    fn test_busy_timeout_configured() {
        let db = Database::new_in_memory().unwrap();
        let timeout: i64 = db
            .get_connection()
            .query_row("PRAGMA busy_timeout", [], |row| row.get(0))
            .unwrap();
        // Should be 5000ms (5 seconds)
        assert_eq!(timeout, 5000);
    }

    #[test]
    fn test_app_state_get_set() {
        let db = Database::new_in_memory().unwrap();

        // Initially should be None
        let value = db.get_app_state("test_key").unwrap();
        assert!(value.is_none());

        // Set a value
        db.set_app_state("test_key", "test_value").unwrap();

        // Get the value back
        let value = db.get_app_state("test_key").unwrap();
        assert_eq!(value, Some("test_value".to_string()));

        // Update the value
        db.set_app_state("test_key", "new_value").unwrap();
        let value = db.get_app_state("test_key").unwrap();
        assert_eq!(value, Some("new_value".to_string()));
    }

    #[test]
    fn test_vcs_mode_defaults_to_jj() {
        let db = Database::new_in_memory().unwrap();
        assert_eq!(db.get_vcs_mode().unwrap(), VcsMode::Jj);
    }

    #[test]
    fn test_vcs_mode_round_trip() {
        let db = Database::new_in_memory().unwrap();
        db.set_vcs_mode(VcsMode::Git).unwrap();
        assert_eq!(db.get_vcs_mode().unwrap(), VcsMode::Git);
        db.set_vcs_mode(VcsMode::Jj).unwrap();
        assert_eq!(db.get_vcs_mode().unwrap(), VcsMode::Jj);
    }

    #[test]
    fn test_with_transaction_commit() {
        let db = Database::new_in_memory().unwrap();

        // Transaction should commit on success
        let result = db.with_transaction(|| {
            db.set_app_state("tx_key", "tx_value")?;
            Ok("success")
        });

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), "success");

        // Value should persist after commit
        let value = db.get_app_state("tx_key").unwrap();
        assert_eq!(value, Some("tx_value".to_string()));
    }

    #[test]
    fn test_with_transaction_rollback() {
        let db = Database::new_in_memory().unwrap();

        // Set initial value
        db.set_app_state("rollback_key", "initial").unwrap();

        // Transaction should rollback on error
        let result: crate::utils::Result<()> = db.with_transaction(|| {
            db.set_app_state("rollback_key", "changed")?;
            Err(crate::utils::TrackError::Cancelled)
        });

        assert!(result.is_err());

        // Value should be rolled back to initial
        let value = db.get_app_state("rollback_key").unwrap();
        assert_eq!(value, Some("initial".to_string()));
    }

    #[test]
    fn test_with_transaction_returns_value() {
        let db = Database::new_in_memory().unwrap();

        // Transaction should return the value from closure
        let result = db.with_transaction(|| {
            db.set_app_state("key", "value")?;
            Ok(42i64)
        });

        assert!(result.is_ok());
        assert_eq!(result.unwrap(), 42);
    }

    #[test]
    fn test_current_task_id() {
        let db = Database::new_in_memory().unwrap();

        // Initially should be None
        let task_id = db.get_current_task_id().unwrap();
        assert!(task_id.is_none());

        // Set a task ID
        db.set_current_task_id(42).unwrap();

        // Get the task ID back
        let task_id = db.get_current_task_id().unwrap();
        assert_eq!(task_id, Some(42));

        // Clear the task ID
        db.clear_current_task_id().unwrap();
        let task_id = db.get_current_task_id().unwrap();
        assert!(task_id.is_none());
    }

    #[test]
    fn test_schema_initialization() {
        let db = Database::new_in_memory().unwrap();
        let conn = db.get_connection();

        // Verify all tables exist
        let tables = vec![
            "app_state",
            "tasks",
            "todos",
            "links",
            "scraps",
            "worktrees",
            "repo_links",
        ];
        for table in tables {
            let result: i64 = conn
                .query_row(
                    "SELECT COUNT(*) FROM sqlite_master WHERE type='table' AND name=?1",
                    [table],
                    |row| row.get(0),
                )
                .unwrap();
            assert_eq!(result, 1, "Table {} should exist", table);
        }
    }

    #[test]
    fn test_increment_rev() {
        let db = Database::new_in_memory().unwrap();

        // Initially should be 0
        let rev = db.get_rev("todos").unwrap();
        assert_eq!(rev, 0);

        // Increment should return 1
        let new_rev = db.increment_rev("todos").unwrap();
        assert_eq!(new_rev, 1);

        // Get should return 1
        let rev = db.get_rev("todos").unwrap();
        assert_eq!(rev, 1);

        // Increment again should return 2
        let new_rev = db.increment_rev("todos").unwrap();
        assert_eq!(new_rev, 2);
    }

    #[test]
    fn test_get_rev_default() {
        let db = Database::new_in_memory().unwrap();

        // Non-existent sections should return 0
        let rev = db.get_rev("nonexistent").unwrap();
        assert_eq!(rev, 0);
    }

    #[test]
    fn test_increment_rev_different_sections() {
        let db = Database::new_in_memory().unwrap();

        // Increment different sections
        db.increment_rev("todos").unwrap();
        db.increment_rev("todos").unwrap();
        db.increment_rev("scraps").unwrap();
        db.increment_rev("links").unwrap();
        db.increment_rev("links").unwrap();
        db.increment_rev("links").unwrap();

        // Verify each section has independent rev
        assert_eq!(db.get_rev("todos").unwrap(), 2);
        assert_eq!(db.get_rev("scraps").unwrap(), 1);
        assert_eq!(db.get_rev("links").unwrap(), 3);
        assert_eq!(db.get_rev("repos").unwrap(), 0);
    }

    #[test]
    fn test_get_all_revs() {
        let db = Database::new_in_memory().unwrap();

        // Initially all should be 0
        let revs = db.get_all_revs().unwrap();
        assert_eq!(revs.task, 0);
        assert_eq!(revs.todos, 0);
        assert_eq!(revs.scraps, 0);
        assert_eq!(revs.links, 0);
        assert_eq!(revs.repos, 0);
        assert_eq!(revs.worktrees, 0);

        // Increment some sections
        db.increment_rev("task").unwrap();
        db.increment_rev("todos").unwrap();
        db.increment_rev("todos").unwrap();
        db.increment_rev("worktrees").unwrap();

        // Verify get_all_revs returns correct values
        let revs = db.get_all_revs().unwrap();
        assert_eq!(revs.task, 1);
        assert_eq!(revs.todos, 2);
        assert_eq!(revs.scraps, 0);
        assert_eq!(revs.links, 0);
        assert_eq!(revs.repos, 0);
        assert_eq!(revs.worktrees, 1);
    }

    #[test]
    fn test_section_revs_equality() {
        let db = Database::new_in_memory().unwrap();

        let revs1 = db.get_all_revs().unwrap();
        let revs2 = db.get_all_revs().unwrap();
        assert_eq!(revs1, revs2);

        db.increment_rev("todos").unwrap();
        let revs3 = db.get_all_revs().unwrap();
        assert_ne!(revs1, revs3);
    }

    #[test]
    fn test_status_check_constraints_enforced() {
        use crate::models::{TaskStatus, TodoStatus};

        let db = Database::new_in_memory().unwrap();
        let conn = db.get_connection();

        let tasks_sql: String = conn
            .query_row(
                "SELECT sql FROM sqlite_master WHERE type='table' AND name='tasks'",
                [],
                |row| row.get(0),
            )
            .unwrap();
        assert!(tasks_sql.contains("CHECK"));

        conn.execute(
            "INSERT INTO tasks (name, status, created_at) VALUES ('Bad', ?1, datetime('now'))",
            rusqlite::params![TaskStatus::ACTIVE],
        )
        .unwrap();

        let err = conn
            .execute(
                "INSERT INTO tasks (name, status, created_at) VALUES ('Bad', 'invalid', datetime('now'))",
                [],
            )
            .unwrap_err();
        assert!(err.to_string().contains("CHECK constraint failed"));

        conn.execute(
            "INSERT INTO todos (task_id, task_index, content, status, worktree_requested, created_at) VALUES (1, 1, 'Ok', ?1, 0, datetime('now'))",
            rusqlite::params![TodoStatus::PENDING],
        )
        .unwrap();

        let err = conn
            .execute(
                "INSERT INTO todos (task_id, task_index, content, status, worktree_requested, created_at) VALUES (1, 2, 'Bad', 'reopened', 0, datetime('now'))",
                [],
            )
            .unwrap_err();
        assert!(err.to_string().contains("CHECK constraint failed"));
    }
}