iomgaa
58c4af28ea
Both reviews landed on the same line independently. _as_role swaps the credentials in the DSN with a regex, and when the pattern does not match it returned the string unchanged. Two shapes miss it: no inline credentials, and a unix socket URL. Either one is a legal DSN. What that costs is not a broken test. The sandbox builds, every assertion still passes, and bare_dsn is now the admin connection, so the worst-case case runs the real script with --apply as a superuser against the shared table. The verifier ran that command as a dry run to see what it would have done: target public.llm_calls, 11 rows to delete. The case would still have gone red on the exit code, after the rows were gone. It raises now. There is also a second check that connects and compares current_user, because a successful string substitution is not the same as connecting as that role -- PGUSER and friends still override. The whole design rests on that connection having no grant on the shared table; a string comparison is too thin a thing to rest it on. That check has to stay inside the try. Past it the cleanup statements have already been merged into the fixture-level stack, and unwinding again runs DROP OWNED BY twice, which has no IF EXISTS. The catalog probe took any SQL and ran it on the admin connection. The design claims withholding the DSN makes the boundary structural; that was only true of the connection string, not of the capability. It takes SELECT now. --table's schema half is restricted to plain identifiers. Not a security fix, since the name goes through a parameter and _quote: the help text says complex identifiers are unsupported and the code was accepting them anyway.
434 lines
14 KiB
Python
434 lines
14 KiB
Python
"""`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
|
|
|
|
# --- --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 TestTableIdentifierWhitelist:
|
|
"""schema 段只收普通标识符: 让 `--help` 说的"不支持复杂标识符"成为事实。
|
|
|
|
这不是安全边界(`to_regclass($1)` 参数化 + `_quote` 转义,注入面本就不存在),
|
|
是**契约边界**: 帮助文本写着不支持,实现却照单全收,受害的是照文档做判断的人。
|
|
"""
|
|
|
|
def test_schema_with_a_space_exits_one(self):
|
|
result = _run(
|
|
"--backend",
|
|
"postgres",
|
|
"--dsn",
|
|
"postgresql://x/y",
|
|
"--older-than-days",
|
|
"7",
|
|
"--table",
|
|
"bad schema.llm_calls",
|
|
)
|
|
|
|
assert result.returncode == 1
|
|
assert "--table" in result.stderr
|
|
|
|
def test_schema_with_a_semicolon_exits_one(self):
|
|
result = _run(
|
|
"--backend",
|
|
"postgres",
|
|
"--dsn",
|
|
"postgresql://x/y",
|
|
"--older-than-days",
|
|
"7",
|
|
"--table",
|
|
"a;b.llm_calls",
|
|
)
|
|
|
|
assert result.returncode == 1
|
|
assert "--table" in result.stderr
|
|
|
|
def test_a_plain_identifier_with_underscores_and_digits_is_accepted(self):
|
|
"""收紧不得误伤正常名字: 这条走到连接阶段才失败(退出 2),说明校验放行了。"""
|
|
result = _run(
|
|
"--backend",
|
|
"postgres",
|
|
"--dsn",
|
|
"postgresql://127.0.0.1:1/nope",
|
|
"--older-than-days",
|
|
"7",
|
|
"--table",
|
|
"pgw_s_a1b2c3.llm_calls",
|
|
)
|
|
|
|
assert result.returncode == 2
|
|
|
|
|
|
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
|
|
|
|
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
|