test: build the sandbox factory the PG tests will run inside
Seven copies of "create a schema, hang it off search_path, drop it in teardown" were spread across two files, each with its own cleanup. Any one of them written wrong leaves the residue on a database shared with real batch runs. This is one implementation, and it makes "the test cannot reach the admin connection" a structural fact rather than a note in a docstring. Three role modes cover every fixture that exists today: none for plain schema isolation, owner for the retention script's own runs, grantee for the least-privilege deployment cases. Owner runs its DDL as itself so it ends up owning the table; grantee is the opposite, since that case only means anything when someone else built it. The schema and the role deliberately get different prefixes. Give them the same name and "$user" resolves to the sandbox, which hides the shared table and quietly turns the worst-case test into a test of nothing. Writing it also turned up a bug in my first version: rolling back a failed sandbox unwound the whole stack, so an earlier sandbox in the same test lost its role mid-use. The test for it fails with a password authentication error, which is what that looks like from the outside. Each call now unwinds only what it created, and cleanup tries every statement before raising, since one failure stranding the rest means global roles left behind by hand.
This commit is contained in:
@@ -0,0 +1,211 @@
|
||||
"""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 的用户名口令段换成沙箱角色的,其余(主机/库/参数)原样保留。"""
|
||||
return re.sub(r"//[^@/]+@", f"//{role}:{_SANDBOX_PASSWORD}@", dsn, count=1)
|
||||
|
||||
|
||||
@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]:
|
||||
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 —— 缺的正是这一项
|
||||
except BaseException:
|
||||
await _unwind(local)
|
||||
raise
|
||||
cleanups.extend(local)
|
||||
|
||||
used = bare if role_name is not None else admin_dsn
|
||||
return PgSandbox(
|
||||
schema=schema,
|
||||
role=role_name,
|
||||
dsn=_with_search_path(used, schema),
|
||||
bare_dsn=bare,
|
||||
)
|
||||
|
||||
yield make
|
||||
await _unwind(cleanups)
|
||||
Reference in New Issue
Block a user