UHD_Media_Tool/scripts/dev-db-clean.ts

// 开发库清理:只针对 55433 开发库(docker umt-dev-postgres-1),清掉测试残留。
//
// 用法:
//   deno run -A scripts/dev-db-clean.ts            # 同 --dry-run:只打印清单,不改库
//   deno run -A scripts/dev-db-clean.ts --yes      # 真正删除
//   DATABASE_URL=postgres://umt:umt@localhost:55433/umt deno run -A scripts/dev-db-clean.ts --dry-run
//
// 清理范围:
//   1. 测试 schema:名字匹配 TEST_SCHEMA_PATTERNS(各测试 harness 建的带 pid+随机后缀 schema)。
//      public / openworkflow / 系统 schema 永不删;有其它会话持有锁的 schema(测试正在跑)跳过。
//   2. public 里名字带测试前缀(TEST_ENTITY_PREFIX:w32- / w33- / fe4- / qa-)的作品、站点、订阅,
//      连同按外键 restrict 挂在它们下面的行(cascade / set null 交给数据库)。
// 不动:admin_user(w31-admin)、pi_upstream、名字不带测试前缀的站点(fake 上游 观众/HDSky/馒头 等)与用户登记的数据。

import pg from "pg";

const DEFAULT_URL = "postgres://umt:umt@localhost:55433/umt";
const ALLOWED_HOSTS = new Set(["localhost", "127.0.0.1", "[::1]", "::1"]);
const DEV_PORT = "55433";

/** 测试 schema 前缀(见 apps/server/test 与 packages 下各 harness) */
export const TEST_SCHEMA_PATTERNS: RegExp[] = [
  /^tasks_test_/,
  /^test_/,
  /^e2e_/,
  /^umt_test_/,
  /^ao_test_/,
  /^fm_test_/,
  /^w2a_test_/,
  /^w2b_/,
  /^w7a_mig_/,
  /^w9g_search_/,
  /^server_test_/,
];
const PROTECTED_SCHEMAS = new Set(["public", "openworkflow", "information_schema"]);

/** 测试实体名前缀(不区分大小写) */
export const TEST_ENTITY_PREFIX = /^(w32|w33|fe4|qa)-/i;
const PREFIX_SQL = "^(w32|w33|fe4|qa)-";

/** 根:表 → 按哪一列匹配前缀 */
const ROOTS: { table: string; column: string }[] = [
  { table: "subscription", column: "title_query" },
  { table: "media_entity", column: "title" },
  { table: "site", column: "name" },
];
/** 即使被外键牵连也绝不删除的表 */
const NEVER_TOUCH = new Set(["admin_user", "admin_session", "api_key", "pi_upstream", "settings", "secret"]);

export function assertDevUrl(url: string): void {
  let u: URL;
  try {
    u = new URL(url);
  } catch {
    throw new Error(`DATABASE_URL 不是合法 URL`);
  }
  if (!/^postgres(ql)?:$/.test(u.protocol)) throw new Error("DATABASE_URL 必须是 postgres://");
  if (!ALLOWED_HOSTS.has(u.hostname) || u.port !== DEV_PORT)
    throw new Error(`拒绝运行:只清理 localhost:${DEV_PORT} 开发库,当前是 ${u.hostname}:${u.port || "5432"}`);
}

export const isTestSchema = (name: string) =>
  !PROTECTED_SCHEMAS.has(name) && !name.startsWith("pg_") && TEST_SCHEMA_PATTERNS.some((p) => p.test(name));

interface Fk {
  child: string;
  column: string;
  parent: string;
  parentColumn: string;
  /** r/a = restrict/no action(需要先删子行);c = cascade;n/d = set null/default */
  rule: string;
  childHasId: boolean;
}

async function loadFks(c: pg.Client): Promise<Fk[]> {
  const r = await c.query<Fk>(`
    select cl.relname as child, a.attname as column, pc.relname as parent, pa.attname as "parentColumn",
           con.confdeltype as rule,
           exists (select 1 from pg_attribute i where i.attrelid = con.conrelid and i.attname = 'id' and not i.attisdropped)
             as "childHasId"
      from pg_constraint con
      join pg_class cl on cl.oid = con.conrelid
      join pg_class pc on pc.oid = con.confrelid
      join pg_attribute a on a.attrelid = con.conrelid and a.attnum = con.conkey[1]
      join pg_attribute pa on pa.attrelid = con.confrelid and pa.attnum = con.confkey[1]
     where con.contype = 'f' and con.connamespace = 'public'::regnamespace and array_length(con.conkey, 1) = 1`);
  return r.rows;
}

type Plan = Map<string, Set<string>>;

