"""Postgres 后端的契约与跨连接测试。 需要一台可连的 Postgres，DSN 从环境变量给出:: CHONGMING_TEST_POSTGRES='postgresql://user@127.0.0.1:5432/tasks' \\ uv run pytest tests/integration/test_postgres_store.py 未设置时整组跳过——它依赖外部服务，与内置实现的契约测试（纯本地）刻意分开。 这里覆盖三类只有真数据库才能证明的事： 1. **完整契约**——与内存 / SQLite 跑的是同一套 ``StoreContract``； 2. **跨连接原子性**——设计说明 §15.4 明确要求后端自建并发用例： N 个连接并发 ``save`` 恰好一个成功、并发 ``claim_commands`` 不重复下发、 持有者被强杀后租约仍按时间过期且连接池自愈； 3. **接线**——``Controller`` 真的能把它当存储用，跨 store 的控制走指令通道。 """ from __future__ import annotations import os import threading import time from datetime import UTC, datetime, timedelta from typing import Any from uuid import uuid4 import pytest psycopg = pytest.importorskip("psycopg") pytest.importorskip("psycopg_pool") from psycopg import sql # noqa: E402 from psycopg.conninfo import make_conninfo # noqa: E402 from chongming_task.controller import Controller # noqa: E402 from chongming_task.errors import ConcurrencyConflict, StoreError # noqa: E402 from chongming_task.states import TaskStatus # noqa: E402 from chongming_task.stores.postgres import PostgresTaskStore # noqa: E402 from chongming_task.stores.protocol import TaskStore # noqa: E402 from chongming_task.testing.store_contract import ( # noqa: E402 StoreContract, command, simple_snapshot, ) _DSN = os.environ.get("CHONGMING_TEST_POSTGRES") _T0 = datetime(2026, 1, 1, tzinfo=UTC) pytestmark = pytest.mark.skipif( not _DSN, reason="未设置 CHONGMING_TEST_POSTGRES（DSN）：Postgres 集成测试需要一台可连的数据库", ) # --------------------------------------------------------------------------- # 隔离：每个存储实例独占一个 schema（与 SQLite 的临时文件等价） # --------------------------------------------------------------------------- def _isolated_dsn() -> tuple[str, str]: """建一个全新 schema，返回 ``(schema 名, 指向它的 DSN)``。""" assert _DSN is not None schema = f"cstore_{uuid4().hex[:12]}" with psycopg.connect(_DSN, autocommit=True) as conn: conn.execute(sql.SQL("CREATE SCHEMA {}").format(sql.Identifier(schema))) return schema, make_conninfo(_DSN, options=f"-csearch_path={schema}") def _drop_schema(schema: str) -> None: assert _DSN is not None with psycopg.connect(_DSN, autocommit=True) as conn: conn.execute(sql.SQL("DROP SCHEMA IF EXISTS {} CASCADE").format(sql.Identifier(schema))) # --------------------------------------------------------------------------- # 1. 完整契约 # --------------------------------------------------------------------------- class TestPostgresStoreContract(StoreContract): """与内存 / SQLite 实现跑**完全相同**的用例。""" _sandboxes: list[str] def setup_method(self) -> None: self._sandboxes = [] super().setup_method() def make_store(self) -> TaskStore: schema, dsn = _isolated_dsn() self._sandboxes.append(schema) return PostgresTaskStore(dsn) def teardown_method(self) -> None: super().teardown_method() for schema in self._sandboxes: _drop_schema(schema) # --------------------------------------------------------------------------- # 2. 跨连接原子性（StoreContract 覆盖不到，设计说明 §15.4 要求） # --------------------------------------------------------------------------- def test_concurrent_saves_have_exactly_one_winner() -> None: """N 个连接并发 save 同一任务：恰好一个成功，其余收到冲突。 乐观并发在**跨连接**下真正成立，而不是只在单连接语义里成立。 """ schema, dsn = _isolated_dsn() try: seed = PostgresTaskStore(dsn) seed.create(simple_snapshot("t1", version=1)) seed.close() workers = 8 barrier = threading.Barrier(workers) outcomes: list[str] = [] lock = threading.Lock() def attempt() -> None: store = PostgresTaskStore(dsn) try: barrier.wait(timeout=10) try: store.save( simple_snapshot("t1", status=TaskStatus.SUCCEEDED, version=2), expect_version=1, ) except ConcurrencyConflict: outcome = "conflict" else: outcome = "ok" finally: store.close() with lock: outcomes.append(outcome) threads = [threading.Thread(target=attempt) for _ in range(workers)] for thread in threads: thread.start() for thread in threads: thread.join(timeout=30) assert sorted(outcomes) == ["conflict"] * (workers - 1) + ["ok"] finally: _drop_schema(schema) def test_concurrent_claims_hand_out_each_command_exactly_once() -> None: """N 个连接并发认领：同一条指令绝不会被两个连接取到。 "挑选 + 打租约"是一条 ``UPDATE ... FOR UPDATE SKIP LOCKED``，所以并发认领互不 阻塞、也不会重复。这是 SQLite 用 ``BEGIN IMMEDIATE`` 换来的同一个保证。 """ schema, dsn = _isolated_dsn() try: total = 20 seed = PostgresTaskStore(dsn) for index in range(total): seed.submit_command(command(f"c{index}", created_offset=float(index))) seed.close() claimed: list[str] = [] lock = threading.Lock() def drain() -> None: store = PostgresTaskStore(dsn) try: while True: batch = store.claim_commands(now=_T0, lease_seconds=300, limit=3) if not batch: return with lock: claimed.extend(item.command_id for item in batch) finally: store.close() threads = [threading.Thread(target=drain) for _ in range(8)] for thread in threads: thread.start() for thread in threads: thread.join(timeout=30) assert sorted(claimed) == sorted(f"c{index}" for index in range(total)) assert len(claimed) == len(set(claimed)), "同一条指令被下发了两次" finally: _drop_schema(schema) def test_a_killed_holder_does_not_wedge_the_lease() -> None: """持有者被强杀：租约仍按时间过期，别的连接可以接管；连接池自愈。 租约是**存储里的数据**，不是连接上的会话状态——所以杀掉持有者的后端连接既不会 提前释放它，也不会永久锁住它。 """ schema, dsn = _isolated_dsn() try: store_a = PostgresTaskStore(dsn, min_size=1, max_size=1) store_a.create(simple_snapshot("t1")) store_a.submit_command(command("c1")) claimed = store_a.claim_commands(now=_T0, lease_seconds=30) assert [item.command_id for item in claimed] == ["c1"] # 模拟持有者进程消失：杀掉它的后端连接（本测试此刻只有 store_a 在用）。 killed: list[bool] = [] try: with psycopg.connect(_DSN, autocommit=True) as admin: rows = admin.execute( "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " "WHERE application_name = 'chongming-task' AND pid <> pg_backend_pid()" ).fetchall() killed = [bool(row[0]) for row in rows] except psycopg.errors.InsufficientPrivilege: # pragma: no cover - 权限受限的环境 pytest.skip("当前角色无权 pg_terminate_backend，跳过") assert any(killed), "没有杀掉任何连接，测试前提不成立" store_b = PostgresTaskStore(dsn) # 租约未到期：接管者取不到。 assert store_b.claim_commands(now=_T0 + timedelta(seconds=10), lease_seconds=30) == [] # 租约到期：可以接管。 reclaimed = store_b.claim_commands(now=_T0 + timedelta(seconds=31), lease_seconds=30) assert [item.command_id for item in reclaimed] == ["c1"] # 被杀掉的持有者：连接池把死连接换掉，操作照常。 assert store_a.get("t1") is not None store_a.close() store_b.close() finally: _drop_schema(schema) # --------------------------------------------------------------------------- # 3. close() 的语义（模块文档承诺过，就要有用例守着） # --------------------------------------------------------------------------- def test_close_is_idempotent_and_operations_fail_afterwards() -> None: schema, dsn = _isolated_dsn() try: store = PostgresTaskStore(dsn) store.close() store.close() # 幂等：重复关闭是空操作 with pytest.raises(StoreError, match="已关闭"): store.get("t1") finally: _drop_schema(schema) def test_a_role_without_ddl_privileges_can_use_an_existing_schema() -> None: """建表是**条件性**的：表已存在时一个 DDL 都不执行。 因此运维控制台可以用只读账号接进来（只给 SELECT 就够看板了）。若实现无条件跑 ``CREATE TABLE IF NOT EXISTS``，这个账号会因为缺 CREATE 权限直接连不上。 """ schema, dsn = _isolated_dsn() role = f"cstore_ro_{uuid4().hex[:8]}" password = uuid4().hex try: seed = PostgresTaskStore(dsn) seed.create(simple_snapshot("t1", name="只读也能看到")) seed.close() try: with psycopg.connect(_DSN, autocommit=True) as admin: admin.execute( sql.SQL("CREATE ROLE {} LOGIN PASSWORD {}").format( sql.Identifier(role), sql.Literal(password) ) ) except psycopg.errors.InsufficientPrivilege: # pragma: no cover - 权限受限的环境 pytest.skip("当前角色无权创建角色，跳过") with psycopg.connect(_DSN, autocommit=True) as admin: admin.execute( sql.SQL("GRANT USAGE ON SCHEMA {} TO {}").format( sql.Identifier(schema), sql.Identifier(role) ) ) admin.execute( sql.SQL( "GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA {} TO {}" ).format(sql.Identifier(schema), sql.Identifier(role)) ) info = psycopg.conninfo.conninfo_to_dict(dsn) if dsn else {} info.update({"user": role, "password": password}) readonly_dsn = psycopg.conninfo.make_conninfo(**info) store = PostgresTaskStore(readonly_dsn) # 没有 CREATE 权限，也必须能起来 try: snapshot = store.get("t1") assert snapshot is not None and snapshot.name == "只读也能看到" assert len(store.query()) == 1 finally: store.close() finally: with psycopg.connect(_DSN, autocommit=True) as admin: admin.execute(sql.SQL("DROP OWNED BY {}").format(sql.Identifier(role))) admin.execute(sql.SQL("DROP ROLE IF EXISTS {}").format(sql.Identifier(role))) _drop_schema(schema) # --------------------------------------------------------------------------- # 4. 接线：Controller 真的能把它当存储用 # --------------------------------------------------------------------------- def _add(ctx: Any, left: int, right: int) -> int: del ctx return left + right def _block(ctx: Any) -> str: for _ in range(1_000_000): ctx.checkpoint() time.sleep(0.005) return "done" def _wait_until(predicate: Any, timeout: float = 10.0) -> bool: deadline = time.monotonic() + timeout while time.monotonic() < deadline: if predicate(): return True time.sleep(0.01) return False def test_state_transitions_persist_and_are_auditable() -> None: schema, dsn = _isolated_dsn() try: store = PostgresTaskStore(dsn) controller = Controller(store=store) controller.register("add", _add) controller.submit("add", 2, 3, task_id="t1") controller.wait_all(timeout=10) snapshot = store.get("t1") assert snapshot is not None assert snapshot.status is TaskStatus.SUCCEEDED assert snapshot.outcome == {"value": 5} assert [item.to_status for item in store.history("t1")] == [ TaskStatus.QUEUED, TaskStatus.STARTING, TaskStatus.RUNNING, TaskStatus.SUCCEEDED, ] store.close() finally: _drop_schema(schema) def test_stop_crosses_stores_through_postgres() -> None: """两个 store = 两个"进程"：worker 在跑，运维从另一套连接停掉它。""" schema, dsn = _isolated_dsn() try: owner_store = PostgresTaskStore(dsn) owner = Controller(store=owner_store, command_poll_interval=0.05, owner="worker-pg") owner.register("block", _block) run = owner.submit("block", task_id="t1") assert _wait_until(lambda: run.status is TaskStatus.RUNNING) client_store = PostgresTaskStore(dsn) client = Controller(store=client_store, owner="ops-pg") assert client.runs() == (), "客户端并不持有这个任务" assert client.stop("t1", reason="跨 store 停止") is True assert _wait_until(lambda: run.status is TaskStatus.STOPPED), "持有者没有取走指令" final = client.get("t1") assert final is not None and final.status is TaskStatus.STOPPED owner.close(timeout=10) client.close(timeout=5) finally: _drop_schema(schema)