mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:21:58 +00:00
新增 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:
@@ -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())
|
||||||
@@ -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
|
||||||
|
|
||||||
@@ -7,6 +7,7 @@
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import contextlib
|
||||||
import re
|
import re
|
||||||
from typing import Any
|
from typing import Any
|
||||||
|
|
||||||
@@ -75,3 +76,27 @@ def refuse_inserts(table: str, condition: str, *, name: str) -> tuple[list[str],
|
|||||||
return create, drop
|
return create, drop
|
||||||
return ([f"CREATE TRIGGER {name} BEFORE INSERT ON {table} WHEN {condition} "
|
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}"])
|
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)')
|
||||||
|
|||||||
@@ -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,)]
|
||||||
Reference in New Issue
Block a user