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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user