zkxw/backend/scripts/migrate_minio_urls.py

#!/usr/bin/env python3
"""批量替换数据库历史 URL 前缀（MinIO 域名迁移脚本）。

用法（在 backend 目录执行）：
1) 先预览：
   python scripts/migrate_minio_urls.py --dry-run

2) 正式执行：
   python scripts/migrate_minio_urls.py

3) 自定义前缀：
   python scripts/migrate_minio_urls.py \
       --old-prefix "https://www.chengyantek.cn:20004" \
       --new-prefix "https://www.chengyantek.cn"
"""
from __future__ import annotations

import argparse
import sys
from datetime import datetime
from pathlib import Path

from sqlalchemy import text

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from app.core.database import engine  # noqa: E402


DEFAULT_OLD_PREFIX = "https://www.chengyantek.cn:20004"
DEFAULT_NEW_PREFIX = "https://www.chengyantek.cn"

TEXT_TYPES = (
    "char",
    "varchar",
    "tinytext",
    "text",
    "mediumtext",
    "longtext",
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="批量替换数据库中的 URL 前缀")
    parser.add_argument("--old-prefix", default=DEFAULT_OLD_PREFIX, help="旧 URL 前缀")
    parser.add_argument("--new-prefix", default=DEFAULT_NEW_PREFIX, help="新 URL 前缀")
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="仅预览受影响记录数，不执行更新",
    )
    return parser.parse_args()


def quote_ident(name: str) -> str:
    """安全引用 MySQL 标识符。"""
    return f"`{name.replace('`', '``')}`"


def get_text_columns(conn) -> list[tuple[str, str]]:
    sql = text(
        """
        SELECT table_name, column_name
        FROM information_schema.columns
        WHERE table_schema = DATABASE()
          AND data_type IN :types
        ORDER BY table_name, ordinal_position
        """
    )
    rows = conn.execute(sql, {"types": TEXT_TYPES}).fetchall()
    return [(r[0], r[1]) for r in rows]


def count_affected(conn, table_name: str, column_name: str, old_prefix: str) -> int:
    sql = text(
        f"""
        SELECT COUNT(*) AS cnt
        FROM {quote_ident(table_name)}
        WHERE {quote_ident(column_name)} LIKE :pattern
        """
    )
    return int(conn.execute(sql, {"pattern": f"%{old_prefix}%"}).scalar() or 0)


def update_column(
    conn,
    table_name: str,
    column_name: str,
    old_prefix: str,
    new_prefix: str,
) -> int:
    sql = text(
        f"""
        UPDATE {quote_ident(table_name)}
        SET {quote_ident(column_name)} = REPLACE({quote_ident(column_name)}, :old_prefix, :new_prefix)
        WHERE {quote_ident(column_name)} LIKE :pattern
        """
    )
    result = conn.execute(
        sql,
        {
            "old_prefix": old_prefix,
            "new_prefix": new_prefix,
            "pattern": f"%{old_prefix}%",
        },
    )
    return int(result.rowcount or 0)


def make_backup_table_name(table_name: str, suffix: str) -> str:
    """生成备份表名，控制在 MySQL 64 字符限制内。"""
    max_table_len = 64
    reserve = len("_bak_") + len(suffix)
    max_base_len = max_table_len - reserve
    base = table_name[:max_base_len]
    return f"{base}_bak_{suffix}"


def backup_table(conn, table_name: str, backup_table_name: str) -> None:
    """备份单张表（结构+数据）。"""
    conn.execute(
        text(
            f"CREATE TABLE {quote_ident(backup_table_name)} LIKE {quote_ident(table_name)}"
        )
    )
    conn.execute(
        text(
            f"INSERT INTO {quote_ident(backup_table_name)} SELECT * FROM {quote_ident(table_name)}"
        )
    )


def main() -> None:
    args = parse_args()
    old_prefix = args.old_prefix.strip().rstrip("/")
    new_prefix = args.new_prefix.strip().rstrip("/")
    dry_run = args.dry_run

    if not old_prefix or not new_prefix:
        raise ValueError("old-prefix/new-prefix 不能为空")
    if old_prefix == new_prefix:
        raise ValueError("old-prefix 与 new-prefix 不能相同")

    print(f"数据库 URL 迁移开始，模式：{'预览(dry-run)' if dry_run else '正式执行'}")
    print(f"旧前缀: {old_prefix}")
    print(f"新前缀: {new_prefix}")

    total_rows = 0
    hit_columns = 0
    affected_items: list[tuple[str, str, int]] = []

    with engine.begin() as conn:
        columns = get_text_columns(conn)
        for table_name, column_name in columns:
            affected = count_affected(conn, table_name, column_name, old_prefix)
            if affected <= 0:
                continue
            affected_items.append((table_name, column_name, affected))

        if not affected_items:
            print("未发现需要替换的数据，已结束。")
            return

        if not dry_run:
            backup_suffix = datetime.now().strftime("%Y%m%d%H%M%S")
            backup_tables: dict[str, str] = {}
            for table_name, _, _ in affected_items:
                if table_name in backup_tables:
                    continue
                backup_name = make_backup_table_name(table_name, backup_suffix)
                backup_table(conn, table_name, backup_name)
                backup_tables[table_name] = backup_name
                print(f"[BACKUP] {table_name} -> {backup_name}")

            print(f"备份完成：共备份 {len(backup_tables)} 张表。")

        for table_name, column_name, affected in affected_items:
            hit_columns += 1
            if dry_run:
                print(f"[DRY-RUN] {table_name}.{column_name} -> {affected} 行")
                total_rows += affected
                continue

            updated = update_column(conn, table_name, column_name, old_prefix, new_prefix)
            print(f"[UPDATED] {table_name}.{column_name} -> {updated} 行")
            total_rows += updated

    if dry_run:
        print(f"预览完成：共 {hit_columns} 个字段命中，预计影响 {total_rows} 行。")
    else:
        print(f"执行完成：共更新 {hit_columns} 个字段，影响 {total_rows} 行。")


if __name__ == "__main__":
    main()