PolyGateway/tests/integration/conftest.py

"""PG 集成测试的一次性沙箱工厂(issue #18)。

**为什么把它收敛成一份**: 在此之前,"建临时 schema → 挂 search_path → teardown
删净"这套样板在两个测试文件里重复了七处,清理逻辑各写各的——任何一处写漏,残留都
落在与真实批跑共用的那个库上。工厂让清理只有一份实现,并让"用例拿不到管理连接"
成为结构事实而不是纪律。

**admin DSN 不做成 fixture**: 它能对共享表执行任何语句。做成 fixture 等于把这个
能力摆在每一条用例面前,"用例不该直接用"就只是一句提醒。故它是模块私有函数,
只被工厂内部调用,`PgSandbox` 也不携带它。
"""

from __future__ import annotations

import os
import re
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal
from uuid import uuid4

import pytest
from dotenv import dotenv_values

if TYPE_CHECKING:
    from collections.abc import Sequence

# 测试专用口令: 这些角色只在单条用例的生命周期内存在,且只对自建 schema 有权。
# 它不是机密,写死在这里比走 .env 更清楚——.env 里的每一项都该是真实部署会用的。
_SANDBOX_PASSWORD = "pgw-sandbox-not-a-secret"  # noqa: S105

_Role = Literal["none", "owner", "grantee"]


@dataclass(frozen=True)
class PgSandbox:
    """一次性 PG 沙箱: 独立 schema + 可选独占登录角色。"""

    schema: str
    role: str | None
    dsn: str
    """已挂 `options=-csearch_path=<schema>`,用例默认用它。"""
    bare_dsn: str | None
    """同角色但**不挂** search_path(回落 `"$user", public`);`role="none"` 时为 None。"""


def _admin_dsn() -> str | None:
    """读 `.env` 的 `PGW_TELEMETRY_PG_DSN` 并剥掉 SQLAlchemy 风格的 `+driver` 后缀。"""
    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}"


def _require_admin_dsn() -> str:
    """取管理连接串;未配置则 skip,连错库则 fail(不是 skip)。

    库名守卫不肯降级成 skip: 这个实例上还有 app/chs_prod 等在用库,把"连错库"
    悄悄跳过,等于让一次配置事故以"没跑那些测试"的形态过关。
    """
    value = _admin_dsn()
    if value is None:
        pytest.skip("PGW_TELEMETRY_PG_DSN 未配置")
    if not value.rstrip("/").endswith("/polygateway"):
        pytest.fail(f"PG 集成测试只允许连 polygateway 专用库,当前 DSN 库名不符: {value!r}")
    return value


def _with_search_path(dsn: str, schema: str) -> str:
    sep = "&" if "?" in dsn else "?"
    return f"{dsn}{sep}options=-csearch_path%3D{schema}"


def _as_role(dsn: str, role: str) -> str:
    """把 DSN 的用户名口令段换成沙箱角色的,其余(主机/库/参数)原样保留。

    **换不掉就报错,绝不原样返回**: `postgresql://h:5432/db`(口令走 PGPASSWORD /
    .pgpass / trust)与 `postgresql:///db?host=/var/run/postgresql`(unix socket)
    都是合法 DSN,却没有可替换的内联凭据段。静默返回原串的后果不是测试报错,而是
    沙箱以**管理身份**建成、用例照常绿,同时 `bare_dsn` 变成超级用户连接——最坏
    情况用例会拿它跑真实 `--apply`,删空共享表之后才在退出码断言上红。
    这正是 P5"严禁默认值掩盖错误"要挡的形态。
    """
    swapped, count = re.subn(r"//[^@/]+@", f"//{role}:{_SANDBOX_PASSWORD}@", dsn, count=1)
    if count != 1:
        raise RuntimeError(
            f"DSN 里没有可替换的内联凭据段,沙箱角色 {role} 无法生效,拒绝以管理身份继续。"
            "请把 PGW_TELEMETRY_PG_DSN 写成 postgresql://<用户>:<口令>@<主机>/<库> 的形态。"
        )
    return swapped


@pytest.fixture
async def pg_catalog_probe():
    """只读地查 PG catalog,**仅供工厂自测核对残留**,不是通用查询入口。

    它拿的是管理连接,故有意只暴露给 `test_pg_sandbox.py` 这一类"验证隔离本身
    是否成立"的用例;业务断言一律走 `PgSandbox.dsn`。
    """
    import asyncpg

    dsn = _require_admin_dsn()

    async def probe(sql: str, *args: object) -> list[tuple]:
        # 只读校验不是形式主义: 这个闭包持的是管理连接,不设限就等于把"用例够不到
        # 管理能力"这句话降格成一句 docstring 里的请求。
        if not sql.lstrip().upper().startswith("SELECT"):
            raise RuntimeError(f"pg_catalog_probe 只接受 SELECT 语句,收到: {sql[:60]!r}")
        conn = await asyncpg.connect(dsn, timeout=10)
        try:
            return [tuple(r) for r in await conn.fetch(sql, *args)]
        finally:
            await conn.close()

    return probe


