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

//! Database migration binary for zesdex-backend.
//!
//! Scans all session directories under the store path and initializes or
//! upgrades the SQLite schema (`messages.sqlite`) for each one. This is
//! a standalone CLI tool invoked as `cargo run --bin migrate`.
//!
//! ## Workflow
//! 1. Resolve the base store directory via `Store::new()`
//! 2. Iterate over each subdirectory under `sessions/`
//! 3. For each session directory, call `migrate_session_msglog()` to
//!    create/upgrade the `messages.sqlite` schema
//! 4. Report count of succeeded and failed migrations
//! 5. Exit with error if any session failed
//!
//! ## Schema
//! - `messages` table — stores conversation message rows
//! - `archives` table — stores session archive metadata
//! - `blobs` table — stores binary blob data per session
//! - Indexes on `session_id`, `created_at`, and `role` columns
//!
//! ## Versioning
//! SQLite `PRAGMA user_version` tracks schema version for incremental upgrades.

use std::path::Path;


/// Entry point: migrate all session databases.
///
/// Flow: load store → iterate sessions → migrate each → summarise.
///
/// Returns an error if any session migration failed.
fn main() -> anyhow::Result<()> {
    tracing::info!("starting database migration");
    let store = zesdex_entities::domain::common::store::Store::new();

    // Resolve the sessions directory under the store base path
    let sessions_dir = store.base_dir.join("sessions");
    if !sessions_dir.exists() {
        tracing::info!("no sessions directory found at {:?}", sessions_dir);
        eprintln!("No sessions directory found, nothing to migrate");
        return Ok(());
    }

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

    // Iterate over all session subdirectories
    for entry in std::fs::read_dir(&sessions_dir)? {
        let entry = entry?;
        let path = entry.path();
        if !path.is_dir() {
            continue; // skip non-directory entries
        }

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

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

/// Open (or create) a session's `messages.sqlite` and ensure its schema is current.
///
/// Flow: resolve path → open/ create DB → set PRAGMAs → create tables → upgrade version.
///
/// ## Parameters
/// - `session_dir`: path to the individual session directory
///
/// ## Returns
/// - `Ok(())` on success
/// - `Err` if file I/O or SQLite operations fail
fn migrate_session_msglog(session_dir: &Path) -> anyhow::Result<()> {
    tracing::debug!("migrating session at {:?}", session_dir);
    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(())
}