mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 04:21:57 +00:00
ThearchyHelios
3e922cbefb
- 先停服务：来源旁边留着 agentlab.db-wal / checkpoints.db-wal（还有连接开着，或者上次是强杀），或者目标库上有 application_name=agentlab 的连接，就拒绝，什么都不写；强杀留下的 -wal 用 --assume-stopped 放过。查不到的（另一台机器上的服务）在说明里写明要自己保证 - 读快照：SQLite 备份接口拷到 --work-dir（缺省系统临时目录），在拷贝上跑一遍启动时的建表补列，老版本的库按当前表结构读；来源按读写打开、只许读（mode=ro 打开 WAL 库会留下它自己建的 -wal，下一次导入就当成服务还开着） - 目标：同启动时的 create_schema 建表；按外键先后逐表、按主键分批 INSERT … ON CONFLICT DO NOTHING——目标上已有同一主键（或撞了别的唯一约束）的行一律不动，重跑只补没有的；整批撞了外键之类的约束就逐行重试，插不进的列出主键和原因；自增主键的序列接到导入的最大值后面 - 值：按模型的列类型读写（时间、布尔），JSON 原文照搬（SQL 的 NULL 和 JSON 的 null 分开、键的先后不变）；PostgreSQL 不收的（NaN、NUL、整数列里的浮点和数字字符串、文本列里的非文本）换掉并计数打印 - 断点：用两个 saver 的公开接口逐条读出写入（含 pending writes，特殊通道和普通写入分开写，序号同来源），目标上已有的断点不重写；按线程核对断点和 pending writes 的条数 - 核对：每张表按主键数来源的行目标上有没有（导入前已有、这次补上分开数），逐表打印；有没导进去的退出码非零。--dry-run 只核对、不建表不写 - 测试：另起进程在 SQLite 数据目录里跑一次停在审批上的运行（事件、审批、待办、工件引用），导入新建的空库、再导一遍什么都不重复，再另起进程连这个 PostgreSQL、同一个数据目录批掉审批，运行从断点接着跑完、事件序号接着往下编、工件读得到、新的变化记录的序号接在导入的后面；dry-run 不写；目标上服务还连着就拒绝；来源、参数、值转换的单测
507 lines
25 KiB
Python
507 lines
25 KiB
Python
"""把 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
|
||
|