fix: PG init backfill rollback scope + audit-loop state upsert registration
- db: translate INSERT OR IGNORE ... SELECT to ON CONFLICT DO NOTHING (VALUES-only regex left SELECT form untranslated -> PG syntax error); register task_audit_loop_state (task_id) in _OR_REPLACE_PK so PG INSERT OR REPLACE becomes a real upsert instead of DO NOTHING. - server: _backfill_default_quotas uses SAVEPOINT containment and no longer commits the caller's transaction, so a backfill failure cannot roll back init_db's migrate ALTERs. - worker: _coord_child_rows rolls back before column-fallback retries so a missing-column failure on PG does not cascade into InFailedSqlTransaction. - tests: test_pg_init_fix.py (translation, SQLite backfill/upsert, real-PG fresh-init + merge + upsert + backfill containment, skipped without CLOUD_TEST_PG_DSN/driver).
This commit is contained in:
host committed
1 parent
d99853b684
commit
b0347b374d
4 files changed
+381
-4
No files matched your search
@@ -0,0 +1,309 @@
|
||||
"""PG init/backfill 回滚与审计状态 upsert 登记修复测试。
|
||||
|
||||
只覆盖两处修复，不跑真实模型回合、不联网：
|
||||
- 问题一：新鲜 PG 库 init 后 tasks 缺 paused_by_user /
|
||||
interrupted_reason，合流随后连环报 InFailedSqlTransaction。
|
||||
根因是两层：(a) INSERT OR IGNORE ... SELECT 形在 PG 适配层
|
||||
下未翻译，直接报 syntax error at or near "OR"；
|
||||
(b) _backfill_default_quotas 的 except 整事务 rollback 把
|
||||
同事务里 migrate 的 ALTER 一并丢掉。修法：翻译 SELECT 形、
|
||||
backfill 改保存点 containment、成功时不再由本函数提交。
|
||||
- 问题二：task_audit_loop_state 未登记进 _OR_REPLACE_PK，
|
||||
PG 下 INSERT OR REPLACE 退化为 ON CONFLICT DO NOTHING，
|
||||
已存在行的 round / last_decision 等更新静默丢弃。
|
||||
|
||||
PG 测试与 test_postgres_integration.py 同款口径：须显式设置
|
||||
CLOUD_TEST_PG_DSN 且装有 psycopg 驱动，否则 skip，绝不 fail；
|
||||
DSN 只从环境变量读，测试不打印它。
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
|
||||
try:
|
||||
from cloud.audit_loop import LoopState, save_loop_state
|
||||
from cloud.db import _translate_sql
|
||||
from cloud.server import _backfill_default_quotas, init_db
|
||||
from cloud.worker import Worker
|
||||
except ImportError:
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||||
from cloud.audit_loop import LoopState, save_loop_state
|
||||
from cloud.db import _translate_sql
|
||||
from cloud.server import _backfill_default_quotas, init_db
|
||||
from cloud.worker import Worker
|
||||
|
||||
try:
|
||||
from cloud.tests.test_postgres_backend import (
|
||||
_has_real_driver,
|
||||
_real_pg_params,
|
||||
)
|
||||
_PG_PARAMS = _real_pg_params()
|
||||
except Exception:
|
||||
_PG_PARAMS = None
|
||||
|
||||
|
||||
def _pg_env():
|
||||
"""返回 PG 测试用的环境变量 dict；无真实 PG 时返回 None。"""
|
||||
if _PG_PARAMS is None or not _has_real_driver():
|
||||
return None
|
||||
host, port, dbname, user, password = _PG_PARAMS
|
||||
return {
|
||||
"CLOUD_DB_BACKEND": "postgres",
|
||||
"CLOUD_DB_HOST": host,
|
||||
"CLOUD_DB_PORT": str(port),
|
||||
"CLOUD_DB_NAME": dbname,
|
||||
"CLOUD_DB_USER": user,
|
||||
"CLOUD_DB_PASSWORD": password,
|
||||
}
|
||||
|
||||
|
||||
def _completed_runner(cmd):
|
||||
return SimpleNamespace(returncode=0, stdout="{}", stderr="")
|
||||
|
||||
|
||||
class TranslateTestCase(unittest.TestCase):
|
||||
def test_ignore_select_form_translated(self):
|
||||
sql, kind = _translate_sql(
|
||||
"INSERT OR IGNORE INTO user_quotas(user_id,"
|
||||
" quota_usd_micros, updated_at)"
|
||||
" SELECT u.id, ?, ? FROM users u"
|
||||
" LEFT JOIN user_quotas q ON q.user_id = u.id"
|
||||
" WHERE q.user_id IS NULL")
|
||||
self.assertEqual(kind, "normal")
|
||||
self.assertNotIn("OR IGNORE", sql)
|
||||
self.assertIn("ON CONFLICT DO NOTHING", sql)
|
||||
self.assertIn("%s", sql)
|
||||
self.assertNotIn("?", sql)
|
||||
|
||||
def test_audit_loop_state_or_replace_updates(self):
|
||||
sql, kind = _translate_sql(
|
||||
"INSERT OR REPLACE INTO task_audit_loop_state"
|
||||
"(task_id, round, passed_gates, criteria, last_decision,"
|
||||
" updated_at) VALUES(?, ?, ?, ?, ?, ?)")
|
||||
self.assertEqual(kind, "normal")
|
||||
self.assertIn("ON CONFLICT (task_id) DO UPDATE SET", sql)
|
||||
self.assertNotIn("DO NOTHING", sql)
|
||||
self.assertIn("round = EXCLUDED.round", sql)
|
||||
self.assertIn("last_decision = EXCLUDED.last_decision", sql)
|
||||
|
||||
|
||||
class SqliteBackfillTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||||
|
||||
def tearDown(self):
|
||||
self.tmp.cleanup()
|
||||
|
||||
def _conn_with_users(self):
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
conn.execute("CREATE TABLE users(id TEXT PRIMARY KEY)")
|
||||
conn.execute("INSERT INTO users(id) VALUES('u1')")
|
||||
conn.commit()
|
||||
return conn
|
||||
|
||||
def test_backfill_does_not_commit_caller_transaction(self):
|
||||
# 旧版成功时 conn.commit()，会把调用方未提交的改动一并
|
||||
# 提交；改后提交由调用方负责，rollback 后调用方的改动
|
||||
# 与回填行应一并消失。
|
||||
conn = self._conn_with_users()
|
||||
try:
|
||||
conn.execute("INSERT INTO users(id) VALUES('u2')")
|
||||
self.assertEqual(_backfill_default_quotas(conn), 2)
|
||||
conn.rollback()
|
||||
users = [r[0] for r in conn.execute(
|
||||
"SELECT id FROM users ORDER BY id").fetchall()]
|
||||
self.assertEqual(users, ["u1"])
|
||||
# 回滚连 _ensure 建的 user_quotas 表一起丢掉（同事务），
|
||||
# 表不在即等于回填行已消失；表若在则行数必为 0。
|
||||
tables = [r[0] for r in conn.execute(
|
||||
"SELECT name FROM sqlite_master"
|
||||
" WHERE name = 'user_quotas'").fetchall()]
|
||||
if tables:
|
||||
quotas = conn.execute(
|
||||
"SELECT COUNT(*) FROM user_quotas").fetchone()[0]
|
||||
self.assertEqual(quotas, 0)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def test_backfill_failure_rolls_back_only_to_savepoint(self):
|
||||
# user_quotas 缺 quota_usd_micros 列（_ensure 为
|
||||
# CREATE TABLE IF NOT EXISTS，不会补列），INSERT 必失败；
|
||||
# 旧版整事务 rollback 会把同事务里调用方的改动也丢掉，
|
||||
# 保存点版只回滚回填本身。
|
||||
conn = sqlite3.connect(self.db_path)
|
||||
try:
|
||||
conn.execute("CREATE TABLE users(id TEXT PRIMARY KEY)")
|
||||
conn.execute(
|
||||
"CREATE TABLE user_quotas(user_id TEXT PRIMARY KEY,"
|
||||
" updated_at DOUBLE PRECISION)")
|
||||
conn.execute("INSERT INTO users(id) VALUES('u1')")
|
||||
conn.commit()
|
||||
conn.execute("INSERT INTO users(id) VALUES('u2')")
|
||||
self.assertEqual(_backfill_default_quotas(conn), 0)
|
||||
users = [r[0] for r in conn.execute(
|
||||
"SELECT id FROM users ORDER BY id").fetchall()]
|
||||
self.assertEqual(users, ["u1", "u2"])
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class SqliteUpsertTestCase(unittest.TestCase):
|
||||
def test_save_loop_state_updates_existing_row(self):
|
||||
conn = sqlite3.connect(":memory:")
|
||||
try:
|
||||
conn.execute(
|
||||
"CREATE TABLE task_audit_loop_state(task_id TEXT"
|
||||
" PRIMARY KEY, round INTEGER, passed_gates TEXT,"
|
||||
" criteria TEXT, last_decision TEXT,"
|
||||
" updated_at DOUBLE PRECISION)")
|
||||
save_loop_state(conn, LoopState(
|
||||
task_id="P1", round=1, passed_gates=["child:a"],
|
||||
criteria={"child:a": "child:a completes"},
|
||||
last_decision="pushback"))
|
||||
save_loop_state(conn, LoopState(
|
||||
task_id="P1", round=2, passed_gates=["child:a"],
|
||||
criteria={"child:a": "child:a completes"},
|
||||
last_decision="escalate"))
|
||||
conn.commit()
|
||||
row = conn.execute(
|
||||
"SELECT round, last_decision FROM task_audit_loop_state"
|
||||
" WHERE task_id = 'P1'").fetchone()
|
||||
self.assertEqual(row, (2, "escalate"))
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
class PgInitFixTestCase(unittest.TestCase):
|
||||
"""真实 PG 测试：无 DSN/驱动时全部 skip。"""
|
||||
|
||||
def setUp(self):
|
||||
env = _pg_env()
|
||||
if env is None:
|
||||
self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN 或未装驱动，"
|
||||
"跳过真实 PG 测试")
|
||||
self._patcher = mock.patch.dict(os.environ, env)
|
||||
self._patcher.start()
|
||||
self.addCleanup(self._patcher.stop)
|
||||
self._old = os.environ.get("CODEX_PUSH_ALLOW_PRIVATE_TARGETS")
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = "1"
|
||||
self.addCleanup(self._restore_push_env)
|
||||
# 一次性库由外部建好并指向 DSN；init_db 直接跑，不做
|
||||
# 任何夹具补偿（这正是本修复要证明的点）。
|
||||
init_db(":memory:")
|
||||
self.worker = Worker(":memory:", "/fake/host.py",
|
||||
runner=_completed_runner,
|
||||
env={"XCM_API_KEY": "test-placeholder"},
|
||||
max_retries=0, retry_backoff_seconds=0)
|
||||
self.worker._audit = lambda p, a, d="": None
|
||||
|
||||
def _restore_push_env(self):
|
||||
if self._old is None:
|
||||
os.environ.pop("CODEX_PUSH_ALLOW_PRIVATE_TARGETS", None)
|
||||
else:
|
||||
os.environ["CODEX_PUSH_ALLOW_PRIVATE_TARGETS"] = self._old
|
||||
|
||||
def _pg(self):
|
||||
from cloud.db import connect
|
||||
return connect(":memory:")
|
||||
|
||||
def test_fresh_pg_init_then_merge_without_failed_transaction(self):
|
||||
conn = self._pg()
|
||||
try:
|
||||
cols = {r[1] for r in conn.execute(
|
||||
"PRAGMA table_info(tasks)").fetchall()}
|
||||
self.assertIn("paused_by_user", cols)
|
||||
self.assertIn("interrupted_reason", cols)
|
||||
now = time.time()
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" idempotency_key, depends_on, priority, created_by,"
|
||||
" created_at, updated_at, last_heartbeat, audit_loop,"
|
||||
" parent_task_id, coordinator_part, result_text,"
|
||||
" error_text, coordinator_parts, started_at)"
|
||||
" VALUES('P1', 'p1', 'do', 'running', NULL, '[]', 0,"
|
||||
" 'tester', ?, ?, NULL, 1, NULL, NULL, NULL, NULL,"
|
||||
" ?, ?)",
|
||||
(now, now,
|
||||
json.dumps([{"name": "a", "on_failure": "fail"}]), now))
|
||||
conn.execute(
|
||||
"INSERT INTO tasks(id, project_id, prompt, status,"
|
||||
" idempotency_key, depends_on, priority, created_by,"
|
||||
" created_at, updated_at, last_heartbeat, audit_loop,"
|
||||
" parent_task_id, coordinator_part, result_text,"
|
||||
" error_text, coordinator_parts, started_at)"
|
||||
" VALUES('C1', 'p1', 'do', 'failed', NULL, '[]', 0,"
|
||||
" 'tester', ?, ?, NULL, 0, 'P1', 'a', NULL, 'boom',"
|
||||
" NULL, ?)",
|
||||
(now, now, now))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
# 不做任何补偿重跑 migrate：直接调合流，旧代码下这里
|
||||
# 会因缺列回落连环报 InFailedSqlTransaction。
|
||||
out = self.worker._coord_audit_loop_merge(
|
||||
{"id": "P1", "project_id": "p1", "audit_loop": 1},
|
||||
[{"name": "a", "on_failure": "fail"}],
|
||||
[{"name": "a", "id": "C1"}])
|
||||
self.assertEqual(out["decision"], "pushback")
|
||||
|
||||
def test_pg_upsert_updates_existing_state_row(self):
|
||||
conn = self._pg()
|
||||
try:
|
||||
save_loop_state(conn, LoopState(
|
||||
task_id="P9", round=1, passed_gates=[],
|
||||
criteria={}, last_decision="pushback"))
|
||||
conn.commit()
|
||||
save_loop_state(conn, LoopState(
|
||||
task_id="P9", round=2, passed_gates=["child:a"],
|
||||
criteria={"child:a": "child:a completes"},
|
||||
last_decision="escalate"))
|
||||
conn.commit()
|
||||
row = conn.execute(
|
||||
"SELECT round, passed_gates, criteria, last_decision"
|
||||
" FROM task_audit_loop_state WHERE task_id = 'P9'"
|
||||
).fetchone()
|
||||
self.assertEqual(row[0], 2)
|
||||
self.assertEqual(json.loads(row[1]), ["child:a"])
|
||||
self.assertEqual(json.loads(row[2]),
|
||||
{"child:a": "child:a completes"})
|
||||
self.assertEqual(row[3], "escalate")
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def test_pg_backfill_failure_preserves_caller_transaction(self):
|
||||
# 与 SQLite 版同款保存点断言，跑在 PG 上：缺列导致
|
||||
# INSERT 失败时，同事务里调用方的改动必须还在。
|
||||
conn = self._pg()
|
||||
try:
|
||||
conn.execute("DROP TABLE IF EXISTS user_quotas")
|
||||
conn.execute(
|
||||
"CREATE TABLE user_quotas(user_id TEXT PRIMARY KEY,"
|
||||
" updated_at DOUBLE PRECISION)")
|
||||
conn.execute("CREATE TABLE IF NOT EXISTS users"
|
||||
"(id TEXT PRIMARY KEY)")
|
||||
conn.execute("INSERT INTO users(id) VALUES('u9')"
|
||||
" ON CONFLICT DO NOTHING")
|
||||
conn.commit()
|
||||
conn.execute("INSERT INTO users(id) VALUES('u10')"
|
||||
" ON CONFLICT DO NOTHING")
|
||||
self.assertEqual(_backfill_default_quotas(conn), 0)
|
||||
users = [r[0] for r in conn.execute(
|
||||
"SELECT id FROM users WHERE id IN ('u9', 'u10')"
|
||||
" ORDER BY id").fetchall()]
|
||||
self.assertEqual(users, sorted(["u9", "u10"]))
|
||||
conn.rollback()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in new issue
Block a user