PolyGateway/src/polygateway/telemetry/postgres.py

"""Postgres 遥测后端(M2 设计 §5): asyncpg lazy 池 + 两级降级。

参考仓无先例(三项目遥测全 SQLite);asyncpg 工程写法取 GovDoc
`taskrun/postgres_store.py`($n 占位、`ON CONFLICT DO NOTHING`),但其
"失败冒泡"方向按遥测铁律**有意反转**:
① 结构性失败 → warning 一次后永久降级(所有写入短路);
② 运行时单条写失败 → 逐条 warning 丢弃,不降级不重试(连接抖动由
asyncpg 池自恢复;避免浸泡开头一次抖动导致后续全程失遥测)。
构造不连库(lazy),22 列 schema 与 SQLite 版同名同序。

**"结构性"的判据是「确定写不进去」,不是「初始化时出过错」**(issue #9):
只有建池失败(重试要在业务路径上内联吞掉 connect 超时)与"表确定不存在
且建不出来"(后续 INSERT 必然全败)才判死;探测失败、补列失败、取连接
失败一律只 warning,让写入照常尝试或下次调用重试。
"""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from loguru import logger

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

# 新列排在 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"),
)

# 探测表是否存在;不需要任何权限,且与 INSERT 走同一套 search_path 解析
_TABLE_EXISTS = "SELECT to_regclass('llm_calls')"

# 探测现有列;尊重 search_path(to_regclass 按当前 search_path 解析)
_EXISTING_COLUMNS = (
    "SELECT attname FROM pg_attribute "
    "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",
)

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


class PostgresRecorder:
    """TelemetryRecorder 端口的 Postgres 实现;asyncpg 原生异步,无线程桥接。"""

    def __init__(self, dsn: str, *, pool: asyncpg.Pool | None = None) -> None:
        try:
            import asyncpg  # noqa: F401 - 仅探测 extra 是否安装
        except ImportError as exc:
            raise ImportError(
                "Postgres 遥测未启用: 安装 pip install 'polygateway[postgres]' 后重试"
            ) from exc
        self._dsn = dsn
        self._pool: asyncpg.Pool | None = pool
        self._external_pool = pool is not None
        self._schema_ready = False
        self._failed = False  # 结构性降级标志: 置位后所有写入短路
        self._init_lock = asyncio.Lock()

    async def _ensure_ready(self) -> asyncpg.Pool | None:
        """lazy 建池+备表;判死只认「确定写不进去」(issue #9),其余失败都留活路。"""
        if self._failed:
            return None
        if self._schema_ready:
            return self._pool
        async with self._init_lock:
            if self._failed:
                return None
            if self._schema_ready:
                return self._pool
            pool = await self._open_pool()
            if pool is None:
                return None
            return await self._prepare_schema(pool)

    async def _open_pool(self) -> asyncpg.Pool | None:
        """建池;失败即永久降级(唯一一处「无条件判死」)。"""
        if self._pool is not None:
            return self._pool
        try:
            import asyncpg

            self._pool = await asyncpg.create_pool(self._dsn, timeout=10)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            # 池建不出来 = 确定写不进去;且每次调用重试都要内联吞掉 connect
            # 超时,而遥测是业务路径上的 await —— 此处必须永久降级
            self._failed = True
            logger.warning("Postgres 遥测建池失败,后续记录降级为 no-op: {}", exc)
            return None
        return self._pool

    async def _prepare_schema(self, pool: asyncpg.Pool) -> asyncpg.Pool | None:
        """备好表并交回可用的池;瞬时失败只跳过本次,确定写不进去才判死。"""
        try:
            async with pool.acquire() as conn:
                writable = await self._prepare_table(conn)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            # 池已在手,取连接/探测失败多为瞬时抖动: 不判死也不标就绪,
            # 只跳过本次记录,下次调用重新准备
            logger.warning("Postgres 遥测建表探测失败(跳过本条,下次重试): {}", exc)
            return None
        if not writable:
            self._failed = True
            return None
        self._schema_ready = True
        return pool

    async def _prepare_table(self, conn: object) -> bool:
        """备好 `llm_calls`;**表存在就绝不发 DDL**。返回 False 仅表示表确定不存在。

        `CREATE TABLE IF NOT EXISTS` 不能无条件发: PostgreSQL 对 schema 的
        CREATE 权限检查**早于** `IF NOT EXISTS` 的存在性判断(PG 16.14 实测:
        只授 `SELECT, INSERT ON llm_calls` 的角色,表明明在、也写得进去,这一句
        照样被拒 `permission denied for schema`)。这与 `_backfill_columns` 撞的
        是同一类问题(issue #3/#9),故守卫也必须同款: 先探测,后 DDL。
        探测走 `to_regclass`,不需要任何权限,且与 INSERT 的 search_path 解析
        口径一致——比裸 DDL 更准(裸 `CREATE TABLE` 落在首个**可建**的 schema,
        可能与 INSERT 命中的不是同一张表)。
        """
        exists = await conn.fetchval(_TABLE_EXISTS) is not None  # type: ignore[attr-defined]
        if exists:
            await self._backfill_columns(conn)  # 旧表可能缺列;失败只逐行降级
            return True
        try:
            await conn.execute(_DDL)  # type: ignore[attr-defined]
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("Postgres 遥测建表失败(表不存在,记录无处可落): {}", exc)
            return False
        return True  # 新建表列已齐全,无需再走补列

    async def _backfill_columns(self, conn: object) -> None:
        """给已存在的旧表补新列(issue #3);**先探测再 ALTER,失败绝不置 `_failed`**。

        两条纪律各有实测理由:
        ① 不置 `_failed`: 应用账号只有 INSERT 权限时,`ALTER TABLE` 的 ownership
           检查早于 `IF NOT EXISTS` 的存在性判断——列明明齐全也会失败。置位会让
           整个 recorder 永久 no-op,与「补列失败只降级为逐行丢弃」的承诺相悖
           (SQLite 侧同款守卫,两侧必须对称)。
        ② 先探测: `ADD COLUMN IF NOT EXISTS` 即便列已存在,也会**先取 ACCESS
           EXCLUSIVE 锁**再判存在性(实测会被一个开着的读事务阻塞)。遥测是内联
           await,让每个进程的首次写入都去抢共享审计表的排他锁,等于用记录基础设施
           拖垮业务调用。探测走 ACCESS SHARE,稳态下一条 ALTER 都不会发。
        """
        try:
            existing = {row["attname"] for row in await conn.fetch(_EXISTING_COLUMNS)}  # type: ignore[attr-defined]
            for column, statement in _BACKFILL:
                if column not in existing:
                    await conn.execute(statement)  # type: ignore[attr-defined]
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("Postgres 遥测补列失败(写入将逐行降级): {}", exc)

    async def record_llm_call(self, **fields: object) -> None:
        """写一行遥测;单条失败逐条 warning 丢弃(两级降级之二),绝不冒泡。"""
        pool = await self._ensure_ready()
        if pool is None:
            return
        row = tuple(fields[col] for col in _COLUMNS)
        try:
            async with pool.acquire() as conn:
                await conn.execute(_INSERT, *row)
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            # 遥测铁律: 丢一条 < 拖垮调用;仅记 warning(非 pass),池自恢复
            logger.warning("Postgres 遥测写入失败(丢弃该行): {}", exc)

    async def aclose(self) -> None:
        """幂等关闭自建池;注入的池归注入方管理。"""
        pool, self._pool = self._pool, None
        self._schema_ready = False
        if pool is not None and not self._external_pool:
            await pool.close()