"""PG init/backfill 回滚与审计状态 upsert 登记修复测试。 演练注：add/add 冲突取并集——以 main 版（真实接口版）为基底， 另并入移植版独有的 2 项纯 _translate_sql 回归（无列清单 SELECT 形、VALUES 形不变）；移植版其余用例均已被 main 版同名/更强 用例覆盖（save_loop_state、worker 合流、PG 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 traceback import unittest from multiprocessing import Barrier, Process, Queue 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_ignore_select_form_without_column_list_translated(self): # 正则里列清单是可选组；无列清单的 SELECT 形同样要翻译， # 不能原样透传到 PG 报 syntax error。 sql, kind = _translate_sql( "INSERT OR IGNORE INTO user_quotas" " 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("SELECT u.id", sql) def test_ignore_values_form_unchanged(self): # 回归：既有 VALUES 形翻译不受新增 SELECT 形分支影响。 sql, kind = _translate_sql( "INSERT OR IGNORE INTO user_quotas(user_id," " quota_usd_micros, updated_at) VALUES(?, ?, ?)") self.assertEqual(kind, "normal") self.assertNotIn("OR IGNORE", sql) self.assertIn("ON CONFLICT DO NOTHING", 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() def _concurrent_init_db_worker(barrier, queue, wid, env): """并发启动子进程：与其它进程同时跑 server.init_db。""" try: os.environ.update(env) barrier.wait(timeout=60) from cloud.server import init_db init_db(":memory:") queue.put((wid, "ok", "")) except Exception: queue.put((wid, "FAIL", traceback.format_exc())) 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() def test_concurrent_init_db_no_unique_violation(self): """并发启动回归：4 进程同时 init_db，须全部成功。 覆盖整条启动路径（init_postgres_schema + _ensure_*）： _ensure_* 建的 task_continue_messages / user_gateway_keys / task_attempts / task_audit_loop_state 不在 POSTGRES_SCHEMA 里，只锁 schema 函数不够；init_db 持会话级 advisory 锁 覆盖全程。空库上跑，未修复时第 1 轮即 3/4 进程报 UniqueViolation（pg_type_typname_nsp_index）。 """ env = _pg_env() if env is None: self.skipTest("未设置可达的 CLOUD_TEST_PG_DSN 或未装驱动，" "跳过真实 PG 测试") from cloud.db import PostgresDatabase db = PostgresDatabase( host=env["CLOUD_DB_HOST"], port=int(env["CLOUD_DB_PORT"]), dbname=env["CLOUD_DB_NAME"], user=env["CLOUD_DB_USER"], password=env["CLOUD_DB_PASSWORD"]) conn = db.connect() try: conn.execute("DROP SCHEMA public CASCADE") conn.execute("CREATE SCHEMA public") conn.commit() finally: conn.close() nworkers = 4 barrier = Barrier(nworkers) queue = Queue() procs = [Process(target=_concurrent_init_db_worker, args=(barrier, queue, i, env)) for i in range(nworkers)] try: for proc in procs: proc.start() for proc in procs: proc.join(timeout=180) finally: for proc in procs: if proc.is_alive(): proc.terminate() results = [queue.get() for _ in range(nworkers)] failures = [r for r in results if r[1] == "FAIL"] self.assertEqual([], failures, "并发 init_db 有进程失败：%s" % (failures,)) # 初始化完整：schema 表与 _ensure_* 表都在。 conn = db.connect() try: rows = db.read_all( conn, "SELECT table_name FROM information_schema.tables" " WHERE table_schema = 'public'", ()) names = {r["table_name"] for r in rows} for table in ("users", "projects", "tasks", "task_continue_messages", "user_gateway_keys", "task_attempts", "task_audit_loop_state"): self.assertIn(table, names) finally: conn.close() if __name__ == "__main__": unittest.main()