scryer-db 0.2.1

use anyhow::Context;

use crate::models::{Project, SchemaVersion, ToolInvocationMetric};
use crate::time::{format_unix_millis, parse_legacy};

/// Current schema version.
///
/// Version 1 is the baseline: the models pushed by `push_schema` plus the idempotent
/// statements in [`TABLE_STATEMENTS`], [`apply_column_migrations`] and [`INDEX_STATEMENTS`],
/// which run on every start. To change the schema afterwards, bump this and append a
/// `(version, statements)` step to [`MIGRATION_STEPS`]; steps run once, in order, in a
/// transaction. `push_schema` only runs for a brand-new database, so a new table or column
/// *must* be added through a step (and, for a new model, also to `models![]` in `db.rs`).
pub const SCHEMA_VERSION: u64 = 1;

/// Ordered schema changes after the baseline: `(version, statements)`.
pub const MIGRATION_STEPS: &[(u64, &[&str])] = &[];

pub const TABLE_STATEMENTS: &[&str] = &[
    "CREATE TABLE IF NOT EXISTS schema_version (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    version INTEGER NOT NULL,
    applied_at TEXT NOT NULL
);",
    "CREATE TABLE IF NOT EXISTS dependency_package (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name TEXT NOT NULL,
    version TEXT NOT NULL,
    source_type TEXT NOT NULL,
    package_hash TEXT NOT NULL,
    root_path TEXT NOT NULL,
    manifest_path TEXT NOT NULL,
    indexed_at TEXT NOT NULL
);",
    "CREATE TABLE IF NOT EXISTS project_dependency (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    project_id INTEGER NOT NULL,
    dependency_package_id INTEGER NOT NULL,
    is_direct INTEGER NOT NULL,
    features TEXT NOT NULL,
    FOREIGN KEY (project_id) REFERENCES project(id) ON DELETE CASCADE,
    FOREIGN KEY (dependency_package_id) REFERENCES dependency_package(id) ON DELETE CASCADE
);",
];

pub const INDEX_STATEMENTS: &[&str] = &[
    "CREATE INDEX IF NOT EXISTS idx_symbol_proj_name ON symbol (project_id, name);",
    "CREATE INDEX IF NOT EXISTS idx_symbol_proj_qualified ON symbol (project_id, qualified_name);",
    "CREATE INDEX IF NOT EXISTS idx_files_proj_path ON source_file (project_id, path);",
    "CREATE INDEX IF NOT EXISTS idx_scope_proj_file ON scope (project_id, file_id, start_line, end_line);",
    "CREATE INDEX IF NOT EXISTS idx_ref_proj_sym ON symbol_reference (project_id, symbol_id);",
    "CREATE INDEX IF NOT EXISTS idx_edge_proj_source ON code_graph_edge (project_id, source_symbol_id);",
    "CREATE INDEX IF NOT EXISTS idx_edge_proj_target ON code_graph_edge (project_id, target_symbol_id);",
    "CREATE UNIQUE INDEX IF NOT EXISTS idx_dep_pkg_name_ver_hash ON dependency_package (name, version, package_hash);",
    "CREATE INDEX IF NOT EXISTS idx_proj_dep_proj_pkg ON project_dependency (project_id, dependency_package_id);",
    "CREATE INDEX IF NOT EXISTS idx_files_dep_pkg ON source_file (dependency_package_id);",
    "CREATE INDEX IF NOT EXISTS idx_symbol_dep_pkg ON symbol (dependency_package_id);",
    "CREATE INDEX IF NOT EXISTS idx_symbol_dep_pkg_name ON symbol (dependency_package_id, name);",
];

/// Add columns that may be missing in existing databases.
pub async fn apply_column_migrations(db: &mut toasty::Db) -> anyhow::Result<()> {
    for (table, col) in [
        ("source_file", "dependency_package_id"),
        ("symbol", "dependency_package_id"),
    ] {
        let stmt = format!("ALTER TABLE {table} ADD COLUMN {col} INTEGER;");
        match toasty::sql::statement(stmt).exec(db).await {
            Ok(_) => {}
            Err(e) => {
                let err_str = e.to_string();
                if err_str.contains("duplicate column name")
                    || err_str.contains("already exists")
                    || err_str.contains("no such table")
                {
                    // Already exists or table not created yet
                } else {
                    return Err(e).with_context(|| format!("Failed to alter {table} add {col}"));
                }
            }
        }
    }
    Ok(())
}

