PolyGateway/tools/telemetry_retention.py

#!/usr/bin/env python3
"""遥测表 `llm_calls` 的保留期清理脚本(issue #12;独立运维工具,库本体不 import 它)。

**为什么是脚本而不是库能力**: 库对下游数据库只做 SELECT/INSERT 加可选建表,一切
改结构与删数据的操作交给下游(ARCHITECTURE D15)。库若持有 DELETE 权限,就与生产
部署模板推荐的 `REVOKE UPDATE, DELETE ON llm_calls FROM app` 直接冲突。

**默认 dry-run**: 本脚本会永久删除审计数据,故不带 `--apply` 时只统计不删,并把
行数、`created_at` 窗口、`tenant_id` 分布三样一并打出——运维据此判断"删掉的是不是
我想删的",判断不了就不该按下 `--apply`。

**失败方向与库相反**: 这是运维工具,缺依赖/连不上/表不存在一律明确报错退出,绝不
静默降级成"删了 0 行"——静默的 0 行会被当成"已清理干净"。

用法见 `--help`。
"""

from __future__ import annotations

import argparse
import asyncio
import re
import sqlite3
import sys
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import TYPE_CHECKING, Any, NoReturn

if TYPE_CHECKING:
    from collections.abc import Sequence

TABLE = "llm_calls"

# --table 的 schema 段白名单。收紧到普通标识符不是为了防注入(目标名走 to_regclass
# 的参数化占位,且用 _quote 转义),而是让 --help 里"不支持复杂标识符"这句话与实现
# 一致——文档说不支持、实现却照单全收,受害的是照文档做判断的人。
_PLAIN_IDENTIFIER = re.compile(r"[A-Za-z_][A-Za-z0-9_$]*")

# 退出码是本脚本对调度器(cron/systemd)的公共契约,改动即破坏下游告警规则
EXIT_OK = 0
EXIT_USAGE = 1
EXIT_BACKEND = 2
EXIT_PARTITIONED = 3

# 库写 SQLite 的 created_at 是 UTC 文本(DEFAULT (datetime('now'))),故截止时刻
# 也必须是同格式文本——该格式定长且高位在前,字符串比较与时间序等价。
# PG 的 created_at 是 TIMESTAMPTZ,直接传 aware datetime,两端口径不可互换。
_SQLITE_TIME_FORMAT = "%Y-%m-%d %H:%M:%S"

_EPILOG = """\
退出码:
  0  正常完成(含 dry-run)
  1  参数错误
  2  连接/权限/目标表不可用(含缺少 asyncpg)
  3  目标是 PostgreSQL 分区表 —— 请改用 DETACH/DROP PARTITION,脚本不会 DELETE

权限: 请用**维护角色**(表属主)跑本脚本,不要用应用账号 —— 生产部署模板已对应用
      账号 REVOKE UPDATE, DELETE ON llm_calls(遥测表按不可变审计表对待)。

推荐路径(本脚本是存量兜底,不是首选):
  PostgreSQL  把 llm_calls 建成按 created_at 的 RANGE 分区表,过期靠
              ALTER TABLE ... DETACH PARTITION + DROP TABLE 做 O(1) 清理。
  SQLite      按天/按实验轮转库文件(如 runs/<date>.db),到期直接删文件。

时间口径: 截止时刻 = 当前 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,只看会删什么
  python tools/telemetry_retention.py --backend postgres --dsn "$DSN" \\
      --older-than-days 90 --apply --batch-size 1000
"""


class _Parser(argparse.ArgumentParser):
    """把 argparse 的参数错误退出码从 2 改成 1。

    2 在本脚本的契约里留给"连接/权限失败",两者混用会让调度器分不清"我写错了参数"
    与"数据库连不上"——后者要告警重试,前者不该重试。
    """

    def error(self, message: str) -> NoReturn:
        self.print_usage(sys.stderr)
        print(f"{self.prog}: 参数错误: {message}", file=sys.stderr)
        raise SystemExit(EXIT_USAGE)


