tradingSystem/scripts/init_db.py

# -*- coding: utf-8 -*-
"""
建表脚本: 把 ddl_pms_v1.sql 应用到 153 代理侧
===============================================
运行:
  docker compose run --rm pms-web python scripts/init_db.py          # 演练, 只打印不执行
  docker compose run --rm pms-web python scripts/init_db.py --yes    # 实际建表
  docker compose run --rm pms-web python scripts/init_db.py --yes --only pms_command,pms_plan

说明:
  * 全部语句为 `CREATE TABLE IF NOT EXISTS`, 重复执行安全 (已存在的表不会被改动)。
  * 直接用引擎执行, 绕过 db.session 的单表守卫 —— DDL 本身就只涉及一张表,
    但语句里的中文注释可能撞上守卫的关键字启发式, 没必要让它误伤。
  * ShardingSphere-Proxy 对 DDL 的支持视配置而定。若某条语句被代理拒绝,
    脚本会把完整语句打出来, 直接拿去物理库 (my_quant_db) 执行即可, 其余语句不受影响。
"""
import argparse
import os
import re
import sys

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

DDL_FILE = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
                        "ddl_pms_v1.sql")


def split_sql(text: str) -> list:
    """按分号切语句, 但**跳过字符串字面量里的分号**。

    DDL 的列注释里有 `COMMENT '参数命令: EFFECTIVE/SUPERSEDED; 任务命令: PENDING/...'`,
    直接 split(";") 会把一条 CREATE TABLE 劈成三段残句 —— 必须带引号状态机。
    """
    out, buf, in_str, i = [], [], False, 0
    while i < len(text):
        ch = text[i]
        if in_str:
            if ch == "\\" and i + 1 < len(text):        # 反斜杠转义
                buf.append(text[i:i + 2])
                i += 2
                continue
            if ch == "'":
                if i + 1 < len(text) and text[i + 1] == "'":   # '' 转义
                    buf.append("''")
                    i += 2
                    continue
                in_str = False
            buf.append(ch)
        elif ch == "'":
            in_str = True
            buf.append(ch)
        elif ch == ";":
            out.append("".join(buf))
            buf = []
        else:
            buf.append(ch)
        i += 1
    out.append("".join(buf))
    return out


def parse_statements(sql_text: str) -> list:
    """去掉整行 `--` 注释后切语句。返回 [(种类, 表名, 语句)]。

    种类: `table` = CREATE TABLE / `seed` = 幂等的初始化 INSERT / `?` = 认不出的残句。

    认 seed 是因为 `pms_ws_state` 是**单行表**, 那一行 (id=1) 本身就属于表结构的一部分,
    跟建表放一起最省事。但只放行 `INSERT ... ON DUPLICATE KEY UPDATE` / `INSERT IGNORE`
    这种可重复执行的写法 —— 建表脚本必须幂等, 一条会重复插数的 INSERT 混进来, 比认不出
    它更危险。

    `?` 一律中止且不执行任何语句: 它通常意味着 split_sql 把某条 CREATE TABLE 劈断了
    (列注释里的分号是老坑), 这时候硬执行残句只会建出半张表。
    """
    body = "\n".join(ln for ln in sql_text.splitlines() if not ln.strip().startswith("--"))
    out = []
    for stmt in split_sql(body):
        s = stmt.strip()
        if not s:
            continue
        m = re.search(r"CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?`?([A-Za-z0-9_]+)`?", s, re.I)
        if m:
            out.append(("table", m.group(1), s))
            continue
        m = re.match(r"INSERT\s+(?:IGNORE\s+)?INTO\s+`?([A-Za-z0-9_]+)`?", s, re.I)
        if m and re.search(r"ON\s+DUPLICATE\s+KEY\s+UPDATE|INSERT\s+IGNORE", s, re.I):
            out.append(("seed", m.group(1), s))
            continue
        out.append(("?", "?", s))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--yes", action="store_true", help="确认执行 (缺省只演练)")
    ap.add_argument("--only", default="", help="只处理指定表, 逗号分隔")
    args = ap.parse_args()

    if not os.path.exists(DDL_FILE):
        print(f"FAIL  找不到 DDL 文件: {DDL_FILE}")
        sys.exit(1)
    with open(DDL_FILE, encoding="utf-8") as f:
        stmts = parse_statements(f.read())
    only = {x.strip() for x in args.only.split(",") if x.strip()}
    if only:
        stmts = [(k, t, s) for k, t, s in stmts if t in only]
    if not stmts:
        print("FAIL  没有可执行的语句 (检查 --only 是否写对)")
        sys.exit(1)
    broken = [s for k, _, s in stmts if k == "?"]
    if broken:
        print(f"FAIL  解析出 {len(broken)} 条既不是 CREATE TABLE 也不是幂等 INSERT 的残句, "
              f"说明 DDL 切分有误, 已中止 (不执行任何语句)。首条残句:\n{broken[0][:200]}")
        sys.exit(1)

    tables = [t for k, t, _ in stmts if k == "table"]
    seeds = [t for k, t, _ in stmts if k == "seed"]
    print(f"DDL 文件: {DDL_FILE}")
    print(f"待建表 {len(tables)} 张: {', '.join(tables)}")
    if seeds:
        print(f"待写初始行 {len(seeds)} 条 (幂等): {', '.join(seeds)}")
    if not args.yes:
        print("\n[演练模式] 未执行任何语句。确认无误后加 --yes 重跑:")
        print("  docker compose run --rm pms-web python scripts/init_db.py --yes")
        return

    # 连库依赖放到演练之后再导入: 演练模式只做语句切分校验, 无依赖环境也能跑
    from sqlalchemy import text
    from app.db.session import get_engine

    eng = get_engine("proxy")
    okc, failed = 0, []
    print()
    for kind, tbl, stmt in stmts:
        label = tbl if kind == "table" else f"{tbl} (初始行)"
        try:
            with eng.begin() as c:
                c.execute(text(stmt))
            print(f"  OK    {label}")
            okc += 1
        except Exception as e:
            print(f"  FAIL  {label}: {type(e).__name__}: {e}")
            failed.append((tbl, stmt, str(e)))

    print("\n[验证] 逐表 COUNT(*)")
    missing = []
    for tbl in dict.fromkeys(t for _, t, _ in stmts):      # 去重且保序
        try:
            with eng.connect() as c:
                n = c.execute(text(f"SELECT COUNT(*) FROM {tbl}")).scalar()
            print(f"  OK    {tbl:<20} {n} 行")
        except Exception as e:
            print(f"  FAIL  {tbl:<20} {type(e).__name__}: {e}")
            missing.append(tbl)

    print("\n" + "-" * 62)
    if failed:
        print(f"以下 {len(failed)} 张表建失败, 完整语句如下 —— 可直接拿到物理库执行:")
        for tbl, stmt, err in failed:
            print(f"\n### {tbl}  ({err.splitlines()[0]})\n{stmt};")
    if missing:
        print(f"\nFAILED: {len(missing)} 张表仍不可用: {', '.join(missing)}")
        sys.exit(1)
    print(f"ALL OK: {len(tables)} 张表就绪 ({okc} 条语句执行成功), "
          f"接着跑 python scripts/check_db.py 做完整自检")


if __name__ == "__main__":
    main()