@pytest.fixture
async def pg_sandbox():
    """一次性沙箱工厂: `await pg_sandbox(ddl=..., role=...)`,清理由 fixture 兜底。

    同一条用例可以要多个沙箱(如"A 的角色去动 B 的表"),它们按后进先出清理。
    """
    import asyncpg

    admin_dsn = _require_admin_dsn()
    # 清理动作栈: 每建成一个对象就入栈一条,setup 中途失败与正常 teardown 共用
    # 同一条退栈路径——两处各写一份的话,失败那条永远是没被测过的那份。
    cleanups: list[str] = []

    async def _run_as_admin(*statements: str) -> None:
        conn = await asyncpg.connect(admin_dsn, timeout=10)
        try:
            for statement in statements:
                await conn.execute(statement)
        finally:
            await conn.close()

    async def _unwind(statements: list[str]) -> None:
        """逆序执行清理并**逐条容错**: 一条失败不该拖累其余对象的清理。

        吞掉异常是不行的(残留会静默累积),但让第一条失败中断整栈更糟——角色是
        全局对象,漏掉的每一个都要人手工去删。故全部试完再抛出第一个异常。
        """
        first: BaseException | None = None
        for statement in reversed(statements):
            try:
                await _run_as_admin(statement)
            except Exception as exc:  # noqa: BLE001 — 见 docstring: 收集而非吞没
                first = first or exc
        statements.clear()
        if first is not None:
            raise first

    async def make(
        *,
        ddl: str | None = None,
        extra: Sequence[str] = (),
        role: _Role = "none",
        grants: Sequence[str] = ("SELECT", "INSERT"),
    ) -> PgSandbox:
        # 权限门在建任何对象**之前**: pytest.skip 抛的是 BaseException,若它在
        # 已建对象之后触发,清理会去 DROP 从未建成的东西并把 skip 盖掉。
        if role != "none":
            conn = await asyncpg.connect(admin_dsn, timeout=10)
            try:
                can_create = await conn.fetchval(
                    "SELECT rolcreaterole OR rolsuper FROM pg_roles WHERE rolname = current_user"
                )
            finally:
                await conn.close()
            if not can_create:
                pytest.skip("当前账号无权建临时角色,跳过需要独占角色的用例")

        # schema 与角色的前缀有意不同: 同名会让 "$user" 命中自有 schema 并遮蔽
        # 共享表,于是"search_path 落到共享表"这个最坏情况就再也构造不出来。
        suffix = uuid4().hex[:12]
        schema = f"pgw_s_{suffix}"
        role_name = f"pgw_r_{suffix}" if role != "none" else None

        # 本次调用自己的清理栈: 失败只回滚**本次**建成的对象。同一条用例常要两个
        # 沙箱(如"A 的角色去动 B 的表"),回滚整栈会把已通过断言依赖的对象也删掉。
        local: list[str] = []
        try:
            if role_name is not None:
                await _run_as_admin(f"CREATE ROLE {role_name} LOGIN PASSWORD '{_SANDBOX_PASSWORD}'")
                # DROP OWNED BY 必须排在 DROP ROLE 之前: 角色仍持有对象时删不掉
                local.append(f"DROP ROLE IF EXISTS {role_name}")
                local.append(f"DROP OWNED BY {role_name}")
            owner_clause = f" AUTHORIZATION {role_name}" if role == "owner" else ""
            await _run_as_admin(f"CREATE SCHEMA {schema}{owner_clause}")
            local.append(f"DROP SCHEMA IF EXISTS {schema} CASCADE")

            bare = _as_role(admin_dsn, role_name) if role_name is not None else None
            # role="owner" 时 DDL 由角色自己执行,表属主才会是它;"grantee" 的现场
            # 恰恰相反——表由别的账号建好,角色只拿到表级权限。
            ddl_dsn = _with_search_path(bare if role == "owner" else admin_dsn, schema)
            if ddl is not None:
                conn = await asyncpg.connect(ddl_dsn, timeout=10)
                try:
                    await conn.execute(ddl)
                    for statement in extra:
                        await conn.execute(statement)
                finally:
                    await conn.close()

            if role == "grantee":
                await _run_as_admin(f"GRANT USAGE ON SCHEMA {schema} TO {role_name}")
                if ddl is not None:
                    await _run_as_admin(
                        f"GRANT {', '.join(grants)} ON ALL TABLES IN SCHEMA {schema} TO {role_name}"
                    )
                # 关键: 绝不 GRANT CREATE ON SCHEMA —— 缺的正是这一项

            used = bare if role_name is not None else admin_dsn
            sandbox = PgSandbox(
                schema=schema,
                role=role_name,
                dsn=_with_search_path(used, schema),
                bare_dsn=bare,
            )
            if role_name is not None:
                # 字符串替换成功不等于连上去就是那个角色(PGUSER 等环境变量仍可能
                # 盖掉 DSN 里的用户名)。这道校验按**实际身份**兜底: 整个设计的价值
                # 都压在"跑脚本的那个连接对共享表无权"上,不值得只用一次字符串比较
                # 来担保。它必须留在 try 之内——出了这个块,清理动作已经并进 fixture
                # 级的栈,再回滚一次就会对同一个角色跑两遍 DROP OWNED BY(它没有
                # IF EXISTS,第二遍必报错)。
                conn = await asyncpg.connect(sandbox.dsn, timeout=10)
                try:
                    actual = await conn.fetchval("SELECT current_user")
                finally:
                    await conn.close()
                if actual != role_name:
                    raise RuntimeError(
                        f"沙箱 DSN 连上去的身份是 {actual!r},不是预期的 {role_name!r};"
                        "权限边界不成立,拒绝把这个沙箱交出去。"
                    )
        except BaseException:
            await _unwind(local)
            raise
        cleanups.extend(local)

        return sandbox

    yield make
    await _unwind(cleanups)