zesdex/crates/zesdex-backend/src/bin/migrate.rs

//! Database migration: creates/upgrades SQLite schemas for all sessions.
use std::path::Path;

fn main() -> anyhow::Result<()> {
    let store = zesdex_entities::seaorm::common::store::Store::new();

    // Find all session directories
    let sessions_dir = store.base_dir.join("sessions");
    if !sessions_dir.exists() {
        eprintln!("No sessions directory found, nothing to migrate");
        return Ok(());
    }

    let mut migrated = 0u32;
    let mut failed = 0u32;

    for entry in std::fs::read_dir(&sessions_dir)? {
        let entry = entry?;
        let path = entry.path();
        if !path.is_dir() {
            continue;
        }

        match migrate_session_msglog(&path) {
            Ok(_) => {
                migrated += 1;
                eprintln!("Migrated session: {:?}", path.file_name());
            }
            Err(e) => {
                failed += 1;
                eprintln!("Failed to migrate session {:?}: {e}", path.file_name());
            }
        }
    }

    eprintln!("Migration complete: {migrated} succeeded, {failed} failed");
    if failed > 0 {
        anyhow::bail!("{failed} session(s) failed to migrate");
    }
    Ok(())
}

/// Open a session's `messages.sqlite` and initialize its schema.
fn migrate_session_msglog(session_dir: &Path) -> anyhow::Result<()> {
    let msglog_path = session_dir.join("messages.sqlite");

    if let Some(parent) = msglog_path.parent() {
        std::fs::create_dir_all(parent)?;
    }

    let conn = rusqlite::Connection::open(&msglog_path)?;
    conn.execute_batch("PRAGMA journal_mode = WAL;")?;
    conn.execute_batch("PRAGMA busy_timeout = 5000;")?;

    // Initialize schema
    conn.execute_batch("PRAGMA foreign_keys = ON;")?;
    conn.execute_batch(
        "
        CREATE TABLE IF NOT EXISTS messages (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            role TEXT NOT NULL,
            content TEXT,
            tool_call_id TEXT,
            tool_name TEXT,
            tool_arguments TEXT,
            created_at INTEGER NOT NULL
        );
        CREATE TABLE IF NOT EXISTS archives (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL UNIQUE,
            title TEXT,
            model TEXT,
            created_at INTEGER NOT NULL,
            updated_at INTEGER NOT NULL,
            message_count INTEGER DEFAULT 0,
            token_count INTEGER DEFAULT 0,
            summary TEXT
        );
        CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages(session_id);
        CREATE INDEX IF NOT EXISTS idx_messages_created_at ON messages(created_at);
        CREATE INDEX IF NOT EXISTS idx_archives_created_at ON archives(created_at);
        CREATE TABLE IF NOT EXISTS blobs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            session_id TEXT NOT NULL,
            blob_key TEXT NOT NULL,
            data BLOB NOT NULL,
            mime_type TEXT,
            created_at INTEGER NOT NULL,
            UNIQUE(session_id, blob_key)
        );
        ",
    )?;

    // Check and upgrade schema version
    let version: i32 = conn
        .pragma_query_value(None, "user_version", |row| row.get(0))
        .unwrap_or(0);

    if version < 1 {
        conn.pragma_update(None, "user_version", 1)?;
    }
    if version < 2 {
        conn.execute_batch(
            "CREATE INDEX IF NOT EXISTS idx_messages_session_role ON messages(session_id, role);",
        )?;
        conn.pragma_update(None, "user_version", 2)?;
    }

    Ok(())
}