host
b0347b374d
- 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).
310 lines
13 KiB
Python
310 lines
13 KiB
Python
"""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()
|