PolyGateway/src/polygateway/telemetry/schema.py

"""遥测表 `llm_calls` 的 schema 单一事实源: 列序、两端 DDL、补列语句与 INSERT 构造。

两个 recorder(`sqlite.py` / `postgres.py`)与公共函数 `telemetry_schema_sql` 共用本模块。
收敛的理由是**正确性**而非整洁: 打印给下游的 SQL 必须与库真正执行的 DDL 同源——常量在
多处各存一份必然漂移,而漂移的表现是"下游照打印的 SQL 建完表,库仍报缺列"。

**`COLUMNS` 是 INSERT 字段序,不是物理列序**: 数据库自填的 `created_at` 不在其中(它带
`DEFAULT now()` / `datetime('now')`,库从不显式写它)。物理表列 = 24 个 INSERT 字段 +
`created_at` = 25;列数断言一律按物理列数写,两套口径混用是最易错处。

本模块只依赖标准库: `telemetry/` 与 `backends/`、`transports/`、`structured/` 同层且
互不依赖(import-linter 契约执法)。
"""

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Sequence

TABLE = "llm_calls"

# 支持的后端;`insert_sql` / `telemetry_schema_sql` 的取值域
_BACKENDS = ("sqlite", "postgres")

SQLITE_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            REAL,
    max_inter_token_ms REAL,
    cache_hit          INTEGER NOT NULL DEFAULT 0,
    error              TEXT,
    cost               REAL,
    created_at         TEXT    NOT NULL DEFAULT (datetime('now')),
    cached_prompt_tokens INTEGER,
    model_reported       TEXT,
    sampling             TEXT,
    reasoning_tokens     INTEGER,
    tenant_id            TEXT    NOT NULL DEFAULT '',
    meta                 TEXT    NOT NULL DEFAULT '{}'
);
"""

PG_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_BACKFILL = (
    ("cached_prompt_tokens", "INTEGER"),
    ("model_reported", "TEXT"),
    ("sampling", "TEXT"),
    ("reasoning_tokens", "INTEGER"),
    # NOT NULL 补列必须带非 NULL 常量默认值,否则 SQLite 直接拒绝该 ALTER
    # ("Cannot add a NOT NULL column with default value NULL"),补列全盘失败。
    ("tenant_id", "TEXT NOT NULL DEFAULT ''"),
    ("meta", "TEXT NOT NULL DEFAULT '{}'"),
)

# PG 补列的列定义。语句由此派生成两份文本(见下),使"库内执行的那份"与"打印给
# 下游的那份"的列集合与列定义**无法分叉**——本模块存在的全部理由就是不许漂移。
_PG_BACKFILL_DECLS = (
    ("cached_prompt_tokens", "INTEGER"),
    ("model_reported", "TEXT"),
    ("sampling", "TEXT"),
    ("reasoning_tokens", "INTEGER"),
    # 两个默认值都是非易失常量,PG 11+ 只改 catalog 不重写全表,故大表补列亦是秒级
    ("tenant_id", "TEXT NOT NULL DEFAULT ''"),
    ("meta", "JSONB NOT NULL DEFAULT '{}'::jsonb"),
)

# 新列排在 created_at 之后: 与旧表 ALTER 追加的位置一致(见 SQLITE_BACKFILL 同款注释)。
# **库内执行的这份有意不带 `IF NOT EXISTS`**: PG 对它即便列已存在也会先取 ACCESS
# EXCLUSIVE 锁,而遥测是业务路径上的内联 await,故库侧一律"先探测后 ALTER"
# (postgres.py `_backfill_columns` 记有实测)。给人执行的那份见 `telemetry_schema_sql`。
PG_BACKFILL = tuple(
    (column, f"ALTER TABLE {TABLE} ADD COLUMN {column} {decl}")
    for column, decl in _PG_BACKFILL_DECLS
)

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",
)

_COLUMN_SET = frozenset(COLUMNS)


def insert_sql(backend: str, columns: Sequence[str]) -> str:
    """按给定列构造 INSERT;列必须是 `COLUMNS` 的非空子集,否则 ValueError。

    子集校验是**注入面的闸**: 列名来自数据库探测结果,不是常量,不校验就等于把外部
    字符串拼进 SQL(占位符只保护值,保护不了列名)。空集同样来自探测结果,而
    `INSERT INTO llm_calls () VALUES ()` 两端都语法非法——本函数自己拒,不把这个
    不变量押在调用方身上。sqlite 用 `?`、postgres 用 `$n`,
    两端的重复键处理都不绑定具体约束名(`INSERT OR IGNORE` / `ON CONFLICT`)。

    **PG 的 `ON CONFLICT` 一律不带冲突目标,不得"顺手"补回 `(call_id)`**: PG 要求
    分区表的唯一约束必须包含分区键,按 `created_at` 分区(issue #12 的保留期方案)后
    主键变成 `(call_id, created_at)`,带目标的语句匹配不到任何约束,PG 直接拒收
    ("there is no unique or exclusion constraint matching the ON CONFLICT
    specification"),而遥测写失败只逐行 warning——分区部署下会全线静默丢数据。
    无目标版本在两种表形态上都合法,普通表上语义逐字等价(表上只有主键一个唯一约束)。

    Args:
        backend: `"sqlite"` 或 `"postgres"`。
        columns: 要写入的列,顺序即占位符顺序(调用方须按同序取值)。

    Returns:
        完整的 INSERT 语句。

    Raises:
        ValueError: backend 不在取值域内,columns 为空,或含 `COLUMNS` 之外的列名。
    """
    if backend not in _BACKENDS:
        raise ValueError(f"未知遥测后端 {backend!r}: 只支持 {list(_BACKENDS)}")
    selected = tuple(columns)
    if not selected:
        raise ValueError("遥测 INSERT 至少需要一列: 空列集合会拼出语法非法的 SQL")
    unknown = [column for column in selected if column not in _COLUMN_SET]
    if unknown:
        raise ValueError(f"列名不在遥测 schema 内(拒绝拼进 SQL): {unknown}")
    names = ", ".join(selected)
    if backend == "sqlite":
        placeholders = ", ".join("?" for _ in selected)
        return f"INSERT OR IGNORE INTO {TABLE} ({names}) VALUES ({placeholders})"
    placeholders = ", ".join(f"${i + 1}" for i in range(len(selected)))
    return f"INSERT INTO {TABLE} ({names}) VALUES ({placeholders}) ON CONFLICT DO NOTHING"


# 缺列告警要打印的补列语句: 库内执行的那份怎么写,打印给人的就怎么写(同源不许漂移)。
# SQLite 侧常量只有列定义,故在此按 TABLE 拼成整条 ALTER;PG 侧常量本就是整条语句。
_ALTER_BY_BACKEND = {
    "sqlite": {
        column: f"ALTER TABLE {TABLE} ADD COLUMN {column} {decl}"
        for column, decl in SQLITE_BACKFILL
    },
    "postgres": dict(PG_BACKFILL),
}

_BACKEND_LABELS = {"sqlite": "SQLite", "postgres": "Postgres"}

# PG 的 ALTER 取 ACCESS EXCLUSIVE 锁,执行时机得由 DBA 自己挑;SQLite 是下游本地文件,无此顾虑
_EXECUTION_NOTES = {"sqlite": "", "postgres": "(建议挑低峰,ALTER 取 ACCESS EXCLUSIVE 锁)"}


def missing_columns_warning(backend: str, missing: Sequence[str], *, alien_table: bool) -> str:
    """拼 manual 档的缺列告警: 逐列点名 + 讲清后果 + 给出可直接执行的 SQL。

    只说"缺列"是不够的: 静默丢维度的后果是多租户账目全归空串且无任何报错,
    看告警的人必须一眼看到丢的是哪几个维度、以及怎么补。

    **住在本模块而不是两个 recorder 里**: 这条消息拼的是给人执行的 DDL,与库自己
    执行的 ALTER 必须同源——本模块存在的全部理由就是不许这两者漂移。

    Args:
        backend: `"sqlite"` 或 `"postgres"`。
        missing: 缺失的列名(按 `COLUMNS` 保序)。
        alien_table: 连主键列 `call_id` 都没有——该表多半不是本库的 `llm_calls`。

    Returns:
        单条 warning 的完整文本(库只在准备期发一次,不逐行发)。

    Raises:
        ValueError: backend 不在取值域内。
    """
    if backend not in _BACKENDS:
        raise ValueError(f"未知遥测后端 {backend!r}: 只支持 {list(_BACKENDS)}")
    alters = _ALTER_BY_BACKEND[backend]
    selected = tuple(missing)
    statements = [f"{alters[column]};" for column in selected if column in alters]
    unknown = [column for column in selected if column not in alters]
    if unknown:
        # 这些列本库从未经 ALTER 补过(建表即有),给不出单条 ALTER,指向完整脚本
        statements.append(
            f"-- 另缺 {', '.join(unknown)};完整建表脚本见 "
            f'polygateway.telemetry_schema_sql("{backend}")'
        )
    label = _BACKEND_LABELS[backend]
    head = (
        f"{label} 遥测表 {TABLE} 缺主键列 call_id,很可能不是本库的遥测表"
        "(库不做二次判定,仍照常尝试写入)"
        if alien_table
        else f"{label} 遥测表 {TABLE} 缺列,且 auto_migrate=False(库不发任何 DDL)"
    )
    return (
        f"{head};以下维度不会被记录: {', '.join(selected)}。"
        f"补列请自行执行{_EXECUTION_NOTES[backend]}:\n" + "\n".join(statements)
    )


def telemetry_schema_sql(backend: str) -> str:
    """返回可直接粘进迁移文件的完整脚本(建表 + 各补列语句 + 注释)。

    给不愿意让库在自己的生产表上发 DDL 的下游用: 输出与库运行时执行的 DDL 同源,
    照它建完表,库探测到的列就是齐的。

    **补列语句与库内执行的那份是两套文本,不是一份**: 这份给人执行,必须可重复执行,
    故 PG 变体带 `ADD COLUMN IF NOT EXISTS`(它会先取 ACCESS EXCLUSIVE 锁,但执行时机
    由 DBA 自己挑,锁风险可控);库内那份不带,靠先探测后 ALTER 规避锁。SQLite 没有
    `ADD COLUMN IF NOT EXISTS` 语法,只能以注释交代"仅当该列不存在时执行"。

    Args:
        backend: `"sqlite"` 或 `"postgres"`。

    Returns:
        含注释的完整 SQL 脚本。

    Raises:
        ValueError: backend 不在取值域内。
    """
    if backend not in _BACKENDS:
        raise ValueError(f"未知遥测后端 {backend!r}: 只支持 {list(_BACKENDS)}")
    if backend == "sqlite":
        ddl = SQLITE_DDL
        notes = (
            f"-- 旧表补列(库升级后新增的列)。SQLite 无 ADD COLUMN IF NOT EXISTS 语法,\n"
            f"-- 以下每条**仅当该列不存在时执行**(先 PRAGMA table_info({TABLE}) 对照)。"
        )
        alters = [
            f"ALTER TABLE {TABLE} ADD COLUMN {column} {decl};" for column, decl in SQLITE_BACKFILL
        ]
    else:
        ddl = PG_DDL
        notes = (
            "-- 旧表补列(库升级后新增的列)。带 IF NOT EXISTS,整段可重复执行;\n"
            "-- 注意它即便列已存在也会先取 ACCESS EXCLUSIVE 锁,请挑低峰执行。"
        )
        alters = [
            f"ALTER TABLE {TABLE} ADD COLUMN IF NOT EXISTS {column} {decl};"
            for column, decl in _PG_BACKFILL_DECLS
        ]
    header = (
        f"-- PolyGateway 遥测表 {TABLE}({backend})\n"
        f'-- 由 polygateway.telemetry_schema_sql("{backend}") 生成,与库运行时执行的 DDL 同源。\n'
        "-- 新建库执行整段;已有旧表则建表语句自动跳过,只需关注下方补列语句。"
    )
    return "\n".join([header, "", ddl.strip(), "", notes, *alters, ""])