"""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=`,用例默认用它。""" 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)