xiaotutongxue/scripts/dev-clean-e2e.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
开发库与本地磁盘的**测试残留**清理。

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

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

跑几十轮之后真正的问题不是"列表里有测试名字"，而是两笔越滚越大的账：

  · **租户垃圾**：`sys_tenant` 里 48 个 `e2e*`（每个还带着用户、角色、组织）；
  · **磁盘垃圾**：`backend/data/uploads` 里几百 MB 的文件 ——
    **删库不删盘**。压测与 E2E 会删掉自己的库行，但文件留在磁盘上，
    而"哪些文件还有主"只有库知道（`collect_file.std_path/raw_path`、
    `collect_export_file.file_path`）。实测：库里 16 行文件、磁盘上 281 个文件。

用法
----
    python scripts/dev-clean-e2e.py                  # 只看会删什么（默认 dry-run）
    python scripts/dev-clean-e2e.py --apply          # 真删：e2e 租户 + 测试组织 + 它们的业务行
    python scripts/dev-clean-e2e.py --apply --demo-tasks      # 连【演示·旧】任务一起删
    python scripts/dev-clean-e2e.py --apply --orphan-files    # 连"没有库行指向"的磁盘文件一起删

⚠️ 只在开发库上跑，不要对生产执行。
⚠️ `--orphan-files` 删的是**磁盘文件**，不可恢复；它靠"库里没有任何一行指向它"来判定，
   所以只要库是完整的，它就不会删到在用的文件。默认 dry-run，先看一眼再 `--apply`。