def _build_parser() -> _Parser:
    """构造 CLI 解析器(参数契约见设计 §6.2)。"""
    parser = _Parser(
        prog="telemetry_retention.py",
        description="按 created_at 清理 PolyGateway 遥测表 llm_calls 的过期行(默认 dry-run)。",
        epilog=_EPILOG,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument("--backend", required=True, choices=("sqlite", "postgres"))
    parser.add_argument("--path", help="SQLite 库文件路径(--backend sqlite 必填)")
    parser.add_argument("--dsn", help="PostgreSQL DSN(--backend postgres 必填)")
    parser.add_argument(
        "--older-than-days",
        type=int,
        required=True,
        metavar="N",
        help="删除 created_at 早于 N 天前的行;N >= 0",
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="真正执行删除;不给则只统计不删(默认)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        metavar="N",
        help="仅 postgres: 每批删除的行数,每批一个事务(默认 1000)",
    )
    parser.add_argument(
        "--vacuum",
        action="store_true",
        help="仅 sqlite: 删除后执行 VACUUM 回收文件空间;须与 --apply 同时给",
    )
    parser.add_argument(
        "--table",
        metavar="SCHEMA.NAME",
        help=f"仅 postgres: 把目标钉死为 <schema>.{TABLE},绕开 search_path 推断",
    )
    return parser


def _validate(parser: _Parser, args: argparse.Namespace) -> None:
    """校验参数组合;任何不合法组合以退出码 1 结束(P5: 不给默认值掩盖错误)。

    **校验链的顺序就是错误消息的优先级**: 先两端通用,再按 backend 分支——同时给出
    多个错误参数时,报出的是链上最先命中的那条。
    """
    _validate_shared(parser, args)
    if args.backend == "sqlite":
        _validate_sqlite(parser, args)
        return
    _validate_postgres(parser, args)


def _validate_shared(parser: _Parser, args: argparse.Namespace) -> None:
    """两端通用的校验。

    `--vacuum` 与 `--apply` 的联动归在这里(而不是 SQLite 分支): 它是"别在只想看看的
    时候重写整个库"这条安全约束,先于"这个参数属于哪个 backend"成立。
    """
    if args.older_than_days < 0:
        parser.error("--older-than-days 必须 >= 0")
    if args.vacuum and not args.apply:
        parser.error("--vacuum 会重写整个库文件,必须与 --apply 同时给")


def _validate_sqlite(parser: _Parser, args: argparse.Namespace) -> None:
    """SQLite 分支: 必须有 --path,且拒绝一切 postgres 专属参数(不静默忽略)。"""
    if args.path is None:
        parser.error("--backend sqlite 需要 --path")
    if args.dsn is not 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 缺省值。

    `--table` 在此解析成 `args.table_schema`(未给则 None): 校验与解析放在同一处,
    后面的执行路径就只面对一个已经合法的 schema 名,不必再重复判断。
    """
    if args.dsn is None:
        parser.error("--backend postgres 需要 --dsn")
    if args.path is not None:
        parser.error("--backend postgres 不接受 --path")
    if args.vacuum:
        parser.error("--vacuum 仅用于 --backend sqlite")
    if args.batch_size is 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 not _PLAIN_IDENTIFIER.fullmatch(schema):
        parser.error(
            f"--table 的 schema 段只接受普通标识符(字母或下划线开头,其后字母/数字/"
            f"下划线/$),当前: {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:
    """打印将删除行数、created_at 窗口与按 tenant_id 的分布。

    tenant_id 用 repr 打: 空串是"未归属"的哨兵(不是 NULL),裸打会与缺失混淆。
    """
    print(f"将删除行数: {total}")
    print(f"created_at 范围: {low} ~ {high}" if total else "created_at 范围: (无匹配行)")
    print("按 tenant_id 分布:")
    if not tenants:
        print("  (无匹配行)")
    for tenant, count in tenants:
        print(f"  {tenant!r}: {count}")


# --------------------------------------------------------------------------- SQLite


def _run_sqlite(path: str, cutoff: str, apply_: bool, vacuum: bool) -> int:
    """SQLite 分支: 单条 DELETE(本地文件无长事务与锁膨胀问题),VACUUM 须显式要。"""
    file = Path(path)
    if not file.is_file():
        print(f"SQLite 库文件不存在: {file}", file=sys.stderr)
        return EXIT_BACKEND
    try:
        conn = sqlite3.connect(f"file:{file}?mode=rw", uri=True)
    except sqlite3.Error as exc:
        print(f"打开 SQLite 库失败: {file}: {exc}", file=sys.stderr)
        return EXIT_BACKEND
    try:
        exists = conn.execute(
            "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", (TABLE,)
        ).fetchone()
        if exists is None:
            print(f"目标库里没有表 {TABLE}: {file}", file=sys.stderr)
            return EXIT_BACKEND
        print(f"目标表: {file}::{TABLE}")
        total, low, high = conn.execute(
            f"SELECT COUNT(*), MIN(created_at), MAX(created_at) FROM {TABLE} WHERE created_at < ?",
            (cutoff,),
        ).fetchone()
        tenants = conn.execute(
            f"SELECT tenant_id, COUNT(*) FROM {TABLE} WHERE created_at < ? "
            "GROUP BY tenant_id ORDER BY COUNT(*) DESC, tenant_id",
            (cutoff,),
        ).fetchall()
        _print_stats(total, low, high, tenants)
        if not apply_:
            print("模式 dry-run: 未删除任何行。确认无误后加 --apply 才会真正删除。")
            return EXIT_OK
        cursor = conn.execute(f"DELETE FROM {TABLE} WHERE created_at < ?", (cutoff,))
        conn.commit()
        print(f"已删除 {cursor.rowcount} 行。")
        if vacuum:
            print("执行 VACUUM(重写整个库文件,需要与库等量的空闲磁盘)…")
            conn.execute("VACUUM")
            conn.commit()
            print("VACUUM 完成。")
    except sqlite3.Error as exc:
        print(f"SQLite 操作失败: {exc}", file=sys.stderr)
        return EXIT_BACKEND
    finally:
        conn.close()
    return EXIT_OK


# --------------------------------------------------------------------------- PostgreSQL


def _quote(identifier: str) -> str:
    """把 catalog 取回的 schema/表名包成合法标识符(库名含大写或特殊字符时必需)。"""
    escaped = identifier.replace('"', '""')
    return f'"{escaped}"'


async def _run_postgres(
    dsn: str, cutoff: datetime, apply_: bool, batch_size: int, table_schema: str | None
) -> int:
    """PostgreSQL 分支: 分区表让路,普通表分批 DELETE(每批一个事务)。"""
    try:
        import asyncpg
    except ImportError as exc:
        print(
            f"--backend postgres 需要 asyncpg,当前不可用({exc});"
            "请 pip install 'polygateway[postgres]' 或 pip install asyncpg 后重试。",
            file=sys.stderr,
        )
        return EXIT_BACKEND
    try:
        conn = await asyncpg.connect(dsn, timeout=10)
    except (OSError, asyncpg.PostgresError) as exc:
        print(f"连接 PostgreSQL 失败: {exc}", file=sys.stderr)
        return EXIT_BACKEND
    try:
        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
    finally:
        await conn.close()


async def _purge_postgres(
    conn: Any, cutoff: datetime, apply_: bool, batch_size: int, table_schema: str | None
) -> int:
    """已连上后的清理主体(conn 是 asyncpg.Connection,不 import 类型以免脚本硬依赖)。"""
    # 给了 --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)",
        lookup,
    )
    if target is None:
        # 两条路的诊断方向不同,消息分开写: 显式指定找不到多半是名字/大小写写错了,
        # 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"
            "请改用 DETACH/DROP PARTITION —— ALTER TABLE ... DETACH PARTITION <子表> 后 "
            "DROP TABLE <子表>(或交给 pg_partman 的 retention)。\n"
            "那是 O(1) 的,而 DELETE 会全表扫描并留下等量膨胀。"
        )
        return EXIT_PARTITIONED

    stats = await conn.fetchrow(
        f"SELECT COUNT(*) AS total, MIN(created_at) AS low, MAX(created_at) AS high "
        f"FROM {qualified} WHERE created_at < $1",
        cutoff,
    )
    tenants = await conn.fetch(
        f"SELECT tenant_id, COUNT(*) AS total FROM {qualified} WHERE created_at < $1 "
        "GROUP BY tenant_id ORDER BY COUNT(*) DESC, tenant_id",
        cutoff,
    )
    _print_stats(
        stats["total"], stats["low"], stats["high"], [(r["tenant_id"], r["total"]) for r in tenants]
    )
    if not apply_:
        print("模式 dry-run: 未删除任何行。确认无误后加 --apply 才会真正删除。")
        return EXIT_OK

    # 分批: 一条大 DELETE 会撑出长事务(阻塞 autovacuum、堆积 WAL、锁膨胀),
    # 中断后还得整批回滚重来。每批独立提交,中断只影响未删批次。
    deleted = 0
    batches = 0
    statement = (
        f"DELETE FROM {qualified} WHERE ctid IN "
        f"(SELECT ctid FROM {qualified} WHERE created_at < $1 ORDER BY created_at LIMIT $2)"
    )
    while True:
        async with conn.transaction():
            status = await conn.execute(statement, cutoff, batch_size)
        count = int(status.rsplit(" ", 1)[-1])
        if count == 0:
            break
        deleted += count
        batches += 1
        print(f"  批次 {batches}: 删除 {count} 行(已提交)")
    print(f"已删除 {deleted} 行,共 {batches} 批。")
    return EXIT_OK


# --------------------------------------------------------------------------- 入口


def main(argv: Sequence[str] | None = None) -> int:
    """解析参数并分派到对应后端;返回值即进程退出码。"""
    parser = _build_parser()
    args = parser.parse_args(argv)
    _validate(parser, args)

    cutoff = datetime.now(UTC) - timedelta(days=args.older_than_days)
    print(f"后端: {args.backend}")
    print(
        f"截止时间(UTC): {cutoff.strftime(_SQLITE_TIME_FORMAT)}"
        f"(--older-than-days {args.older_than_days};删除 created_at 早于该时刻的行)"
    )
    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, args.table_schema)
    )


if __name__ == "__main__":
    sys.exit(main())