tradingSystem/scripts/init_db.py

146 lines
5.5 KiB
Python
Raw Permalink Normal View History

# -*- 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:
"""去掉整行 `--` 注释后切语句。返回 [(表名, 语句)]; 非 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)
out.append((m.group(1) if m else "?", 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 = [(t, s) for t, s in stmts if t in only]
if not stmts:
print("FAIL 没有可执行的语句 (检查 --only 是否写对)")
sys.exit(1)
broken = [s for t, s in stmts if t == "?"]
if broken:
print(f"FAIL 解析出 {len(broken)} 条非 CREATE TABLE 残句, 说明 DDL 切分有误, "
f"已中止 (不执行任何语句)。首条残句:\n{broken[0][:200]}")
sys.exit(1)
print(f"DDL 文件: {DDL_FILE}")
print(f"待建表 {len(stmts)} 张: {', '.join(t for t, _ in stmts)}")
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 tbl, stmt in stmts:
try:
with eng.begin() as c:
c.execute(text(stmt))
print(f" OK {tbl}")
okc += 1
except Exception as e:
print(f" FAIL {tbl}: {type(e).__name__}: {e}")
failed.append((tbl, stmt, str(e)))
print("\n[验证] 逐表 COUNT(*)")
missing = []
for tbl, _ 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: {okc} 张表就绪, 接着跑 python scripts/check_db.py 做完整自检")
if __name__ == "__main__":
main()