fix(transfer): support PostgreSQL checkpoint CAS

This commit is contained in:
jxxghp
2026-08-28 22:27:59 +08:00
parent 92e745fc97
commit 0445c655de
2 changed files with 70 additions and 4 deletions

View File

@@ -10,6 +10,7 @@ from sqlalchemy import (
Text,
UniqueConstraint,
and_,
cast as sql_cast,
column,
delete,
exists,
@@ -19,7 +20,9 @@ from sqlalchemy import (
table,
update,
)
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, Session, mapped_column
from sqlalchemy.sql.elements import ColumnElement
from app.db.base import Base, execute_dml, get_id_column
@@ -389,7 +392,11 @@ class TransferPending(Base):
now_utc: str,
updated_at: str,
) -> int:
"""仅以有效租约和完整冻结计划把任务推进或保持为 running。"""
"""仅以有效租约和完整冻结计划把任务推进或保持为 running。
PostgreSQL 的 JSON 类型不支持直接等值比较,计划匹配条件会在该方言下
临时转换为 JSONB,避免执行 CAS 时触发数据库操作符错误。
"""
if not all((
task_id,
lease_token,
@@ -400,6 +407,10 @@ class TransferPending(Base):
updated_at,
)):
return 0
checkpoint_payload_match = cls._checkpoint_payload_match(
db,
checkpoint_payload,
)
return execute_dml(
db,
update(cls)
@@ -408,7 +419,7 @@ class TransferPending(Base):
cls.state == admission_state,
cls.state.in_(("planned", "provider_pending")),
cls.checkpoint_version == checkpoint_version,
cls.checkpoint_payload == checkpoint_payload,
checkpoint_payload_match,
cls.planned_at.is_not(None),
cls.lease_token == lease_token,
cls.lease_expires_at.is_not(None),
@@ -424,6 +435,20 @@ class TransferPending(Base):
execution_options={"synchronize_session": False},
)
@classmethod
def _checkpoint_payload_match(
cls,
db: Session,
checkpoint_payload: dict[str, Any],
) -> ColumnElement[bool]:
"""构造跨数据库的冻结计划匹配条件。"""
if db.get_bind(mapper=cls).dialect.name == "postgresql":
return sql_cast(cls.checkpoint_payload, JSONB) == sql_cast(
checkpoint_payload,
JSONB,
)
return cls.checkpoint_payload == checkpoint_payload
@classmethod
def defer_execution(
cls,

View File

@@ -4,8 +4,9 @@ from dataclasses import replace
from datetime import datetime, timezone
import pytest
from sqlalchemy import create_engine, select
from sqlalchemy.orm import sessionmaker
from sqlalchemy import create_engine, create_mock_engine, select
from sqlalchemy.dialects import postgresql
from sqlalchemy.orm import Session, sessionmaker
from app.application.transfer.execution import (
TransferExecutionCheckpoint,
@@ -337,6 +338,46 @@ def test_stage_execution_running_cas_binds_exact_plan_identity(execution_store):
assert current == 1
def test_stage_execution_running_casts_postgresql_json_before_comparison(
monkeypatch,
):
"""PostgreSQL 的 JSON 计划比较必须转为 JSONB,避免缺少 json 等值操作符。"""
engine = create_mock_engine(
"postgresql+psycopg://",
lambda *_args, **_kwargs: None,
)
statements = []
def capture_statement(_db, statement, execution_options=None):
"""捕获模型生成的 DML,检查生产 PostgreSQL 方言的 SQL。"""
del execution_options
statements.append(statement)
return 1
monkeypatch.setattr(
"app.db.models.transferpending.execute_dml",
capture_statement,
)
payload = {"schema_version": 1, "items": []}
with Session(engine) as session:
updated = TransferPending.stage_execution_running(
session,
task_id="task-1",
lease_token="lease-1",
admission_state="planned",
checkpoint_version=1,
checkpoint_payload=payload,
now_utc="2026-08-27 01:30:00.000000",
updated_at="2026-08-27 09:30:00",
)
assert updated == 1
assert len(statements) == 1
sql = str(statements[0].compile(dialect=postgresql.dialect()))
assert "CAST(transferpending.checkpoint_payload AS JSONB)" in sql
assert "CAST(%(param_1)s::JSONB AS JSONB)" in sql
def test_success_path_persists_steps_and_execution_checkpoint(execution_store):
"""成功路径应保留每步证据,并提交可供唯一 durable writer 结算的检查点。"""
_seed_pending(execution_store)