380 lines
14 KiB
Python
380 lines
14 KiB
Python
|
"""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)
|