mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:21:58 +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 不写；目标上服务还连着就拒绝；来源、参数、值转换的单测
297 lines
14 KiB
Python
297 lines
14 KiB
Python
"""python -m app.db copy-to-postgres：SQLite 数据目录导入 PostgreSQL（db/to_postgres.py）。
|
||
|
||
端到端那几项要 PostgreSQL（scripts/test-postgres.sh）：另起一个进程在 SQLite 数据目录里跑一次停在审批上的运行（带事件、
|
||
审批、工件引用），导入一个新建的空库，再另起一个进程连这个 PostgreSQL、同一个数据目录，把审批批了，运行接着跑完。
|
||
不连库的几项（来源认不认、服务停没停、参数）两种库都跑。全是虚构数据。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import pytest
|
||
|
||
from app.core.config import settings
|
||
from app.db import to_postgres
|
||
from app.db.to_postgres import CopyRefused, copy_to_postgres
|
||
|
||
PG = bool(os.environ.get("AGENTLAB_TEST_DATABASE_URL", "").strip())
|
||
needs_pg = pytest.mark.skipif(not PG, reason="要 AGENTLAB_TEST_DATABASE_URL（scripts/test-postgres.sh）")
|
||
BACKEND = Path(__file__).resolve().parents[1]
|
||
#: 不会去连的地址：导入前的检查先于连目标库
|
||
NOWHERE = "postgresql+asyncpg://nobody@127.0.0.1:1/none"
|
||
|
||
|
||
def node(nid, ntype, **config):
|
||
return {"id": nid, "type": ntype, "position": {"x": 0, "y": 0}, "data": {"label": nid, "config": config}}
|
||
|
||
|
||
GRAPH = {
|
||
"nodes": [node("start", "input", fields=[{"name": "question"}]),
|
||
node("prep", "transform", mode="template", template="第一段"),
|
||
node("gate", "human", mode="approve", title="看一眼"),
|
||
node("after", "transform", mode="template", template="第二段"),
|
||
node("out", "output", fields=[{"name": "结果", "value": "{{ nodes.after }}"}])],
|
||
"edges": [{"source": a, "target": b} for a, b in
|
||
(("start", "prep"), ("prep", "gate"), ("gate", "after"), ("after", "out"))],
|
||
}
|
||
|
||
MAKE_SOURCE = r'''
|
||
import asyncio, json, sys
|
||
|
||
async def main():
|
||
from sqlalchemy import text
|
||
from app.core import artifact_store
|
||
from app.db.base import SessionLocal, engine, init_db
|
||
from app.db.models import DataSource, Run
|
||
from app.db.models_sources import SourceChange
|
||
from app.engine.runner import run_manager
|
||
|
||
await init_db()
|
||
await run_manager.setup()
|
||
run = await run_manager.start(graph=json.loads(sys.argv[1]), input_payload={"question": "虚构的问题"})
|
||
status = None
|
||
for _ in range(400):
|
||
async with SessionLocal() as s:
|
||
status = (await s.get(Run, run.id)).status
|
||
if status in ("interrupted", "failed", "succeeded"):
|
||
break
|
||
await asyncio.sleep(0.05)
|
||
await run_manager.wait_idle(run.id)
|
||
artifact = await artifact_store.put_json({"区域": "虚构", "数": 1}, kind="test_copy", run_id=run.id)
|
||
async with SessionLocal() as s:
|
||
source = DataSource(name="虚构的源", kind="sqlite", database="/tmp/虚构的源.db")
|
||
s.add(source)
|
||
await s.flush()
|
||
s.add(SourceChange(source_id=source.id, kind="column_added", fingerprint="f" * 64, table_name="订单"))
|
||
long = Run(workflow_name="长" * 300, last_seq=5_000_000_000, origin={"b": 2, "a": 1})
|
||
sql_null = Run(workflow_name="SQL 的空", origin={"x": 1})
|
||
json_null = Run(workflow_name="JSON 的空", origin=None)
|
||
s.add_all([long, sql_null, json_null])
|
||
await s.commit()
|
||
await s.execute(text("UPDATE runs SET origin = NULL WHERE id = :i"), {"i": sql_null.id})
|
||
await s.commit()
|
||
out = {"long": long.id, "sql_null": sql_null.id, "json_null": json_null.id, "source": source.id}
|
||
async with SessionLocal() as s:
|
||
last_seq = (await s.get(Run, run.id)).last_seq
|
||
await run_manager.shutdown()
|
||
await engine.dispose()
|
||
print(json.dumps({"run": run.id, "status": status, "last_seq": last_seq, "artifact": artifact, **out}))
|
||
|
||
asyncio.run(main())
|
||
'''
|
||
|
||
RESUME = r'''
|
||
import asyncio, json, sys
|
||
|
||
async def main():
|
||
from sqlalchemy import select
|
||
from app.core import artifact_store
|
||
from app.db.base import SessionLocal, engine, init_db
|
||
from app.db.models import Approval, Run, RunEvent
|
||
from app.db.models_sources import SourceChange
|
||
from app.engine.runner import run_manager
|
||
|
||
run_id, source_id, artifact = sys.argv[1], sys.argv[2], sys.argv[3]
|
||
await init_db()
|
||
await run_manager.setup()
|
||
await run_manager.resume(run_id, {"approved": True})
|
||
status = None
|
||
for _ in range(400):
|
||
async with SessionLocal() as s:
|
||
status = (await s.get(Run, run_id)).status
|
||
if status in ("failed", "succeeded"):
|
||
break
|
||
await asyncio.sleep(0.05)
|
||
await run_manager.wait_idle(run_id)
|
||
async with SessionLocal() as s:
|
||
seqs = list((await s.execute(select(RunEvent.seq).where(RunEvent.run_id == run_id)
|
||
.order_by(RunEvent.seq))).scalars())
|
||
approvals = list((await s.execute(select(Approval.status).where(Approval.run_id == run_id))).scalars())
|
||
row = await s.get(Run, run_id)
|
||
change = SourceChange(source_id=source_id, kind="column_removed", fingerprint="e" * 64, table_name="订单")
|
||
s.add(change)
|
||
await s.commit()
|
||
out = {"status": status, "error": row.error, "output": row.output, "seqs": seqs, "approvals": approvals,
|
||
"new_change_seq": change.seq, "artifact": artifact_store.load(artifact)}
|
||
await run_manager.shutdown()
|
||
await engine.dispose()
|
||
print(json.dumps(out, ensure_ascii=False))
|
||
|
||
asyncio.run(main())
|
||
'''
|
||
|
||
|
||
def _run(script: str, *args: str, data_dir: Path, database_url: str = "") -> dict:
|
||
env = {**os.environ, "AGENTLAB_DATA_DIR": str(data_dir), "AGENTLAB_DATABASE_URL": database_url,
|
||
"AGENTLAB_CHECKPOINT_URL": "", "AGENTLAB_SCHEDULER": "0", "PYTHONPATH": str(BACKEND)}
|
||
env.pop("AGENTLAB_DB_POOL_SIZE", None)
|
||
done = subprocess.run([sys.executable, "-c", script, *args], env=env, capture_output=True, text=True,
|
||
timeout=180, cwd=str(data_dir))
|
||
assert done.returncode == 0, done.stderr[-4000:]
|
||
return json.loads(done.stdout.strip().splitlines()[-1])
|
||
|
||
|
||
def _sql(url: str, sql: str, *params) -> list[tuple]:
|
||
import psycopg
|
||
|
||
from app.db.urls import libpq_conninfo
|
||
|
||
with psycopg.connect(libpq_conninfo(url)) as conn:
|
||
return [tuple(r) for r in conn.execute(sql, params).fetchall()]
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 不连库：来源、服务停没停、参数
|
||
# --------------------------------------------------------------------------
|
||
|
||
|
||
def test_the_source_must_be_an_sqlite_app_database(tmp_path):
|
||
with pytest.raises(CopyRefused, match="找不到应用库"):
|
||
to_postgres.resolve_source(tmp_path)
|
||
(tmp_path / "agentlab.db").write_bytes(b"not a database at all")
|
||
with pytest.raises(CopyRefused, match="不是 SQLite 库文件"):
|
||
to_postgres.resolve_source(tmp_path)
|
||
|
||
|
||
async def test_a_source_still_open_is_refused_before_anything_is_written(tmp_path):
|
||
import sqlite3
|
||
|
||
sqlite3.connect(tmp_path / "agentlab.db").execute("CREATE TABLE t (x)").connection.close()
|
||
(tmp_path / "agentlab.db-wal").write_bytes(b"")
|
||
with pytest.raises(CopyRefused, match="先停服务"):
|
||
await copy_to_postgres(tmp_path, NOWHERE, progress=lambda _: None)
|
||
|
||
|
||
def test_the_command_needs_a_target(tmp_path, monkeypatch, capsys):
|
||
monkeypatch.setattr(settings, "database_url", "")
|
||
with pytest.raises(SystemExit) as exited:
|
||
to_postgres.main([str(tmp_path)])
|
||
assert exited.value.code == 2 and "没有目标地址" in capsys.readouterr().err
|
||
assert to_postgres.main([str(tmp_path), "mysql://x@h/db"]) == 2
|
||
assert "只认 PostgreSQL" in capsys.readouterr().err
|
||
|
||
|
||
def test_the_dispatcher_knows_its_commands(capsys):
|
||
from app.db.__main__ import main
|
||
|
||
assert main([]) == 2 and "copy-to-postgres" in capsys.readouterr().out
|
||
assert main(["drop-everything"]) == 2 and "认不出的命令" in capsys.readouterr().err
|
||
|
||
|
||
def test_pending_writes_keep_their_order_and_special_channels_apart():
|
||
from langgraph.checkpoint.base import WRITES_IDX_MAP
|
||
|
||
special = next(iter(WRITES_IDX_MAP))
|
||
groups = to_postgres._split_writes([("t1", special, 1), ("t1", "a", 2), ("t2", "b", 3), ("t1", "c", 4)])
|
||
assert groups == {"t1": ([(special, 1)], [("a", 2), ("c", 4)]), "t2": ([], [("b", 3)])}
|
||
|
||
|
||
def test_values_sqlite_kept_but_postgresql_refuses_are_fixed_and_counted():
|
||
from sqlalchemy import Column, Integer, MetaData, String, Table
|
||
|
||
from app.db.base import PortableJSON
|
||
|
||
table = Table("t", MetaData(), Column("n", Integer), Column("s", String(5)), Column("j", PortableJSON))
|
||
fixed: dict[str, int] = {}
|
||
fix = to_postgres._Fixer(fixed)
|
||
assert fix.value(table.c.n, 3.0) == 3 and fix.value(table.c.n, " 12 ") == 12 and fix.value(table.c.n, 7) == 7
|
||
assert fix.value(table.c.s, 42) == "42" and fix.value(table.c.s, "a\x00b") == "ab"
|
||
assert fix.value(table.c.j, '{"b": 1, "a": NaN}') == '{"b": 1, "a": null}'
|
||
assert fix.value(table.c.j, '{"b": 1, "a": 2}') == '{"b": 1, "a": 2}', "平常的 JSON 原文照搬，键的先后不变"
|
||
assert fix.value(table.c.j, None) is None
|
||
assert set(fixed) == {"整数列里的浮点写法", "整数列里的数字字符串", "文本列里的非文本值", "文本里的 NUL 字符",
|
||
"JSON 里的 NaN / Infinity / NUL"}
|
||
|
||
|
||
# --------------------------------------------------------------------------
|
||
# 端到端：SQLite 上停在审批的运行，导入 PostgreSQL 后接着跑完
|
||
# --------------------------------------------------------------------------
|
||
|
||
|
||
@pytest.fixture
|
||
def source_dir(tmp_path):
|
||
"""另起一个进程，在 SQLite 数据目录里留下一次停在审批上的运行和几行只看值的数据。"""
|
||
data = tmp_path / "data"
|
||
data.mkdir()
|
||
made = _run(MAKE_SOURCE, json.dumps(GRAPH), data_dir=data)
|
||
assert made["status"] == "interrupted"
|
||
assert not Path(f"{data}/agentlab.db-wal").exists(), "来源进程正常退出后不该留着 -wal"
|
||
return data, made
|
||
|
||
|
||
@needs_pg
|
||
async def test_a_waiting_run_is_imported_and_finishes_on_postgresql(source_dir):
|
||
from tests.appdb import fresh_postgres_database
|
||
|
||
data, made = source_dir
|
||
lines: list[str] = []
|
||
with fresh_postgres_database("copy") as target:
|
||
report = await copy_to_postgres(data, target, progress=lines.append)
|
||
assert report.ok, "\n".join(to_postgres.render(report))
|
||
by_table = {t.name: t for t in report.tables}
|
||
assert by_table["runs"].source == by_table["runs"].present_after == 4
|
||
assert by_table["run_events"].source > 0 and by_table["approvals"].present_after == 1
|
||
assert by_table["artifacts"].present_after >= 1 and by_table["source_changes"].present_after == 1
|
||
checkpoints, writes = report.checkpoints
|
||
assert checkpoints > 0 and report.checkpoints_target == (checkpoints, writes)
|
||
assert report.threads == 1 and report.checkpoints_written == checkpoints
|
||
|
||
# 值：超长的名字、过 21 亿的整数照存；SQL 的 NULL 和 JSON 的 null 分开；JSON 原文（键的先后）不变
|
||
[(name, seq)] = _sql(target, "SELECT workflow_name, last_seq FROM runs WHERE id = %s", made["long"])
|
||
assert name == "长" * 300 and seq == 5_000_000_000
|
||
assert _sql(target, "SELECT CAST(origin AS TEXT) FROM runs WHERE id = %s", made["long"]) == \
|
||
[('{"b": 2, "a": 1}',)]
|
||
assert _sql(target, "SELECT origin IS NULL FROM runs WHERE id = %s", made["sql_null"]) == [(True,)]
|
||
assert _sql(target, "SELECT CAST(origin AS TEXT) FROM runs WHERE id = %s", made["json_null"]) == [("null",)]
|
||
|
||
# 再导一遍：什么都不重复，目标上已有的不动
|
||
again = await copy_to_postgres(data, target, progress=lambda _: None)
|
||
assert again.ok and again.checkpoints_written == 0
|
||
assert all(t.present_before == t.source for t in again.tables)
|
||
assert {t.name: t.target_after for t in again.tables} == {t.name: t.target_after for t in report.tables}
|
||
|
||
# PostgreSQL 上、同一个数据目录：批了审批，运行从断点接着跑完
|
||
done = _run(RESUME, made["run"], made["source"], made["artifact"], data_dir=data, database_url=target)
|
||
assert done["status"] == "succeeded", done["error"]
|
||
assert done["output"]["结果"] == "第二段"
|
||
assert done["approvals"] == ["resolved"], "导入的那条审批就是被批掉的那条，没有另起一条"
|
||
seqs = done["seqs"]
|
||
assert seqs == sorted(set(seqs)) and seqs[-1] > made["last_seq"], "接着导入前的序号往下编"
|
||
assert done["artifact"] == {"区域": "虚构", "数": 1}
|
||
assert done["new_change_seq"] > 1, "自增主键的序列接在导入的最大值后面"
|
||
|
||
|
||
@needs_pg
|
||
async def test_dry_run_counts_and_writes_nothing(source_dir):
|
||
from tests.appdb import fresh_postgres_database
|
||
|
||
data, _ = source_dir
|
||
with fresh_postgres_database("dry") as target:
|
||
report = await copy_to_postgres(data, target, dry_run=True, progress=lambda _: None)
|
||
assert report.ok and report.checkpoints[0] > 0
|
||
runs = next(t for t in report.tables if t.name == "runs")
|
||
assert runs.source == 4 and runs.present_before == runs.present_after == 0
|
||
assert _sql(target, "SELECT count(*) FROM information_schema.tables WHERE table_name IN "
|
||
"('runs', 'checkpoints')") == [(0,)]
|
||
assert "只核对，什么都没写" in to_postgres.render(report)[-1]
|
||
|
||
|
||
@needs_pg
|
||
async def test_a_target_with_the_service_still_connected_is_refused(source_dir):
|
||
import psycopg
|
||
|
||
from app.db.urls import libpq_conninfo
|
||
from tests.appdb import fresh_postgres_database
|
||
|
||
data, _ = source_dir
|
||
with fresh_postgres_database("busy") as target:
|
||
with psycopg.connect(libpq_conninfo(target), application_name="agentlab"):
|
||
with pytest.raises(CopyRefused, match="服务的连接"):
|
||
await copy_to_postgres(data, target, progress=lambda _: None)
|
||
assert _sql(target, "SELECT count(*) FROM information_schema.tables WHERE table_name = 'runs'") == [(0,)]
|