Files
cma-management/backend/scripts/treasury_migrate.py
T
Hermes CI Fix f1fa781e08
ci/woodpecker/push/main Pipeline was canceled
feat(cma): 司库阶段1 DDL 并入特权迁移步 + cma_api 授权
- db_migrate.py: create_all 后调用 run_structural() → 复用 treasury_migrate(幂等、特权连接)
- treasury_migrate.py: 换特权连接(CMA_MIGRATE_DB_USER 或本机 root socket,禁 cma_api);up() 增 grant_app_account()(cma_api 对 6 张 treasury 表 DML,双 host)
- 效果:deploy 链由特权账号一次性完成 建表+补列(cash_plans 4列/1索引)+种子(8规则/6限额)+授权

Ref: 0-P0-project-20260915-cma-treasury-stage1
2026-09-16 14:39:54 +08:00

439 lines
20 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""CMA 司库(资金)阶段一 —— 迁移脚本(阶段 1 落地版)
落地路径:/root/cma-management/backend/scripts/treasury_migrate.py
评审稿 :wecom-fullstack/pending-tasks/_reports/treasury-stage1-migration-draft-20260915.py(md5 ceaa3b3bf563b8d664ddd7224f01a5bf,归档不回改)
G1 修正项 :P0-1 up() 语句过滤器 / P0-2 计数口径 / P1-3 down() 部分态 / P1-4 down() dump 凭据 /
P1-5 种子展开 / P2-6 data_status fail-safe(DDL 侧)/ P2-7 import_logs 指纹行 / P2-8 金额纪律
用法:
python3 treasury_migrate.py --parse-only # 零 DB:只解析 DDL,打印命中语句与期望计数(P0-1 复算证据)
python3 treasury_migrate.py --check # 只读预检:缺表/缺列/缺索引/种子行数(零写入,默认动作)
python3 treasury_migrate.py --apply # 建表 + 加列 + 索引 + 种子(幂等)
python3 treasury_migrate.py --down # 回滚(先 dump 后删;有数据时要求 --yes)
幂等契约(P0-2 计数口径 = 新增对象数,不是执行语句数):
首跑 --apply 期望:created_new=6 / added_cols=4 / added_indexes=1 / rules=8 行 / limits=6 行
二跑 --apply 期望:created_new=0 / added_cols=0 / added_indexes=0 / 种子行数不变
指纹行:import_logs(batch=treasury-ddl-<sha1[:10]>),按 (import_type,batch) 先查后写,
不依赖表上是否有唯一键(P2-7:entity_id=1 为约定——迁移指纹属实体无关元数据)
边界:只动 6 张 treasury_* 新表 + cash_plans 加列/加索引;不碰 cash.py / 既有端点契约 / 其他模块表结构。
"""
import argparse
import hashlib
import json
import os
import subprocess
import sys
import tempfile
from datetime import datetime
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from sqlalchemy import create_engine, text # noqa: E402
HERE = os.path.dirname(os.path.abspath(__file__))
BACKUP_DIR = os.path.join(HERE, "backups")
DDL_FILE = os.path.join(HERE, "treasury_ddl.sql")
NEW_TABLES = [
"treasury_accounts",
"treasury_balances",
"treasury_control_rules",
"treasury_payments",
"treasury_control_hits",
"treasury_limits",
]
CASH_PLAN_NEW_COLS = ["account_id", "tier", "pool_id", "is_locked"]
CASH_PLAN_NEW_INDEX = "ix_cashplans_entity_acct"
EXPECTED_RULE_ROWS = 8 # e1 4 + e2 4
EXPECTED_LIMIT_ROWS = 6 # e1 3 + e2 3
CASH_PLAN_ALTER = """ALTER TABLE cash_plans
ADD COLUMN account_id INT NULL COMMENT '资金来源/去向账户 treasury_accounts.id(软链,无物理外键)',
ADD COLUMN tier VARCHAR(20) NULL COMMENT 'hq/platform/grassroots(三位一体视角;可由 org_nodes.level 派生)',
ADD COLUMN pool_id INT NULL COMMENT '归属资金池(阶段三启用,本阶段恒 NULL)',
ADD COLUMN is_locked TINYINT DEFAULT 0 COMMENT '是否已被预算/内控锁定'"""
# ─────────────────────────────────────────────────────────────
# 特权迁移连接(2026-09-16 并入):与应用运行账号 cma_api 剥离,走有 DDL 权限的通道
# ─────────────────────────────────────────────────────────────
DEFAULT_SOCKET = "/var/run/mysqld/mysqld.sock"
FORBIDDEN_RUN_ACCOUNTS = {"cma_api"}
GRANT_APP_ACCOUNT = "cma_api"
def get_engine():
"""特权引擎:CMA_MIGRATE_DB_USER/PASS(TCP)优先,否则本机 root unix socket(无口令)。
绝不使用应用运行账号(无 DDL)——与 db_migrate.py 同口径。"""
from app import database as db
db_name = (os.getenv("CMA_MIGRATE_DB_NAME") or getattr(db, "DB_NAME", None) or "cma")
user = os.getenv("CMA_MIGRATE_DB_USER", "").strip()
if user:
if user in FORBIDDEN_RUN_ACCOUNTS:
raise SystemExit(
"[FAIL] CMA_MIGRATE_DB_USER=cma_api 是应用运行账号(无 DDL 权限),已拒绝执行")
pw = os.getenv("CMA_MIGRATE_DB_PASS", "")
host = os.getenv("CMA_MIGRATE_DB_HOST", "127.0.0.1")
port = int(os.getenv("CMA_MIGRATE_DB_PORT", "3306"))
return create_engine(
f"mysql+pymysql://{user}:{pw}@{host}:{port}/{db_name}?charset=utf8mb4",
pool_pre_ping=True)
sock = os.getenv("CMA_MIGRATE_DB_SOCKET", DEFAULT_SOCKET)
if not os.path.exists(sock):
raise SystemExit(
f"[FAIL] 未设 CMA_MIGRATE_DB_USER 且 unix socket 不存在:{sock}")
return create_engine(
f"mysql+pymysql://root@localhost/{db_name}?unix_socket={sock}&charset=utf8mb4",
pool_pre_ping=True)
def grant_app_account(conn):
"""授权应用运行账号 cma_api 对 6 张 treasury 新表的 DML(幂等 GRANT,双 host)。
必须由特权连接执行;否则应用端点读不到/写不进新表(1142/1044)。"""
from app import database as db
dbn = (os.getenv("CMA_MIGRATE_DB_NAME") or getattr(db, "DB_NAME", None) or "cma")
granted = []
for host in ("localhost", "127.0.0.1"):
for t in NEW_TABLES:
conn.execute(text(
f"GRANT SELECT, INSERT, UPDATE, DELETE ON `{dbn}`.`{t}` "
f"TO '{GRANT_APP_ACCOUNT}'@{host}"))
granted.append(f"{t}@{host}")
return granted
# ─────────────────────────────────────────────────────────────
# DDL 解析(P0-1:先剔除注释行再判定语句类型,不再用 startswith 裸判)
# ─────────────────────────────────────────────────────────────
def _strip_comment_lines(sql_text: str) -> str:
"""剔除整行注释(-- 开头)+ 行尾注释,避免注释行带动导致语句被跳过/误判。"""
out = []
for raw in sql_text.splitlines():
line = raw
idx = line.find("--")
if idx >= 0:
line = line[:idx]
if not line.strip():
continue
out.append(line)
return "\n".join(out)
def _split_statements(sql_text: str) -> list:
return [s.strip() for s in _strip_comment_lines(sql_text).split(";") if s.strip()]
def _stmt_table_name(stmt: str, keyword: str):
"""取 `KEYWORD TABLE [IF NOT EXISTS] <name>` 中的表名。"""
toks = stmt.replace("(", " ").replace("\n", " ").split()
up = [t.upper() for t in toks]
try:
i = up.index(keyword)
except ValueError:
return None
j = i + 1
while j < len(up) and up[j] in ("IF", "NOT", "EXISTS"):
j += 1
return toks[j].strip("`") if j < len(toks) else None
def parse_ddl(path=DDL_FILE):
"""解析 DDL → {creates: [(table, stmt)], seeds: [stmt], alter: str|None, index: str|None}"""
sql = open(path, encoding="utf-8").read()
creates, seeds, alter, index_stmt = [], [], None, None
for stmt in _split_statements(sql):
head = stmt.upper()
if head.startswith("CREATE TABLE"):
creates.append((_stmt_table_name(stmt, "TABLE"), stmt))
elif head.startswith("INSERT INTO"):
seeds.append(stmt)
elif head.startswith("ALTER TABLE CASH_PLANS ADD KEY"):
index_stmt = stmt
elif head.startswith("ALTER TABLE CASH_PLANS"):
alter = stmt
return {"creates": creates, "seeds": seeds, "alter": alter, "index": index_stmt,
"ddl_text": sql}
def assert_ddl_parseable(p):
"""P0-1 断言:6 条 CREATE 全命中且表名与 NEW_TABLES 一致;种子 2 条;ALTER/B1 各 1 条。"""
names = [t for t, _ in p["creates"]]
errs = []
if len(p["creates"]) != len(NEW_TABLES):
errs.append(f"CREATE TABLE 命中数 {len(p['creates'])} != {len(NEW_TABLES)}({names})")
if set(names) != set(NEW_TABLES):
errs.append(f"CREATE TABLE 表名集合不一致: {sorted(names)} vs {sorted(NEW_TABLES)}")
if len(p["seeds"]) != 2:
errs.append(f"种子 INSERT 语句数 {len(p['seeds'])} != 2")
if not p["alter"]:
errs.append("DDL §B 未解析到 ALTER TABLE cash_plans 加列语句")
if not p["index"]:
errs.append("DDL §B1 未解析到 ALTER TABLE cash_plans ADD KEY 语句")
if p["alter"] and p["alter"].strip() not in CASH_PLAN_ALTER.strip():
errs.append("DDL §B ALTER 文本与脚本 CASH_PLAN_ALTER 常量漂移(指纹口径会不一致)")
return errs
def _fingerprint(texts):
canon = json.dumps(sorted(texts), ensure_ascii=False, separators=(",", ":"))
return hashlib.sha1(canon.encode("utf-8")).hexdigest()[:10]
# ─────────────────────────────────────────────────────────────
# information_schema 内省(只读)
# ─────────────────────────────────────────────────────────────
def _introspect(conn):
rows = conn.execute(text(
"SELECT TABLE_NAME FROM information_schema.TABLES "
"WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME LIKE 'treasury%'"
)).fetchall()
tables = {r[0] for r in rows}
cols = conn.execute(text(
"SELECT COLUMN_NAME FROM information_schema.COLUMNS "
"WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='cash_plans'"
)).fetchall()
colset = {r[0] for r in cols}
idx = conn.execute(text(
"SELECT DISTINCT INDEX_NAME FROM information_schema.STATISTICS "
"WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='cash_plans'"
)).fetchall()
return tables, colset, {r[0] for r in idx}
def _seed_counts(conn):
counts = {}
for tbl, key in (("treasury_control_rules", "rules_rows"), ("treasury_limits", "limits_rows")):
try:
counts[key] = conn.execute(text(f"SELECT COUNT(*) FROM {tbl}")).scalar()
except Exception:
counts[key] = None
return counts
def check(conn):
"""只读预检 —— 缺表 / 缺列 / 缺索引 / 存量行数 / 种子行数;零写入。"""
tables, cols, idx = _introspect(conn)
cash_rows = conn.execute(text("SELECT COUNT(*) FROM cash_plans")).scalar()
report = {
"missing_tables": [t for t in NEW_TABLES if t not in tables],
"existing_tables": sorted(tables),
"missing_cash_plan_cols": [c for c in CASH_PLAN_NEW_COLS if c not in cols],
"missing_cash_plan_index": [] if CASH_PLAN_NEW_INDEX in idx else [CASH_PLAN_NEW_INDEX],
"cash_plans_rows": cash_rows,
}
report.update(_seed_counts(conn))
print(json.dumps(report, ensure_ascii=False, indent=2))
return report
# ─────────────────────────────────────────────────────────────
# UP / DOWN
# ─────────────────────────────────────────────────────────────
def _write_fingerprint(conn, batch, tables_created):
"""P2-7:指纹行按 (import_type,batch) 先查后写(不依赖唯一键;entity_id=1 为约定)。"""
cols = {r[0] for r in conn.execute(text(
"SELECT COLUMN_NAME FROM information_schema.COLUMNS "
"WHERE TABLE_SCHEMA=DATABASE() AND TABLE_NAME='import_logs'"
)).fetchall()}
need = {"entity_id", "filename", "batch", "total_rows", "success_rows", "import_type"}
if not need <= cols:
raise RuntimeError(f"import_logs 结构不符,缺列 {sorted(need - cols)}(fail-closed,不静默跳过)")
row = conn.execute(text(
"SELECT id FROM import_logs WHERE import_type='treasury-ddl' AND batch=:b ORDER BY id LIMIT 1"
), {"b": batch}).fetchone()
if row:
conn.execute(text(
"UPDATE import_logs SET total_rows=:t, success_rows=:s WHERE id=:i"
), {"t": tables_created, "s": tables_created, "i": row[0]})
return "updated"
conn.execute(text(
"INSERT INTO import_logs (entity_id, filename, batch, total_rows, success_rows, import_type, created_by) "
"VALUES (1, 'treasury_ddl.sql', :b, :t, :s, 'treasury-ddl', 'treasury_migrate')"
), {"b": batch, "t": tables_created, "s": tables_created})
return "inserted"
def up(conn):
p = parse_ddl()
errs = assert_ddl_parseable(p)
if errs:
for e in errs:
print("❌ DDL 解析失败:", e, file=sys.stderr)
return 3
before_tables, before_cols, before_idx = _introspect(conn)
# ① 建表(6 条 CREATE,IF NOT EXISTS 幂等)
for _table, stmt in p["creates"]:
conn.execute(text(stmt))
# ② 条件加列(列级幂等)
missing_cols = [c for c in CASH_PLAN_NEW_COLS if c not in before_cols]
if missing_cols:
conn.execute(text(p["alter"]))
# ③ 条件加索引(D-7)
if CASH_PLAN_NEW_INDEX not in before_idx:
conn.execute(text(p["index"]))
# ③b 授权应用账号 cma_api 对本 6 表 DML(幂等 GRANT,须特权连接)
granted = grant_app_account(conn)
# ④ 种子(ON DUPLICATE KEY UPDATE,数据级幂等)
for stmt in p["seeds"]:
conn.execute(text(stmt))
conn.commit()
after_tables, after_cols, after_idx = _introspect(conn)
created_new = len(after_tables - before_tables)
added_cols = len(after_cols - before_cols)
added_indexes = len(after_idx - before_idx)
counts = _seed_counts(conn)
# ⑤ 自检断言:缺表即非零退出(P0-1,禁静默成功)
post = check(conn)
if post["missing_tables"]:
print(f"❌ up() 后仍有缺表 {post['missing_tables']}(非零退出)", file=sys.stderr)
return 4
# ⑥ 指纹行
batch = "treasury-ddl-" + _fingerprint([p["ddl_text"], CASH_PLAN_ALTER])
action = _write_fingerprint(conn, batch, created_new or len(NEW_TABLES))
conn.commit()
print(json.dumps({
"phase": "up",
"created_new": created_new,
"added_cols": added_cols,
"added_indexes": added_indexes,
"seeded_statements": len(p["seeds"]),
"rules_rows": counts["rules_rows"],
"limits_rows": counts["limits_rows"],
"import_logs_row": action,
"batch": batch,
"granted": granted,
"missing_tables": post["missing_tables"],
"cash_plans_rows": post["cash_plans_rows"],
}, ensure_ascii=False))
return 0
def _dump_cash_plans(out_path):
"""P1-4:写路径账号 + 临时 defaults-file 注入口令(不硬编码、不回显);失败即抛异常(fail-closed)。"""
from app import database as db
fd, cnf = tempfile.mkstemp(prefix=".cma-dump-", suffix=".cnf")
try:
with os.fdopen(fd, "w") as f:
f.write("[client]\nuser=%s\npassword=%s\nhost=%s\nport=%s\n"
% (db.DB_USER, db.DB_PASS, db.DB_HOST, db.DB_PORT))
os.chmod(cnf, 0o600)
with open(out_path, "w") as fh:
proc = subprocess.run(
["mysqldump", f"--defaults-extra-file={cnf}", "--single-transaction",
"--skip-comments", db.DB_NAME, "cash_plans"],
stdout=fh, stderr=subprocess.PIPE, text=True)
if proc.returncode != 0:
raise RuntimeError("mysqldump 失败(rc=%s): %s" % (proc.returncode, (proc.stderr or "").strip()[:300]))
finally:
try:
os.remove(cnf)
except OSError:
pass
def down(conn):
"""回滚:先 dump(fail-closed)→ 逐列删扩展列(P1-3)→ 删索引 → 删表(IF EXISTS)。"""
os.makedirs(BACKUP_DIR, exist_ok=True)
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
dump = os.path.join(BACKUP_DIR, f"cash_plans-before-rollback-{stamp}.sql")
try:
_dump_cash_plans(dump)
except Exception as exc: # fail-closed:dump 不成功就不允许删列
print(f"❌ 备份失败,回滚中止(未做任何结构变更): {exc}", file=sys.stderr)
return 5
print("backup:", dump, os.path.getsize(dump), "bytes")
tables, cols, idx = _introspect(conn)
dropped = []
for col in CASH_PLAN_NEW_COLS: # P1-3:按列存在性逐列 DROP(部分态不再 1091)
if col in cols:
conn.execute(text(f"ALTER TABLE cash_plans DROP COLUMN {col}"))
dropped.append(col)
if CASH_PLAN_NEW_INDEX in idx:
conn.execute(text(f"ALTER TABLE cash_plans DROP INDEX {CASH_PLAN_NEW_INDEX}"))
dropped.append(CASH_PLAN_NEW_INDEX)
for t in reversed(NEW_TABLES):
conn.execute(text(f"DROP TABLE IF EXISTS {t}"))
conn.commit()
print(json.dumps({"phase": "down", "dropped_cash_plan_objects": dropped,
"dropped_tables": list(reversed(NEW_TABLES)), "backup": dump},
ensure_ascii=False))
return 0
# ─────────────────────────────────────────────────────────────
# 入口
# ─────────────────────────────────────────────────────────────
def parse_only():
"""零 DB 的 DDL 解析复算(P0-1 证据):命中 6 条 CREATE + 表名清单 + 种子语句。"""
p = parse_ddl()
errs = assert_ddl_parseable(p)
out = {
"ddl_file": DDL_FILE,
"ddl_md5": hashlib.md5(open(DDL_FILE, "rb").read()).hexdigest(),
"statements_total": len(_split_statements(p["ddl_text"])),
"create_table_hits": len(p["creates"]),
"create_table_names": [t for t, _ in p["creates"]],
"seed_statements": len(p["seeds"]),
"has_cash_plan_alter": bool(p["alter"]),
"has_cash_plan_index": bool(p["index"]),
"expected_first_run": {"created_new": 6, "added_cols": 4, "added_indexes": 1,
"rules_rows": EXPECTED_RULE_ROWS, "limits_rows": EXPECTED_LIMIT_ROWS},
"expected_second_run": {"created_new": 0, "added_cols": 0, "added_indexes": 0,
"rules_rows": EXPECTED_RULE_ROWS, "limits_rows": EXPECTED_LIMIT_ROWS},
"errors": errs,
}
print(json.dumps(out, ensure_ascii=False, indent=2))
return 1 if errs else 0
def main():
ap = argparse.ArgumentParser()
g = ap.add_mutually_exclusive_group()
g.add_argument("--parse-only", action="store_true", help="零 DB:只解析 DDL(P0-1 复算证据)")
g.add_argument("--check", action="store_true", help="只读预检(默认)")
g.add_argument("--apply", action="store_true")
g.add_argument("--down", action="store_true")
ap.add_argument("--yes", action="store_true", help="回滚确认(表内有数据时必需)")
args = ap.parse_args()
if args.parse_only:
return parse_only()
engine = get_engine()
with engine.connect() as conn:
if args.apply:
return up(conn)
if args.down:
tables, _c, _i = _introspect(conn)
payload = 0
if {"treasury_payments", "treasury_balances"} <= tables:
payload = conn.execute(text(
"SELECT (SELECT COUNT(*) FROM treasury_payments) + "
"(SELECT COUNT(*) FROM treasury_balances)")).scalar()
if payload and not args.yes:
print(f"REFUSE: 已有 {payload} 行业务数据,回滚需显式 --yes", file=sys.stderr)
return 2
return down(conn)
return 0 if check(conn) is not None else 1
if __name__ == "__main__":
sys.exit(main())