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:
@@ -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"]
|
||||
Reference in New Issue
Block a user