From d55ae68e7d9afa0d32372f52ce13e7f0fa58a8e0 Mon Sep 17 00:00:00 2001 From: ThearchyHelios Date: Mon, 19 Oct 2026 14:42:33 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A4=8D=E6=A0=B8=20R-b12=EF=BC=88=E7=9B=B4?= =?UTF-8?q?=E8=BF=9E=E5=9C=B0=E5=9D=80=E5=85=B6=E5=AE=9E=E8=BF=98=E6=98=AF?= =?UTF-8?q?=20PgBouncer=E3=80=81=E6=90=AC=E6=96=AD=E7=82=B9=E8=B5=B0?= =?UTF-8?q?=E4=BA=86=E4=BA=8B=E5=8A=A1=E6=A8=A1=E5=BC=8F=EF=BC=89=EF=BC=9A?= =?UTF-8?q?check=5Fpgbouncer=20=E9=99=A4=E4=BA=86=E6=B2=A1=E7=BB=99?= =?UTF-8?q?=E7=9B=B4=E8=BF=9E=E5=9C=B0=E5=9D=80=EF=BC=8C=E7=9B=B4=E8=BF=9E?= =?UTF-8?q?=E5=9C=B0=E5=9D=80=E5=92=8C=E5=BA=94=E7=94=A8=E5=BA=93=E5=9C=B0?= =?UTF-8?q?=E5=9D=80=E8=BF=9E=E7=9A=84=E6=98=AF=E5=90=8C=E4=B8=80=E4=B8=AA?= =?UTF-8?q?=E4=B8=BB=E6=9C=BA=E7=AB=AF=E5=8F=A3=EF=BC=88=E8=A7=84=E8=8C=83?= =?UTF-8?q?=E5=8C=96=E4=B9=8B=E5=90=8E=E6=AF=94=EF=BC=8C=E7=AB=AF=E5=8F=A3?= =?UTF-8?q?=E4=B8=8D=E5=86=99=E6=8C=89=205432=EF=BC=89=E4=B9=9F=E6=8B=92?= =?UTF-8?q?=E7=BB=9D=E5=90=AF=E5=8A=A8=EF=BC=8C=E6=8A=A5=E9=94=99=E8=AF=B4?= =?UTF-8?q?=E6=B8=85=E7=9B=B4=E8=BF=9E=E8=A6=81=E7=BB=95=E8=BF=87=20PgBoun?= =?UTF-8?q?cer=EF=BC=9Bcopy-to-postgres=20=E6=90=AC=E6=96=AD=E7=82=B9?= =?UTF-8?q?=E6=8A=BD=E6=88=90=20checkpoint=5Fconninfo=EF=BC=9A=E7=BB=99?= =?UTF-8?q?=E4=BA=86=E6=96=AD=E7=82=B9=E5=BA=93=E5=9C=B0=E5=9D=80=E7=94=A8?= =?UTF-8?q?=E5=AE=83=EF=BC=8C=E5=90=A6=E5=88=99=E8=BF=9B=E7=9B=AE=E6=A0=87?= =?UTF-8?q?=E5=BA=93=E3=80=81=E7=BB=99=E4=BA=86=E7=9B=B4=E8=BF=9E=E5=9C=B0?= =?UTF-8?q?=E5=9D=80=E5=B0=B1=E8=B5=B0=E7=9B=B4=E8=BF=9E=EF=BC=8C=E5=91=BD?= =?UTF-8?q?=E4=BB=A4=E8=A1=8C=E5=8A=A0=20--direct-url=EF=BC=8C=E4=B8=8D?= =?UTF-8?q?=E7=BB=99=E8=80=8C=E7=9B=AE=E6=A0=87=E5=B0=B1=E6=98=AF=E5=BA=94?= =?UTF-8?q?=E7=94=A8=E5=BA=93=E6=97=B6=E7=94=A8=20AGENTLAB=5FDATABASE=5FDI?= =?UTF-8?q?RECT=5FURL=E3=80=82=E8=A1=A5=E6=B5=8B=E8=AF=95=EF=BC=9A?= =?UTF-8?q?=E5=90=8C=E4=B8=80=E4=B8=AA=E4=B8=BB=E6=9C=BA=E7=AB=AF=E5=8F=A3?= =?UTF-8?q?=EF=BC=88=E5=A4=A7=E5=B0=8F=E5=86=99=E3=80=81=E7=94=A8=E6=88=B7?= =?UTF-8?q?=E3=80=81=E5=8F=82=E6=95=B0=E4=B8=8D=E5=90=8C=EF=BC=89=E6=8B=92?= =?UTF-8?q?=E7=BB=9D=E3=80=81=E5=90=8C=E4=B8=BB=E6=9C=BA=E5=8F=A6=E4=B8=80?= =?UTF-8?q?=E4=B8=AA=E7=AB=AF=E5=8F=A3=E6=94=BE=E8=A1=8C=EF=BC=9B=E6=90=AC?= =?UTF-8?q?=E6=96=AD=E7=82=B9=E7=BC=BA=E7=9C=81=E8=B5=B0=E7=9B=B4=E8=BF=9E?= =?UTF-8?q?=E3=80=81=E5=8D=95=E7=BB=99=E7=9A=84=E6=96=AD=E7=82=B9=E5=BA=93?= =?UTF-8?q?=E5=9C=B0=E5=9D=80=E4=BC=98=E5=85=88=E3=80=81=E7=9B=AE=E6=A0=87?= =?UTF-8?q?=E6=98=AF=E5=88=AB=E7=9A=84=E5=BA=93=E4=B8=8D=E5=A5=97=E7=94=A8?= =?UTF-8?q?=E5=BA=94=E7=94=A8=E5=BA=93=E7=9A=84=E7=9B=B4=E8=BF=9E=E5=9C=B0?= =?UTF-8?q?=E5=9D=80?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/app/db/base.py | 22 ++++++++++++++--- backend/app/db/to_postgres.py | 27 +++++++++++++++++---- backend/tests/test_db_direct.py | 42 +++++++++++++++++++++++++++++++++ 3 files changed, 84 insertions(+), 7 deletions(-) diff --git a/backend/app/db/base.py b/backend/app/db/base.py index b839df41..b76ba890 100644 --- a/backend/app/db/base.py +++ b/backend/app/db/base.py @@ -12,6 +12,7 @@ from sqlalchemy import JSON, BigInteger, Column, DateTime, Integer, MetaData, ev from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column from sqlalchemy.pool import NullPool +from sqlalchemy.engine import make_url from sqlalchemy.types import TypeDecorator from app.core.config import settings @@ -300,11 +301,26 @@ async def dispose_session_engine() -> None: def check_pgbouncer() -> None: - """AGENTLAB_DB_PGBOUNCER=1 却没给直连地址：拒绝启动（事务模式下 LISTEN、会话级咨询锁不成立，跨进程的锁会锁不住、放不掉）。""" - if settings.db_pgbouncer and settings.db_dialect == "postgresql" and not settings.database_direct_url: + """AGENTLAB_DB_PGBOUNCER=1 时直连地址要真的是直连：没给、或者给的和应用库地址指着同一个主机端口（还是那个 PgBouncer），都拒绝 + 启动（事务模式下 LISTEN、会话级咨询锁不成立，跨进程的锁会锁不住、放不掉；复核 R-b12）。""" + if not (settings.db_pgbouncer and settings.db_dialect == "postgresql"): + return + if not settings.database_direct_url: raise RuntimeError("AGENTLAB_DB_PGBOUNCER=1（应用库地址是 PgBouncer 的事务模式）时要另给 AGENTLAB_DATABASE_DIRECT_URL：" "进程间通知（LISTEN）、跨进程的锁（会话级咨询锁）、断点库要直连 PostgreSQL（或 PgBouncer 的会话模式）。" - "见 README「用 PostgreSQL」里的三种部署") + "见 README「多进程」里的「直连与 PgBouncer」") + if _endpoint(settings.database_direct_url) == _endpoint(settings.database_url): + from app.db.urls import describe + + raise RuntimeError(f"AGENTLAB_DATABASE_DIRECT_URL 和 AGENTLAB_DATABASE_URL 指着同一个地址（{describe(settings.database_url)}）：" + "AGENTLAB_DB_PGBOUNCER=1 说应用库地址是 PgBouncer 的事务模式，直连地址要绕过它、直接连 PostgreSQL" + "（或者 PgBouncer 的会话模式那个端口）。见 README「多进程」里的「直连与 PgBouncer」") + + +def _endpoint(url: str) -> tuple[str, int]: + """地址连的是哪个主机端口（端口不写按 5432）。""" + parsed = make_url(url) + return (parsed.host or "").lower(), int(parsed.port or 5432) async def get_session() -> AsyncIterator[AsyncSession]: diff --git a/backend/app/db/to_postgres.py b/backend/app/db/to_postgres.py index bb6093cc..f4c7a6a0 100644 --- a/backend/app/db/to_postgres.py +++ b/backend/app/db/to_postgres.py @@ -481,13 +481,23 @@ async def copy_checkpoints(snapshot_ckpt: Path, conninfo: str, *, dry_run: bool, # -------------------------------------------------------------------------- +def checkpoint_conninfo(target_url: str, *, checkpoint_url: str | None = None, direct_url: str | None = None) -> str: + """断点搬到哪儿（psycopg 的连接串）：给了断点库地址用它；否则进目标库——给了直连地址就走直连（复核 R-b12：目标地址是 + PgBouncer 的事务模式时，psycopg 的预备语句、建表时取的会话级锁都会出问题）。""" + if checkpoint_url: + return normalize_checkpoint_url(checkpoint_url) + return libpq_conninfo(normalize_database_url(direct_url) if direct_url else normalize_database_url(target_url)) + + 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: + direct_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: + """direct_url：目标库的直连地址（目标地址是 PgBouncer 时给；断点走它）。""" 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) + conninfo = checkpoint_conninfo(target_url, checkpoint_url=checkpoint_url, direct_url=direct_url) db, ckpt = resolve_source(source) if (wal := source_in_use(db, ckpt)) and not assume_stopped: raise CopyRefused(f"{wal} 还在：服务看起来还开着这个库。先停服务再导入；确认服务已停（上次是强杀，留下了 " @@ -567,6 +577,8 @@ def main(argv: list[str] | None = None) -> int: 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("--direct-url", help="目标库的直连地址（目标地址是 PgBouncer 时给，断点走它）；不给、目标又是 " + "AGENTLAB_DATABASE_URL 时用 AGENTLAB_DATABASE_DIRECT_URL") parser.add_argument("--dry-run", action="store_true", help="只核对：数来源和目标各有多少行，不建表、不写") parser.add_argument("--assume-stopped", action="store_true", help="不查服务停没停（来源旁边留着强杀后的 -wal 时用）") @@ -576,9 +588,16 @@ def main(argv: list[str] | None = None) -> int: target = args.target or settings.database_url if not target: parser.error("没有目标地址：给第二个参数，或者设 AGENTLAB_DATABASE_URL") + direct = args.direct_url or None + if direct is None and settings.database_direct_url: + try: + same = normalize_database_url(target) == settings.database_url + except ValueError: + same = False + direct = settings.database_direct_url if same else None try: report = asyncio.run(copy_to_postgres( - args.source, target, checkpoint_url=args.checkpoint_url or settings.checkpoint_url or None, + args.source, target, checkpoint_url=args.checkpoint_url or settings.checkpoint_url or None, direct_url=direct, 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: diff --git a/backend/tests/test_db_direct.py b/backend/tests/test_db_direct.py index 552a0c40..b7b152f8 100644 --- a/backend/tests/test_db_direct.py +++ b/backend/tests/test_db_direct.py @@ -67,3 +67,45 @@ def test_transaction_mode_without_a_direct_url_refuses_to_start(monkeypatch): monkeypatch.setattr(settings, "database_direct_url", Settings(database_url=APP, database_direct_url=DIRECT, _env_file=None).database_direct_url) base.check_pgbouncer() + + +def test_transaction_mode_with_a_direct_url_that_is_the_same_endpoint_refuses_to_start(monkeypatch): + """复核 R-b12：直连地址和应用库地址规范化之后连的是同一个主机端口（还是那个 PgBouncer）——同样拒绝，说清楚。""" + from app.db import base + + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "db_pgbouncer", True) + for same in (APP, "postgresql+asyncpg://other:pw@PGBOUNCER.internal:6432/agentlab?ssl=require"): + monkeypatch.setattr(settings, "database_direct_url", + Settings(database_url=APP, database_direct_url=same, _env_file=None).database_direct_url) + with pytest.raises(RuntimeError, match="指着同一个地址"): + base.check_pgbouncer() + monkeypatch.setattr(settings, "database_direct_url", Settings( + database_url=APP, database_direct_url="postgresql://agentlab:secret@pgbouncer.internal:5432/agentlab", + _env_file=None).database_direct_url) + base.check_pgbouncer() # 同一台机器上的另一个端口（直连 PostgreSQL 或会话模式）：放行 + + +def test_copying_checkpoints_goes_through_the_direct_url(monkeypatch, tmp_path): + """复核 R-b12：copy-to-postgres 搬断点缺省走直连地址（设了的话）；单独给了断点库地址就用它。""" + from app.db import to_postgres + + assert "db.internal:5432" in to_postgres.checkpoint_conninfo(APP, direct_url=DIRECT) + assert "pgbouncer.internal:6432" in to_postgres.checkpoint_conninfo(APP) + assert "ckpt.internal" in to_postgres.checkpoint_conninfo(APP, checkpoint_url="postgresql://u:p@ckpt.internal/c", + direct_url=DIRECT) + seen: dict = {} + + async def fake_copy(source, target, **kw): + seen.update(kw, target=target) + raise to_postgres.CopyRefused("测试：到这里为止") + + monkeypatch.setattr(to_postgres, "copy_to_postgres", fake_copy) + monkeypatch.setattr(settings, "database_url", Settings(database_url=APP, _env_file=None).database_url) + monkeypatch.setattr(settings, "database_direct_url", + Settings(database_url=APP, database_direct_url=DIRECT, _env_file=None).database_direct_url) + monkeypatch.setattr(settings, "checkpoint_url", "") + assert to_postgres.main([str(tmp_path)]) == 2 + assert seen["direct_url"] == settings.database_direct_url, "目标就是应用库：断点走它的直连地址" + assert to_postgres.main([str(tmp_path), "postgresql://u:p@elsewhere.internal:5432/x"]) == 2 + assert seen["direct_url"] is None, "目标是别的库：不套用应用库的直连地址"