yggdrasil/db/

migrate.rs

//! 数据库迁移运行器。
//!
//! 在服务器开始监听之前自动执行迁移。
//! 设计要点：
//! - 迁移 SQL 通过 `include_str!` 内联进二进制，部署只需单个二进制。
//! - `schema_migrations` 表记录已应用版本，避免重复执行。
//! - 每个迁移在独立事务里执行，失败自动回滚，版本行不会写入。
//! - 咨询锁（`pg_advisory_lock`）保证多实例启动时只有一个进程执行迁移。
//!
//! 前置条件：目标数据库的存在性由 [`crate::db::pool::ensure_database`] 在本模块
//! 被调用前保证（它先连 `postgres` 维护库做 `CREATE DATABASE IF NOT EXISTS` 等价逻辑）。
//! 本模块负责 schema 与一次性存储格式迁移，不再关心"库不存在"。
//!
//! 仅在 `feature = "server"` 时编译。

use std::collections::HashSet;

/// 咨询锁的固定 key。Postgres 咨询锁是数据库级唯一的；
/// 这里用一个项目专属的大整数，避免与同库其它应用冲突。
/// 该值无语义，仅要求唯一性。
const ADVISORY_LOCK_KEY: i64 = 0x5947_4752_4153_494C;

/// 所有迁移的 (version, sql) 列表，按 version 升序排列。
///
/// 新增迁移时：
/// 1. 在 `migrations/` 下创建 `NNN_描述.sql`
/// 2. 在本数组末尾追加一行 `("NNN", include_str!("../../migrations/NNN_描述.sql"))`
///
/// version 用字符串而非整数，与文件名前缀直接对应，且未来可支持日期戳/语义版本。
const MIGRATIONS: &[(&str, &str)] = &[
    ("001", include_str!("../../migrations/001_init.sql")),
    ("002", include_str!("../../migrations/002_posts.sql")),
    ("003", include_str!("../../migrations/003_indexes.sql")),
    ("004", include_str!("../../migrations/004_search_trgm.sql")),
    ("005", include_str!("../../migrations/005_comments.sql")),
    ("006", include_str!("../../migrations/006_add_toc_html.sql")),
    ("007", include_str!("../../migrations/007_settings.sql")),
    (
        "008",
        include_str!("../../migrations/008_comments_cascade.sql"),
    ),
    (
        "009",
        include_str!("../../migrations/009_cleanup_duplicate_indexes.sql"),
    ),
    (
        "010",
        include_str!("../../migrations/010_post_word_counts.sql"),
    ),
    ("011", include_str!("../../migrations/011_perf_indexes.sql")),
    (
        "012",
        include_str!("../../migrations/012_session_generation.sql"),
    ),
    (
        "013",
        include_str!("../../migrations/013_comment_content_hash_index.sql"),
    ),
    (
        "014",
        include_str!("../../migrations/014_drop_ineffective_trgm_index.sql"),
    ),
    ("015", include_str!("../../migrations/015_assets.sql")),
    (
        "016",
        include_str!("../../migrations/016_assets_content_hash.sql"),
    ),
    ("017", include_str!("../../migrations/017_mcp_tokens.sql")),
    (
        "018",
        include_str!("../../migrations/018_session_generation_trigger.sql"),
    ),
    (
        "019",
        include_str!("../../migrations/019_restore_search_trgm_index.sql"),
    ),
    ("020", include_str!("../../migrations/020_friend_links.sql")),
    (
        "021",
        include_str!("../../migrations/021_site_settings.sql"),
    ),
    (
        "022",
        include_str!("../../migrations/022_backup_settings.sql"),
    ),
    ("023", include_str!("../../migrations/023_user_profile.sql")),
    (
        "024",
        include_str!("../../migrations/024_comments_user_id.sql"),
    ),
    ("025", include_str!("../../migrations/025_logs.sql")),
    (
        "026",
        include_str!("../../migrations/026_katex_html_classes.sql"),
    ),
    ("027", include_str!("../../migrations/027_notes.sql")),
    (
        "028",
        include_str!("../../migrations/028_notebook_archive.sql"),
    ),
    (
        "029",
        include_str!("../../migrations/029_content_cache_revision.sql"),
    ),
    (
        "030",
        include_str!("../../migrations/030_note_asset_refs_asset_index.sql"),
    ),
    (
        "031",
        include_str!("../../migrations/031_session_generation_snapshot.sql"),
    ),
    (
        "032",
        include_str!("../../migrations/032_comment_avatar_asset_refs.sql"),
    ),
    // 新增迁移在此追加，同时在 migrations/ 下创建对应 .sql 文件。
];

