mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 00:21:58 +00:00
复核 R-b12(直连地址其实还是 PgBouncer、搬断点走了事务模式):check_pgbouncer 除了没给直连地址,直连地址和应用库地址连的是同一个主机端口(规范化之后比,端口不写按 5432)也拒绝启动,报错说清直连要绕过 PgBouncer;copy-to-postgres 搬断点抽成 checkpoint_conninfo:给了断点库地址用它,否则进目标库、给了直连地址就走直连,命令行加 --direct-url,不给而目标就是应用库时用 AGENTLAB_DATABASE_DIRECT_URL。补测试:同一个主机端口(大小写、用户、参数不同)拒绝、同主机另一个端口放行;搬断点缺省走直连、单给的断点库地址优先、目标是别的库不套用应用库的直连地址
This commit is contained in:
+19
-3
@@ -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.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker, create_async_engine
|
||||||
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
|
||||||
from sqlalchemy.pool import NullPool
|
from sqlalchemy.pool import NullPool
|
||||||
|
from sqlalchemy.engine import make_url
|
||||||
from sqlalchemy.types import TypeDecorator
|
from sqlalchemy.types import TypeDecorator
|
||||||
|
|
||||||
from app.core.config import settings
|
from app.core.config import settings
|
||||||
@@ -300,11 +301,26 @@ async def dispose_session_engine() -> None:
|
|||||||
|
|
||||||
|
|
||||||
def check_pgbouncer() -> None:
|
def check_pgbouncer() -> None:
|
||||||
"""AGENTLAB_DB_PGBOUNCER=1 却没给直连地址:拒绝启动(事务模式下 LISTEN、会话级咨询锁不成立,跨进程的锁会锁不住、放不掉)。"""
|
"""AGENTLAB_DB_PGBOUNCER=1 时直连地址要真的是直连:没给、或者给的和应用库地址指着同一个主机端口(还是那个 PgBouncer),都拒绝
|
||||||
if settings.db_pgbouncer and settings.db_dialect == "postgresql" and not settings.database_direct_url:
|
启动(事务模式下 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:"
|
raise RuntimeError("AGENTLAB_DB_PGBOUNCER=1(应用库地址是 PgBouncer 的事务模式)时要另给 AGENTLAB_DATABASE_DIRECT_URL:"
|
||||||
"进程间通知(LISTEN)、跨进程的锁(会话级咨询锁)、断点库要直连 PostgreSQL(或 PgBouncer 的会话模式)。"
|
"进程间通知(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]:
|
async def get_session() -> AsyncIterator[AsyncSession]:
|
||||||
|
|||||||
@@ -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,
|
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,
|
direct_url: str | None = None, dry_run: bool = False, assume_stopped: bool = False,
|
||||||
work_dir: str | Path | None = None, progress: Callable[[str], None] = print) -> Report:
|
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
|
from app.db.base import create_schema, make_engine
|
||||||
|
|
||||||
target_url = normalize_database_url(target_url)
|
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)
|
db, ckpt = resolve_source(source)
|
||||||
if (wal := source_in_use(db, ckpt)) and not assume_stopped:
|
if (wal := source_in_use(db, ckpt)) and not assume_stopped:
|
||||||
raise CopyRefused(f"{wal} 还在:服务看起来还开着这个库。先停服务再导入;确认服务已停(上次是强杀,留下了 "
|
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("source", help="数据目录,或者 agentlab.db 的路径(断点库 checkpoints.db 在同一个目录里找)")
|
||||||
parser.add_argument("target", nargs="?", help="目标地址 postgresql+asyncpg://…;不给就用 AGENTLAB_DATABASE_URL")
|
parser.add_argument("target", nargs="?", help="目标地址 postgresql+asyncpg://…;不给就用 AGENTLAB_DATABASE_URL")
|
||||||
parser.add_argument("--checkpoint-url", help="断点放到另一个 PostgreSQL 库;不给就用 AGENTLAB_CHECKPOINT_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("--dry-run", action="store_true", help="只核对:数来源和目标各有多少行,不建表、不写")
|
||||||
parser.add_argument("--assume-stopped", action="store_true",
|
parser.add_argument("--assume-stopped", action="store_true",
|
||||||
help="不查服务停没停(来源旁边留着强杀后的 -wal 时用)")
|
help="不查服务停没停(来源旁边留着强杀后的 -wal 时用)")
|
||||||
@@ -576,9 +588,16 @@ def main(argv: list[str] | None = None) -> int:
|
|||||||
target = args.target or settings.database_url
|
target = args.target or settings.database_url
|
||||||
if not target:
|
if not target:
|
||||||
parser.error("没有目标地址:给第二个参数,或者设 AGENTLAB_DATABASE_URL")
|
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:
|
try:
|
||||||
report = asyncio.run(copy_to_postgres(
|
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),
|
dry_run=args.dry_run, assume_stopped=args.assume_stopped, batch=max(args.batch, 1),
|
||||||
work_dir=args.work_dir))
|
work_dir=args.work_dir))
|
||||||
except (CopyRefused, ValueError) as e:
|
except (CopyRefused, ValueError) as e:
|
||||||
|
|||||||
@@ -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,
|
monkeypatch.setattr(settings, "database_direct_url", Settings(database_url=APP, database_direct_url=DIRECT,
|
||||||
_env_file=None).database_direct_url)
|
_env_file=None).database_direct_url)
|
||||||
base.check_pgbouncer()
|
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, "目标是别的库:不套用应用库的直连地址"
|
||||||
|
|||||||
Reference in New Issue
Block a user