/** 从根行出发,沿 restrict/no action 外键收集必须先删的子行(只收集,不删) */
async function collect(c: pg.Client, fks: Fk[], table: string, ids: string[], plan: Plan): Promise<void> {
  const seen = plan.get(table) ?? new Set<string>();
  const fresh = ids.filter((id) => !seen.has(id));
  if (fresh.length === 0) return;
  if (NEVER_TOUCH.has(table)) throw new Error(`会牵连受保护表 ${table},中止`);
  for (const id of fresh) seen.add(id);
  plan.set(table, seen);
  for (const fk of fks.filter((f) => f.parent === table && (f.rule === "r" || f.rule === "a"))) {
    if (!fk.childHasId) throw new Error(`${fk.child} 没有 id 列且以 restrict 引用 ${table},需手工处理`);
    const r = await c.query<{ id: string }>(
      `select id::text as id from public."${fk.child}" where "${fk.column}"::text = any($1::text[])`,
      [fresh],
    );
    if (r.rows.length)
      await collect(
        c,
        fks,
        fk.child,
        r.rows.map((x) => x.id),
        plan,
      );
  }
}

/** 删除顺序:子表先于父表(按 restrict 依赖拓扑;自引用表内一次删完) */
function deleteOrder(plan: Plan, fks: Fk[]): string[] {
  const tables = [...plan.keys()];
  const out: string[] = [];
  const visiting = new Set<string>();
  const visit = (t: string) => {
    if (out.includes(t) || visiting.has(t)) return;
    visiting.add(t);
    for (const fk of fks)
      if (fk.parent === t && fk.child !== t && plan.has(fk.child) && (fk.rule === "r" || fk.rule === "a"))
        visit(fk.child);
    visiting.delete(t);
    out.push(t);
  };
  for (const t of tables) visit(t);
  return out;
}

async function main() {
  const args = new Set(process.argv.slice(2));
  const yes = args.has("--yes");
  if (yes && args.has("--dry-run")) throw new Error("--yes 与 --dry-run 不能同时给");
  const url = process.env.DATABASE_URL || DEFAULT_URL;
  assertDevUrl(url);
  const c = new pg.Client({ connectionString: url, connectionTimeoutMillis: 3000 });
  await c.connect();
  try {
    // ── 1. 测试 schema ──
    const schemas = (
      await c.query<{ name: string; tables: number; locked: boolean }>(`
        select n.nspname as name,
               (select count(*)::int from pg_class k where k.relnamespace = n.oid and k.relkind = 'r') as tables,
               exists (select 1 from pg_locks l join pg_class k on k.oid = l.relation
                        where k.relnamespace = n.oid and l.pid <> pg_backend_pid()) as locked
          from pg_namespace n order by n.nspname`)
    ).rows.filter((s) => isTestSchema(s.name));
    const busy = schemas.filter((s) => s.locked);
    const drop = schemas.filter((s) => !s.locked);
    const groups = new Map<string, number>();
    for (const s of drop) {
      const g = TEST_SCHEMA_PATTERNS.find((p) => p.test(s.name))!.source.replace(/^\^/, "");
      groups.set(g, (groups.get(g) ?? 0) + 1);
    }
    console.log(`测试 schema:${drop.length} 个待删${busy.length ? `,${busy.length} 个正被使用(跳过)` : ""}`);
    for (const [g, n] of [...groups].sort((a, b) => b[1] - a[1])) console.log(`  ${g.padEnd(14)} ${n}`);
    if (args.has("--verbose")) for (const s of drop) console.log(`    ${s.name} (${s.tables} 表)`);
    for (const s of busy) console.log(`  跳过(有锁):${s.name}`);

    // ── 2. public 里的测试实体 ──
    const fks = await loadFks(c);
    const plan: Plan = new Map();
    const roots: string[] = [];
    for (const r of ROOTS) {
      const rows = await c.query<{ id: string; label: string }>(
        `select id::text as id, "${r.column}" as label from public."${r.table}" where "${r.column}" ~* $1`,
        [PREFIX_SQL],
      );
      for (const x of rows.rows) roots.push(`  ${r.table.padEnd(13)} ${x.label}`);
      if (rows.rows.length)
        await collect(
          c,
          fks,
          r.table,
          rows.rows.map((x) => x.id),
          plan,
        );
    }
    const order = deleteOrder(plan, fks);
    const total = [...plan.values()].reduce((a, s) => a + s.size, 0);
    console.log(`测试实体(前缀 w32-/w33-/fe4-/qa-):${roots.length} 个根,连同依赖共 ${total} 行`);
    for (const line of roots) console.log(line);
    for (const t of order) console.log(`  删除 ${t}: ${plan.get(t)!.size} 行`);

    if (!yes) {
      console.log("\n(dry-run)未改动。确认后加 --yes 执行。");
      return;
    }
    for (const s of drop) await c.query(`drop schema if exists "${s.name}" cascade`);
    await c.query("begin");
    try {
      for (const t of order)
        await c.query(`delete from public."${t}" where id::text = any($1::text[])`, [[...plan.get(t)!]]);
      await c.query("commit");
    } catch (e) {
      await c.query("rollback");
      throw e;
    }
    console.log(`\n已删除 ${drop.length} 个 schema、${total} 行测试实体。`);
  } finally {
    await c.end();
  }
}

if (import.meta.main) {
  main().catch((e) => {
    console.error(e instanceof Error ? e.message : e);
    process.exitCode = 1;
  });
}