feat: add a retention script downstreams can schedule

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.
This commit is contained in:
2026-08-19 14:11:22 -04:00
parent c26b34e854
commit 511aa4899c
3 changed files with 902 additions and 0 deletions
+298
View File
@@ -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"]
+273
View File
@@ -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
+331
View File
@@ -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/<date>.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())