feat: let the retention script be told which table it may delete from

Until now the target came from whatever search_path resolved to. The script printed what it found, but that print and the DELETE happen in the same run with nobody in between, so it only ever helped the person who ran a dry-run first. Swap the role that runs it and "$user" can resolve somewhere else entirely. --table takes the whole qualified name and resolves it directly. The table half has to be llm_calls: a version that accepts any name turns one typo into a general purpose row deleter, and any table with a created_at and a tenant_id would go through the same batched DELETE without complaint. The tests that run it now run as a role that owns its own scratch table and holds nothing on the shared one, so the row-count snapshot could go. What replaced it is a case that lets the script fall through to the shared table on purpose and asserts it exits 2 having deleted nothing. That one has no red-first path, since making it red means running it as the superuser, which is the thing being prevented; the finding's probe covers it instead. Five of the new usage tests passed before the flag existed, because argparse rejects an unknown --table with exit 1 and the word --table in stderr, which is exactly what they asserted. They now also assert the error is not "unrecognized", which is the difference between testing the validation and testing argparse.
2026-08-26 10:38:21 -04:00
parent 064f22a0a0
commit 503c06327e
4 changed files with 393 additions and 116 deletions
@@ -59,6 +59,11 @@ _EPILOG = """\
 时间口径: 截止时刻 = 当前 UTC 时刻 - N 天,删除 created_at < 截止时刻 的行;
          --older-than-days 0 即"删除此刻之前的全部行"。

+--table: 本脚本只清理表 llm_calls,故 --table 只有 schema 一段可变(写成
+         --table <schema>.llm_calls)。给了它,目标就由参数精确解析、不再经
+         search_path 推断。含点或引号的复杂标识符不支持,此时请不给 --table,
+         退回 search_path 解析那条路径。
+
 示例:
  python tools/telemetry_retention.py --backend sqlite --path runs/telemetry.db \\
      --older-than-days 90                      # dry-run,只看会删什么
@@ -114,6 +119,11 @@ def _build_parser() -> _Parser:
        action="store_true",
        help="仅 sqlite: 删除后执行 VACUUM 回收文件空间;须与 --apply 同时给",
    )
+    parser.add_argument(
+        "--table",
+        metavar="SCHEMA.NAME",
+        help=f"仅 postgres: 把目标钉死为 <schema>.{TABLE},绕开 search_path 推断",
+    )
    return parser


@@ -150,10 +160,16 @@ def _validate_sqlite(parser: _Parser, args: argparse.Namespace) -> None:
        parser.error("--backend sqlite 不接受 --dsn")
    if args.batch_size is not None:
        parser.error("--batch-size 仅用于 --backend postgres")
+    if args.table is not None:
+        parser.error("--table 仅用于 --backend postgres:SQLite 的库文件即目标,无 schema 可消歧")


 def _validate_postgres(parser: _Parser, args: argparse.Namespace) -> None:
-    """Postgres 分支: 必须有 --dsn,拒绝 sqlite 专属参数,并在此落 --batch-size 缺省值。"""
+    """Postgres 分支: 必须有 --dsn,拒绝 sqlite 专属参数,并在此落 --batch-size 缺省值。
+
+    `--table` 在此解析成 `args.table_schema`(未给则 None): 校验与解析放在同一处,
+    后面的执行路径就只面对一个已经合法的 schema 名,不必再重复判断。
+    """
    if args.dsn is None:
        parser.error("--backend postgres 需要 --dsn")
    if args.path is not None:
@@ -164,6 +180,35 @@ def _validate_postgres(parser: _Parser, args: argparse.Namespace) -> None:
        args.batch_size = 1000
    elif args.batch_size < 1:
        parser.error("--batch-size 必须 >= 1")
+    args.table_schema = None if args.table is None else _parse_table(parser, args.table)
+
+
+def _parse_table(parser: _Parser, value: str) -> str:
+    """校验 `--table SCHEMA.NAME` 并返回 schema 段;任何不合法形态退出 1。
+
+    **表名段为什么不可变**: 只校验"两段、非空"的话,一次手误 `--table audit.events`
+    就会让本脚本对一张恰好也有 `created_at` / `tenant_id` 的业务表跑同一套分批 DELETE。
+    脚本的名字、退出码 3 的分区提示、README 的定位全都围绕遥测表写,它从未声称自己
+    是通用清理器;把这条校验去掉等于在一个拿 DELETE 权限跑的脚本上开静默的口子。
+    """
+    segments = value.split(".")
+    if len(segments) != 2:
+        parser.error(f"--table 必须是 <schema>.{TABLE} 这样的两段限定名,当前: {value!r}")
+    schema, name = segments
+    # 段内不可能再含 "." (上面按 "." 切成恰好两段),故此处只需查引号
+    if not schema or not name:
+        parser.error(f"--table 的 schema 段与表名段都不得为空,当前: {value!r}")
+    if '"' in schema or '"' in name:
+        parser.error(
+            f"--table 不支持含引号的复杂标识符,当前: {value!r};"
+            "这种情形请不给 --table,退回 search_path 解析那条路径。"
+        )
+    if name != TABLE:
+        parser.error(
+            f"--table 的表名段必须逐字等于 {TABLE}:本脚本只清理遥测表 {TABLE},"
+            f"不是通用清理器,当前: {value!r}"
+        )
+    return schema


 def _print_stats(total: int, low: object, high: object, tenants: Sequence[tuple[str, int]]) -> None:
@@ -240,7 +285,9 @@ def _quote(identifier: str) -> str:
    return f'"{escaped}"'


-async def _run_postgres(dsn: str, cutoff: datetime, apply_: bool, batch_size: int) -> int:
+async def _run_postgres(
+    dsn: str, cutoff: datetime, apply_: bool, batch_size: int, table_schema: str | None
+) -> int:
    """PostgreSQL 分支: 分区表让路,普通表分批 DELETE(每批一个事务)。"""
    try:
        import asyncpg
@@ -257,7 +304,7 @@ async def _run_postgres(dsn: str, cutoff: datetime, apply_: bool, batch_size: in
        print(f"连接 PostgreSQL 失败: {exc}", file=sys.stderr)
        return EXIT_BACKEND
    try:
-        return await _purge_postgres(conn, cutoff, apply_, batch_size)
+        return await _purge_postgres(conn, cutoff, apply_, batch_size, table_schema)
    except asyncpg.PostgresError as exc:
        print(f"PostgreSQL 操作失败: {exc}", file=sys.stderr)
        return EXIT_BACKEND
@@ -265,23 +312,41 @@ async def _run_postgres(dsn: str, cutoff: datetime, apply_: bool, batch_size: in
        await conn.close()


-async def _purge_postgres(conn: Any, cutoff: datetime, apply_: bool, batch_size: int) -> int:
+async def _purge_postgres(
+    conn: Any, cutoff: datetime, apply_: bool, batch_size: int, table_schema: str | None
+) -> int:
    """已连上后的清理主体(conn 是 asyncpg.Connection,不 import 类型以免脚本硬依赖)。"""
-    # 先解析目标: to_regclass 走连接自己的 search_path,故必须把解析结果打出来——
-    # "我删的到底是哪张表"是这个脚本唯一不能猜的事(共享库里另有同名表的场景常见)。
+    # 给了 --table 就用引号限定名精确解析(绕开 search_path),否则维持裸表名解析——
+    # 后者走连接自己的 search_path,故无论哪条路都必须把解析结果打出来:"我删的到底是
+    # 哪张表"是这个脚本唯一不能猜的事(共享库里另有同名表的场景常见)。
+    lookup = TABLE if table_schema is None else f"{_quote(table_schema)}.{_quote(TABLE)}"
    target = await conn.fetchrow(
        "SELECT n.nspname AS schema, c.relname AS name, "
        "EXISTS (SELECT 1 FROM pg_partitioned_table p WHERE p.partrelid = c.oid) AS partitioned "
        "FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace "
        "WHERE c.oid = to_regclass($1)",
-        TABLE,
+        lookup,
    )
    if target is None:
-        print(f"目标库的 search_path 下找不到表 {TABLE}", file=sys.stderr)
+        # 两条路的诊断方向不同,消息分开写: 显式指定找不到多半是名字/大小写写错了,
+        # search_path 找不到则是连接配置的事。
+        if table_schema is None:
+            print(f"目标库的 search_path 下找不到表 {TABLE}", file=sys.stderr)
+        else:
+            print(
+                f"显式指定的表 {table_schema}.{TABLE} 不存在或当前角色不可见。"
+                "注意: PG 中未加引号建的标识符在 catalog 里是小写。",
+                file=sys.stderr,
+            )
        return EXIT_BACKEND
    schema, name = target["schema"], target["name"]
    qualified = f"{_quote(schema)}.{_quote(name)}"
    print(f"目标表: {schema}.{name}")
+    if apply_ and table_schema is None:
+        # 只在 --apply 时提示: dry-run 不可逆性为零,且它本就以"看清楚再决定"为用途。
+        print(
+            "注意: 目标表由连接的 search_path 推断得到。要把目标钉死,请加 --table <schema>.<表名>。"
+        )
    if target["partitioned"]:
        print(
            f"{schema}.{name} 是分区表: 本脚本拒绝对分区表执行 DELETE。\n"
@@ -347,7 +412,9 @@ def main(argv: Sequence[str] | None = None) -> int:
    print(f"模式: {'apply(将真正删除)' if args.apply else 'dry-run(只统计,不删除)'}")
    if args.backend == "sqlite":
        return _run_sqlite(args.path, cutoff.strftime(_SQLITE_TIME_FORMAT), args.apply, args.vacuum)
-    return asyncio.run(_run_postgres(args.dsn, cutoff, args.apply, args.batch_size))
+    return asyncio.run(
+        _run_postgres(args.dsn, cutoff, args.apply, args.batch_size, args.table_schema)
+    )


 if __name__ == "__main__":