From 3e922cbefbb3d2bbc181d3e2492fca3114060354 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 19 Oct 2026 10:07:06 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=20python=20-m=20app.db=20cop?= =?UTF-8?q?y-to-postgres=EF=BC=9A=E6=8A=8A=20SQLite=20=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E7=9B=AE=E5=BD=95=E9=87=8C=E7=9A=84=E5=BA=94=E7=94=A8=E5=BA=93?= =?UTF-8?q?=E5=92=8C=E6=96=AD=E7=82=B9=E5=AF=BC=E5=85=A5=20PostgreSQL?= =?UTF-8?q?=EF=BC=88db/to=5Fpostgres.py=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 先停服务：来源旁边留着 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 不写；目标上服务还连着就拒绝；来源、参数、值转换的单测 --- backend/app/db/__main__.py | 31 ++ backend/app/db/to_postgres.py | 506 +++++++++++++++++++++++++ backend/tests/appdb.py | 25 ++ backend/tests/test_copy_to_postgres.py | 296 +++++++++++++++ 4 files changed, 858 insertions(+) create mode 100644 backend/app/db/__main__.py create mode 100644 backend/app/db/to_postgres.py create mode 100644 backend/tests/test_copy_to_postgres.py diff --git a/backend/app/db/__main__.py b/backend/app/db/__main__.py new file mode 100644 index 00000000..661cc3b9 --- /dev/null +++ b/backend/app/db/__main__.py @@ -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()) diff --git a/backend/app/db/to_postgres.py b/backend/app/db/to_postgres.py new file mode 100644 index 00000000..4dec6a25 --- /dev/null +++ b/backend/app/db/to_postgres.py @@ -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 + diff --git a/backend/tests/appdb.py b/backend/tests/appdb.py index bf217064..fa3c643a 100644 --- a/backend/tests/appdb.py +++ b/backend/tests/appdb.py @@ -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)') diff --git a/backend/tests/test_copy_to_postgres.py b/backend/tests/test_copy_to_postgres.py new file mode 100644 index 00000000..b1391432 --- /dev/null +++ b/backend/tests/test_copy_to_postgres.py @@ -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,)]