iomgaa/PolyGateway
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases 9 Wiki Activity

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

Browse Source 
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.
This commit is contained in:
iomgaa
2026-08-26 10:38:21 -04:00
parent 064f22a0a0
commit 503c06327e
4 changed files with 393 additions and 116 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/integration/test_retention_tool_pg.py
+209 -106
View File
@@ -1,11 +1,12 @@
"""`tools/telemetry_retention.py` 的 PostgreSQL 分支测试(issue #12 Task 3,真实 PG)。
DSN 走 .env `PGW_TELEMETRY_PG_DSN`,缺则 skip。
隔离纪律(issue #18): 本文件跑的是一个**会删数据的脚本**,而实例上的共享表 `llm_calls`
与真实批跑共用。故**凡启动脚本的用例一律用 `pg_sandbox(role="owner")` 的临时角色跑**:
该角色对共享表一无所有,越界不是"会被发现",而是数据库层面做不到。
隔离纪律(M4 事故教训): `public.llm_calls` 是与真实批跑共享的表,而本测试跑的是
一个**会删数据的脚本**——一律在自建的临时 schema 里操作(DSN 挂 search_path),
teardown 只 `DROP SCHEMA ... CASCADE`;分批删除那例另行断言 `public.llm_calls`
的行数前后不变,把"search_path 没生效"这种最坏情况钉成红灯而不是静默删库。
这条纪律取代了此前那条"跑完对比共享表行数"的安全网——行数快照守的是安全属性,却把它
编码成对全局可变量的观测: 外部进程一写就假红,外部插入与脚本误删互相抵消则假阴。
权限边界两个方向都没有。
"""
from __future__ import annotations
@@ -16,10 +17,8 @@ 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
@@ -55,20 +54,6 @@ def _partitioned_ddl() -> str:
)
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
@@ -84,43 +69,6 @@ def _run(*args: str, env: dict[str, str] | None = None) -> subprocess.CompletedP
)
@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
@@ -142,50 +90,31 @@ async def _call_ids(schema_dsn: str) -> list[str]:
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):
async def partitioned_sandbox(pg_sandbox):
"""临时 schema 内的**分区表**: 脚本必须认出它并让路给 DROP PARTITION。"""
name = await _make_schema(
dsn,
"part",
_partitioned_ddl(),
return await pg_sandbox(
ddl=_partitioned_ddl(),
extra=(
"CREATE TABLE llm_calls_all PARTITION OF llm_calls "
"FOR VALUES FROM ('2000-01-01') TO ('2100-01-01')",
),
role="owner",
)
yield _search_path_dsn(dsn, name), name
await _drop_schema(dsn, name)
@pytest.fixture
async def plain_schema(dsn):
async def plain_sandbox(pg_sandbox):
"""临时 schema 内的普通表: 存量场景,脚本的分批 DELETE 兜底路径。"""
name = await _make_schema(dsn, "plain", PG_DDL)
yield _search_path_dsn(dsn, name), name
await _drop_schema(dsn, name)
return await pg_sandbox(ddl=PG_DDL, role="owner")
class TestPartitionedTarget:
async def test_partitioned_table_exits_three_without_deleting_anything(
self, partitioned_schema
self, partitioned_sandbox
):
schema_dsn, schema = partitioned_schema
await _seed(
schema_dsn,
partitioned_sandbox.dsn,
[
("part-old-1", "", _stamp(timedelta(days=-30))),
("part-old-2", "acme", _stamp(timedelta(days=-20))),
@@ -194,24 +123,28 @@ class TestPartitionedTarget:
# 带 --apply 跑: 危险的那条路径必须在真正删之前就被分区探测拦住
result = _run(
"--backend", "postgres", "--dsn", schema_dsn, "--older-than-days", "7", "--apply"
"--backend",
"postgres",
"--dsn",
partitioned_sandbox.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 await _call_ids(partitioned_sandbox.dsn) == ["part-old-1", "part-old-2"]
# 脚本必须报出它解析到的**限定表名**: 这是"我删的到底是哪张表"的唯一凭据
assert f"{schema}.llm_calls" in result.stdout
assert f"{partitioned_sandbox.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)
async def test_apply_deletes_only_expired_rows_in_batches(self, plain_sandbox):
await _seed(
schema_dsn,
plain_sandbox.dsn,
[
("old-1", "", _stamp(timedelta(days=-40))),
("old-2", "acme", _stamp(timedelta(days=-30))),
@@ -227,7 +160,7 @@ class TestPlainTableBatches:
"--backend",
"postgres",
"--dsn",
schema_dsn,
plain_sandbox.dsn,
"--older-than-days",
"7",
"--apply",
@@ -236,8 +169,8 @@ class TestPlainTableBatches:
)
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 await _call_ids(plain_sandbox.dsn) == ["fresh-1", "fresh-2"]
assert f"{plain_sandbox.schema}.llm_calls" in result.stdout
assert "将删除行数: 5" in result.stdout
assert "'acme': 3" in result.stdout
# 5 行 / 每批 2 行 = 3 批,每批各自提交;批次行必须真的出现三条
@@ -245,30 +178,200 @@ class TestPlainTableBatches:
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)))])
async def test_dry_run_on_a_plain_table_deletes_nothing(self, plain_sandbox):
await _seed(plain_sandbox.dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))])
result = _run("--backend", "postgres", "--dsn", schema_dsn, "--older-than-days", "7")
result = _run("--backend", "postgres", "--dsn", plain_sandbox.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"]
assert await _call_ids(plain_sandbox.dsn) == ["old-1"]
class TestExplicitTable:
"""`--table SCHEMA.NAME` 的真实解析行为(issue #18 设计 §4.1;判据 1b)。
单测只能验参数分类,验不了 `to_regclass` 的语义——schema 不存在返 NULL 而非抛错、
引号限定名区分大小写、无权限落在 `COUNT` 而非解析,这三条都必须真连库才成立。
"""
async def test_explicit_table_deletes_exactly_like_the_implicit_path(self, plain_sandbox):
await _seed(
plain_sandbox.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",
plain_sandbox.dsn,
"--older-than-days",
"7",
"--apply",
"--batch-size",
"2",
"--table",
f"{plain_sandbox.schema}.llm_calls",
)
# 与不给 --table 的那条用例逐条同款: 显式声明只改"怎么找到表",不改任何行为
assert result.returncode == 0, (result.stdout, result.stderr)
assert await _call_ids(plain_sandbox.dsn) == ["fresh-1", "fresh-2"]
assert f"{plain_sandbox.schema}.llm_calls" in result.stdout
assert "将删除行数: 5" in result.stdout
assert "'acme': 3" in result.stdout
assert "批次 1" in result.stdout
assert "批次 3" in result.stdout
assert "批次 4" not in result.stdout
assert "已删除 5 行" in result.stdout
async def test_table_in_a_nonexistent_schema_exits_two(self, plain_sandbox):
"""schema 不存在时 `to_regclass` 返 NULL(不抛错),故落进既有的"目标不可用"。"""
await _seed(plain_sandbox.dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))])
missing = "pgw_s_nosuchxxxxxxxx"
result = _run(
"--backend",
"postgres",
"--dsn",
plain_sandbox.dsn,
"--older-than-days",
"7",
"--apply",
"--table",
f"{missing}.llm_calls",
)
assert result.returncode == 2, (result.stdout, result.stderr)
assert missing in result.stderr
assert "llm_calls" in result.stderr
assert await _call_ids(plain_sandbox.dsn) == ["old-1"]
async def test_table_owned_by_another_role_exits_two(self, pg_sandbox):
"""拿 A 的连接指 B 的表: 权限拒绝,两张表都不能少一行。"""
sandbox_a = await pg_sandbox(ddl=PG_DDL, role="owner")
sandbox_b = await pg_sandbox(ddl=PG_DDL, role="owner")
await _seed(sandbox_a.dsn, [("a-old", "acme", _stamp(timedelta(days=-30)))])
await _seed(sandbox_b.dsn, [("b-old", "acme", _stamp(timedelta(days=-30)))])
result = _run(
"--backend",
"postgres",
"--dsn",
sandbox_a.dsn,
"--older-than-days",
"7",
"--apply",
"--table",
f"{sandbox_b.schema}.llm_calls",
)
assert result.returncode == 2, (result.stdout, result.stderr)
assert await _call_ids(sandbox_a.dsn) == ["a-old"]
assert await _call_ids(sandbox_b.dsn) == ["b-old"]
async def test_explicit_partitioned_table_still_exits_three(self, partitioned_sandbox):
await _seed(partitioned_sandbox.dsn, [("part-old-1", "acme", _stamp(timedelta(days=-30)))])
result = _run(
"--backend",
"postgres",
"--dsn",
partitioned_sandbox.dsn,
"--older-than-days",
"7",
"--apply",
"--table",
f"{partitioned_sandbox.schema}.llm_calls",
)
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(partitioned_sandbox.dsn) == ["part-old-1"]
class TestInferredTargetHint:
async def test_apply_without_table_warns_that_the_target_was_inferred(self, plain_sandbox):
"""未钉死目标时必须当场说清"这张表是猜出来的"(设计 §4.4;判据 2)。
该提示行只在 PG 分支打印,不连库的单测触发不到它,故验收落在集成层。
"""
await _seed(plain_sandbox.dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))])
result = _run(
"--backend",
"postgres",
"--dsn",
plain_sandbox.dsn,
"--older-than-days",
"7",
"--apply",
)
assert result.returncode == 0, (result.stdout, result.stderr)
assert "search_path" in result.stdout
assert "--table" in result.stdout
async def test_dry_run_does_not_print_the_hint(self, plain_sandbox):
"""dry-run 不可逆性为零,它本就以"看清楚再决定"为用途,多一行提示是噪音。"""
await _seed(plain_sandbox.dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))])
result = _run("--backend", "postgres", "--dsn", plain_sandbox.dsn, "--older-than-days", "7")
assert result.returncode == 0, (result.stdout, result.stderr)
assert "--table" not in result.stdout
class TestSearchPathFallsThrough:
async def test_bare_search_path_cannot_touch_the_shared_table(self, plain_sandbox):
"""最坏情况: `search_path` 没生效,脚本落到共享表 `llm_calls` 上(设计 §5.3)。
用沙箱角色的**裸** DSN 跑(search_path 回落 `"$user", public`,而角色名与 schema
名有意错开,故 `"$user"` 命不中沙箱),不给 `--table`,带 `--apply`。角色对共享表
无任何权限,于是两条可能的路都收敛到退出码 2: 库里有那张表则 `COUNT` 被权限拒绝,
没有则解析不到。**不断言 PG 的英文原文**——服务端 `lc_messages` 不由测试掌握。
"""
await _seed(plain_sandbox.dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))])
result = _run(
"--backend",
"postgres",
"--dsn",
plain_sandbox.bare_dsn,
"--older-than-days",
"7",
"--apply",
)
assert result.returncode == 2, (result.stdout, result.stderr)
assert result.stderr.strip()
assert "llm_calls" in result.stderr
# 沙箱表一行不少: 脚本既没删共享表,也没绕回来删自己这张
assert await _call_ids(plain_sandbox.dsn) == ["old-1"]
class TestMissingAsyncpg:
async def test_missing_asyncpg_exits_two_without_touching_rows(self, plain_schema, tmp_path):
async def test_missing_asyncpg_exits_two_without_touching_rows(self, plain_sandbox, 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)))])
await _seed(plain_sandbox.dsn, [("old-1", "acme", _stamp(timedelta(days=-30)))])
stub = tmp_path / "stub"
stub.mkdir()
(stub / "asyncpg.py").write_text(
@@ -285,7 +388,7 @@ class TestMissingAsyncpg:
"--backend",
"postgres",
"--dsn",
schema_dsn,
plain_sandbox.dsn,
"--older-than-days",
"7",
"--apply",
@@ -295,4 +398,4 @@ class TestMissingAsyncpg:
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"]
assert await _call_ids(plain_sandbox.dsn) == ["old-1"]
tests/unit/test_retention_tool.py
+107
View File
@@ -261,6 +261,104 @@ class TestUsageErrors:
assert result.returncode == 1
# --- --table 的参数分类(issue #18 设计 §4.2);真实解析行为在集成层验 ---
def test_sqlite_with_table_exits_one(self, tmp_path):
"""SQLite 库文件即目标,无 schema 概念,故 `--table` 在该分支无歧义可消。"""
result = _run(
"--backend",
"sqlite",
"--path",
str(tmp_path / "x.db"),
"--older-than-days",
"7",
"--table",
"some_schema.llm_calls",
)
assert result.returncode == 1
assert "--table" in result.stderr
# 单看退出码与 "--table" 字样会被 argparse 的 "unrecognized arguments" 蒙混
# 过去(它也是退出 1、也回显参数名)。必须钉住"参数已被识别、因规则被拒"。
assert "unrecognized" not in result.stderr
def test_table_without_schema_qualifier_exits_one(self):
"""单段等于没声明: 目标仍由 `search_path` 决定,隐式性原样保留,故拒绝。"""
result = _run(
"--backend",
"postgres",
"--dsn",
"postgresql://x/y",
"--older-than-days",
"7",
"--table",
"llm_calls",
)
assert result.returncode == 1
assert "--table" in result.stderr
# 单看退出码与 "--table" 字样会被 argparse 的 "unrecognized arguments" 蒙混
# 过去(它也是退出 1、也回显参数名)。必须钉住"参数已被识别、因规则被拒"。
assert "unrecognized" not in result.stderr
def test_table_with_empty_segment_exits_one(self):
result = _run(
"--backend",
"postgres",
"--dsn",
"postgresql://x/y",
"--older-than-days",
"7",
"--table",
".llm_calls",
)
assert result.returncode == 1
assert "--table" in result.stderr
# 单看退出码与 "--table" 字样会被 argparse 的 "unrecognized arguments" 蒙混
# 过去(它也是退出 1、也回显参数名)。必须钉住"参数已被识别、因规则被拒"。
assert "unrecognized" not in result.stderr
def test_table_with_quote_in_a_segment_exits_one(self):
"""含引号的复杂标识符不支持: 此时退回不给 `--table` 的路径(见 epilog)。"""
result = _run(
"--backend",
"postgres",
"--dsn",
"postgresql://x/y",
"--older-than-days",
"7",
"--table",
'sch"ema.llm_calls',
)
assert result.returncode == 1
assert "--table" in result.stderr
# 单看退出码与 "--table" 字样会被 argparse 的 "unrecognized arguments" 蒙混
# 过去(它也是退出 1、也回显参数名)。必须钉住"参数已被识别、因规则被拒"。
assert "unrecognized" not in result.stderr
def test_table_naming_another_table_exits_one(self):
"""表名段锁死: 不加这条,`--table` 会把本脚本扩成"任意同形表删除工具"。"""
result = _run(
"--backend",
"postgres",
"--dsn",
"postgresql://x/y",
"--older-than-days",
"7",
"--table",
"audit.events",
)
assert result.returncode == 1
assert "--table" in result.stderr
# 单看退出码与 "--table" 字样会被 argparse 的 "unrecognized arguments" 蒙混
# 过去(它也是退出 1、也回显参数名)。必须钉住"参数已被识别、因规则被拒"。
assert "unrecognized" not in result.stderr
# 错误消息要当场把边界说清: 本脚本的作用域到 llm_calls 为止
assert "llm_calls" in result.stderr
class TestHelp:
def test_help_names_the_maintenance_role_and_the_recommended_path(self):
@@ -271,3 +369,12 @@ class TestHelp:
assert "维护角色" in result.stdout
assert "REVOKE" in result.stdout
assert "PARTITION" in result.stdout
def test_help_states_the_table_name_is_fixed(self):
"""`--table` 只有 schema 一段可变,这条边界必须写在运维会读到的地方。"""
result = _run("--help")
assert result.returncode == 0
assert "--table" in result.stdout
assert "只清理" in result.stdout
assert "llm_calls" in result.stdout