PolyGateway/tests/integration/test_postgres_telemetry.py

"""PostgresRecorder 集成测试(M2 设计 §5;真实实验室 Postgres,polygateway 专用库)。

DSN 走 .env `PGW_TELEMETRY_PG_DSN`,缺则 skip。该实例上有 app/chs_prod 等
在用库——本测试只允许连 polygateway 专用库(fixture 里守卫)。

隔离纪律(M4 事故教训): `llm_calls` 是与真实批跑/迁移项目共享的表,
**严禁 DROP/TRUNCATE**——本测试以 run 级 call_id 前缀隔离,断言只看
自己写入的行,teardown 只删自己的行。
"""

from __future__ import annotations

import asyncio
import os
from uuid import uuid4

import pytest
from dotenv import dotenv_values

from polygateway.telemetry.postgres import PostgresRecorder

_EXPECTED_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",
    "created_at",
    "cached_prompt_tokens",
    "model_reported",
    "sampling",
]

# run 级前缀: 同库并存的其他运行(迁移批跑/另一开发机)互不可见
_RUN_PREFIX = f"pgwtest-{uuid4().hex[:8]}"


def _cid(suffix: str) -> str:
    return f"{_RUN_PREFIX}-{suffix}"


def _dsn() -> str | None:
    merged = {**dotenv_values(".env"), **os.environ}
    raw = merged.get("PGW_TELEMETRY_PG_DSN")
    if not raw:
        return None
    scheme, sep, rest = raw.partition("://")
    return f"{scheme.partition('+')[0]}{sep}{rest}"


@pytest.fixture
async def dsn():
    value = _dsn()
    if value is None:
        pytest.skip("PGW_TELEMETRY_PG_DSN 未配置")
    # 隔离守卫: 该实例有 app/chs_prod/mimiciv 等在用库,只许打 polygateway 专用库
    if not value.rstrip("/").endswith("/polygateway"):
        pytest.fail(f"遥测测试只允许连 polygateway 专用库,当前 DSN 库名不符: {value!r}")
    yield value
    # teardown: 只删本 run 写入的行;表可能尚不存在(全新库)则忽略
    import asyncpg

    conn = await asyncpg.connect(value, timeout=10)
    try:
        if await conn.fetchval("SELECT to_regclass('llm_calls')") is not None:
            await conn.execute("DELETE FROM llm_calls WHERE call_id LIKE $1", f"{_RUN_PREFIX}-%")
    finally:
        await conn.close()


async def _record_minimal(
    recorder: PostgresRecorder, call_id: str | None = None, **overrides
) -> None:
    fields = {
        "call_id": call_id if call_id is not None else _cid("c1"),
        "parent_call_id": None,
        "session_id": "sess-1",
        "model": "m",
        "provider": "p",
        "source_name": "s1",
        "messages": "[]",
        "response": "ok",
        "thinking": "",
        "prompt_tokens": 1,
        "completion_tokens": 2,
        "usage_source": "measured",
        "latency_ms": 10,
        "ttft_ms": None,
        "max_inter_token_ms": None,
        "cache_hit": False,
        "error": None,
        "cost": None,
        "cached_prompt_tokens": None,
        "model_reported": None,
        "sampling": None,
    }
    fields.update(overrides)
    await recorder.record_llm_call(**fields)


async def _fetch(dsn: str, sql: str, *args):
    import asyncpg

    conn = await asyncpg.connect(dsn, timeout=10)
    try:
        return await conn.fetch(sql, *args)
    finally:
        await conn.close()


_LEGACY_DDL = """
CREATE TABLE {schema}.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()
)
"""


@pytest.fixture
async def legacy_schema(dsn):
    """在**自建的临时 schema** 里造一张 18 列旧表,验证补列(issue #3)。

    绝不碰共享的 public.llm_calls: 用 search_path 把 recorder 指向临时 schema,
    teardown 只 DROP 自己建的 schema。
    """
    import asyncpg

    name = f"pgwtest_{uuid4().hex[:8]}"
    conn = await asyncpg.connect(dsn, timeout=10)
    try:
        await conn.execute(f"CREATE SCHEMA {name}")
        await conn.execute(_LEGACY_DDL.format(schema=name))
    finally:
        await conn.close()
    sep = "&" if "?" in dsn else "?"
    yield f"{dsn}{sep}options=-csearch_path%3D{name}", name
    conn = await asyncpg.connect(dsn, timeout=10)
    try:
        await conn.execute(f"DROP SCHEMA {name} CASCADE")
    finally:
        await conn.close()