/// Highest schema version recorded in the database (0 when none is, e.g. a legacy database).
async fn stored_schema_version(db: &mut toasty::Db) -> anyhow::Result<u64> {
    Ok(SchemaVersion::all()
        .exec(&mut *db)
        .await?
        .into_iter()
        .map(|v| v.version)
        .max()
        .unwrap_or(0))
}

/// Fail if the database was written by a newer Scryer than this one.
///
/// Opening it anyway would query tables or columns this build doesn't know about, or worse,
/// write rows it doesn't understand.
pub async fn check_schema_version(db: &mut toasty::Db) -> anyhow::Result<()> {
    // The table is created here, not only by `apply_indexes`, so legacy databases can be read.
    toasty::sql::statement(TABLE_STATEMENTS[0])
        .exec(&mut *db)
        .await
        .context("Failed to create schema_version table")?;
    let stored = stored_schema_version(db).await?;
    if stored > SCHEMA_VERSION {
        anyhow::bail!(
            "This database has schema version {stored}, but this Scryer only understands up to \
             version {SCHEMA_VERSION}. Upgrade Scryer, or point it at another database with \
             `--db-url` / `SCRYER_DB_URL`."
        );
    }
    Ok(())
}

/// Run pending [`MIGRATION_STEPS`] and record the resulting version. Idempotent.
pub async fn record_schema_version(db: &mut toasty::Db) -> anyhow::Result<()> {
    let stored = stored_schema_version(db).await?;
    if stored >= SCHEMA_VERSION {
        return Ok(());
    }
    for (version, statements) in MIGRATION_STEPS.iter().filter(|(v, _)| *v > stored) {
        let mut tx = db.transaction().await?;
        for stmt in *statements {
            toasty::sql::statement(*stmt)
                .exec(&mut tx)
                .await
                .with_context(|| format!("Schema migration {version} failed: {stmt}"))?;
        }
        tx.commit().await?;
        tracing::info!("Applied schema migration {version}");
    }
    SchemaVersion::create()
        .version(SCHEMA_VERSION)
        .applied_at(crate::time::now_rfc3339())
        .exec(&mut *db)
        .await?;
    Ok(())
}

/// Applies required multi-tenant composite indexes and tables to Turso database.
pub async fn apply_indexes(db: &mut toasty::Db) -> anyhow::Result<()> {
    for stmt in TABLE_STATEMENTS {
        toasty::sql::statement(*stmt)
            .exec(db)
            .await
            .with_context(|| format!("Failed to execute table migration: {stmt}"))?;
    }

    apply_column_migrations(db).await?;

    for stmt in INDEX_STATEMENTS {
        toasty::sql::statement(*stmt)
            .exec(db)
            .await
            .with_context(|| format!("Failed to execute index migration: {stmt}"))?;
    }
    Ok(())
}

/// Rewrite timestamps stored in the legacy `"<secs>.<millis>Z"` format as RFC 3339.
///
/// Idempotent: rows already in RFC 3339 are left alone.
pub async fn upgrade_legacy_timestamps(db: &mut toasty::Db) -> anyhow::Result<()> {
    let mut statements = Vec::new();
    for project in Project::all().exec(&mut *db).await? {
        let created = parse_legacy(&project.created_at);
        let updated = parse_legacy(&project.updated_at);
        if created.is_none() && updated.is_none() {
            continue;
        }
        let created = created.map_or(project.created_at, format_unix_millis);
        let updated = updated.map_or(project.updated_at, format_unix_millis);
        statements.push(format!(
            "UPDATE project SET created_at = '{created}', updated_at = '{updated}' WHERE id = {};",
            project.id
        ));
    }
    for metric in ToolInvocationMetric::all().exec(&mut *db).await? {
        if let Some(millis) = parse_legacy(&metric.created_at) {
            statements.push(format!(
                "UPDATE tool_invocation_metric SET created_at = '{}' WHERE id = {};",
                format_unix_millis(millis),
                metric.id
            ));
        }
    }
    if statements.is_empty() {
        return Ok(());
    }

    tracing::info!(
        "Upgrading {} legacy timestamps to RFC 3339",
        statements.len()
    );
    let mut tx = db.transaction().await?;
    for stmt in &statements {
        toasty::sql::statement(stmt).exec(&mut tx).await?;
    }
    tx.commit().await?;
    Ok(())
}