xiaotutongxue/scripts/dev-clean-e2e.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
清理 e2e-m0.py 产生的测试数据。

E2E 每跑一次会：
  1. 新建一个租户（编码 e2e######）并给它建角色与账号；
  2. **在 demo 租户下**建三个组织（学校/院系/班级，编码带时间戳）。

第 2 点容易被忽略——这些组织挂在正常租户下，不清掉会让组织树里
堆满 "E2E 示范学校 586923" 这种名字。

用法
----
    python scripts/dev-clean-e2e.py            # 只看会删什么（默认 dry-run）
    python scripts/dev-clean-e2e.py --apply    # 真删

⚠️ 只在开发库上跑，不要对生产执行。
"""
from __future__ import annotations

import argparse
import subprocess
import sys

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")

# 子表在前、主表在后（无外键约束，顺序只为可读性与排查方便）
TENANT_CHILD_TABLES = [
    ("角色授权", "sys_role_permission"),
    ("用户角色", "sys_user_role"),
    ("登录日志", "sys_login_log"),
    ("组织", "sys_org"),
    ("角色", "sys_role"),
    ("用户", "sys_user"),
]

TENANT_FILTER = "tenant_code LIKE 'e2e%'"
SUB_FILTER = f"tenant_id IN (SELECT id FROM sys_tenant WHERE {TENANT_FILTER})"

# demo 租户下由 E2E 建的组织：编码都是 ROOT/DEPT/CLS + 6 位时间戳
DEMO_ORG_FILTER = (
    "org_code ~ '^(ROOT|DEPT|CLS|BAD)[0-9]{6}$' "
    "OR org_name LIKE 'E2E %' OR org_name LIKE '重复编码%' OR org_name LIKE '层级颠倒%'"
)


class Db:
    def __init__(self, container: str, user: str, name: str):
        self.container = container
        self.user = user
        self.name = name

    def rows(self, sql: str) -> list[str]:
        cmd = ["docker", "exec", self.container, "psql", "-U", self.user, "-d", self.name,
               "-v", "ON_ERROR_STOP=1", "-t", "-A", "-F", "|", "-c", sql]
        r = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8")
        if r.returncode != 0:
            print(r.stderr.strip(), file=sys.stderr)
            sys.exit(1)
        return [x for x in r.stdout.split("\n") if x.strip()]


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--apply", action="store_true", help="真正执行删除（默认只预览）")
    ap.add_argument("--container", default="xt-pg")
    ap.add_argument("--user", default="fill")
    ap.add_argument("--db", default="filldb")
    args = ap.parse_args()

    db = Db(args.container, args.user, args.db)

    tenants = db.rows(
        f"SELECT id || '|' || tenant_code || '|' || tenant_name "
        f"FROM sys_tenant WHERE {TENANT_FILTER} ORDER BY id")
    demo_orgs = db.rows(
        f"SELECT id || '|' || org_name FROM sys_org "
        f"WHERE deleted = 0 AND ({DEMO_ORG_FILTER}) ORDER BY id")

    if not tenants and not demo_orgs:
        print("没有 e2e 测试数据，无需清理")
        return 0

    if tenants:
        print(f"① e2e 测试租户 {len(tenants)} 个：")
        for r in tenants[:15]:
            p = r.split("|")
            print("   id=%-5s code=%-14s %s" % (p[0], p[1], p[2] if len(p) > 2 else ""))
        if len(tenants) > 15:
            print(f"   ... 另有 {len(tenants) - 15} 个")

    if demo_orgs:
        print(f"\n② 其它租户下的 e2e 测试组织 {len(demo_orgs)} 个：")
        for r in demo_orgs[:15]:
            p = r.split("|")
            print("   id=%-5s %s" % (p[0], p[1] if len(p) > 1 else ""))
        if len(demo_orgs) > 15:
            print(f"   ... 另有 {len(demo_orgs) - 15} 个")

    if not args.apply:
        print("\n这是预览。加 --apply 才会真正删除。")
        return 0

    print("\n开始清理：")
    total = 0

    for label, table in TENANT_CHILD_TABLES:
        n = len(db.rows(f"DELETE FROM {table} WHERE {SUB_FILTER} RETURNING id"))
        total += n
        print("   %-10s %-22s 删除 %d 行" % (label, table, n))

    n = len(db.rows(f"DELETE FROM sys_tenant WHERE {TENANT_FILTER} RETURNING id"))
    total += n
    print("   %-10s %-22s 删除 %d 行" % ("租户", "sys_tenant", n))

    n = len(db.rows(f"DELETE FROM sys_org WHERE {DEMO_ORG_FILTER} RETURNING id"))
    total += n
    print("   %-10s %-22s 删除 %d 行" % ("测试组织", "sys_org", n))

    print(f"\n完成，共删除 {total} 行。")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())