refactor: make the telemetry schema a single source of truth

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.
This commit is contained in:
2026-08-19 11:18:06 -04:00
parent e9adb36577
commit 1471e0a2c6
5 changed files with 367 additions and 182 deletions
+6 -82
View File
@@ -21,56 +21,11 @@ from typing import TYPE_CHECKING
from loguru import logger
from polygateway.telemetry.schema import COLUMNS, PG_BACKFILL, PG_DDL, insert_sql
if TYPE_CHECKING:
import asyncpg
_DDL = """
CREATE TABLE IF NOT EXISTS 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,
tenant_id TEXT NOT NULL DEFAULT '',
meta JSONB NOT NULL DEFAULT '{}'::jsonb
);
"""
# 新列排在 created_at 之后: 与旧表 ALTER 追加的位置一致(见 sqlite.py 同款注释)
_BACKFILL = (
("cached_prompt_tokens", "ALTER TABLE llm_calls ADD COLUMN cached_prompt_tokens INTEGER"),
("model_reported", "ALTER TABLE llm_calls ADD COLUMN model_reported TEXT"),
("sampling", "ALTER TABLE llm_calls ADD COLUMN sampling TEXT"),
("reasoning_tokens", "ALTER TABLE llm_calls ADD COLUMN reasoning_tokens INTEGER"),
# 两个默认值都是非易失常量,PG 11+ 只改 catalog 不重写全表,故大表补列亦是秒级
(
"tenant_id",
"ALTER TABLE llm_calls ADD COLUMN tenant_id TEXT NOT NULL DEFAULT ''",
),
(
"meta",
"ALTER TABLE llm_calls ADD COLUMN meta JSONB NOT NULL DEFAULT '{}'::jsonb",
),
)
# 探测表是否存在;不需要任何权限,且与 INSERT 走同一套 search_path 解析
_TABLE_EXISTS = "SELECT to_regclass('llm_calls')"
@@ -80,38 +35,7 @@ _EXISTING_COLUMNS = (
"WHERE attrelid = to_regclass('llm_calls') AND attnum > 0 AND NOT attisdropped"
)
_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",
"cached_prompt_tokens",
"model_reported",
"sampling",
"reasoning_tokens",
"tenant_id",
"meta",
)
_INSERT = (
f"INSERT INTO llm_calls ({', '.join(_COLUMNS)}) "
f"VALUES ({', '.join(f'${i + 1}' for i in range(len(_COLUMNS)))}) "
"ON CONFLICT (call_id) DO NOTHING"
)
_INSERT = insert_sql("postgres", COLUMNS)
class PostgresRecorder:
@@ -200,7 +124,7 @@ class PostgresRecorder:
await self._backfill_columns(conn) # 旧表可能缺列;失败只逐行降级
return True
try:
await conn.execute(_DDL) # type: ignore[attr-defined]
await conn.execute(PG_DDL) # type: ignore[attr-defined]
except asyncio.CancelledError:
raise
except Exception as exc:
@@ -223,7 +147,7 @@ class PostgresRecorder:
"""
try:
existing = {row["attname"] for row in await conn.fetch(_EXISTING_COLUMNS)} # type: ignore[attr-defined]
for column, statement in _BACKFILL:
for column, statement in PG_BACKFILL:
if column not in existing:
await conn.execute(statement) # type: ignore[attr-defined]
except asyncio.CancelledError:
@@ -236,7 +160,7 @@ class PostgresRecorder:
pool = await self._ensure_ready()
if pool is None:
return
row = tuple(fields[col] for col in _COLUMNS)
row = tuple(fields[col] for col in COLUMNS)
try:
async with pool.acquire() as conn:
await conn.execute(_INSERT, *row)