/// 迁移执行错误。
#[derive(Debug)]
pub enum MigrateError {
    /// 无法从连接池获取连接。
    Pool(deadpool_postgres::PoolError),
    /// 执行查询（建表、查版本、咨询锁等）失败。
    Query(tokio_postgres::Error),
    /// 某个具体迁移执行失败（包含版本号便于定位）。
    Apply {
        version: String,
        source: tokio_postgres::Error,
    },
    /// 迁移 026 的 HTML 解析/改写失败；整笔迁移回滚，不丢弃正文。
    KatexHtml {
        table: &'static str,
        id: i64,
        source: lol_html::errors::RewritingError,
    },
}

impl From<deadpool_postgres::PoolError> for MigrateError {
    fn from(e: deadpool_postgres::PoolError) -> Self {
        MigrateError::Pool(e)
    }
}

impl From<tokio_postgres::Error> for MigrateError {
    fn from(e: tokio_postgres::Error) -> Self {
        MigrateError::Query(e)
    }
}

impl std::fmt::Display for MigrateError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        // 主动展开 source 链：tokio_postgres::Error 的 Display 对 DB 侧错误只打印
        // 无信息量的 `db error`，真正的 message/SQLSTATE/约束名藏在 source() 里。
        match self {
            MigrateError::Pool(e) => {
                write!(
                    f,
                    "database pool error: {}",
                    crate::db::format_with_sources(e)
                )
            }
            MigrateError::Query(e) => {
                write!(
                    f,
                    "database query error: {}",
                    crate::db::format_with_sources(e)
                )
            }
            MigrateError::Apply { version, source } => {
                write!(
                    f,
                    "migration {} failed: {}",
                    version,
                    crate::db::format_with_sources(source)
                )
            }
            MigrateError::KatexHtml { table, id, source } => {
                write!(
                    f,
                    "migration 026 failed rewriting {table} row {id}: {source}"
                )
            }
        }
    }
}

impl std::error::Error for MigrateError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        match self {
            MigrateError::Pool(e) => Some(e),
            MigrateError::Query(e) => Some(e),
            MigrateError::Apply { source, .. } => Some(source),
            MigrateError::KatexHtml { source, .. } => Some(source),
        }
    }
}

/// 在**已获取**的连接上执行迁移主体逻辑（咨询锁 + 建表 + 应用迁移 + 解锁）。
///
/// 调用方负责自行控制连接获取策略——例如 `main.rs` 启动时用
/// [`get_conn_for_startup`](crate::db::pool::get_conn_for_startup)（长重试窗口）
/// 拿到连接后再调用本函数，以应对 DB 尚未就绪的场景。
///
/// 流程：
/// 1. 抢咨询锁（多实例启动时串行化）
/// 2. 确保 `schema_migrations` 表存在
/// 3. 查询已应用版本集合
/// 4. 按序应用未应用的迁移（每个一个事务）
/// 5. 释放咨询锁
///
/// 失败时返回错误；调用方（`main.rs`）应让进程退出，避免启动半残服务。
pub async fn run_on_conn(conn: &mut deadpool_postgres::Object) -> Result<(), MigrateError> {
    // 抢咨询锁：多实例滚动发布时只有一个进程能进入迁移循环，
    // 其余实例在此等待；锁释放后它们查版本表发现已全部应用，直接返回。
    conn.execute("SELECT pg_advisory_lock($1)", &[&ADVISORY_LOCK_KEY])
        .await?;

    // 在已持锁连接上执行迁移主体逻辑。
    // 锁释放策略：
    // - 正常返回 / 返回 Err：下面的显式 pg_advisory_unlock 释放锁。
    // - 进程被强杀（SIGKILL 等）：连接断开，Postgres 在检测到会话断开后释放
    //   session 级咨询锁。
    // - `main.rs` 在迁移失败时用 `std::process::exit(1)` 终止进程（不再 panic）：
    //   `exit(1)` 不会 unwind，但同样会关闭进程持有的所有 socket / 池连接，
    //   效果等价于会话断开——Postgres 会释放 session 级咨询锁。
    //   因此把 `.expect()` 改成 `exit(1)` 不破坏原有的锁安全保证。
    let result = run_inner(conn).await;

    // 无论成功失败都尝试显式释放锁；释放失败不应掩盖原始错误，仅记录告警。
    if let Err(unlock_err) = conn
        .execute("SELECT pg_advisory_unlock($1)", &[&ADVISORY_LOCK_KEY])
        .await
    {
        tracing::warn!("failed to release migration advisory lock: {}", unlock_err);
    }

    result
}

