mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:11:59 +00:00
ThearchyHelios
2e7d6103ea
- 配置：database_url、checkpoint_url（启动时校验写法，只认 PostgreSQL）、连接池四项；启动日志写用的是哪种库（主机、库名、用户，不写密码） - 引擎（db/base.make_engine）：SQLite 照旧 WAL、外键、忙等；PostgreSQL 用连接池、探活、会话时区 UTC，装上方言调整（db/postgres.py）：建表时 String(n) 是不限长的 VARCHAR（SQLite 从不查长度，导过来的存量也可能超长）、ORDER BY 里可空列的 NULL 位置照 SQLite 补写（升序在前、降序在后，窗口函数里同样）；Integer 列在 PostgreSQL 上是 BIGINT（列挂到表上时换成 with_variant，绑定参数跟着列按 BIGINT 转换）；JSON 列两边都存原文（PostgreSQL 的 json，不用 JSONB：JSONB 重排对象里的键，快照的列序、表的先后读回来就变了）；JSON 里的 NaN 写成 null、NUL 去掉（PostgreSQL 不收） - 迁移按方言：已有的列和索引用 inspector 认；PostgreSQL 上补列的类型取模型里这一列的、缺省值取迁移表里的（布尔 0/1 换成 false/true），ADD COLUMN IF NOT EXISTS；去重语句里的布尔按方言写；部分唯一索引补 postgresql_where；init_db 在 PostgreSQL 上拿咨询锁（几个进程同时启动不撞） - 应用库上的 SQLite 专用写法改成按方言编译的表达式（db/portable.py）：json_extract 的十几处（口径召回、出具、对账、个人应用、幂等键、访问控制、数据源变化通知、剖析任务、待办超时）换成 json_eq / json_flag / json_text，json_type 换成 json_is_string（两种库同一套带类型的比较）；能力登记的 sqlite insert 换成按会话方言的 insert_ignore；能力使用次数的 json_each 按方言分支；附件「没绑记录」的条件不靠 json 的等号；按报错文字认唯一约束的地方换成 unique_violation（按表和列认，两种库的报错都认得） - 断点（engine/checkpoints.py）：PostgreSQL 时开 psycopg 连接池、在咨询锁里 setup；runner 只认 BaseCheckpointSaver，读断点仍只走 saver 的公开接口 - 依赖：新 extra postgres（asyncpg、psycopg[binary] 3、psycopg-pool、langgraph-checkpoint-postgres 3.1.2）；langgraph-checkpoint 卡在 <4.3（断点格式就在这个包里，postgres 包新版要 4.3，不动已装的 4.2.0）；锁文件补上新包 - 运行收尾：记 action.unknown、安排核实做完才把运行摘掉（以前先摘掉，等「不在跑了」再读事件的一方在库慢时读不到这一条） - 测试：AGENTLAB_TEST_DATABASE_URL 时跑在 PostgreSQL 上（每个 xdist 进程建自己的库、按事件循环分的连接池 tests/pgpool.py），不设时盖掉环境里的数据库地址照旧用 SQLite；scripts/test-postgres.sh 起临时的 postgres:16 容器跑全量；sqlite_only 标记；另开连接改库的几个竞态测试、触发器测试改成两种库都能跑（tests/appdb.py）；新增 test_db_portable 守着方言调整、JSON 键序、超长字符串、大整数、NULL 排序、时区、PostgreSQL 老库迁移
78 lines
3.3 KiB
Python
78 lines
3.3 KiB
Python
"""测试里另开一条连接直接改应用库：模拟别处（另一个进程、恢复、继续运行）抢在中间提交。
|
||
|
||
应用库是 SQLite 时用 sqlite3 连数据目录里的 agentlab.db（同以前各处的写法）；是 PostgreSQL 时（AGENTLAB_TEST_DATABASE_URL，
|
||
见 conftest）用 psycopg 连同一个库。SQL 一律用 ? 占位，两种库写成同一句；锁等 5 秒（同以前 sqlite3 的 timeout=5）：
|
||
另一条连接要等的锁正好被调用方拿着的话，等满报错，不会把测试卡死。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
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}"])
|