starpod_db/

lib.rs

//! Unified SQLite database for Starpod's transactional data.
//!
//! Consolidates sessions, cron scheduling, and authentication into a single
//! `core.db` file, replacing the previous `session.db` + `cron.db` + `users.db`
//! setup. A single shared connection pool (WAL mode, foreign keys enabled)
//! serves all three domains.
//!
//! # Architecture
//!
//! ```text
//! ┌──────────┐
//! │  CoreDb   │  owns SqlitePool (max 10 conns, WAL, FK ON)
//! └────┬─────┘
//!      │ pool.clone()
//!      ├──────────────► SessionManager::from_pool(pool)
//!      ├──────────────► CronStore::from_pool(pool)
//!      └──────────────► AuthStore::from_pool(pool)
//! ```
//!
//! # Databases kept separate
//!
//! - **memory.db** — FTS5 + vector blobs, bulk reindex I/O, different access pattern
//! - **vault.db** — AES-256-GCM encrypted, optional (needs `.vault_key`), isolated security boundary
//!
//! # Legacy migration
//!
//! On first open, [`CoreDb::new`] detects old `session.db` / `cron.db` / `users.db`
//! in the same directory and migrates their data into `core.db`, then renames the
//! old files to `*.db.migrated`.
//!
//! # Usage
//!
//! ```no_run
//! # async fn example() -> starpod_core::Result<()> {
//! use starpod_db::CoreDb;
//!
//! let db = CoreDb::new(std::path::Path::new(".starpod/db")).await?;
//! // Pass db.pool().clone() to SessionManager, CronStore, AuthStore
//! # Ok(())
//! # }
//! ```

mod migrate;

use std::path::Path;
use std::str::FromStr;

use sqlx::sqlite::{SqliteConnectOptions, SqlitePoolOptions};
use sqlx::SqlitePool;
use tracing::{debug, info};

use starpod_core::{StarpodError, Result};

/// Unified database for sessions, cron, and auth.
///
/// Owns a single `SqlitePool` backed by `core.db`. Individual stores
/// (`SessionManager`, `CronStore`, `AuthStore`) receive a clone of the
/// pool via `from_pool()` instead of opening their own connections.
///
/// The pool is configured with:
/// - **WAL journal mode** — concurrent readers don't block writers
/// - **Foreign keys ON** — referential integrity across all tables
/// - **10 max connections** — shared across all stores
pub struct CoreDb {
    pool: SqlitePool,
}

impl CoreDb {
    /// Open (or create) `core.db` inside `db_dir`.
    ///
    /// Runs all migrations, then checks for legacy database files
    /// (`session.db`, `cron.db`, `users.db`) and migrates their data
    /// into the unified database if found.
    pub async fn new(db_dir: &Path) -> Result<Self> {
        std::fs::create_dir_all(db_dir)?;

        let db_path = db_dir.join("core.db");
        let opts = SqliteConnectOptions::from_str(
            &format!("sqlite://{}?mode=rwc", db_path.display()),
        )
        .map_err(|e| StarpodError::Database(format!("Invalid DB path: {}", e)))?
        .pragma("journal_mode", "WAL")
        .pragma("foreign_keys", "ON");

        let pool = SqlitePoolOptions::new()
            .max_connections(10)
            .connect_with(opts)
            .await
            .map_err(|e| StarpodError::Database(format!("Failed to open core db: {}", e)))?;

        sqlx::migrate!("./migrations")
            .run(&pool)
            .await
            .map_err(|e| StarpodError::Database(format!("Core migration failed: {}", e)))?;

        debug!("core.db ready at {}", db_path.display());

        // Migrate legacy databases if present
        if migrate::has_legacy_dbs(db_dir) {
            info!("Legacy database files detected — migrating to core.db");
            migrate::migrate_legacy_dbs(&pool, db_dir).await?;
        }

        Ok(Self { pool })
    }

    /// Create an in-memory `CoreDb` for testing.
    ///
    /// Runs all migrations on a shared in-memory database. Each call
    /// returns a fresh, empty database.
    pub async fn in_memory() -> Result<Self> {
        let opts = SqliteConnectOptions::from_str("sqlite::memory:")
            .map_err(|e| StarpodError::Database(format!("Invalid memory DB: {}", e)))?
            .pragma("foreign_keys", "ON");

        let pool = SqlitePoolOptions::new()
            .max_connections(1)
            .connect_with(opts)
            .await
            .map_err(|e| StarpodError::Database(format!("Failed to open in-memory db: {}", e)))?;

        sqlx::migrate!("./migrations")
            .run(&pool)
            .await
            .map_err(|e| StarpodError::Database(format!("Core migration failed: {}", e)))?;

        Ok(Self { pool })
    }