/// 在已持有咨询锁的连接上执行迁移主体逻辑。
async fn run_inner(conn: &mut deadpool_postgres::Object) -> Result<(), MigrateError> {
    // 确保版本表存在（独立语句，不在事务里，否则建表失败无法记录）。
    ensure_versions_table(conn).await?;

    // 查询已应用的版本集合。
    let applied = applied_versions(conn).await?;

    // 按序应用未应用的迁移。
    let mut applied_count = 0usize;
    for (version, sql) in MIGRATIONS {
        if applied.contains(*version) {
            continue;
        }
        tracing::info!("applying migration {}", version);
        apply_one(conn, version, sql).await?;
        applied_count += 1;
    }

    if applied_count == 0 {
        tracing::info!("database is up to date, 0 migrations applied");
    } else {
        tracing::info!("successfully applied {} migration(s)", applied_count);
    }
    Ok(())
}

/// 创建 `schema_migrations` 表（若不存在）。
async fn ensure_versions_table(conn: &deadpool_postgres::Object) -> Result<(), MigrateError> {
    conn.batch_execute(
        "CREATE TABLE IF NOT EXISTS schema_migrations (
            version    TEXT PRIMARY KEY,
            applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
        )",
    )
    .await?;
    Ok(())
}

/// 查询已应用的版本集合。
async fn applied_versions(
    conn: &deadpool_postgres::Object,
) -> Result<HashSet<String>, MigrateError> {
    let rows = conn
        .query("SELECT version FROM schema_migrations", &[])
        .await?;
    let mut set = HashSet::with_capacity(rows.len());
    for row in rows {
        set.insert(row.get::<_, String>(0));
    }
    Ok(set)
}

/// 在一个事务内应用单个迁移：执行 SQL + 写入版本行，失败则回滚。
async fn apply_one(
    conn: &mut deadpool_postgres::Object,
    version: &str,
    sql: &str,
) -> Result<(), MigrateError> {
    let tx = conn.transaction().await.map_err(MigrateError::Query)?;

    // batch_execute 执行整段 SQL（可含多条语句）。
    if let Err(e) = tx.batch_execute(sql).await {
        // 显式回滚以尽早释放事务（而非等 Transaction drop 的隐式回滚）；
        // 回滚本身的错误丢弃，因为已有更具信息量的 apply 错误要上报。
        let _ = tx.rollback().await;
        return Err(MigrateError::Apply {
            version: version.to_string(),
            source: e,
        });
    }

    if version == "026" {
        if let Err(error) = migrate_katex_html(&tx).await {
            let _ = tx.rollback().await;
            return Err(error);
        }
    }

    // 记录版本行。显式构造 Apply 错误（不能用 ?，否则会被 blanket From 映射成 Query）。
    if let Err(e) = tx
        .execute(
            "INSERT INTO schema_migrations (version) VALUES ($1)",
            &[&version],
        )
        .await
    {
        let _ = tx.rollback().await;
        return Err(MigrateError::Apply {
            version: version.to_string(),
            source: e,
        });
    }

    tx.commit().await.map_err(|e| MigrateError::Apply {
        version: version.to_string(),
        source: e,
    })?;
    Ok(())
}

