fix: probe for the telemetry table before creating it
PostgreSQL checks the schema CREATE privilege before the IF NOT EXISTS existence test, so an account with only table-level INSERT was denied on CREATE TABLE IF NOT EXISTS even though the table was right there and writable. The denial set _failed and the whole recorder went no-op for the process lifetime, silently: 150+ calls downstream lost their latency, token and cost rows with nothing but one warning to show for it. The probe is the direct fix. The larger fix is the criterion: structural degradation now means "provably cannot write" (pool creation failed, or the table is absent and cannot be created), not "something threw during init" -- a probe or acquire failure just skips the row and retries on the next call. SQLite stays as it is on purpose. Measured: it short-circuits the statement at parse time, so it passes even under another connection's EXCLUSIVE lock or on a read-only file. A probe there would buy nothing; the docstring now says so to keep symmetry-minded future edits away.
This commit is contained in:
@@ -13,6 +13,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
@@ -299,3 +300,88 @@ class TestDegradation:
|
||||
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.postgres import _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(_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()
|
||||
|
||||
@@ -257,17 +257,41 @@ class TestSQLiteColumnBackfill:
|
||||
|
||||
|
||||
class _FakePgConn:
|
||||
"""记录执行过的语句;可让 ALTER 抛错以模拟权限不足。"""
|
||||
"""记录执行过的语句;可让 ALTER/CREATE/探测抛错以模拟权限不足与抖动。
|
||||
|
||||
def __init__(self, existing: list[str], *, fail_alter: bool = False):
|
||||
`existing` 为空列表即表示**表不存在**(与真实 PG 一致: `to_regclass` 为 NULL
|
||||
时列探测必然零行),故 `fetchval` 与 `fetch` 共用同一份事实。
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
existing: list[str],
|
||||
*,
|
||||
fail_alter: bool = False,
|
||||
fail_create: bool = False,
|
||||
probe_errors: int = 0,
|
||||
):
|
||||
self.existing = existing
|
||||
self.fail_alter = fail_alter
|
||||
self.fail_create = fail_create
|
||||
self.probe_errors = probe_errors
|
||||
self.statements: list[str] = []
|
||||
|
||||
async def execute(self, sql, *args):
|
||||
self.statements.append(sql)
|
||||
if sql.startswith("ALTER TABLE") and self.fail_alter:
|
||||
raise RuntimeError("must be owner of table llm_calls")
|
||||
if sql.lstrip().startswith("CREATE TABLE"):
|
||||
if self.fail_create:
|
||||
raise RuntimeError("permission denied for schema public")
|
||||
self.existing = list(_EXPECTED_COLUMNS)
|
||||
|
||||
async def fetchval(self, sql, *args):
|
||||
self.statements.append(sql)
|
||||
if self.probe_errors > 0:
|
||||
self.probe_errors -= 1
|
||||
raise RuntimeError("connection was closed in the middle of operation")
|
||||
return "llm_calls" if self.existing else None
|
||||
|
||||
async def fetch(self, sql, *args):
|
||||
self.statements.append(sql)
|
||||
@@ -338,6 +362,76 @@ class TestPostgresBackfillDiscipline:
|
||||
assert all("IF NOT EXISTS" not in s for s in altered) # 探测已确认缺列,无需再判
|
||||
|
||||
|
||||
class TestPostgresTableProbe:
|
||||
"""建表必须先探测,且"判死"只认"确定写不进去"(issue #9)。
|
||||
|
||||
实测(PostgreSQL 16.14,只有表级 SELECT/INSERT 的角色): `CREATE TABLE IF NOT
|
||||
EXISTS` 被拒 permission denied for schema,而同一连接的 `INSERT` 通过——
|
||||
PG 对 schema 的 CREATE 权限检查早于 `IF NOT EXISTS` 的存在性判断。无条件发
|
||||
DDL 会让这类最小权限部署的整个进程静默失遥测。
|
||||
"""
|
||||
|
||||
_CURRENT = [
|
||||
"call_id",
|
||||
"cost",
|
||||
"created_at",
|
||||
"cached_prompt_tokens",
|
||||
"model_reported",
|
||||
"sampling",
|
||||
"reasoning_tokens",
|
||||
]
|
||||
|
||||
def _recorder(self, conn):
|
||||
from polygateway.telemetry.postgres import PostgresRecorder
|
||||
|
||||
return PostgresRecorder("postgresql://u:p@h:5432/polygateway", pool=_FakePgPool(conn))
|
||||
|
||||
def _created(self, conn):
|
||||
return [s for s in conn.statements if s.lstrip().startswith("CREATE TABLE")]
|
||||
|
||||
async def test_existing_table_is_never_recreated(self):
|
||||
"""表已存在就一条 DDL 都不发——这是权限被拒的唯一根治办法。"""
|
||||
conn = _FakePgConn(self._CURRENT)
|
||||
await _record_minimal(self._recorder(conn))
|
||||
assert not self._created(conn)
|
||||
|
||||
async def test_create_denied_on_existing_table_keeps_recording(self):
|
||||
"""就算 DDL 仍被发出并被拒,表存在时也不得判死整个 recorder。"""
|
||||
conn = _FakePgConn(self._CURRENT, fail_create=True)
|
||||
recorder = self._recorder(conn)
|
||||
await _record_minimal(recorder) # 不得抛
|
||||
assert recorder._failed is False
|
||||
assert any(s.startswith("INSERT INTO llm_calls") for s in conn.statements)
|
||||
|
||||
async def test_missing_table_is_created_and_not_backfilled(self):
|
||||
"""表不存在→建表;新建表列已齐全,不得再发补列 ALTER。"""
|
||||
conn = _FakePgConn([])
|
||||
recorder = self._recorder(conn)
|
||||
await _record_minimal(recorder)
|
||||
assert len(self._created(conn)) == 1
|
||||
assert not [s for s in conn.statements if s.startswith("ALTER TABLE")]
|
||||
assert recorder._failed is False
|
||||
assert any(s.startswith("INSERT INTO llm_calls") for s in conn.statements)
|
||||
|
||||
async def test_create_failure_on_missing_table_degrades_to_noop(self):
|
||||
"""表确定不存在且建不出来 = 确定写不进去: 此时才允许永久 no-op。"""
|
||||
conn = _FakePgConn([], fail_create=True)
|
||||
recorder = self._recorder(conn)
|
||||
await _record_minimal(recorder) # 不得抛
|
||||
assert recorder._failed is True
|
||||
assert not [s for s in conn.statements if s.startswith("INSERT INTO llm_calls")]
|
||||
|
||||
async def test_probe_failure_is_transient_not_terminal(self):
|
||||
"""探测失败多为连接抖动: 跳过本次,下次调用必须重试,绝不永久判死。"""
|
||||
conn = _FakePgConn(self._CURRENT, probe_errors=1)
|
||||
recorder = self._recorder(conn)
|
||||
await _record_minimal(recorder, call_id="first") # 不得抛
|
||||
assert recorder._failed is False
|
||||
assert not [s for s in conn.statements if s.startswith("INSERT INTO llm_calls")]
|
||||
await _record_minimal(recorder, call_id="second")
|
||||
assert [s for s in conn.statements if s.startswith("INSERT INTO llm_calls")]
|
||||
|
||||
|
||||
class _MemoryRecorder:
|
||||
def __init__(self):
|
||||
self.rows = []
|
||||
|
||||
Reference in New Issue
Block a user