    /// Get a reference to the shared connection pool.
    pub fn pool(&self) -> &SqlitePool {
        &self.pool
    }
}

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

    // ── Basic lifecycle ─────────────────────────────────────────────

    #[tokio::test]
    async fn in_memory_creates_all_tables() {
        let db = CoreDb::in_memory().await.unwrap();
        let pool = db.pool();

        // Auth tables
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM api_keys")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM telegram_links")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM auth_audit_log")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        // Session tables
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM session_metadata")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM session_messages")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM usage_stats")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM compaction_log")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        // Cron tables
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM cron_jobs")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM cron_runs")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);
    }

    #[tokio::test]
    async fn on_disk_creates_core_db() {
        let tmp = tempfile::tempdir().unwrap();
        let db = CoreDb::new(tmp.path()).await.unwrap();

        assert!(tmp.path().join("core.db").exists());

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users")
            .fetch_one(db.pool()).await.unwrap();
        assert_eq!(row.0, 0);
    }

    #[tokio::test]
    async fn on_disk_creates_parent_dirs() {
        let tmp = tempfile::tempdir().unwrap();
        let nested = tmp.path().join("deep").join("nested").join("db");
        let db = CoreDb::new(&nested).await.unwrap();

        assert!(nested.join("core.db").exists());
        drop(db);
    }

    #[tokio::test]
    async fn reopen_is_idempotent() {
        let tmp = tempfile::tempdir().unwrap();

        // First open — creates the DB
        let db1 = CoreDb::new(tmp.path()).await.unwrap();
        sqlx::query(
            "INSERT INTO users (id, email, display_name, role, is_active, created_at, updated_at) \
             VALUES ('u1', 'a@b.com', 'A', 'admin', 1, '2024-01-01', '2024-01-01')"
        ).execute(db1.pool()).await.unwrap();
        drop(db1);

        // Second open — should find existing data, not recreate
        let db2 = CoreDb::new(tmp.path()).await.unwrap();
        let row: (String,) = sqlx::query_as("SELECT email FROM users WHERE id = 'u1'")
            .fetch_one(db2.pool()).await.unwrap();
        assert_eq!(row.0, "a@b.com");
    }

    // ── Foreign key enforcement ─────────────────────────────────────

    #[tokio::test]
    async fn fk_rejects_invalid_api_key_user() {
        let db = CoreDb::in_memory().await.unwrap();

        let result = sqlx::query(
            "INSERT INTO api_keys (id, user_id, prefix, key_hash, created_at) \
             VALUES ('k1', 'nonexistent', 'sp_', 'hash', '2024-01-01')"
        ).execute(db.pool()).await;

        assert!(result.is_err(), "FK should reject api_key with invalid user_id");
    }

    #[tokio::test]
    async fn fk_rejects_invalid_telegram_link_user() {
        let db = CoreDb::in_memory().await.unwrap();

        let result = sqlx::query(
            "INSERT INTO telegram_links (telegram_id, user_id, username, linked_at) \
             VALUES (123, 'nonexistent', 'bob', '2024-01-01')"
        ).execute(db.pool()).await;

        assert!(result.is_err(), "FK should reject telegram_link with invalid user_id");
    }

    #[tokio::test]
    async fn fk_rejects_invalid_session_message() {
        let db = CoreDb::in_memory().await.unwrap();

        let result = sqlx::query(
            "INSERT INTO session_messages (session_id, role, content, timestamp) \
             VALUES ('nonexistent', 'user', 'hello', '2024-01-01')"
        ).execute(db.pool()).await;

        assert!(result.is_err(), "FK should reject message with invalid session_id");
    }

    #[tokio::test]
    async fn fk_rejects_invalid_cron_run_job() {
        let db = CoreDb::in_memory().await.unwrap();

        let result = sqlx::query(
            "INSERT INTO cron_runs (id, job_id, started_at, status) \
             VALUES ('r1', 'nonexistent', 1000, 'pending')"
        ).execute(db.pool()).await;

        assert!(result.is_err(), "FK should reject cron_run with invalid job_id");
    }

    // ── CASCADE deletes ─────────────────────────────────────────────

    #[tokio::test]
    async fn cascade_delete_user_removes_api_keys() {
        let db = CoreDb::in_memory().await.unwrap();
        let pool = db.pool();

        sqlx::query(
            "INSERT INTO users (id, email, role, is_active, created_at, updated_at) \
             VALUES ('u1', 'a@b.com', 'admin', 1, '2024-01-01', '2024-01-01')"
        ).execute(pool).await.unwrap();

        sqlx::query(
            "INSERT INTO api_keys (id, user_id, prefix, key_hash, created_at) \
             VALUES ('k1', 'u1', 'sp_', 'hash1', '2024-01-01')"
        ).execute(pool).await.unwrap();

        sqlx::query(
            "INSERT INTO api_keys (id, user_id, prefix, key_hash, created_at) \
             VALUES ('k2', 'u1', 'sp_', 'hash2', '2024-01-01')"
        ).execute(pool).await.unwrap();

        // Delete user
        sqlx::query("DELETE FROM users WHERE id = 'u1'")
            .execute(pool).await.unwrap();

        // API keys should be gone
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM api_keys")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);
    }

    #[tokio::test]
    async fn cascade_delete_user_removes_telegram_links() {
        let db = CoreDb::in_memory().await.unwrap();
        let pool = db.pool();

        sqlx::query(
            "INSERT INTO users (id, role, is_active, created_at, updated_at) \
             VALUES ('u1', 'admin', 1, '2024-01-01', '2024-01-01')"
        ).execute(pool).await.unwrap();

        sqlx::query(
            "INSERT INTO telegram_links (telegram_id, user_id, username, linked_at) \
             VALUES (999, 'u1', 'bob', '2024-01-01')"
        ).execute(pool).await.unwrap();

        sqlx::query("DELETE FROM users WHERE id = 'u1'")
            .execute(pool).await.unwrap();

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM telegram_links")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);
    }

    #[tokio::test]
    async fn cascade_delete_session_removes_messages_and_compaction() {
        let db = CoreDb::in_memory().await.unwrap();
        let pool = db.pool();

        sqlx::query(
            "INSERT INTO session_metadata (id, created_at, last_message_at) \
             VALUES ('s1', '2024-01-01', '2024-01-01')"
        ).execute(pool).await.unwrap();

        sqlx::query(
            "INSERT INTO session_messages (session_id, role, content, timestamp) \
             VALUES ('s1', 'user', 'hi', '2024-01-01')"
        ).execute(pool).await.unwrap();

        sqlx::query(
            "INSERT INTO compaction_log (session_id, timestamp, trigger, pre_tokens, summary) \
             VALUES ('s1', '2024-01-01', 'auto', 1000, 'summary')"
        ).execute(pool).await.unwrap();

        sqlx::query("DELETE FROM session_metadata WHERE id = 's1'")
            .execute(pool).await.unwrap();

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM session_messages")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM compaction_log")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);
    }

    #[tokio::test]
    async fn cascade_delete_cron_job_removes_runs() {
        let db = CoreDb::in_memory().await.unwrap();
        let pool = db.pool();

        sqlx::query(
            "INSERT INTO cron_jobs (id, name, prompt, schedule_type, schedule_value, created_at) \
             VALUES ('j1', 'test', 'do stuff', 'interval', '60000', 1000)"
        ).execute(pool).await.unwrap();

        sqlx::query(
            "INSERT INTO cron_runs (id, job_id, started_at, status) \
             VALUES ('r1', 'j1', 2000, 'success')"
        ).execute(pool).await.unwrap();

        sqlx::query("DELETE FROM cron_jobs WHERE id = 'j1'")
            .execute(pool).await.unwrap();

        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM cron_runs")
            .fetch_one(pool).await.unwrap();
        assert_eq!(row.0, 0);
    }

    // ── Cross-domain queries (the whole point of consolidation) ─────

    #[tokio::test]
    async fn cross_domain_join_sessions_with_usage_by_user() {
        let db = CoreDb::in_memory().await.unwrap();
        let pool = db.pool();

        // Create a user
        sqlx::query(
            "INSERT INTO users (id, email, role, is_active, created_at, updated_at) \
             VALUES ('u1', 'alice@test.com', 'admin', 1, '2024-01-01', '2024-01-01')"
        ).execute(pool).await.unwrap();

        // Create sessions for this user
        sqlx::query(
            "INSERT INTO session_metadata (id, created_at, last_message_at, user_id) \
             VALUES ('s1', '2024-01-01', '2024-01-01', 'u1')"
        ).execute(pool).await.unwrap();

        // Record usage
        sqlx::query(
            "INSERT INTO usage_stats (session_id, turn, input_tokens, output_tokens, cost_usd, timestamp, user_id) \
             VALUES ('s1', 1, 100, 200, 0.01, '2024-01-01', 'u1')"
        ).execute(pool).await.unwrap();

        // Cross-domain query: total cost per user (joins users + usage_stats)
        let row: (String, f64) = sqlx::query_as(
            "SELECT u.email, SUM(us.cost_usd) as total_cost \
             FROM users u \
             JOIN usage_stats us ON us.user_id = u.id \
             GROUP BY u.id"
        ).fetch_one(pool).await.unwrap();

        assert_eq!(row.0, "alice@test.com");
        assert!((row.1 - 0.01).abs() < 0.001);
    }

    #[tokio::test]
    async fn pool_clone_shares_state() {
        let db = CoreDb::in_memory().await.unwrap();

        // Insert on original pool
        sqlx::query(
            "INSERT INTO users (id, role, is_active, created_at, updated_at) \
             VALUES ('u1', 'admin', 1, '2024-01-01', '2024-01-01')"
        ).execute(db.pool()).await.unwrap();

        // Read from cloned pool (simulates what stores do)
        let pool2 = db.pool().clone();
        let row: (i64,) = sqlx::query_as("SELECT COUNT(*) FROM users")
            .fetch_one(&pool2).await.unwrap();
        assert_eq!(row.0, 1);
    }
}