"""`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"]