iomgaa
503c06327e
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.
402 lines
15 KiB
Python
402 lines
15 KiB
Python
"""`tools/telemetry_retention.py` 的 PostgreSQL 分支测试(issue #12 Task 3,真实 PG)。
|
|
|
|
隔离纪律(issue #18): 本文件跑的是一个**会删数据的脚本**,而实例上的共享表 `llm_calls`
|
|
与真实批跑共用。故**凡启动脚本的用例一律用 `pg_sandbox(role="owner")` 的临时角色跑**:
|
|
该角色对共享表一无所有,越界不是"会被发现",而是数据库层面做不到。
|
|
|
|
这条纪律取代了此前那条"跑完对比共享表行数"的安全网——行数快照守的是安全属性,却把它
|
|
编码成对全局可变量的观测: 外部进程一写就假红,外部插入与脚本误删互相抵消则假阴。
|
|
权限边界两个方向都没有。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
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 _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,
|
|
)
|
|
|
|
|
|
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]
|
|
|
|
|
|
@pytest.fixture
|
|
async def partitioned_sandbox(pg_sandbox):
|
|
"""临时 schema 内的**分区表**: 脚本必须认出它并让路给 DROP PARTITION。"""
|
|
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",
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
async def plain_sandbox(pg_sandbox):
|
|
"""临时 schema 内的普通表: 存量场景,脚本的分批 DELETE 兜底路径。"""
|
|
return await pg_sandbox(ddl=PG_DDL, role="owner")
|
|
|
|
|
|
class TestPartitionedTarget:
|
|
async def test_partitioned_table_exits_three_without_deleting_anything(
|
|
self, partitioned_sandbox
|
|
):
|
|
await _seed(
|
|
partitioned_sandbox.dsn,
|
|
[
|
|
("part-old-1", "", _stamp(timedelta(days=-30))),
|
|
("part-old-2", "acme", _stamp(timedelta(days=-20))),
|
|
],
|
|
)
|
|
|
|
# 带 --apply 跑: 危险的那条路径必须在真正删之前就被分区探测拦住
|
|
result = _run(
|
|
"--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(partitioned_sandbox.dsn) == ["part-old-1", "part-old-2"]
|
|
# 脚本必须报出它解析到的**限定表名**: 这是"我删的到底是哪张表"的唯一凭据
|
|
assert f"{partitioned_sandbox.schema}.llm_calls" in result.stdout
|
|
|
|
|
|
class TestPlainTableBatches:
|
|
async def test_apply_deletes_only_expired_rows_in_batches(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",
|
|
)
|
|
|
|
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
|
|
# 5 行 / 每批 2 行 = 3 批,每批各自提交;批次行必须真的出现三条
|
|
assert "批次 1" in result.stdout
|
|
assert "批次 3" in result.stdout
|
|
assert "批次 4" not in result.stdout
|
|
assert "已删除 5 行" in result.stdout
|
|
|
|
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", 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(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_sandbox, tmp_path):
|
|
"""缺 asyncpg 必须明确报错退出(码 2),不静默降级——这是运维工具不是库路径。
|
|
|
|
用一个只 `raise ImportError` 的临时 `asyncpg.py` 挂进子进程的 PYTHONPATH 构造该
|
|
场景: 脚本跑在子进程里,monkeypatch 对它无效。DSN 用**真实可连**的临时 schema,
|
|
这样"没有导入守卫"的实现会走通并退出 0,而不是碰巧也退出 2 而假绿。
|
|
"""
|
|
await _seed(plain_sandbox.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",
|
|
plain_sandbox.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(plain_sandbox.dsn) == ["old-1"]
|