iomgaa
1471e0a2c6
DDL, column order and backfill statements lived twice, once in each recorder. A public telemetry_schema_sql() would have made three copies, and the drift shows up downstream as "I ran the printed SQL and the library still reports a missing column". Move both DDLs, both backfill lists and the 24 INSERT fields into telemetry/schema.py verbatim; the recorders now import them and build _INSERT through insert_sql(backend, COLUMNS) at import time. The generated statements are byte-identical to the previous constants, so runtime behaviour is unchanged (the postgres conflict target stays bound to call_id for now). insert_sql() validates its columns against COLUMNS: from the next task on those names come from database probing, not from a constant, so the subset check is the gate on the only injection surface. The new telemetry_schema_sql() prints a paste-ready migration script; its postgres backfill deliberately uses ADD COLUMN IF NOT EXISTS while the library's own statements do not, because that form takes an ACCESS EXCLUSIVE lock even when the column exists. Both variants are derived from one declaration list so their column sets cannot drift.
631 lines
24 KiB
Python
631 lines
24 KiB
Python
"""PostgresRecorder 集成测试(M2 设计 §5;真实实验室 Postgres,polygateway 专用库)。
|
|
|
|
DSN 走 .env `PGW_TELEMETRY_PG_DSN`,缺则 skip。该实例上有 app/chs_prod 等
|
|
在用库——本测试只允许连 polygateway 专用库(fixture 里守卫)。
|
|
|
|
隔离纪律(M4 事故教训): `llm_calls` 是与真实批跑/迁移项目共享的表,
|
|
**严禁 DROP/TRUNCATE**——本测试以 run 级 call_id 前缀隔离,断言只看
|
|
自己写入的行,teardown 只删自己的行。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import re
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from dotenv import dotenv_values
|
|
|
|
from polygateway.telemetry.postgres import PostgresRecorder
|
|
|
|
_EXPECTED_COLUMNS = [
|
|
"call_id",
|
|
"parent_call_id",
|
|
"session_id",
|
|
"model",
|
|
"provider",
|
|
"source_name",
|
|
"messages",
|
|
"response",
|
|
"thinking",
|
|
"prompt_tokens",
|
|
"completion_tokens",
|
|
"usage_source",
|
|
"latency_ms",
|
|
"ttft_ms",
|
|
"max_inter_token_ms",
|
|
"cache_hit",
|
|
"error",
|
|
"cost",
|
|
"created_at",
|
|
"cached_prompt_tokens",
|
|
"model_reported",
|
|
"sampling",
|
|
"reasoning_tokens",
|
|
"tenant_id",
|
|
"meta",
|
|
]
|
|
|
|
# run 级前缀: 同库并存的其他运行(迁移批跑/另一开发机)互不可见
|
|
_RUN_PREFIX = f"pgwtest-{uuid4().hex[:8]}"
|
|
|
|
|
|
def _cid(suffix: str) -> str:
|
|
return f"{_RUN_PREFIX}-{suffix}"
|
|
|
|
|
|
def _dsn() -> 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}"
|
|
|
|
|
|
@pytest.fixture
|
|
async def dsn():
|
|
value = _dsn()
|
|
if value is None:
|
|
pytest.skip("PGW_TELEMETRY_PG_DSN 未配置")
|
|
# 隔离守卫: 该实例有 app/chs_prod/mimiciv 等在用库,只许打 polygateway 专用库
|
|
if not value.rstrip("/").endswith("/polygateway"):
|
|
pytest.fail(f"遥测测试只允许连 polygateway 专用库,当前 DSN 库名不符: {value!r}")
|
|
yield value
|
|
# teardown: 只删本 run 写入的行;表可能尚不存在(全新库)则忽略
|
|
import asyncpg
|
|
|
|
conn = await asyncpg.connect(value, timeout=10)
|
|
try:
|
|
if await conn.fetchval("SELECT to_regclass('llm_calls')") is not None:
|
|
await conn.execute("DELETE FROM llm_calls WHERE call_id LIKE $1", f"{_RUN_PREFIX}-%")
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
async def _record_minimal(
|
|
recorder: PostgresRecorder, call_id: str | None = None, **overrides
|
|
) -> None:
|
|
fields = {
|
|
"call_id": call_id if call_id is not None else _cid("c1"),
|
|
"parent_call_id": None,
|
|
"session_id": "sess-1",
|
|
"model": "m",
|
|
"provider": "p",
|
|
"source_name": "s1",
|
|
"messages": "[]",
|
|
"response": "ok",
|
|
"thinking": "",
|
|
"prompt_tokens": 1,
|
|
"completion_tokens": 2,
|
|
"usage_source": "measured",
|
|
"latency_ms": 10,
|
|
"ttft_ms": None,
|
|
"max_inter_token_ms": None,
|
|
"cache_hit": False,
|
|
"error": None,
|
|
"cost": None,
|
|
"cached_prompt_tokens": None,
|
|
"model_reported": None,
|
|
"sampling": None,
|
|
"reasoning_tokens": None,
|
|
# 到达 recorder 时已由 emitter 归一化: None → '',空 dict → '{}'
|
|
"tenant_id": "",
|
|
"meta": "{}",
|
|
}
|
|
fields.update(overrides)
|
|
await recorder.record_llm_call(**fields)
|
|
|
|
|
|
async def _fetch(dsn: str, sql: str, *args):
|
|
import asyncpg
|
|
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
return await conn.fetch(sql, *args)
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
_LEGACY_DDL = """
|
|
CREATE TABLE {schema}.llm_calls (
|
|
call_id TEXT PRIMARY KEY,
|
|
parent_call_id TEXT,
|
|
session_id TEXT,
|
|
model TEXT NOT NULL,
|
|
provider TEXT NOT NULL,
|
|
source_name TEXT NOT NULL,
|
|
messages TEXT NOT NULL,
|
|
response TEXT NOT NULL,
|
|
thinking TEXT NOT NULL DEFAULT '',
|
|
prompt_tokens INTEGER NOT NULL,
|
|
completion_tokens INTEGER NOT NULL,
|
|
usage_source TEXT NOT NULL,
|
|
latency_ms INTEGER NOT NULL,
|
|
ttft_ms DOUBLE PRECISION,
|
|
max_inter_token_ms DOUBLE PRECISION,
|
|
cache_hit BOOLEAN NOT NULL DEFAULT FALSE,
|
|
error TEXT,
|
|
cost DOUBLE PRECISION,
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
|
)
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
async def legacy_schema(dsn):
|
|
"""在**自建的临时 schema** 里造一张 18 列旧表,验证补列(issue #3)。
|
|
|
|
绝不碰共享的 public.llm_calls: 用 search_path 把 recorder 指向临时 schema,
|
|
teardown 只 DROP 自己建的 schema。
|
|
"""
|
|
import asyncpg
|
|
|
|
name = f"pgwtest_{uuid4().hex[:8]}"
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await conn.execute(f"CREATE SCHEMA {name}")
|
|
await conn.execute(_LEGACY_DDL.format(schema=name))
|
|
finally:
|
|
await conn.close()
|
|
sep = "&" if "?" in dsn else "?"
|
|
yield f"{dsn}{sep}options=-csearch_path%3D{name}", name
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await conn.execute(f"DROP SCHEMA {name} CASCADE")
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
class TestObservabilityColumns:
|
|
"""issue #3: 两列写入可回读,且已存在的 18 列旧表会被自动补列。"""
|
|
|
|
async def test_values_round_trip(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await _record_minimal(recorder, call_id=_cid("hit"), cached_prompt_tokens=64)
|
|
await _record_minimal(recorder, call_id=_cid("zero"), cached_prompt_tokens=0)
|
|
await _record_minimal(recorder, call_id=_cid("model"), model_reported="MiniMax-01")
|
|
await _record_minimal(
|
|
recorder, call_id=_cid("samp"), sampling='{"seed": 42, "temperature": 0}'
|
|
)
|
|
rows = await _fetch(
|
|
dsn,
|
|
"SELECT call_id, cached_prompt_tokens, model_reported, sampling FROM llm_calls "
|
|
"WHERE call_id LIKE $1",
|
|
f"{_RUN_PREFIX}-%",
|
|
)
|
|
by_id = {r["call_id"]: r for r in rows}
|
|
assert by_id[_cid("hit")]["cached_prompt_tokens"] == 64
|
|
assert by_id[_cid("zero")]["cached_prompt_tokens"] == 0 # 真实零命中 ≠ NULL
|
|
assert by_id[_cid("model")]["cached_prompt_tokens"] is None
|
|
assert by_id[_cid("model")]["model_reported"] == "MiniMax-01"
|
|
# issue #4: PG 侧也须验非空 sampling 能读回原值(不只是列存在)
|
|
assert json.loads(by_id[_cid("samp")]["sampling"]) == {"seed": 42, "temperature": 0}
|
|
assert by_id[_cid("hit")]["sampling"] is None
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_legacy_table_is_upgraded_in_place(self, legacy_schema):
|
|
"""18 列旧表不补列的话,每行写入都会被逐行 warning 丢弃(遥测静默全失)。"""
|
|
schema_dsn, schema = legacy_schema
|
|
recorder = PostgresRecorder(schema_dsn)
|
|
try:
|
|
await _record_minimal(
|
|
recorder, call_id=_cid("legacy"), cached_prompt_tokens=7, model_reported="m-real"
|
|
)
|
|
cols = await _fetch(
|
|
schema_dsn,
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_schema = $1 AND table_name = 'llm_calls' ORDER BY ordinal_position",
|
|
schema,
|
|
)
|
|
# ALTER 只能追加到末尾: 与新建库的列序一致才不会分叉
|
|
assert [r["column_name"] for r in cols] == _EXPECTED_COLUMNS
|
|
rows = await _fetch(
|
|
schema_dsn,
|
|
"SELECT cached_prompt_tokens, model_reported FROM llm_calls WHERE call_id = $1",
|
|
_cid("legacy"),
|
|
)
|
|
assert (rows[0]["cached_prompt_tokens"], rows[0]["model_reported"]) == (7, "m-real")
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
|
|
class TestSchema:
|
|
async def test_schema_has_frozen_columns_in_order(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await _record_minimal(recorder)
|
|
rows = await _fetch(
|
|
dsn,
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_name='llm_calls' ORDER BY ordinal_position",
|
|
)
|
|
assert [r["column_name"] for r in rows] == _EXPECTED_COLUMNS
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_call_id_idempotent(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await _record_minimal(recorder, call_id=_cid("dup"))
|
|
await _record_minimal(recorder, call_id=_cid("dup"), response="second")
|
|
rows = await _fetch(
|
|
dsn, "SELECT response FROM llm_calls WHERE call_id = $1", _cid("dup")
|
|
)
|
|
assert [r["response"] for r in rows] == ["ok"] # ON CONFLICT DO NOTHING
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_concurrent_writes_all_land(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await asyncio.gather(
|
|
*(_record_minimal(recorder, call_id=_cid(f"c{i}")) for i in range(50))
|
|
)
|
|
rows = await _fetch(
|
|
dsn,
|
|
"SELECT count(*) AS n FROM llm_calls WHERE call_id LIKE $1",
|
|
f"{_RUN_PREFIX}-c%",
|
|
)
|
|
assert rows[0]["n"] == 50
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
|
|
class TestDegradation:
|
|
async def test_unreachable_server_degrades_silently(self):
|
|
"""结构性失败(建池不通)→ warning 一次后永久降级,业务零感知。"""
|
|
recorder = PostgresRecorder("postgresql://u:p@127.0.0.1:1/x")
|
|
await _record_minimal(recorder) # 不抛
|
|
await _record_minimal(recorder, call_id=_cid("c2")) # 已降级短路,同样不抛
|
|
await recorder.aclose()
|
|
|
|
async def test_row_failure_does_not_poison_later_rows(self, dsn):
|
|
"""运行时单条写失败(NUL 字节文本被 PG 拒)→ 丢该行,后续行照常落库。"""
|
|
recorder = PostgresRecorder(dsn)
|
|
try:
|
|
await _record_minimal(recorder, call_id=_cid("bad"), response="nul\x00byte")
|
|
await _record_minimal(recorder, call_id=_cid("good"))
|
|
rows = await _fetch(
|
|
dsn,
|
|
"SELECT call_id FROM llm_calls WHERE call_id = ANY($1::text[]) ORDER BY call_id",
|
|
[_cid("bad"), _cid("good")],
|
|
)
|
|
assert [r["call_id"] for r in rows] == [_cid("good")]
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_aclose_idempotent(self, dsn):
|
|
recorder = PostgresRecorder(dsn)
|
|
await _record_minimal(recorder)
|
|
await recorder.aclose()
|
|
await recorder.aclose()
|
|
|
|
|
|
_PROBE_PASSWORD = "pgw_issue9_probe" # 临时角色,teardown 删除;非任何真实凭据
|
|
|
|
|
|
@pytest.fixture
|
|
async def least_privilege_dsn(dsn):
|
|
"""临时 schema + 临时角色: 只授表级 SELECT/INSERT,**不授 schema CREATE**。
|
|
|
|
这是 issue #9 的现场——最小权限部署的标准形态。fixture 建的一切
|
|
(schema、表、角色)都在 teardown 里删净,共享的 public.llm_calls 不受影响;
|
|
连不上或无权建角色(非超级用户)时 skip,不让 CI 假绿。
|
|
"""
|
|
import asyncpg
|
|
|
|
from polygateway.telemetry.schema import PG_DDL
|
|
|
|
name = f"pgwtest_lp_{uuid4().hex[:8]}"
|
|
admin = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
if not await admin.fetchval(
|
|
"SELECT rolcreaterole OR rolsuper FROM pg_roles WHERE rolname = current_user"
|
|
):
|
|
pytest.skip("当前账号无权建临时角色,跳过最小权限用例")
|
|
await admin.execute(f"CREATE ROLE {name} LOGIN PASSWORD '{_PROBE_PASSWORD}'")
|
|
await admin.execute(f"CREATE SCHEMA {name}")
|
|
await admin.execute(f"SET search_path = {name}")
|
|
await admin.execute(PG_DDL) # 表由**别的账号**建好,与现场一致
|
|
await admin.execute(f"GRANT USAGE ON SCHEMA {name} TO {name}")
|
|
await admin.execute(f"GRANT SELECT, INSERT ON {name}.llm_calls TO {name}")
|
|
# 关键: 绝不 GRANT CREATE ON SCHEMA —— 缺的正是这一项
|
|
finally:
|
|
await admin.close()
|
|
low = re.sub(r"//[^@/]+@", f"//{name}:{_PROBE_PASSWORD}@", dsn, count=1)
|
|
sep = "&" if "?" in low else "?"
|
|
yield f"{low}{sep}options=-csearch_path%3D{name}", name
|
|
admin = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await admin.execute(f"DROP SCHEMA IF EXISTS {name} CASCADE")
|
|
await admin.execute(f"DROP OWNED BY {name}")
|
|
await admin.execute(f"DROP ROLE IF EXISTS {name}")
|
|
finally:
|
|
await admin.close()
|
|
|
|
|
|
class TestLeastPrivilegeDeployment:
|
|
"""issue #9: 只有表级写权限的账号,遥测必须照常落库而不是整体判死。"""
|
|
|
|
async def test_create_table_if_not_exists_is_denied_for_this_role(self, least_privilege_dsn):
|
|
"""库外事实先钉死: 表存在、写得进去,DDL 仍被拒——PG 的权限检查早于 IF NOT EXISTS。
|
|
|
|
修复依赖的是这条 PG 语义;若某天它变了,这里先红,而不是让下面那条
|
|
用例悄悄变成"永远通过"的空断言。
|
|
"""
|
|
import asyncpg
|
|
|
|
low_dsn, _ = least_privilege_dsn
|
|
conn = await asyncpg.connect(low_dsn, timeout=10)
|
|
try:
|
|
assert await conn.fetchval("SELECT to_regclass('llm_calls')") is not None
|
|
with pytest.raises(asyncpg.exceptions.InsufficientPrivilegeError):
|
|
await conn.execute("CREATE TABLE IF NOT EXISTS llm_calls (call_id TEXT)")
|
|
finally:
|
|
await conn.close()
|
|
|
|
async def test_records_land_without_schema_create_privilege(self, least_privilege_dsn):
|
|
"""修复前: 建表被拒 → _failed → 整个进程一条不落(下游 150 次调用全丢)。"""
|
|
low_dsn, schema = least_privilege_dsn
|
|
recorder = PostgresRecorder(low_dsn)
|
|
try:
|
|
await _record_minimal(recorder, call_id=_cid("lp1"))
|
|
await _record_minimal(recorder, call_id=_cid("lp2"), cost=1.5)
|
|
assert recorder._failed is False # 判死开关不得被建表权限触发
|
|
rows = await _fetch(
|
|
low_dsn,
|
|
"SELECT call_id, cost FROM llm_calls WHERE call_id LIKE $1 ORDER BY call_id",
|
|
f"{_RUN_PREFIX}-lp%",
|
|
)
|
|
assert [(r["call_id"], r["cost"]) for r in rows] == [
|
|
(_cid("lp1"), None),
|
|
(_cid("lp2"), 1.5),
|
|
]
|
|
assert schema # teardown 会连表带角色删净
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
|
|
# issue #11 之前的表形态: 22 个 recorder 字段 + created_at = 23 个物理列,没有任何租户维度
|
|
_PRE_TENANT_DDL = """
|
|
CREATE TABLE {schema}.llm_calls (
|
|
call_id TEXT PRIMARY KEY,
|
|
parent_call_id TEXT,
|
|
session_id TEXT,
|
|
model TEXT NOT NULL,
|
|
provider TEXT NOT NULL,
|
|
source_name TEXT NOT NULL,
|
|
messages TEXT NOT NULL,
|
|
response TEXT NOT NULL,
|
|
thinking TEXT NOT NULL DEFAULT '',
|
|
prompt_tokens INTEGER NOT NULL,
|
|
completion_tokens INTEGER NOT NULL,
|
|
usage_source TEXT NOT NULL,
|
|
latency_ms INTEGER NOT NULL,
|
|
ttft_ms DOUBLE PRECISION,
|
|
max_inter_token_ms DOUBLE PRECISION,
|
|
cache_hit BOOLEAN NOT NULL DEFAULT FALSE,
|
|
error TEXT,
|
|
cost DOUBLE PRECISION,
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
|
cached_prompt_tokens INTEGER,
|
|
model_reported TEXT,
|
|
sampling TEXT,
|
|
reasoning_tokens INTEGER
|
|
)
|
|
"""
|
|
|
|
_PRE_TENANT_INSERT = (
|
|
"INSERT INTO {schema}.llm_calls (call_id, model, provider, source_name, messages, response, "
|
|
"prompt_tokens, completion_tokens, usage_source, latency_ms) "
|
|
"VALUES ($1, 'm', 'p', 's1', '[]', 'old body', 1, 2, 'measured', 10)"
|
|
)
|
|
|
|
|
|
def _search_path_dsn(dsn: str, schema: str) -> str:
|
|
sep = "&" if "?" in dsn else "?"
|
|
return f"{dsn}{sep}options=-csearch_path%3D{schema}"
|
|
|
|
|
|
@pytest.fixture
|
|
async def captured_warnings():
|
|
"""捕获库发出的 WARNING;loguru 不经标准 logging,pytest 的 caplog 抓不到。
|
|
|
|
名字避开裸 `warnings`: 那会遮蔽标准库模块名,本文件将来任何一次
|
|
`import warnings` 都会与它静默互相顶掉,而报错点离真因很远。
|
|
"""
|
|
from loguru import logger
|
|
|
|
messages: list[str] = []
|
|
sink_id = logger.add(messages.append, level="WARNING")
|
|
yield messages
|
|
logger.remove(sink_id)
|
|
|
|
|
|
@pytest.fixture
|
|
async def pre_tenant_schema(dsn):
|
|
"""自建临时 schema 里造一张 **22 字段的 issue #11 之前的表**,并留一行历史数据。
|
|
|
|
绝不碰共享的 public.llm_calls——本机那张表早已被 `_BACKFILL` 真实补过列,
|
|
指望它还是旧形态的测试第二次跑就会空转。schema 名带 uuid,可重复运行。
|
|
"""
|
|
import asyncpg
|
|
|
|
name = f"pgwtest_pre_{uuid4().hex[:8]}"
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await conn.execute(f"CREATE SCHEMA {name}")
|
|
await conn.execute(_PRE_TENANT_DDL.format(schema=name))
|
|
await conn.execute(_PRE_TENANT_INSERT.format(schema=name), _cid("old"))
|
|
finally:
|
|
await conn.close()
|
|
yield _search_path_dsn(dsn, name), name
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await conn.execute(f"DROP SCHEMA {name} CASCADE")
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
@pytest.fixture
|
|
async def fresh_schema(dsn):
|
|
"""空 schema: recorder 自己建表,验"新建库"这条路径而不依赖共享表的历史状态。"""
|
|
import asyncpg
|
|
|
|
name = f"pgwtest_new_{uuid4().hex[:8]}"
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await conn.execute(f"CREATE SCHEMA {name}")
|
|
finally:
|
|
await conn.close()
|
|
yield _search_path_dsn(dsn, name), name
|
|
conn = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await conn.execute(f"DROP SCHEMA {name} CASCADE")
|
|
finally:
|
|
await conn.close()
|
|
|
|
|
|
@pytest.fixture
|
|
async def least_privilege_pre_tenant_dsn(dsn):
|
|
"""22 字段旧表 + 只有 `SELECT, INSERT` 权限的角色: 补列必然失败的现场。
|
|
|
|
与 `least_privilege_dsn` 分开而非复用: 那个 fixture 建的是列已齐全的当前表
|
|
(测的是 CREATE 被拒),这里必须是缺列的旧表,才能让 `ALTER TABLE` 真的发出去
|
|
并撞上 ownership 检查(该检查早于 `IF NOT EXISTS` 的存在性判断)。
|
|
"""
|
|
import asyncpg
|
|
|
|
name = f"pgwtest_lppre_{uuid4().hex[:8]}"
|
|
admin = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
if not await admin.fetchval(
|
|
"SELECT rolcreaterole OR rolsuper FROM pg_roles WHERE rolname = current_user"
|
|
):
|
|
pytest.skip("当前账号无权建临时角色,跳过最小权限用例")
|
|
await admin.execute(f"CREATE ROLE {name} LOGIN PASSWORD '{_PROBE_PASSWORD}'")
|
|
await admin.execute(f"CREATE SCHEMA {name}")
|
|
await admin.execute(_PRE_TENANT_DDL.format(schema=name)) # 表属主是 admin,不是应用账号
|
|
await admin.execute(f"GRANT USAGE ON SCHEMA {name} TO {name}")
|
|
await admin.execute(f"GRANT SELECT, INSERT ON {name}.llm_calls TO {name}")
|
|
finally:
|
|
await admin.close()
|
|
low = re.sub(r"//[^@/]+@", f"//{name}:{_PROBE_PASSWORD}@", dsn, count=1)
|
|
yield _search_path_dsn(low, name)
|
|
admin = await asyncpg.connect(dsn, timeout=10)
|
|
try:
|
|
await admin.execute(f"DROP SCHEMA IF EXISTS {name} CASCADE")
|
|
await admin.execute(f"DROP OWNED BY {name}")
|
|
await admin.execute(f"DROP ROLE IF EXISTS {name}")
|
|
finally:
|
|
await admin.close()
|
|
|
|
|
|
class TestCallerDimensionsAcceptance:
|
|
"""issue #11 的机械化验收(PG 侧,真实实例): 新建库 / 旧表补列 / 补列失败方向。"""
|
|
|
|
async def test_fresh_schema_round_trips_the_dimensions(self, fresh_schema):
|
|
"""新建库: 列齐全,且维度值原样读回——只验列存在会漏掉写错列位的错。"""
|
|
fresh_dsn, schema = fresh_schema
|
|
recorder = PostgresRecorder(fresh_dsn)
|
|
try:
|
|
await _record_minimal(
|
|
recorder, call_id=_cid("dim"), tenant_id="tenant-a", meta='{"batch": "b7"}'
|
|
)
|
|
cols = await _fetch(
|
|
fresh_dsn,
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_schema = $1 AND table_name = 'llm_calls' ORDER BY ordinal_position",
|
|
schema,
|
|
)
|
|
assert [r["column_name"] for r in cols] == _EXPECTED_COLUMNS
|
|
rows = await _fetch(
|
|
fresh_dsn,
|
|
"SELECT tenant_id, meta FROM llm_calls WHERE call_id = $1",
|
|
_cid("dim"),
|
|
)
|
|
assert rows[0]["tenant_id"] == "tenant-a"
|
|
assert json.loads(rows[0]["meta"]) == {"batch": "b7"}
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_pre_tenant_table_gains_columns_and_old_rows_stay_auditable(
|
|
self, pre_tenant_schema
|
|
):
|
|
"""22 字段旧表补列后,新行带维度,而**老行的 tenant_id 是空串而非 NULL**。
|
|
|
|
这条直接验收 issue #11 的核心论点(先启用落库、后加列,补列之前的行没有
|
|
租户归属)。断言方向必须是空串: PG 的 RLS `USING` 表达式对返回 false **或
|
|
NULL** 的行一律隐藏且不报错,故 NULL 的 `tenant_id` 不是"未归属",而是对
|
|
所有人永久不可见的黑洞;哨兵空串则能被一条 `COUNT(*) WHERE tenant_id = ''`
|
|
审计出来,历史欠账是可见、可量化、可补录的。
|
|
"""
|
|
schema_dsn, schema = pre_tenant_schema
|
|
recorder = PostgresRecorder(schema_dsn)
|
|
try:
|
|
await _record_minimal(
|
|
recorder, call_id=_cid("new"), tenant_id="tenant-a", meta='{"k": 1}'
|
|
)
|
|
cols = await _fetch(
|
|
schema_dsn,
|
|
"SELECT column_name FROM information_schema.columns "
|
|
"WHERE table_schema = $1 AND table_name = 'llm_calls' ORDER BY ordinal_position",
|
|
schema,
|
|
)
|
|
# 22 → 24 个 recorder 字段(加 created_at 共 25 个物理列),且新列追加在末尾
|
|
assert [r["column_name"] for r in cols] == _EXPECTED_COLUMNS
|
|
rows = await _fetch(
|
|
schema_dsn,
|
|
"SELECT call_id, tenant_id, meta FROM llm_calls "
|
|
"WHERE call_id = ANY($1::text[]) ORDER BY call_id",
|
|
[_cid("new"), _cid("old")],
|
|
)
|
|
by_id = {r["call_id"]: r for r in rows}
|
|
assert by_id[_cid("new")]["tenant_id"] == "tenant-a"
|
|
assert json.loads(by_id[_cid("new")]["meta"]) == {"k": 1}
|
|
assert by_id[_cid("old")]["tenant_id"] == "" # 不是 None: NULL 会被 RLS 静默吞掉
|
|
assert json.loads(by_id[_cid("old")]["meta"]) == {}
|
|
finally:
|
|
await recorder.aclose()
|
|
|
|
async def test_alter_is_denied_for_a_role_that_can_still_insert(
|
|
self, least_privilege_pre_tenant_dsn
|
|
):
|
|
"""库外事实先钉死: 表存在、写得进去,补列的 ALTER 仍被拒(ownership 检查早于存在性判断)。
|
|
|
|
没有这条,下面那个降级用例可能因为 ALTER 其实成功了而变成"永远通过"的空断言。
|
|
"""
|
|
import asyncpg
|
|
|
|
conn = await asyncpg.connect(least_privilege_pre_tenant_dsn, timeout=10)
|
|
try:
|
|
assert await conn.fetchval("SELECT to_regclass('llm_calls')") is not None
|
|
with pytest.raises(asyncpg.exceptions.InsufficientPrivilegeError):
|
|
await conn.execute("ALTER TABLE llm_calls ADD COLUMN IF NOT EXISTS tenant_id TEXT")
|
|
finally:
|
|
await conn.close()
|
|
|
|
async def test_backfill_failure_degrades_per_row_not_wholesale(
|
|
self, least_privilege_pre_tenant_dsn, captured_warnings
|
|
):
|
|
"""补列失败的降级方向: 记 warning、不置 `_failed`、后续 INSERT 仍照发。
|
|
|
|
置 `_failed` 会让整个进程从此一条遥测都不写(比逐行丢弃严重得多),
|
|
且一旦 DBA 补上列也不会自愈——必须等重启。
|
|
"""
|
|
recorder = PostgresRecorder(least_privilege_pre_tenant_dsn)
|
|
try:
|
|
await _record_minimal(recorder, call_id=_cid("lpp1")) # 不得抛
|
|
assert recorder._failed is False
|
|
assert any("补列失败" in m for m in captured_warnings)
|
|
# 缺列的表上 INSERT 必然失败;逐行 warning 正是"INSERT 照发了"的证据
|
|
assert any("写入失败" in m for m in captured_warnings)
|
|
finally:
|
|
await recorder.aclose()
|