host
8b2686548c
ci / test (push) In progress
Startup race: concurrent init_db (cloud-server + worker) ran CREATE TABLE / CREATE INDEX IF NOT EXISTS against a DB where the object did not exist yet; the check-then-create is not atomic, so both sessions created it and one died with UniqueViolation on pg_type_typname_nsp_index (tables) or pg_class_relname_nsp_index (indexes). Hit on the37ba687and7983567deploys (7983567added idx_tasks_parent); systemd restart healed it every time. Fix (advisory lock, not catch-duplicate_object): init_postgres_schema takes pg_advisory_xact_lock (key 20261108, released at its commit), and server.init_db takes the session-level lock for the whole PG branch so the _ensure_* objects outside POSTGRES_SCHEMA (task_continue_messages, user_gateway_keys, task_attempts, task_audit_loop_state) are covered too. Same idiom as the existing RATE_LIMIT_SWEEP_ADVISORY_KEY sweep lock; key is distinct. Tests: two new real-PG regression tests (skipped without CLOUD_TEST_PG_DSN), 4 processes racing on an empty schema, for init_postgres_schema and for full init_db; both fail on the unfixed code and pass with the fix.
412 lines
17 KiB
Python
412 lines
17 KiB
Python
"""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()
|