tradingSystem/scripts/migrate_archived_at.py

# -*- coding: utf-8 -*-
"""
迁移: 给四张表加 archived_at 列 (软归档)
==========================================
运行:
  docker compose run --rm pms-web python scripts/migrate_archived_at.py          # 演练, 只打印不执行
  docker compose run --rm pms-web python scripts/migrate_archived_at.py --yes     # 实际加列

为什么单独一个脚本, 不放进 init_db:
  init_db.py 只认 `CREATE TABLE IF NOT EXISTS` 与幂等 INSERT, 任何 ALTER 语句会被它判成
  残句并整批中止。而 `CREATE TABLE IF NOT EXISTS` 对已存在的表一个字都不会改 —— 给旧表加列
  只能靠 ALTER。所以: 已有库的 archived_at 列由本脚本补; 建表语句 (ddl_pms_v1.sql) 也已带上
  该列, **新库直接就有**, 不需要跑这个脚本。

幂等: 已经有该列的表自动跳过, 重复跑安全。绕过 db.session 的单表守卫 (直接用引擎), 与 init_db 一致
  —— ALTER 本身只涉及一张表, 但语句里的中文注释可能撞上守卫的关键字启发式, 没必要让它误伤。
若某条 ALTER 被 ShardingSphere-Proxy 拒绝, 脚本会把完整语句打出来, 直接拿去物理库执行即可。
"""
import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

TABLES = ["pms_command", "pms_instruction", "pms_proposal", "pms_strategy"]
COL = "archived_at"
COMMENT = "软归档时间: 非空=已从在办视图移除, 行仍在库供审计; 只归档终态"


def _has_column(conn, table: str, col: str) -> bool:
    from sqlalchemy import text
    rows = conn.execute(text(f"SHOW COLUMNS FROM {table}")).fetchall()
    return any(r[0] == col for r in rows)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--yes", action="store_true", help="确认执行 (缺省只演练)")
    args = ap.parse_args()

    from sqlalchemy import text
    from app.db.session import TABLE_PREFIX, get_engine, map_tables
    eng = get_engine("proxy")
    if TABLE_PREFIX:
        print(f"[双实例] 表名前缀 {TABLE_PREFIX!r} 生效: 实际迁移的是 {TABLE_PREFIX}pms_* 这一套表")

    todo, already, unknown = [], [], []
    with eng.connect() as c:
        for t in TABLES:
            try:
                (already if _has_column(c, map_tables(t), COL) else todo).append(t)
            except Exception as e:
                print(f"  ?     {t}: 查列失败 {type(e).__name__}: {e} (表不存在就先跑 init_db)")
                unknown.append(t)

    print(f"已有 {COL} 列 ({len(already)}): {', '.join(already) or '无'}")
    print(f"待加 {COL} 列 ({len(todo)}): {', '.join(todo) or '无 —— 都已就位'}")

    if not todo:
        print("\n无需迁移: 四张表都已有该列 (或表还没建, 先 init_db)。")
        sys.exit(1 if unknown else 0)

    stmts = {t: map_tables(f"ALTER TABLE {t} ADD COLUMN {COL} DATETIME NULL DEFAULT NULL "
                           f"COMMENT '{COMMENT}'") for t in todo}

    if not args.yes:
        print("\n[演练模式] 未执行任何语句。确认无误后加 --yes 重跑。将执行:")
        for t in todo:
            print("  " + stmts[t])
        return

    failed = []
    print()
    for t in todo:
        try:
            with eng.begin() as c:
                c.execute(text(stmts[t]))
            print(f"  OK    {t} 加列成功")
        except Exception as e:
            print(f"  FAIL  {t}: {type(e).__name__}: {e}")
            failed.append((t, stmts[t], str(e)))

    print("\n[验证] 逐表确认列已在")
    miss = []
    with eng.connect() as c:
        for t in todo:
            try:
                has = _has_column(c, t, COL)
                print(f"  {'OK  ' if has else 'FAIL'}  {t}.{COL} {'存在' if has else '仍缺'}")
                if not has:
                    miss.append(t)
            except Exception as e:
                print(f"  FAIL  {t}: {e}")
                miss.append(t)

    print("\n" + "-" * 62)
    if failed:
        print(f"以下 {len(failed)} 张表加列失败, 完整语句如下 —— 可直接拿到物理库 (my_quant_db) 执行:")
        for t, s, err in failed:
            print(f"\n### {t}  ({err.splitlines()[0]})\n{s};")
    if failed or miss or unknown:
        # unknown (查列失败, 多为表不存在) 也算未就绪 (2026-08-28 修): 原来 todo 非空时
        # 它被忘掉, 四张表迁了三张也报 ALL OK —— 缺列的表要等页面报 SQL 错才暴露
        print(f"\nFAILED: {len(failed) + len(miss) + len(unknown)} 张表仍未就绪"
              + (f" (含查列失败 {len(unknown)} 张: {', '.join(unknown)})" if unknown else ""))
        sys.exit(1)
    print(f"ALL OK: {len(todo)} 张表已加 {COL} 列。页面的「移除 / 显示已完成」现在可用。")


if __name__ == "__main__":
    main()