host
2dc35bd331
- db: translate INSERT OR IGNORE ... SELECT to ON CONFLICT DO NOTHING; register task_audit_loop_state (task_id) in _OR_REPLACE_PK (direct port). - server: _backfill_default_quotas uses SAVEPOINT containment, no longer commits the caller's transaction (direct port; callers are init_db only). - worker: NOT ported — integration _coord_child_rows has no column-fallback retry loop, so the InFailedSqlTransaction cascade site does not exist here. - tests: rewritten for the integration base (no cloud.audit_loop / _coord_audit_loop_merge here); real-PG cases skip without DSN/driver. (cherry picked from commit 79e6843610fcd275af874e7b4afc57af7ffa579e)
282 lines
12 KiB
Python
282 lines
12 KiB
Python
"""PG init/backfill 回滚与审计状态 upsert 登记修复测试（集成基底重写版）。
|
||
|
||
移植来源：feat-pg-init-fix 分支提交 854f810 的
|
||
cloud/tests/test_pg_init_fix.py。原测试不可直接带，原因：
|
||
|
||
- 原测试 import main 上不存在的 cloud.audit_loop（LoopState、
|
||
save_loop_state），并调 worker._coord_audit_loop_merge；集成基底
|
||
（02835f6）没有这两个接口，照搬会 import 失败。
|
||
- 重写版只用集成基底上存在的接口：cloud.db._translate_sql、
|
||
cloud.db.connect、cloud.server._backfill_default_quotas、
|
||
cloud.server.init_db，以及 cloud/tests/test_postgres_backend.py
|
||
的 _has_real_driver / _real_pg_params（真实 PG 守卫口径与原
|
||
测试同款：须显式设置 CLOUD_TEST_PG_DSN 且装有 psycopg 驱动，
|
||
否则 skip，绝不 fail；DSN 只从环境变量读，测试不打印它）。
|
||
|
||
覆盖的移植内容（与原测试一一对应，落点不同处已注明）：
|
||
|
||
- 问题一 (a)：INSERT OR IGNORE ... SELECT 形在 PG 适配层下未
|
||
翻译，直接报 syntax error at or near "OR"。重写版直接断言
|
||
_translate_sql 的翻译结果（VALUES 形、SELECT 形含列 / 不含列
|
||
三种），不依赖真实 PG 即可回归。
|
||
- 问题一 (b)：_backfill_default_quotas 的 except 整事务 rollback
|
||
把同事务里 migrate 的 ALTER 一并丢掉。SQLite 版断言保存点
|
||
containment（失败只回滚回填本身、成功不提前提交调用方事务）；
|
||
真实 PG 版断言同款保存点语义 + 新鲜 init 后 tasks 仍有
|
||
paused_by_user / interrupted_reason（即 migrate 的 ALTER 未被
|
||
回填失败丢掉；原测试用 Worker 合流证明同一链条，基底无该
|
||
接口，改以 init 后列存在性直接证明）。
|
||
- 问题二：task_audit_loop_state 未登记进 _OR_REPLACE_PK，PG 下
|
||
INSERT OR REPLACE 退化为 ON CONFLICT DO NOTHING，已存在行的
|
||
round / last_decision 更新静默丢弃。重写版断言 _translate_sql
|
||
对该表产出 ON CONFLICT (task_id) DO UPDATE SET（含 round /
|
||
last_decision），即登记本身；SQLite 原生 OR REPLACE 的行更新
|
||
语义只作对照（证明“应有语义”长什么样，不证明 PG 行为）；
|
||
PG 侧的等价行为由真实 PG 测试覆盖（无 DSN/驱动时 skip）。
|
||
|
||
Worker 部分说明：源提交的 worker 改动是给 _coord_child_rows 的
|
||
列回落重试（7 列 -> 6 列 -> 4 列）之间加 rollback，防 PG
|
||
InFailedSqlTransaction 连环；集成基底的 _coord_child_rows 只有
|
||
单次 4 列 SELECT、无回落重试、无 except 重查，级联发生的落点
|
||
不存在，故本移植不带 worker 改动（见摘要）。本文件不测 worker。
|
||
"""
|
||
|
||
import os
|
||
import sqlite3
|
||
import sys
|
||
import tempfile
|
||
import unittest
|
||
from pathlib import Path
|
||
from unittest import mock
|
||
|
||
try:
|
||
from cloud.db import _translate_sql, connect
|
||
from cloud.server import _backfill_default_quotas, init_db
|
||
except ImportError:
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
|
||
from cloud.db import _translate_sql, connect
|
||
from cloud.server import _backfill_default_quotas, init_db
|
||
|
||
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,
|
||
}
|
||
|
||
|
||
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 SqliteUpsertSemanticsTestCase(unittest.TestCase):
|
||
def test_sqlite_native_replace_updates_existing_row(self):
|
||
# 对照用：SQLite 原生 OR REPLACE 的应有语义——同主键再写
|
||
# 一行，round / last_decision 被新值覆盖。PG 侧的等价行
|
||
# 为由 _OR_REPLACE_PK 登记保证，见 TranslateTestCase 与
|
||
# 真实 PG 测试；本用例不证明 PG 行为。
|
||
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)")
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO task_audit_loop_state"
|
||
" VALUES('P1', 1, '[\"child:a\"]', '{}', 'pushback', 1.0)")
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO task_audit_loop_state"
|
||
" VALUES('P1', 2, '[\"child:a\"]', '{}', 'escalate', 2.0)")
|
||
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)
|
||
# 一次性库由外部建好并指向 DSN；init_db 直接跑，不做
|
||
# 任何夹具补偿（这正是本修复要证明的点）。
|
||
init_db(":memory:")
|
||
|
||
def _pg(self):
|
||
return connect(":memory:")
|
||
|
||
def test_pg_fresh_init_keeps_migrate_columns(self):
|
||
# 新鲜 PG 库 init：migrate_tasks_table 的 ALTER 与回填同
|
||
# 事务；旧代码下回填 SELECT 形未翻译报错、整事务 rollback
|
||
# 会把 ALTER 一并丢掉，tasks 缺 paused_by_user /
|
||
# interrupted_reason。改后两列必须在。
|
||
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)
|
||
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()
|