/// 一次性格式迁移：不重渲染 Markdown、不触碰原文/目录/时间戳/审核状态。
/// SQL 已锁住两张表，批次仅限制内存；所有批次与版本行仍属于同一事务。
async fn migrate_katex_html(tx: &deadpool_postgres::Transaction<'_>) -> Result<(), MigrateError> {
    let db_error = |source| MigrateError::Apply {
        version: "026".to_string(),
        source,
    };
    // 只暂停时间戳触发器；外键等完整性约束照常执行。失败时 DDL 一同回滚。
    tx.batch_execute("ALTER TABLE comments DISABLE TRIGGER trg_comments_updated_at")
        .await
        .map_err(db_error)?;

    for table in ["posts", "comments"] {
        // table 仅来自上面的常量列表，绝不接受外部输入。
        let select = tx
            .prepare(&format!(
                "SELECT id::bigint, content_html FROM {table} \
                 WHERE ($1::bigint IS NULL OR id > $1) AND content_html LIKE '%katex%' \
                 ORDER BY id LIMIT 100"
            ))
            .await
            .map_err(db_error)?;
        let update = tx
            .prepare(&format!(
                "UPDATE {table} SET content_html = $1 WHERE id = $2::bigint"
            ))
            .await
            .map_err(db_error)?;
        let mut last_id: Option<i64> = None;
        loop {
            let rows = tx.query(&select, &[&last_id]).await.map_err(db_error)?;
            if rows.is_empty() {
                break;
            }
            for row in rows {
                let id: i64 = row.get(0);
                let html: &str = row.get(1);
                if let Some(rewritten) = migrate_katex_classes(html)
                    .map_err(|source| MigrateError::KatexHtml { table, id, source })?
                {
                    tx.execute(&update, &[&rewritten, &id])
                        .await
                        .map_err(db_error)?;
                }
                last_id = Some(id);
            }
        }
    }

    tx.batch_execute("ALTER TABLE comments ENABLE TRIGGER trg_comments_updated_at")
        .await
        .map_err(db_error)?;
    Ok(())
}

/// KaTeX 0.18 的精确 class 改名表（上游 commit 6f5c44f，katex-rs 0.3）。
/// `hbox` 的未使用 CSS 被上游删除，没有对应的旧渲染节点需要迁移。
fn is_old_katex_class(class: &str) -> bool {
    matches!(
        class,
        "accent"
            | "base"
            | "fix"
            | "hdashline"
            | "hline"
            | "inner"
            | "newline"
            | "overlay"
            | "overline"
            | "root"
            | "rule"
            | "sizing"
            | "smash"
            | "sout"
            | "stretchy"
            | "strut"
            | "tag"
            | "thinbox"
            | "underline"
            | "vbox"
    )
}

/// 仅修改公式视觉层 span 的完整 class token，保留文字、其它属性和注释。
/// lol_html 不重新序列化未改动的节点；无旧 class 返回 None，避免无效 UPDATE。
fn migrate_katex_classes(html: &str) -> Result<Option<String>, lol_html::errors::RewritingError> {
    let mut changed = false;
    let settings = lol_html::RewriteStrSettings::new().append_element_content_handler(
        lol_html::element!("span.katex > span.katex-html span[class]", |el| {
            let Some(classes) = el.get_attribute("class") else {
                return Ok(());
            };
            if !classes.split_ascii_whitespace().any(is_old_katex_class) {
                return Ok(());
            }
            let mut rewritten = String::with_capacity(classes.len() + 6);
            for class in classes.split_ascii_whitespace() {
                if !rewritten.is_empty() {
                    rewritten.push(' ');
                }
                if is_old_katex_class(class) {
                    rewritten.push_str("katex-");
                }
                rewritten.push_str(class);
            }
            el.set_attribute("class", &rewritten)?;
            changed = true;
            Ok(())
        }),
    );
    let rewritten = lol_html::rewrite_str(html, settings)?;
    Ok(changed.then_some(rewritten))
}

#[cfg(all(test, feature = "server"))]
mod tests {
    use super::*;