class TestObservabilityColumns:
    """issue #3: 两列写入可回读,且已存在的 18 列旧表会被自动补列。"""

    async def test_values_round_trip(self, dsn):
        recorder = PostgresRecorder(dsn)
        try:
            await _record_minimal(recorder, call_id=_cid("hit"), cached_prompt_tokens=64)
            await _record_minimal(recorder, call_id=_cid("zero"), cached_prompt_tokens=0)
            await _record_minimal(recorder, call_id=_cid("model"), model_reported="MiniMax-01")
            rows = await _fetch(
                dsn,
                "SELECT call_id, cached_prompt_tokens, model_reported FROM llm_calls "
                "WHERE call_id LIKE $1",
                f"{_RUN_PREFIX}-%",
            )
            by_id = {r["call_id"]: r for r in rows}
            assert by_id[_cid("hit")]["cached_prompt_tokens"] == 64
            assert by_id[_cid("zero")]["cached_prompt_tokens"] == 0  # 真实零命中 ≠ NULL
            assert by_id[_cid("model")]["cached_prompt_tokens"] is None
            assert by_id[_cid("model")]["model_reported"] == "MiniMax-01"
        finally:
            await recorder.aclose()

    async def test_legacy_table_is_upgraded_in_place(self, legacy_schema):
        """18 列旧表不补列的话,每行写入都会被逐行 warning 丢弃(遥测静默全失)。"""
        schema_dsn, schema = legacy_schema
        recorder = PostgresRecorder(schema_dsn)
        try:
            await _record_minimal(
                recorder, call_id=_cid("legacy"), cached_prompt_tokens=7, model_reported="m-real"
            )
            cols = await _fetch(
                schema_dsn,
                "SELECT column_name FROM information_schema.columns "
                "WHERE table_schema = $1 AND table_name = 'llm_calls' ORDER BY ordinal_position",
                schema,
            )
            # ALTER 只能追加到末尾: 与新建库的列序一致才不会分叉
            assert [r["column_name"] for r in cols] == _EXPECTED_COLUMNS
            rows = await _fetch(
                schema_dsn,
                "SELECT cached_prompt_tokens, model_reported FROM llm_calls WHERE call_id = $1",
                _cid("legacy"),
            )
            assert (rows[0]["cached_prompt_tokens"], rows[0]["model_reported"]) == (7, "m-real")
        finally:
            await recorder.aclose()


class TestSchema:
    async def test_schema_has_frozen_columns_in_order(self, dsn):
        recorder = PostgresRecorder(dsn)
        try:
            await _record_minimal(recorder)
            rows = await _fetch(
                dsn,
                "SELECT column_name FROM information_schema.columns "
                "WHERE table_name='llm_calls' ORDER BY ordinal_position",
            )
            assert [r["column_name"] for r in rows] == _EXPECTED_COLUMNS
        finally:
            await recorder.aclose()

    async def test_call_id_idempotent(self, dsn):
        recorder = PostgresRecorder(dsn)
        try:
            await _record_minimal(recorder, call_id=_cid("dup"))
            await _record_minimal(recorder, call_id=_cid("dup"), response="second")
            rows = await _fetch(
                dsn, "SELECT response FROM llm_calls WHERE call_id = $1", _cid("dup")
            )
            assert [r["response"] for r in rows] == ["ok"]  # ON CONFLICT DO NOTHING
        finally:
            await recorder.aclose()

    async def test_concurrent_writes_all_land(self, dsn):
        recorder = PostgresRecorder(dsn)
        try:
            await asyncio.gather(
                *(_record_minimal(recorder, call_id=_cid(f"c{i}")) for i in range(50))
            )
            rows = await _fetch(
                dsn,
                "SELECT count(*) AS n FROM llm_calls WHERE call_id LIKE $1",
                f"{_RUN_PREFIX}-c%",
            )
            assert rows[0]["n"] == 50
        finally:
            await recorder.aclose()


class TestDegradation:
    async def test_unreachable_server_degrades_silently(self):
        """结构性失败(建池不通)→ warning 一次后永久降级,业务零感知。"""
        recorder = PostgresRecorder("postgresql://u:p@127.0.0.1:1/x")
        await _record_minimal(recorder)  # 不抛
        await _record_minimal(recorder, call_id=_cid("c2"))  # 已降级短路,同样不抛
        await recorder.aclose()

    async def test_row_failure_does_not_poison_later_rows(self, dsn):
        """运行时单条写失败(NUL 字节文本被 PG 拒)→ 丢该行,后续行照常落库。"""
        recorder = PostgresRecorder(dsn)
        try:
            await _record_minimal(recorder, call_id=_cid("bad"), response="nul\x00byte")
            await _record_minimal(recorder, call_id=_cid("good"))
            rows = await _fetch(
                dsn,
                "SELECT call_id FROM llm_calls WHERE call_id = ANY($1::text[]) ORDER BY call_id",
                [_cid("bad"), _cid("good")],
            )
            assert [r["call_id"] for r in rows] == [_cid("good")]
        finally:
            await recorder.aclose()

    async def test_aclose_idempotent(self, dsn):
        recorder = PostgresRecorder(dsn)
        await _record_minimal(recorder)
        await recorder.aclose()
        await recorder.aclose()