mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:11:59 +00:00
ThearchyHelios
3e922cbefb
- 先停服务：来源旁边留着 agentlab.db-wal / checkpoints.db-wal（还有连接开着，或者上次是强杀），或者目标库上有 application_name=agentlab 的连接，就拒绝，什么都不写；强杀留下的 -wal 用 --assume-stopped 放过。查不到的（另一台机器上的服务）在说明里写明要自己保证 - 读快照：SQLite 备份接口拷到 --work-dir（缺省系统临时目录），在拷贝上跑一遍启动时的建表补列，老版本的库按当前表结构读；来源按读写打开、只许读（mode=ro 打开 WAL 库会留下它自己建的 -wal，下一次导入就当成服务还开着） - 目标：同启动时的 create_schema 建表；按外键先后逐表、按主键分批 INSERT … ON CONFLICT DO NOTHING——目标上已有同一主键（或撞了别的唯一约束）的行一律不动，重跑只补没有的；整批撞了外键之类的约束就逐行重试，插不进的列出主键和原因；自增主键的序列接到导入的最大值后面 - 值：按模型的列类型读写（时间、布尔），JSON 原文照搬（SQL 的 NULL 和 JSON 的 null 分开、键的先后不变）；PostgreSQL 不收的（NaN、NUL、整数列里的浮点和数字字符串、文本列里的非文本）换掉并计数打印 - 断点：用两个 saver 的公开接口逐条读出写入（含 pending writes，特殊通道和普通写入分开写，序号同来源），目标上已有的断点不重写；按线程核对断点和 pending writes 的条数 - 核对：每张表按主键数来源的行目标上有没有（导入前已有、这次补上分开数），逐表打印；有没导进去的退出码非零。--dry-run 只核对、不建表不写 - 测试：另起进程在 SQLite 数据目录里跑一次停在审批上的运行（事件、审批、待办、工件引用），导入新建的空库、再导一遍什么都不重复，再另起进程连这个 PostgreSQL、同一个数据目录批掉审批，运行从断点接着跑完、事件序号接着往下编、工件读得到、新的变化记录的序号接在导入的后面；dry-run 不写；目标上服务还连着就拒绝；来源、参数、值转换的单测
103 lines
4.3 KiB
Python
103 lines
4.3 KiB
Python
"""测试里另开一条连接直接改应用库：模拟别处（另一个进程、恢复、继续运行）抢在中间提交。
|
||
|
||
应用库是 SQLite 时用 sqlite3 连数据目录里的 agentlab.db（同以前各处的写法）；是 PostgreSQL 时（AGENTLAB_TEST_DATABASE_URL，
|
||
见 conftest）用 psycopg 连同一个库。SQL 一律用 ? 占位，两种库写成同一句；锁等 5 秒（同以前 sqlite3 的 timeout=5）：
|
||
另一条连接要等的锁正好被调用方拿着的话，等满报错，不会把测试卡死。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import contextlib
|
||
import re
|
||
from typing import Any
|
||
|
||
from app.core.config import settings
|
||
|
||
|
||
def is_postgres() -> bool:
|
||
return settings.db_dialect == "postgresql"
|
||
|
||
|
||
def _pg_conninfo() -> str:
|
||
from app.db.urls import libpq_conninfo
|
||
|
||
return libpq_conninfo(settings.db_url)
|
||
|
||
|
||
def execute(*statements: str | tuple[str, tuple[Any, ...]]) -> None:
|
||
"""在另一条连接上依次执行、一起提交。每一项是 SQL，或者 (SQL, 参数)。"""
|
||
items = [(s, ()) if isinstance(s, str) else (s[0], tuple(s[1])) for s in statements]
|
||
if is_postgres():
|
||
import psycopg
|
||
|
||
with psycopg.connect(_pg_conninfo(), options="-c lock_timeout=5000") as conn:
|
||
for sql, params in items:
|
||
conn.execute(re.sub(r"\?", "%s", sql), params)
|
||
conn.commit()
|
||
return
|
||
import sqlite3
|
||
|
||
db = sqlite3.connect(settings.db_path, timeout=5)
|
||
try:
|
||
for sql, params in items:
|
||
db.execute(sql, params)
|
||
db.commit()
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
def query(sql: str, *params: Any) -> list[tuple[Any, ...]]:
|
||
"""在另一条连接上读：看别人这时提交了什么（不是调用方事务里的）。"""
|
||
if is_postgres():
|
||
import psycopg
|
||
|
||
with psycopg.connect(_pg_conninfo(), options="-c lock_timeout=5000") as conn:
|
||
return [tuple(r) for r in conn.execute(re.sub(r"\?", "%s", sql), params).fetchall()]
|
||
import sqlite3
|
||
|
||
db = sqlite3.connect(settings.db_path, timeout=5)
|
||
try:
|
||
return [tuple(r) for r in db.execute(sql, params).fetchall()]
|
||
finally:
|
||
db.close()
|
||
|
||
|
||
def refuse_inserts(table: str, condition: str, *, name: str) -> tuple[list[str], list[str]]:
|
||
"""建一个触发器：往 table 插入满足 condition（用 NEW.列 写）的行时报错，报的是完整性错误（IntegrityError），同 SQLite
|
||
的 RAISE(ABORT)。返回 (建的语句, 删的语句)，各是一串语句，逐条在应用库的连接上执行。
|
||
|
||
PostgreSQL 要先建一个触发器函数（plpgsql）再挂上去，错误码给 23000（integrity_constraint_violation）。"""
|
||
if is_postgres():
|
||
create = [f"CREATE FUNCTION {name}_fn() RETURNS trigger AS $$ BEGIN IF {condition} THEN "
|
||
f"RAISE EXCEPTION 'refused by test' USING ERRCODE = 'integrity_constraint_violation'; END IF; "
|
||
f"RETURN NEW; END $$ LANGUAGE plpgsql",
|
||
f"CREATE TRIGGER {name} BEFORE INSERT ON {table} FOR EACH ROW EXECUTE FUNCTION {name}_fn()"]
|
||
drop = [f"DROP TRIGGER IF EXISTS {name} ON {table}", f"DROP FUNCTION IF EXISTS {name}_fn()"]
|
||
return create, drop
|
||
return ([f"CREATE TRIGGER {name} BEFORE INSERT ON {table} WHEN {condition} "
|
||
f"BEGIN SELECT RAISE(ABORT, 'refused by test'); END"], [f"DROP TRIGGER IF EXISTS {name}"])
|
||
|
||
|
||
def admin_conninfo() -> str:
|
||
"""测试用的 PostgreSQL 服务器上能建库删库的连接串（就是这个 xdist 进程自己的库，建删别的库用）。"""
|
||
return _pg_conninfo()
|
||
|
||
|
||
@contextlib.contextmanager
|
||
def fresh_postgres_database(tag: str):
|
||
"""同一台测试服务器上另建一个空库，交回它的地址（postgresql+asyncpg://…），用完删掉。只在 PostgreSQL 时能用。"""
|
||
import uuid
|
||
|
||
import psycopg
|
||
from sqlalchemy.engine import make_url
|
||
|
||
url = make_url(settings.db_url)
|
||
name = f"{url.database}_{tag}_{uuid.uuid4().hex[:8]}"
|
||
with psycopg.connect(_pg_conninfo(), autocommit=True) as conn:
|
||
conn.execute(f'CREATE DATABASE "{name}"')
|
||
try:
|
||
yield url.set(database=name).render_as_string(hide_password=False)
|
||
finally:
|
||
with psycopg.connect(_pg_conninfo(), autocommit=True) as conn:
|
||
conn.execute(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)')
|