    #[test]
    fn katex_migration_preserves_non_formula_html() {
        let before = r#"<!-- base class="overline" --><p class="base underline" title="overline">base &amp; underline</p><pre><code>&lt;span class="base"&gt;</code></pre>"#;
        let formula = r#"<span class="katex"><span class="katex-html" aria-hidden="true"><span class="base"><span class="mord overline" title="base &amp; overline">overline &amp; base</span></span></span></span>"#;
        let after = r#"<span class="overline">unchanged</span><span class="katex"><span class="base">not a rendered visual layer</span></span>"#;
        let input = format!("{before}{formula}{after}");
        let output = migrate_katex_classes(&input).unwrap().unwrap();
        assert!(output.starts_with(before));
        assert!(output.ends_with(after));
        assert!(output.contains(r#"class="katex-base""#));
        assert!(output.contains(
            r#"class="mord katex-overline" title="base &amp; overline">overline &amp; base</span>"#
        ));
        assert_eq!(migrate_katex_classes(&output).unwrap(), None);
    }

    #[test]
    fn katex_migration_renames_only_complete_layout_tokens() {
        let input = r#"<span class="katex"><span class="katex-html"><span class="accent base fix hdashline hline inner newline overlay overline root rule sizing smash sout stretchy strut tag thinbox underline vbox mord overline-line underline-line accent-body reset-size6 size3 mtight my-base katex-base">x</span></span></span>"#;
        let expected = r#"<span class="katex"><span class="katex-html"><span class="katex-accent katex-base katex-fix katex-hdashline katex-hline katex-inner katex-newline katex-overlay katex-overline katex-root katex-rule katex-sizing katex-smash katex-sout katex-stretchy katex-strut katex-tag katex-thinbox katex-underline katex-vbox mord overline-line underline-line accent-body reset-size6 size3 mtight my-base katex-base">x</span></span></span>"#;
        assert_eq!(
            migrate_katex_classes(input).unwrap().as_deref(),
            Some(expected)
        );
    }

    #[test]
    fn migrations_are_sorted_ascending() {
        let mut sorted = MIGRATIONS.iter().map(|(v, _)| *v).collect::<Vec<_>>();
        sorted.sort_unstable();
        let original: Vec<&str> = MIGRATIONS.iter().map(|(v, _)| *v).collect();
        assert_eq!(
            original, sorted,
            "MIGRATIONS must be in ascending version order"
        );
    }

    #[test]
    fn migrations_have_unique_versions() {
        let mut versions: Vec<&str> = MIGRATIONS.iter().map(|(v, _)| *v).collect();
        let total = versions.len();
        versions.sort_unstable();
        versions.dedup();
        assert_eq!(
            versions.len(),
            total,
            "MIGRATIONS has duplicate version strings"
        );
    }

    /// 防止"新建了 .sql 但忘记在 MIGRATIONS 加行"的脚枪。
    /// 扫描 migrations/ 目录，断言每个 .sql 文件都在 MIGRATIONS 里有对应版本行。
    /// 仅在 server feature + test 下运行（WASM 无文件系统）。
    #[test]
    fn migrations_match_files_on_disk() {
        use std::collections::HashSet;
        use std::fs;

        // CARGO_MANIFEST_DIR 指向 crate 根目录（yggdrasil/）。
        let manifest_dir = env!("CARGO_MANIFEST_DIR");
        let migrations_dir = std::path::Path::new(manifest_dir).join("migrations");

        let mut files_on_disk: HashSet<String> = HashSet::new();
        for entry in fs::read_dir(&migrations_dir)
            .unwrap_or_else(|e| panic!("failed to read {}: {}", migrations_dir.display(), e))
        {
            let entry = entry.unwrap();
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) == Some("sql") {
                let filename = path
                    .file_name()
                    .and_then(|n| n.to_str())
                    .unwrap_or_else(|| panic!("non-utf8 filename: {}", path.display()));
                // 文件名形如 "001_init.sql"，取前 3 位数字作为 version。
                let version = filename
                    .split('_')
                    .next()
                    .unwrap_or_else(|| panic!("filename has no '_' separator: {}", filename));
                files_on_disk.insert(version.to_string());
            }
        }

        let versions_in_array: HashSet<String> =
            MIGRATIONS.iter().map(|(v, _)| v.to_string()).collect();

        // 磁盘上有但数组里没有 → 忘记加行（会静默不执行该迁移）。
        let missing_in_array: Vec<&String> = files_on_disk.difference(&versions_in_array).collect();
        assert!(
            missing_in_array.is_empty(),
            "migrations/*.sql files not registered in MIGRATIONS: {:?}. \
             Add a row for each in src/db/migrate.rs.",
            missing_in_array
        );

        // 数组里有但磁盘上没有 → include_str! 本就会编译失败，这里只是双保险。
        let missing_on_disk: Vec<&String> = versions_in_array.difference(&files_on_disk).collect();
        assert!(
            missing_on_disk.is_empty(),
            "MIGRATIONS rows without a corresponding .sql file: {:?}",
            missing_on_disk
        );
    }
}