From 511aa4899ce6dbd351a3a7a7fa2e0c1c075c759d Mon Sep 17 00:00:00 2001 From: iomgaa Date: Wed, 19 Aug 2026 14:11:22 -0400 Subject: [PATCH] feat: add a retention script downstreams can schedule MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The library only ever SELECTs/INSERTs into llm_calls (D15), so expiring rows has to live outside it — holding DELETE would contradict the REVOKE UPDATE, DELETE the deployment template recommends. tools/telemetry_retention.py is dry-run by default and prints the row count, the created_at window and the tenant_id spread so an operator can tell whether the rows about to go are the intended ones. The Postgres branch refuses partitioned targets with exit code 3 (DETACH/DROP PARTITION is O(1); DELETE is not) and otherwise deletes in per-batch transactions. Missing asyncpg exits 2 rather than degrading quietly: this is an ops tool, and a silent "0 rows" reads as "already clean". Exit codes are the contract with the scheduler, so argparse errors were moved off 2 (now 1) to keep "bad flags" distinguishable from "cannot reach the database". The Postgres cases run against the real instance in throwaway schemas — never public.llm_calls — and the batch case asserts the shared table's row count is unchanged, so a search_path that failed to apply lands as a red test instead of a deletion. --- tests/integration/test_retention_tool_pg.py | 298 ++++++++++++++++++ tests/unit/test_retention_tool.py | 273 ++++++++++++++++ tools/telemetry_retention.py | 331 ++++++++++++++++++++ 3 files changed, 902 insertions(+) create mode 100644 tests/integration/test_retention_tool_pg.py create mode 100644 tests/unit/test_retention_tool.py create mode 100644 tools/telemetry_retention.py diff --git a/tests/integration/test_retention_tool_pg.py b/tests/integration/test_retention_tool_pg.py new file mode 100644 index 0000000..7c95ba4 --- /dev/null +++ b/tests/integration/test_retention_tool_pg.py @@ -0,0 +1,298 @@ +"""`tools/telemetry_retention.py` 的 PostgreSQL 分支测试(issue #12 Task 3,真实 PG)。 + +DSN 走 .env `PGW_TELEMETRY_PG_DSN`,缺则 skip。 + +隔离纪律(M4 事故教训): `public.llm_calls` 是与真实批跑共享的表,而本测试跑的是 +一个**会删数据的脚本**——一律在自建的临时 schema 里操作(DSN 挂 search_path), +teardown 只 `DROP SCHEMA ... CASCADE`;分批删除那例另行断言 `public.llm_calls` +的行数前后不变,把"search_path 没生效"这种最坏情况钉成红灯而不是静默删库。 +""" + +from __future__ import annotations + +import os +import re +import subprocess +import sys +from datetime import UTC, datetime, timedelta +from pathlib import Path +from uuid import uuid4 + +import pytest +from dotenv import dotenv_values + +from polygateway.telemetry.schema import PG_DDL + +_ROOT = Path(__file__).resolve().parents[2] +_SCRIPT = _ROOT / "tools" / "telemetry_retention.py" + +_INSERT = ( + "INSERT INTO llm_calls (call_id, model, provider, source_name, messages, response, " + "prompt_tokens, completion_tokens, usage_source, latency_ms, tenant_id, created_at) " + "VALUES ($1, 'm', 'p', 's1', '[]', 'ok', 1, 2, 'measured', 10, $2, $3)" +) + + +def _partitioned_ddl() -> str: + """由库的真实 `PG_DDL` 派生一份 RANGE 分区版建表语句。 + + 不另抄一份 DDL: 抄的那份与库的 schema 必然漂移,而漂移后本测试验的就不再是 + "库建的表被做成分区后脚本认不认得"。两处改动都是分区表的**硬性要求**—— + 分区表上的唯一约束必须包含分区键,故 `call_id` 单列主键不再合法。 + """ + body, count = re.subn( + r"call_id(\s+)TEXT PRIMARY KEY", r"call_id\1TEXT NOT NULL", PG_DDL, count=1 + ) + if count != 1: + raise AssertionError("PG_DDL 的 call_id 主键声明形态已变,分区版 DDL 需同步") + body = body.strip().rstrip(";").strip() + if not body.endswith(")"): + raise AssertionError("PG_DDL 结尾形态已变,分区版 DDL 需同步") + return ( + f"{body[:-1].rstrip()},\n" + " PRIMARY KEY (call_id, created_at)\n" + ") PARTITION BY RANGE (created_at)" + ) + + +def _dsn_value() -> str | None: + merged = {**dotenv_values(".env"), **os.environ} + raw = merged.get("PGW_TELEMETRY_PG_DSN") + if not raw: + return None + scheme, sep, rest = raw.partition("://") + return f"{scheme.partition('+')[0]}{sep}{rest}" + + +def _search_path_dsn(dsn: str, schema: str) -> str: + sep = "&" if "?" in dsn else "?" + return f"{dsn}{sep}options=-csearch_path%3D{schema}" + + +def _stamp(delta: timedelta) -> datetime: + return datetime.now(UTC) + delta + + +def _run(*args: str, env: dict[str, str] | None = None) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(_SCRIPT), *args], + capture_output=True, + text=True, + cwd=_ROOT, + env=env, + timeout=120, + ) + + +@pytest.fixture +async def dsn(): + value = _dsn_value() + if value is None: + pytest.skip("PGW_TELEMETRY_PG_DSN 未配置") + # 隔离守卫: 该实例有 app/chs_prod 等在用库,只许打 polygateway 专用库 + if not value.rstrip("/").endswith("/polygateway"): + pytest.fail(f"保留期脚本测试只允许连 polygateway 专用库,当前 DSN 库名不符: {value!r}") + return value + + +async def _make_schema(dsn_value: str, prefix: str, ddl: str, extra: tuple[str, ...] = ()) -> str: + import asyncpg + + name = f"pgwret_{prefix}_{uuid4().hex[:8]}" + conn = await asyncpg.connect(dsn_value, timeout=10) + try: + await conn.execute(f"CREATE SCHEMA {name}") + await conn.execute(f"SET search_path = {name}") + await conn.execute(ddl) + for statement in extra: + await conn.execute(statement) + finally: + await conn.close() + return name + + +async def _drop_schema(dsn_value: str, name: str) -> None: + import asyncpg + + conn = await asyncpg.connect(dsn_value, timeout=10) + try: + await conn.execute(f"DROP SCHEMA {name} CASCADE") + finally: + await conn.close() + + +async def _seed(schema_dsn: str, rows: list[tuple[str, str, datetime]]) -> None: + import asyncpg + + conn = await asyncpg.connect(schema_dsn, timeout=10) + try: + await conn.executemany(_INSERT, rows) + finally: + await conn.close() + + +async def _call_ids(schema_dsn: str) -> list[str]: + import asyncpg + + conn = await asyncpg.connect(schema_dsn, timeout=10) + try: + rows = await conn.fetch("SELECT call_id FROM llm_calls ORDER BY call_id") + finally: + await conn.close() + return [r["call_id"] for r in rows] + + +async def _public_count(dsn_value: str) -> int: + """共享表的行数;本测试全程不得让它变动一行。""" + import asyncpg + + conn = await asyncpg.connect(dsn_value, timeout=10) + try: + if await conn.fetchval("SELECT to_regclass('public.llm_calls')") is None: + return -1 + return await conn.fetchval("SELECT COUNT(*) FROM public.llm_calls") + finally: + await conn.close() + + +@pytest.fixture +async def partitioned_schema(dsn): + """临时 schema 内的**分区表**: 脚本必须认出它并让路给 DROP PARTITION。""" + name = await _make_schema( + dsn, + "part", + _partitioned_ddl(), + extra=( + "CREATE TABLE llm_calls_all PARTITION OF llm_calls " + "FOR VALUES FROM ('2000-01-01') TO ('2100-01-01')", + ), + ) + yield _search_path_dsn(dsn, name), name + await _drop_schema(dsn, name) + + +@pytest.fixture +async def plain_schema(dsn): + """临时 schema 内的普通表: 存量场景,脚本的分批 DELETE 兜底路径。""" + name = await _make_schema(dsn, "plain", PG_DDL) + yield _search_path_dsn(dsn, name), name + await _drop_schema(dsn, name) + + +class TestPartitionedTarget: + async def test_partitioned_table_exits_three_without_deleting_anything( + self, partitioned_schema + ): + schema_dsn, schema = partitioned_schema + await _seed( + schema_dsn, + [ + ("part-old-1", "", _stamp(timedelta(days=-30))), + ("part-old-2", "acme", _stamp(timedelta(days=-20))), + ], + ) + + # 带 --apply 跑: 危险的那条路径必须在真正删之前就被分区探测拦住 + result = _run( + "--backend", "postgres", "--dsn", schema_dsn, "--older-than-days", "7", "--apply" + ) + + assert result.returncode == 3, (result.stdout, result.stderr) + combined = result.stdout + result.stderr + assert "DROP PARTITION" in combined + assert "DETACH" in combined + assert await _call_ids(schema_dsn) == ["part-old-1", "part-old-2"] + # 脚本必须报出它解析到的**限定表名**: 这是"我删的到底是哪张表"的唯一凭据 + assert f"{schema}.llm_calls" in result.stdout + + +class TestPlainTableBatches: + async def test_apply_deletes_only_expired_rows_in_batches(self, plain_schema, dsn): + schema_dsn, schema = plain_schema + before_public = await _public_count(dsn) + await _seed( + schema_dsn, + [ + ("old-1", "", _stamp(timedelta(days=-40))), + ("old-2", "acme", _stamp(timedelta(days=-30))), + ("old-3", "acme", _stamp(timedelta(days=-20))), + ("old-4", "acme", _stamp(timedelta(days=-15))), + ("old-5", "", _stamp(timedelta(days=-10))), + ("fresh-1", "acme", _stamp(timedelta(days=-1))), + ("fresh-2", "", _stamp(timedelta(hours=-1))), + ], + ) + + result = _run( + "--backend", + "postgres", + "--dsn", + schema_dsn, + "--older-than-days", + "7", + "--apply", + "--batch-size", + "2", + ) + + assert result.returncode == 0, (result.stdout, result.stderr) + assert await _call_ids(schema_dsn) == ["fresh-1", "fresh-2"] + assert f"{schema}.llm_calls" in result.stdout + assert "将删除行数: 5" in result.stdout + assert "'acme': 3" in result.stdout + # 5 行 / 每批 2 行 = 3 批,每批各自提交;批次行必须真的出现三条 + assert "批次 1" in result.stdout + assert "批次 3" in result.stdout + assert "批次 4" not in result.stdout + assert "已删除 5 行" in result.stdout + assert await _public_count(dsn) == before_public + + async def test_dry_run_on_a_plain_table_deletes_nothing(self, plain_schema): + schema_dsn, _ = plain_schema + await _seed(schema_dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))]) + + result = _run("--backend", "postgres", "--dsn", schema_dsn, "--older-than-days", "7") + + assert result.returncode == 0, (result.stdout, result.stderr) + assert "将删除行数: 1" in result.stdout + assert "dry-run" in result.stdout + assert await _call_ids(schema_dsn) == ["old-1"] + + +class TestMissingAsyncpg: + async def test_missing_asyncpg_exits_two_without_touching_rows(self, plain_schema, tmp_path): + """缺 asyncpg 必须明确报错退出(码 2),不静默降级——这是运维工具不是库路径。 + + 用一个只 `raise ImportError` 的临时 `asyncpg.py` 挂进子进程的 PYTHONPATH 构造该 + 场景: 脚本跑在子进程里,monkeypatch 对它无效。DSN 用**真实可连**的临时 schema, + 这样"没有导入守卫"的实现会走通并退出 0,而不是碰巧也退出 2 而假绿。 + """ + schema_dsn, _ = plain_schema + await _seed(schema_dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))]) + stub = tmp_path / "stub" + stub.mkdir() + (stub / "asyncpg.py").write_text( + 'raise ImportError("asyncpg 未安装(测试构造)")\n', encoding="utf-8" + ) + env = { + **os.environ, + "PYTHONPATH": os.pathsep.join( + [str(stub), *([p] if (p := os.environ.get("PYTHONPATH")) else [])] + ), + } + + result = _run( + "--backend", + "postgres", + "--dsn", + schema_dsn, + "--older-than-days", + "7", + "--apply", + env=env, + ) + + assert result.returncode == 2, (result.stdout, result.stderr) + assert "asyncpg" in result.stderr + assert "pip install" in result.stderr + assert await _call_ids(schema_dsn) == ["old-1"] diff --git a/tests/unit/test_retention_tool.py b/tests/unit/test_retention_tool.py new file mode 100644 index 0000000..249ea5a --- /dev/null +++ b/tests/unit/test_retention_tool.py @@ -0,0 +1,273 @@ +"""`tools/telemetry_retention.py` 的 SQLite 分支测试(issue #12 Task 3)。 + +一律经 `subprocess` 跑真实脚本 + 真实临时 SQLite 库文件: 脚本是独立运维工具、 +不被库 import,用 monkeypatch 或直接 import 私有函数测出来的"通过"与运维实际 +执行的那条路径不是同一条(退出码、argparse 行为、stdout 全都测不到)。 +""" + +from __future__ import annotations + +import sqlite3 +import subprocess +import sys +from datetime import UTC, datetime, timedelta +from pathlib import Path + +from polygateway.telemetry.schema import SQLITE_DDL + +_ROOT = Path(__file__).resolve().parents[2] +_SCRIPT = _ROOT / "tools" / "telemetry_retention.py" +_TIME_FORMAT = "%Y-%m-%d %H:%M:%S" + +# 库写入 SQLite 的 created_at 是 UTC 的 'YYYY-MM-DD HH:MM:SS' 文本(schema 的 +# DEFAULT (datetime('now'))),测试数据必须同款,否则字符串比较的口径就假了 +_INSERT = ( + "INSERT INTO llm_calls (call_id, model, provider, source_name, messages, response, " + "prompt_tokens, completion_tokens, usage_source, latency_ms, tenant_id, created_at) " + "VALUES (?, 'm', 'p', 's1', '[]', 'ok', 1, 2, 'measured', 10, ?, ?)" +) + + +def _stamp(delta: timedelta) -> str: + return (datetime.now(UTC) + delta).strftime(_TIME_FORMAT) + + +def _make_db(tmp_path: Path, rows: list[tuple[str, str, str]]) -> Path: + """按库的真实 DDL 建临时库并灌入 (call_id, tenant_id, created_at) 三元组。""" + path = tmp_path / "telemetry.db" + conn = sqlite3.connect(path) + try: + conn.executescript(SQLITE_DDL) + conn.executemany(_INSERT, rows) + conn.commit() + finally: + conn.close() + return path + + +def _run(*args: str) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(_SCRIPT), *args], + capture_output=True, + text=True, + cwd=_ROOT, + timeout=120, + ) + + +def _rows(path: Path) -> list[str]: + conn = sqlite3.connect(path) + try: + return [r[0] for r in conn.execute("SELECT call_id FROM llm_calls ORDER BY call_id")] + finally: + conn.close() + + +def _aged_db(tmp_path: Path) -> Path: + return _make_db( + tmp_path, + [ + ("old-1", "", _stamp(timedelta(days=-30))), + ("old-2", "acme", _stamp(timedelta(days=-20))), + ("old-3", "acme", _stamp(timedelta(days=-10))), + ("fresh-1", "acme", _stamp(timedelta(days=-1))), + ("fresh-2", "", _stamp(timedelta(hours=-1))), + ], + ) + + +class TestSqliteDryRun: + def test_dry_run_deletes_nothing_and_reports_counts_range_and_tenants(self, tmp_path): + """缺省(不带 --apply)是 dry-run: 一行不删,且报出足以判断"删的是不是我想删的"的三样。""" + path = _aged_db(tmp_path) + + result = _run("--backend", "sqlite", "--path", str(path), "--older-than-days", "7") + + assert result.returncode == 0, result.stderr + assert _rows(path) == ["fresh-1", "fresh-2", "old-1", "old-2", "old-3"] + assert "将删除行数: 3" in result.stdout + assert "created_at 范围:" in result.stdout + assert "按 tenant_id 分布" in result.stdout + # 空串是"未归属"的哨兵而非 NULL,repr 让它在输出里不被误读成缺失 + assert "'acme': 2" in result.stdout + assert "'': 1" in result.stdout + assert "dry-run" in result.stdout + + def test_dry_run_reports_the_actual_created_at_window(self, tmp_path): + """时间范围报的必须是**命中行**的窗口,不是全表的。""" + path = _aged_db(tmp_path) + + result = _run("--backend", "sqlite", "--path", str(path), "--older-than-days", "7") + + conn = sqlite3.connect(path) + try: + low, high = conn.execute( + "SELECT MIN(created_at), MAX(created_at) FROM llm_calls WHERE call_id LIKE 'old-%'" + ).fetchone() + finally: + conn.close() + assert f"{low} ~ {high}" in result.stdout + + +class TestSqliteApply: + def test_apply_removes_only_expired_rows(self, tmp_path): + path = _aged_db(tmp_path) + + result = _run( + "--backend", "sqlite", "--path", str(path), "--older-than-days", "7", "--apply" + ) + + assert result.returncode == 0, result.stderr + assert _rows(path) == ["fresh-1", "fresh-2"] + assert "已删除 3 行" in result.stdout + + def test_older_than_days_zero_deletes_everything_before_now(self, tmp_path): + """N=0 的边界: 截止时刻即"此刻",此刻之前的全删、之后的(未来戳)留下。""" + path = _make_db( + tmp_path, + [ + ("past", "", _stamp(timedelta(seconds=-5))), + ("future", "", _stamp(timedelta(hours=1))), + ], + ) + + result = _run( + "--backend", "sqlite", "--path", str(path), "--older-than-days", "0", "--apply" + ) + + assert result.returncode == 0, result.stderr + assert _rows(path) == ["future"] + + def test_vacuum_with_apply_rewrites_the_file(self, tmp_path): + path = _aged_db(tmp_path) + + result = _run( + "--backend", + "sqlite", + "--path", + str(path), + "--older-than-days", + "7", + "--apply", + "--vacuum", + ) + + assert result.returncode == 0, result.stderr + assert "VACUUM" in result.stdout + assert _rows(path) == ["fresh-1", "fresh-2"] + + def test_deleting_from_a_db_without_the_table_is_a_backend_failure(self, tmp_path): + """连得上但没有 llm_calls: 属"目标不可用",退出码 2 且**不**静默当成 0 行。""" + path = tmp_path / "empty.db" + sqlite3.connect(path).close() + + result = _run("--backend", "sqlite", "--path", str(path), "--older-than-days", "7") + + assert result.returncode == 2 + assert "llm_calls" in result.stderr + + def test_missing_db_file_exits_two(self, tmp_path): + result = _run( + "--backend", "sqlite", "--path", str(tmp_path / "nope.db"), "--older-than-days", "7" + ) + + assert result.returncode == 2 + assert "nope.db" in result.stderr + + +class TestUsageErrors: + """参数层的一切错误都是退出码 1(argparse 默认的 2 已被本脚本改写,2 留给连接失败)。""" + + def test_sqlite_with_dsn_exits_one(self, tmp_path): + result = _run( + "--backend", + "sqlite", + "--path", + str(tmp_path / "x.db"), + "--dsn", + "postgresql://x/y", + "--older-than-days", + "7", + ) + + assert result.returncode == 1 + assert "--dsn" in result.stderr + + def test_sqlite_without_path_exits_one(self): + result = _run("--backend", "sqlite", "--older-than-days", "7") + + assert result.returncode == 1 + assert "--path" in result.stderr + + def test_sqlite_with_batch_size_exits_one(self, tmp_path): + result = _run( + "--backend", + "sqlite", + "--path", + str(tmp_path / "x.db"), + "--older-than-days", + "7", + "--batch-size", + "10", + ) + + assert result.returncode == 1 + assert "--batch-size" in result.stderr + + def test_postgres_with_vacuum_exits_one(self): + result = _run( + "--backend", + "postgres", + "--dsn", + "postgresql://x/y", + "--older-than-days", + "7", + "--apply", + "--vacuum", + ) + + assert result.returncode == 1 + assert "--vacuum" in result.stderr + + def test_vacuum_without_apply_exits_one(self, tmp_path): + result = _run( + "--backend", + "sqlite", + "--path", + str(tmp_path / "x.db"), + "--older-than-days", + "7", + "--vacuum", + ) + + assert result.returncode == 1 + assert "--apply" in result.stderr + + def test_missing_older_than_days_exits_one(self, tmp_path): + result = _run("--backend", "sqlite", "--path", str(tmp_path / "x.db")) + + assert result.returncode == 1 + + def test_negative_older_than_days_exits_one(self, tmp_path): + result = _run( + "--backend", "sqlite", "--path", str(tmp_path / "x.db"), "--older-than-days", "-1" + ) + + assert result.returncode == 1 + assert "--older-than-days" in result.stderr + + def test_unknown_backend_exits_one(self, tmp_path): + result = _run("--backend", "mysql", "--path", str(tmp_path / "x.db")) + + assert result.returncode == 1 + + +class TestHelp: + def test_help_names_the_maintenance_role_and_the_recommended_path(self): + """帮助文本是运维唯一会读的文档,权限口径与"推荐不是 DELETE"必须在里面。""" + result = _run("--help") + + assert result.returncode == 0 + assert "维护角色" in result.stdout + assert "REVOKE" in result.stdout + assert "PARTITION" in result.stdout diff --git a/tools/telemetry_retention.py b/tools/telemetry_retention.py new file mode 100644 index 0000000..2e13015 --- /dev/null +++ b/tools/telemetry_retention.py @@ -0,0 +1,331 @@ +#!/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 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" + +# 退出码是本脚本对调度器(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/.db),到期直接删文件。 + +时间口径: 截止时刻 = 当前 UTC 时刻 - N 天,删除 created_at < 截止时刻 的行; + --older-than-days 0 即"删除此刻之前的全部行"。 + +示例: + 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 同时给", + ) + return parser + + +def _validate(parser: _Parser, args: argparse.Namespace) -> None: + """校验参数组合;任何不合法组合以退出码 1 结束(P5: 不给默认值掩盖错误)。""" + if args.older_than_days < 0: + parser.error("--older-than-days 必须 >= 0") + if args.vacuum and not args.apply: + parser.error("--vacuum 会重写整个库文件,必须与 --apply 同时给") + if args.backend == "sqlite": + 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") + return + 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") + + +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) -> 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) + 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) -> int: + """已连上后的清理主体(conn 是 asyncpg.Connection,不 import 类型以免脚本硬依赖)。""" + # 先解析目标: to_regclass 走连接自己的 search_path,故必须把解析结果打出来—— + # "我删的到底是哪张表"是这个脚本唯一不能猜的事(共享库里另有同名表的场景常见)。 + 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, + ) + if target is None: + print(f"目标库的 search_path 下找不到表 {TABLE}", file=sys.stderr) + return EXIT_BACKEND + schema, name = target["schema"], target["name"] + qualified = f"{_quote(schema)}.{_quote(name)}" + print(f"目标表: {schema}.{name}") + 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)) + + +if __name__ == "__main__": + sys.exit(main())