新增 python -m app.db copy-to-postgres:把 SQLite 数据目录里的应用库和断点导入 PostgreSQL(db/to_postgres.py)

- 先停服务:来源旁边留着 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 不写;目标上服务还连着就拒绝;来源、参数、值转换的单测
This commit is contained in:
ThearchyHelios
2026-10-19 10:07:06 +08:00
parent f63d037b2c
commit 3e922cbefb
4 changed files with 858 additions and 0 deletions
+31
View File
@@ -0,0 +1,31 @@
"""应用库的维护命令:python -m app.db <命令>。
copy-to-postgres 把 SQLite 数据目录里的应用库和断点导入 PostgreSQL(db/to_postgres.py;先停服务)
"""
from __future__ import annotations
import sys
COMMANDS = {"copy-to-postgres": "把 SQLite 数据目录里的应用库和断点导入 PostgreSQL(先停服务)"}
def main(argv: list[str] | None = None) -> int:
args = list(sys.argv[1:] if argv is None else argv)
if not args or args[0] in ("-h", "--help"):
print("用法:python -m app.db <命令> [参数]\n")
for name, what in COMMANDS.items():
print(f" {name:<18}{what}")
print("\n各命令加 --help 看参数")
return 0 if args else 2
command, rest = args[0], args[1:]
if command == "copy-to-postgres":
from app.db.to_postgres import main as copy_main
return copy_main(rest)
print(f"认不出的命令:{command}(可用的:{'、'.join(COMMANDS)})", file=sys.stderr)
return 2
if __name__ == "__main__":
raise SystemExit(main())
+506
View File
@@ -0,0 +1,506 @@
"""把 SQLite 数据目录里的应用库和断点导入 PostgreSQL:`python -m app.db copy-to-postgres`。
python -m app.db copy-to-postgres <数据目录或 agentlab.db> [目标地址] [--checkpoint-url 地址] [--dry-run]
目标地址缺省取 AGENTLAB_DATABASE_URL;断点缺省进同一个库(AGENTLAB_CHECKPOINT_URL 或 --checkpoint-url 另给)。
**先停服务。** 来源库旁边有 agentlab.db-wal(WAL 模式下还有连接开着,或者上次是强杀),或者目标库上有 application_name
是 agentlab 的连接(服务已经连着 PostgreSQL 跑起来了),就拒绝导入。确认是强杀留下的 -wal 可以加 --assume-stopped;
这只是能查到的两样,查不到的(比如服务连着别的数据目录的拷贝、另一台机器上的服务)要自己保证。
**读的是一份快照。** 先用 SQLite 的备份接口把 agentlab.db、checkpoints.db 拷到临时目录(拷贝过程中别人在写也拿到
一致的一份;放在 --work-dir,缺省系统临时目录,要留出和两个库一样大的空间),在拷贝上跑一遍和启动时一样的建表补列
(老版本的库升级到当前的表结构),从拷贝读,来源的内容不改。
**幂等。** 目标上建表(同启动时的 create_schema),按外键先后逐表、按主键分批插入,`ON CONFLICT DO NOTHING`:
目标上已经有同一主键(或撞了别的唯一约束)的行一律不动,不覆盖、不合并——重跑只补上还没有的行,不重复。所以
导入之后又在 SQLite 上改过的行,重跑不会带过去;要整份重来就清空目标库再导。断点同理:目标上已有的断点不重写,
pending writes 按(断点、任务、序号)插入,已有的不动。
**核对。** 每张表按主键核对:来源的每一行,目标上有没有同一主键的行(导入前就有的、这次补上的分开数),断点和
pending writes 按条数核对,逐项打印;有来源的行目标上没有(撞了约束没插进去的会列出主键和原因),退出码非零。
目标上另有来源里没有的行不算错(同一个目标导过别的、或者服务已经在上面跑过)。
**值的转换。** 读写都按模型的列类型走:时间是不带时区的 UTC、布尔是真布尔、JSON 原文照搬(SQL 的 NULL 和 JSON 的
null 分开,不混)。SQLite 不查类型、PostgreSQL 查:整数列里的整数值的浮点、数字字符串换成整数,文本列里的非文本换成
字符串;文本里的 NUL 字符(PostgreSQL 不收)去掉,JSON 里的 NaN / Infinity 写成 null,各自计数打印。
"""
from __future__ import annotations
import argparse
import asyncio
import json
import sqlite3
import sys
import tempfile
from collections.abc import Callable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
from sqlalchemy import JSON, Boolean, Float, Integer, LargeBinary, String, Text, bindparam, cast, func, select, text
from sqlalchemy import create_engine as create_sync_engine
from sqlalchemy.dialects import postgresql
from sqlalchemy.exc import IntegrityError
from app.db.urls import describe, libpq_conninfo, normalize_checkpoint_url, normalize_database_url
COPY_APPLICATION_NAME = "agentlab-copy"
DEFAULT_BATCH = 1000
@dataclass
class TableResult:
name: str
#: 来源的行数
source: int = 0
#: 来源的这些主键,目标上导入前已经有的(--dry-run 时就是「不用补的」)、导入后有的
present_before: int = 0
present_after: int = 0
#: 目标整张表的行数(导入前、导入后):比来源多出来的是目标上原有的
target_before: int = 0
target_after: int = 0
refused: list[str] = field(default_factory=list)
@property
def missing(self) -> int:
"""来源有、目标上还没有的行(按主键数)。"""
return self.source - self.present_after
@dataclass
class Report:
tables: list[TableResult] = field(default_factory=list)
checkpoints: tuple[int, int] = (0, 0) # 来源(断点、pending writes)
checkpoints_target: tuple[int, int] = (0, 0) # 目标上同一批线程的
checkpoints_written: int = 0
threads: int = 0
fixed: dict[str, int] = field(default_factory=dict)
dry_run: bool = False
@property
def ok(self) -> bool:
if self.dry_run:
return True
tables_ok = all(t.present_after == t.source and not t.refused for t in self.tables)
return tables_ok and self.checkpoints_target[0] >= self.checkpoints[0] \
and self.checkpoints_target[1] >= self.checkpoints[1]
class CopyRefused(RuntimeError):
"""导入前的检查没过(服务还开着、来源不对):照实说明,什么都没写。"""
# --------------------------------------------------------------------------
# 来源
# --------------------------------------------------------------------------
def resolve_source(raw: str | Path) -> tuple[Path, Path | None]:
"""数据目录或 agentlab.db → (应用库文件, 断点库文件或 None)。"""
path = Path(raw).expanduser()
db = path / "agentlab.db" if path.is_dir() else path
if not db.is_file():
raise CopyRefused(f"找不到应用库:{db}(给数据目录,或者直接给 agentlab.db 的路径)")
with open(db, "rb") as f:
if f.read(16) != b"SQLite format 3\x00":
raise CopyRefused(f"{db} 不是 SQLite 库文件")
ckpt = db.parent / "checkpoints.db"
return db, (ckpt if ckpt.is_file() else None)
def source_in_use(db: Path, ckpt: Path | None = None) -> Path | None:
"""来源库是不是还有连接开着:WAL 模式下最后一个连接正常关掉时 -wal 会被删掉,还在就是有人开着(或者上次是强杀)。
交回还在的那个 -wal(应用库、断点库各看一个),都没有是 None。"""
for path in (db, ckpt):
if path is not None and (wal := Path(f"{path}-wal")).exists():
return wal
return None
def snapshot(db: Path, into: Path) -> Path:
"""用 SQLite 的备份接口拷一份一致的快照(别人正在写也是某一刻的完整一份)。
来源按读写打开、只许读(query_only):WAL 模式的库用 mode=ro 打开,关掉时删不了它自己建出来的 -wal、-shm,
下一次导入就会当成服务还开着;读写打开的连接是最后一个关的,关掉时 SQLite 把日志并回主库、删掉这两个文件
(内容不变)。"""
out = into / db.name
src = sqlite3.connect(db)
src.execute("PRAGMA query_only = 1")
try:
dst = sqlite3.connect(out)
try:
src.backup(dst)
finally:
dst.close()
finally:
src.close()
return out
def upgrade_snapshot(path: Path) -> None:
"""在快照上跑一遍启动时的建表补列:老版本的库按当前的表结构读。"""
from app.db import models # noqa: F401
from app.db.base import Base, _migrate
engine = create_sync_engine(f"sqlite:///{path}")
try:
with engine.begin() as conn:
Base.metadata.create_all(conn)
_migrate(conn)
finally:
engine.dispose()
# --------------------------------------------------------------------------
# 值
# --------------------------------------------------------------------------
def _is_json(column: Any) -> bool:
return isinstance(column.type, JSON)
class _Fixer:
"""SQLite 里存得进、PostgreSQL 不收的值,换掉并计数。"""
def __init__(self, fixed: dict[str, int]) -> None:
self.fixed = fixed
def _count(self, what: str) -> None:
self.fixed[what] = self.fixed.get(what, 0) + 1
def value(self, column: Any, value: Any) -> Any:
if value is None:
return None
if _is_json(column):
return self.json(value)
kind = column.type
if isinstance(kind, Boolean):
return bool(value)
if isinstance(kind, Integer):
if isinstance(value, bool) or isinstance(value, int):
return int(value)
if isinstance(value, float) and value.is_integer():
self._count("整数列里的浮点写法")
return int(value)
if isinstance(value, str) and value.strip().lstrip("-").isdigit():
self._count("整数列里的数字字符串")
return int(value.strip())
return value
if isinstance(kind, Float):
return float(value) if isinstance(value, (int, float)) else value
if isinstance(kind, LargeBinary):
return bytes(value) if isinstance(value, (bytes, bytearray, memoryview)) else str(value).encode()
if isinstance(kind, (String, Text)) and not isinstance(value, str):
self._count("文本列里的非文本值")
value = str(value)
if isinstance(value, str) and "\x00" in value:
self._count("文本里的 NUL 字符")
value = value.replace("\x00", "")
return value
def json(self, raw: Any) -> str:
"""JSON 原文 → 写进 PostgreSQL 的原文。平常的原样照搬(键的先后不变);有 NaN / NUL 的才重写一遍。"""
from app.db.base import pg_json_dumps
textual = raw if isinstance(raw, str) else json.dumps(raw)
if "NaN" in textual or "Infinity" in textual or "\\u0000" in textual:
redone = pg_json_dumps(json.loads(textual))
if redone != textual:
self._count("JSON 里的 NaN / Infinity / NUL")
return redone
return textual
# --------------------------------------------------------------------------
# 应用库
# --------------------------------------------------------------------------
def _pk(table: Any) -> list[Any]:
return list(table.primary_key.columns) or list(table.columns)
def _insert(table: Any, columns: list[Any]) -> Any:
"""目标上的插入:JSON 列按原文绑定再转 json(SQL 的 NULL 还是 NULL),撞了任何唯一约束都不动目标上的行。"""
values = {c.name: (cast(bindparam(c.name, type_=Text()), c.type) if _is_json(c) else bindparam(c.name, type_=c.type))
for c in columns}
return postgresql.insert(table).values(values).on_conflict_do_nothing()
async def _count(conn: Any, table: Any) -> int:
return int((await conn.execute(select(func.count()).select_from(table))).scalar() or 0)
async def _write_batch(target: Any, statement: Any, rows: list[dict[str, Any]], table: Any,
refused: list[str]) -> None:
"""一批一个事务。整批撞了约束(外键对不上之类)就逐行重试,插不进的记下主键和原因、跳过。"""
try:
async with target.begin() as conn:
await conn.execute(statement, rows)
return
except IntegrityError:
pass
keys = [c.name for c in _pk(table)]
for row in rows:
try:
async with target.begin() as conn:
await conn.execute(statement, [row])
except IntegrityError as e:
from app.core.errors import first_line
ident = ", ".join(f"{k}={row.get(k)!r}" for k in keys)
refused.append(f"{ident}:{first_line(e.orig or e)}")
def _present(table: Any, keys: list[tuple[Any, ...]]) -> Any:
"""目标上这些主键有几行。"""
from sqlalchemy import tuple_
pk = _pk(table)
where = pk[0].in_([k[0] for k in keys]) if len(pk) == 1 else tuple_(*pk).in_(keys)
return select(func.count()).select_from(table).where(where)
async def copy_tables(snapshot_db: Path, target: Any, *, batch: int, dry_run: bool, report: Report,
progress: Callable[[str], None]) -> None:
"""按外键先后逐表、按主键分批:插入(目标上已有的不动),每批插完按主键数一遍目标上有了几行。"""
from sqlalchemy import type_coerce
from app.db import models # noqa: F401
from app.db.base import Base
fixer = _Fixer(report.fixed)
source = create_sync_engine(f"sqlite:///{snapshot_db}")
try:
with source.connect() as src:
for table in Base.metadata.sorted_tables:
result = TableResult(table.name)
report.tables.append(result)
result.source = int(src.execute(select(func.count()).select_from(table)).scalar() or 0)
async with target.connect() as conn:
exists = await conn.run_sync(lambda sync, name=table.name: _has_table(sync, name))
result.target_before = await _count(conn, table) if exists else 0
columns = list(table.columns)
pk_at = [columns.index(c) for c in _pk(table)]
# JSON 列按原文取(type_coerce 成 Text:SQL 的 NULL 是 None,JSON 的 null 是 'null'),别的按列类型取
reads = [type_coerce(c, Text()).label(c.name) if _is_json(c) else c for c in columns]
statement = _insert(table, columns)
rows = src.execution_options(yield_per=batch).execute(select(*reads).order_by(*_pk(table)))
done = 0
for n, part in enumerate(rows.partitions(batch), 1):
keys = [tuple(row[i] for i in pk_at) for row in part]
if exists:
async with target.connect() as conn:
result.present_before += int((await conn.execute(_present(table, keys))).scalar() or 0)
if not dry_run:
payload = [{c.name: fixer.value(c, row[i]) for i, c in enumerate(columns)} for row in part]
await _write_batch(target, statement, payload, table, result.refused)
async with target.connect() as conn:
result.present_after += int((await conn.execute(_present(table, keys))).scalar() or 0)
done += len(part)
if n % 20 == 0:
progress(f" {table.name}:{done}/{result.source}")
if dry_run:
result.present_after, result.target_after = result.present_before, result.target_before
continue
await _reset_sequences(target, table)
async with target.connect() as conn:
result.target_after = await _count(conn, table)
finally:
source.dispose()
def _has_table(sync: Any, name: str) -> bool:
from sqlalchemy import inspect
return inspect(sync).has_table(name)
async def _reset_sequences(target: Any, table: Any) -> None:
"""自增主键的序列接到已有的最大值后面:导入时写的是来源的主键值,序列自己没动过,不改的话下一条撞主键。"""
for column in table.primary_key.columns:
if column is not table.autoincrement_column:
continue
async with target.begin() as conn:
await conn.execute(text(
f"SELECT setval(pg_get_serial_sequence(:t, :c), GREATEST(COALESCE(MAX({column.name}), 0), 1), "
f"MAX({column.name}) IS NOT NULL) FROM {table.name}"), {"t": table.name, "c": column.name})
async def services_on(target: Any) -> int:
"""目标库上服务的连接数(application_name 是 agentlab 的,不算本命令自己的)。"""
from app.db.base import APPLICATION_NAME
async with target.connect() as conn:
return int((await conn.execute(text(
"SELECT count(*) FROM pg_stat_activity WHERE datname = current_database() AND application_name = :n "
"AND pid <> pg_backend_pid()"), {"n": APPLICATION_NAME})).scalar() or 0)
# --------------------------------------------------------------------------
# 断点
# --------------------------------------------------------------------------
def _split_writes(writes: list[tuple[str, str, Any]]) -> dict[str, tuple[list[tuple[str, Any]], list[tuple[str, Any]]]]:
"""pending writes 按任务分组,各分成「特殊通道」(错误、中断、恢复……,序号是固定的负数)和普通的两拨:普通的分开
单独写,序号才和来源里的一样从 0 排起(来源里一个任务在一个断点上的普通写入是一次写下的)。"""
from langgraph.checkpoint.base import WRITES_IDX_MAP
out: dict[str, tuple[list[tuple[str, Any]], list[tuple[str, Any]]]] = {}
for task_id, channel, value in writes:
special, regular = out.setdefault(task_id, ([], []))
(special if channel in WRITES_IDX_MAP else regular).append((channel, value))
return out
async def copy_checkpoints(snapshot_ckpt: Path, conninfo: str, *, dry_run: bool, report: Report,
progress: Callable[[str], None]) -> None:
"""逐条读出来源的断点(saver 的公开接口,含 pending writes),写进 PostgreSQL 的断点库,再按线程核对条数。"""
from app.engine.checkpoints import count, open_postgres_saver, open_sqlite_saver
threads: set[str] = set()
async with open_sqlite_saver(snapshot_ckpt) as src:
checkpoints = writes = 0
async for tup in src.alist(None):
checkpoints += 1
writes += len(tup.pending_writes or ())
threads.add(str(tup.config["configurable"]["thread_id"]))
report.checkpoints, report.threads = (checkpoints, writes), len(threads)
if dry_run:
return
async with open_postgres_saver(conninfo, max_size=4, application_name=COPY_APPLICATION_NAME) as dst:
done = 0
async for tup in src.alist(None):
conf = tup.config["configurable"]
if await dst.aget_tuple(tup.config) is None:
parent = (tup.parent_config or {}).get("configurable", {}).get("checkpoint_id")
put_config = {"configurable": {"thread_id": conf["thread_id"],
"checkpoint_ns": conf.get("checkpoint_ns", ""),
**({"checkpoint_id": parent} if parent else {})}}
await dst.aput(put_config, tup.checkpoint, tup.metadata,
dict(tup.checkpoint.get("channel_versions") or {}))
report.checkpoints_written += 1
for task_id, (special, regular) in _split_writes(list(tup.pending_writes or ())).items():
for group in (special, regular):
if group:
await dst.aput_writes(tup.config, group, task_id)
done += 1
if done % 500 == 0:
progress(f" 断点:{done}/{checkpoints}")
target = [0, 0]
for thread in sorted(threads):
c, w = await count(dst, thread)
target[0] += c
target[1] += w
report.checkpoints_target = (target[0], target[1])
# --------------------------------------------------------------------------
# 入口
# --------------------------------------------------------------------------
async def copy_to_postgres(source: str | Path, target_url: str, *, checkpoint_url: str | None = None,
dry_run: bool = False, assume_stopped: bool = False, batch: int = DEFAULT_BATCH,
work_dir: str | Path | None = None, progress: Callable[[str], None] = print) -> Report:
from app.db.base import create_schema, make_engine
target_url = normalize_database_url(target_url)
conninfo = normalize_checkpoint_url(checkpoint_url) if checkpoint_url else libpq_conninfo(target_url)
db, ckpt = resolve_source(source)
if (wal := source_in_use(db, ckpt)) and not assume_stopped:
raise CopyRefused(f"{wal} 还在:服务看起来还开着这个库。先停服务再导入;确认服务已停(上次是强杀,留下了 "
"-wal)就加 --assume-stopped")
report = Report(dry_run=dry_run)
target = make_engine(target_url, application_name=COPY_APPLICATION_NAME)
try:
if not assume_stopped and (on := await services_on(target)):
raise CopyRefused(f"目标库上已经有 {on} 条服务的连接(application_name=agentlab):先停掉连着它的服务再导入")
progress(f"来源:{db}{f',断点 {ckpt}' if ckpt else '(没有断点库)'}")
progress(f"目标:PostgreSQL {describe(target_url)}{'(只核对,不写)' if dry_run else ''}")
with tempfile.TemporaryDirectory(prefix="agentlab-copy-", dir=work_dir) as tmp:
snap = snapshot(db, Path(tmp))
upgrade_snapshot(snap)
if not dry_run:
await create_schema(target)
await copy_tables(snap, target, batch=batch, dry_run=dry_run, report=report, progress=progress)
if ckpt is not None:
ckpt_snap = snapshot(ckpt, Path(tmp))
await copy_checkpoints(ckpt_snap, conninfo, dry_run=dry_run, report=report, progress=progress)
finally:
await target.dispose()
return report
def render(report: Report) -> list[str]:
lines = ["", f"{'表':<28}{'来源':>10}{'导入前已有':>10}{'导入后已有':>10}{'目标总行数':>10} 结论"]
for t in report.tables:
lacking = t.source - t.present_after
if not t.source:
verdict = "空表"
elif report.dry_run:
verdict = f"要补 {lacking} 行" if lacking else "都已经有了"
elif t.refused:
verdict = f"少 {lacking} 行:撞了约束,没插进去(见下)"
elif lacking:
verdict = f"少 {lacking} 行"
else:
verdict = "一致" + (f"(这次补了 {t.present_after - t.present_before} 行)"
if t.present_after > t.present_before else "(都已经有了)")
lines.append(f"{t.name:<28}{t.source:>10}{t.present_before:>10}{t.present_after:>10}{t.target_after:>10}"
f" {verdict}")
for t in report.tables:
for item in t.refused[:20]:
lines.append(f" {t.name} 没插进去:{item}")
if len(t.refused) > 20:
lines.append(f" {t.name} 另有 {len(t.refused) - 20} 行没插进去")
c, w = report.checkpoints
if report.dry_run:
lines.append(f"断点:来源 {report.threads} 次运行、{c} 个断点、{w} 条 pending writes")
else:
tc, tw = report.checkpoints_target
lines.append(f"断点:来源 {report.threads} 次运行、{c} 个断点、{w} 条 pending writes;目标上这些运行有 {tc} 个断点、"
f"{tw} 条 pending writes(这次新写 {report.checkpoints_written} 个断点)"
+ ("" if tc >= c and tw >= w else "——对不上"))
for what, n in report.fixed.items():
lines.append(f"转换:{what} {n} 处")
if report.dry_run:
lines.append("只核对,什么都没写")
else:
lines.append("核对通过" if report.ok else "核对没过:见上面标出的表和断点")
return lines
def main(argv: list[str] | None = None) -> int:
from app.core.config import settings
parser = argparse.ArgumentParser(prog="python -m app.db copy-to-postgres",
description="把 SQLite 数据目录里的应用库和断点导入 PostgreSQL(先停服务)。")
parser.add_argument("source", help="数据目录,或者 agentlab.db 的路径(断点库 checkpoints.db 在同一个目录里找)")
parser.add_argument("target", nargs="?", help="目标地址 postgresql+asyncpg://…;不给就用 AGENTLAB_DATABASE_URL")
parser.add_argument("--checkpoint-url", help="断点放到另一个 PostgreSQL 库;不给就用 AGENTLAB_CHECKPOINT_URL,再没有就进目标库")
parser.add_argument("--dry-run", action="store_true", help="只核对:数来源和目标各有多少行,不建表、不写")
parser.add_argument("--assume-stopped", action="store_true",
help="不查服务停没停(来源旁边留着强杀后的 -wal 时用)")
parser.add_argument("--batch", type=int, default=DEFAULT_BATCH, help=f"每批插入的行数,缺省 {DEFAULT_BATCH}")
parser.add_argument("--work-dir", help="放快照的目录(要留出和两个库一样大的空间),缺省系统临时目录")
args = parser.parse_args(argv)
target = args.target or settings.database_url
if not target:
parser.error("没有目标地址:给第二个参数,或者设 AGENTLAB_DATABASE_URL")
try:
report = asyncio.run(copy_to_postgres(
args.source, target, checkpoint_url=args.checkpoint_url or settings.checkpoint_url or None,
dry_run=args.dry_run, assume_stopped=args.assume_stopped, batch=max(args.batch, 1),
work_dir=args.work_dir))
except (CopyRefused, ValueError) as e:
print(f"没有导入:{e}", file=sys.stderr)
return 2
for line in render(report):
print(line)
return 0 if report.ok else 1
+25
View File
@@ -7,6 +7,7 @@
from __future__ import annotations
import contextlib
import re
from typing import Any
@@ -75,3 +76,27 @@ def refuse_inserts(table: str, condition: str, *, name: str) -> tuple[list[str],
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)')
+296
View File
@@ -0,0 +1,296 @@
"""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,)]