"""
from __future__ import annotations

import argparse
import os
import subprocess
import sys

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

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
UPLOADS = os.path.join(ROOT, "backend", "data", "uploads")

# 子表在前、主表在后（无外键约束，顺序只为可读性与排查方便）
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 '层级颠倒%'"
)

# 演示残留：给客户演示时反复摆过的任务，为了不污染演示态，已发布的那批被**改名挪开**
# 而不是删掉（F-14：已发布的任务删不掉）。它们属于"当时的演示数据"，不是需求数据。
#
# ⚠️ 只匹配带「旧」的那一批：`【演示】审核台` 是**当前**演示态在用的任务，
#    删了它 = 把演示态弄坏（而且它的班级就是下面那批 "E2E 24数控1班"，同理不能删）。
DEMO_TASK_FILTER = "task_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 tables_with(db: Db, column: str) -> list[str]:
    """
    所有带这一列的 `collect_*` 表。

    ⚠️ **动态枚举，不写死清单**：写死的话，以后新增一张业务表就会漏清，
    而"漏清"的表现是下一次跑 E2E 时某个统计数字莫名偏大 —— 排查方向完全错。
    """
    sql = ("SELECT c.table_name FROM information_schema.columns c "
           "JOIN information_schema.tables t "
           "  ON t.table_name = c.table_name AND t.table_schema = c.table_schema "
           "WHERE c.table_schema = 'public' AND t.table_type = 'BASE TABLE' "
           f"  AND c.table_name LIKE 'collect\\_%' AND c.column_name = '{column}' "
           "ORDER BY c.table_name")
    return [r.split("|")[0] for r in db.rows(sql)]


def db_file_paths(db: Db) -> set[str]:
    """库里所有被引用的文件路径（相对 uploads 目录）。"""
    known: set[str] = set()
    for table, cols in (("collect_file", ("raw_path", "std_path")),
                        ("collect_export_file", ("file_path",))):
        for col in cols:
            for r in db.rows(f"SELECT {col} FROM {table} WHERE {col} IS NOT NULL"):
                p = r.split("|")[0].strip().replace("\\", "/")
                if not p:
                    continue
                known.add(p)
                if os.path.isabs(p):
                    known.add(os.path.relpath(p, UPLOADS).replace("\\", "/"))
    return known


def scan_orphans(db: Db) -> tuple[list[str], int]:
    """磁盘上有、而库里没有任何一行指向的文件。"""
    known = db_file_paths(db)
    orphans, size = [], 0
    for dirpath, _dirs, files in os.walk(UPLOADS):
        for fn in files:
            full = os.path.join(dirpath, fn)
            rel = os.path.relpath(full, UPLOADS).replace("\\", "/")
            if rel not in known:
                orphans.append(rel)
                try:
                    size += os.path.getsize(full)
                except OSError:
                    pass
    return sorted(orphans), size


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--apply", action="store_true", help="真正执行删除（默认只预览）")
    ap.add_argument("--demo-tasks", action="store_true",
                    help="连【演示·旧】任务（及其名册/材料/填报单/文件）一起删")
    ap.add_argument("--orphan-files", action="store_true",
                    help="连磁盘上「没有库行指向」的文件一起删（不可恢复）")
    ap.add_argument("--only-orphan-files", 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()
    if args.only_orphan_files:
        args.orphan_files = True

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

    tenants = [] if args.only_orphan_files else db.rows(
        f"SELECT id || '|' || tenant_code || '|' || tenant_name "
        f"FROM sys_tenant WHERE {TENANT_FILTER} ORDER BY id")
    demo_orgs = [] if args.only_orphan_files else db.rows(
        f"SELECT id || '|' || org_name FROM sys_org "
        f"WHERE deleted = 0 AND ({DEMO_ORG_FILTER}) ORDER BY id")
    demo_tasks = (db.rows(f"SELECT id || '|' || task_name FROM collect_task "
                          f"WHERE deleted = 0 AND {DEMO_TASK_FILTER} ORDER BY id")
                  if (args.demo_tasks and not args.only_orphan_files) else [])
    orphans, orphan_size = scan_orphans(db) if args.orphan_files else ([], 0)

    if not any((tenants, demo_orgs, demo_tasks, orphans)):
        print("没有需要清理的测试数据")
        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 demo_tasks:
        print(f"\n③ 演示残留任务 {len(demo_tasks)} 个（--demo-tasks）：")
        for r in demo_tasks[:10]:
            p = r.split("|")
            print("   id=%-5s %s" % (p[0], p[1] if len(p) > 1 else ""))
        if len(demo_tasks) > 10:
            print(f"   ... 另有 {len(demo_tasks) - 10} 个")

    if args.orphan_files:
        print(f"\n④ 磁盘上没有库行指向的文件 {len(orphans)} 个 / {orphan_size / 1048576.0:.1f} MB"
              f"（--orphan-files）")
        by_dir: dict[str, int] = {}
        for rel in orphans:
            by_dir[rel.split("/")[0]] = by_dir.get(rel.split("/")[0], 0) + 1
        for k, v in sorted(by_dir.items(), key=lambda x: -x[1]):
            print("   %-10s %d 个" % (k, v))
        for rel in orphans[:5]:
            print("   例：%s" % rel)

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

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

    if not args.only_orphan_files:
        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))

        # e2e 租户的**业务行**：不删的话，租户没了、任务还在，
        # 于是"每个请求多一次租户解析"的旧账会一直留着（F-5）
        for table in tables_with(db, "tenant_id"):
            n = len(db.rows(f"DELETE FROM {table} WHERE tenant_id IN "
                            f"(SELECT id FROM sys_tenant WHERE {TENANT_FILTER}) RETURNING id"))
            if n:
                total += n
                print("   %-10s %-22s 删除 %d 行" % ("业务行", 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))

        if args.demo_tasks:
            # 先按任务 id 删子表（同样动态枚举"带 task_id 的表"）
            for table in tables_with(db, "task_id"):
                n = len(db.rows(f"DELETE FROM {table} WHERE task_id IN "
                                f"(SELECT id FROM collect_task WHERE {DEMO_TASK_FILTER}) RETURNING id"))
                if n:
                    total += n
                    print("   %-10s %-22s 删除 %d 行" % ("演示子表", table, n))
            n = len(db.rows(f"DELETE FROM collect_task WHERE {DEMO_TASK_FILTER} RETURNING id"))
            total += n
            print("   %-10s %-22s 删除 %d 行" % ("演示任务", "collect_task", n))

    removed = 0
    if args.orphan_files:
        # ⚠️ 删文件前**重算一次**：上面刚删过库行，此时"还有主"的文件集合变小了，
        #    但那些刚被删掉主的文件属于"任务数据"，本该跟着一起清 —— 重算正好覆盖它们
        orphans, orphan_size = scan_orphans(db)
        for rel in orphans:
            try:
                os.remove(os.path.join(UPLOADS, rel))
                removed += 1
            except OSError:
                pass
        print("   %-10s %-22s 删除 %d 个文件 / %.1f MB"
              % ("磁盘文件", "backend/data/uploads", removed, orphan_size / 1048576.0))

    print(f"\n完成，共删除 {total} 行" + (f"、{removed} 个文件。" if args.orphan_files else "。"))
    return 0


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