diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py new file mode 100644 index 0000000..3388b62 --- /dev/null +++ b/tests/integration/conftest.py @@ -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=`,用例默认用它。""" + 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) diff --git a/tests/integration/test_pg_sandbox.py b/tests/integration/test_pg_sandbox.py new file mode 100644 index 0000000..4fd6447 --- /dev/null +++ b/tests/integration/test_pg_sandbox.py @@ -0,0 +1,163 @@ +"""`conftest.py` 沙箱工厂自身的行为测试(issue #18 Task 1)。 + +工厂是本次一切隔离的地基: 它若在 setup 中途失败时漏掉清理、或让角色名与 +schema 名撞上,受害的不是这一个文件,而是此后每一条 PG 用例。故它必须先被测。 + +**这里的断言全部只看自建对象与 PG catalog**,不读任何共享数据。 +""" + +from __future__ import annotations + +import pytest + +_DDL = "CREATE TABLE llm_calls (call_id TEXT PRIMARY KEY, created_at TIMESTAMPTZ DEFAULT now())" + + +@pytest.fixture +async def assert_no_leftovers(pg_catalog_probe): + """收集沙箱名,在 `pg_sandbox` 清理之后回查它们是否真的没了。 + + 必须比 `pg_sandbox` **先** setup: pytest 的 finalizer 是后进先出,先 setup + 的后 teardown——本 fixture 的检查因此发生在沙箱清理之后,而不是之前。 + """ + seen: list[tuple[str, str | None]] = [] + yield seen + for schema, role in seen: + left = await pg_catalog_probe("SELECT nspname FROM pg_namespace WHERE nspname = $1", schema) + assert left == [], f"沙箱 schema 未清理: {schema}" + if role is not None: + left = await pg_catalog_probe("SELECT rolname FROM pg_roles WHERE rolname = $1", role) + assert left == [], f"沙箱角色未清理: {role}" + + +async def _oid_of_llm_calls(dsn: str) -> int | None: + import asyncpg + + conn = await asyncpg.connect(dsn, timeout=10) + try: + return await conn.fetchval("SELECT to_regclass('llm_calls')::oid") + finally: + await conn.close() + + +class TestSchemaOnlySandbox: + async def test_table_lands_in_the_sandbox_schema_and_bare_dsn_is_absent(self, pg_sandbox): + """`role="none"`: 表落在自建 schema 下;不发角色,故没有裸 DSN 可给。""" + sandbox = await pg_sandbox(ddl=_DDL) + + import asyncpg + + conn = await asyncpg.connect(sandbox.dsn, timeout=10) + try: + where = await conn.fetchval( + "SELECT n.nspname FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace " + "WHERE c.oid = to_regclass('llm_calls')" + ) + finally: + await conn.close() + + assert where == sandbox.schema + assert sandbox.role is None + assert sandbox.bare_dsn is None + + +class TestOwnerRoleSandbox: + async def test_the_role_owns_its_own_table(self, pg_sandbox): + """`role="owner"`: 表由角色自己建,故属主是它——与"用维护角色跑"的现场一致。""" + sandbox = await pg_sandbox(ddl=_DDL, role="owner") + + import asyncpg + + conn = await asyncpg.connect(sandbox.dsn, timeout=10) + try: + owner = await conn.fetchval( + "SELECT pg_get_userbyid(relowner) FROM pg_class WHERE oid = to_regclass('llm_calls')" + ) + finally: + await conn.close() + + assert owner == sandbox.role + # 名字必须错开: 同名会让 "$user" 命中自有 schema 并遮蔽真表, + # 最坏情况用例就再也走不到那条真实路径上(设计 §5.1 实测) + assert sandbox.role != sandbox.schema + assert not sandbox.role.startswith("pgw_s_") + assert not sandbox.schema.startswith("pgw_r_") + + async def test_bare_dsn_falls_through_to_the_default_search_path(self, pg_sandbox): + """裸 DSN 必须真的回落到 `"$user", public`——最坏情况用例全靠它构造现场。""" + sandbox = await pg_sandbox(ddl=_DDL, role="owner") + + import asyncpg + + conn = await asyncpg.connect(sandbox.bare_dsn, timeout=10) + try: + path = await conn.fetchval("SHOW search_path") + finally: + await conn.close() + + assert path == '"$user", public' + # 裸 DSN 解析到的绝不能是沙箱里那张表(否则"落到共享表"的现场是假的) + assert await _oid_of_llm_calls(sandbox.bare_dsn) != await _oid_of_llm_calls(sandbox.dsn) + + +class TestGranteeRoleSandbox: + async def test_grantee_can_write_but_cannot_create(self, pg_sandbox): + """`role="grantee"`: 表属主是 admin,角色只拿表级权限——最小权限部署的现场。""" + import asyncpg + + sandbox = await pg_sandbox(ddl=_DDL, role="grantee") + + conn = await asyncpg.connect(sandbox.dsn, timeout=10) + try: + await conn.execute("INSERT INTO llm_calls (call_id) VALUES ('g1')") + assert await conn.fetchval("SELECT count(*) FROM llm_calls") == 1 + with pytest.raises(asyncpg.exceptions.InsufficientPrivilegeError): + await conn.execute("CREATE TABLE another (x TEXT)") + finally: + await conn.close() + + +class TestCleanup: + async def test_setup_failure_leaves_nothing_behind(self, pg_sandbox, pg_catalog_probe): + """建到一半失败时也必须删净——角色是**全局**对象,残留不随库消失。""" + before_schemas = await pg_catalog_probe( + "SELECT nspname FROM pg_namespace WHERE nspname LIKE 'pgw!_%' ESCAPE '!'" + ) + before_roles = await pg_catalog_probe( + "SELECT rolname FROM pg_roles WHERE rolname LIKE 'pgw!_%' ESCAPE '!'" + ) + + with pytest.raises(Exception): # noqa: B017 — 工厂原样抛出 PG 的 DDL 错误 + await pg_sandbox(ddl="CREATE TABLE llm_calls (bad NOT_A_REAL_TYPE)", role="owner") + + assert ( + await pg_catalog_probe( + "SELECT nspname FROM pg_namespace WHERE nspname LIKE 'pgw!_%' ESCAPE '!'" + ) + == before_schemas + ) + assert ( + await pg_catalog_probe( + "SELECT rolname FROM pg_roles WHERE rolname LIKE 'pgw!_%' ESCAPE '!'" + ) + == before_roles + ) + + async def test_a_failure_does_not_roll_back_earlier_sandboxes(self, pg_sandbox): + """一次失败只回滚它自己建的东西——同一条用例里先建成的沙箱必须毫发无损。 + + "A 的角色去动 B 的表"这类用例一条要两个沙箱;若失败回滚把整栈清空,受害的 + 是那些**已经通过**的断言所依赖的对象,而症状会以"表不见了"的形态出现在 + 与真因无关的地方。 + """ + good = await pg_sandbox(ddl=_DDL, role="owner") + + with pytest.raises(Exception): # noqa: B017 — 工厂原样抛出 PG 的 DDL 错误 + await pg_sandbox(ddl="CREATE TABLE llm_calls (bad NOT_A_REAL_TYPE)", role="owner") + + assert await _oid_of_llm_calls(good.dsn) is not None, "先前建成的沙箱被误清理" + + async def test_teardown_removes_schema_and_role(self, assert_no_leftovers, pg_sandbox): + """正常路径的清理: 断言发生在 `pg_sandbox` teardown **之后**(见 fixture 说明)。""" + sandbox = await pg_sandbox(ddl=_DDL, role="owner") + assert_no_leftovers.append((sandbox.schema, sandbox.role))