AgentLab/backend/tests/test_db_direct.py

"""直连地址和 PgBouncer 事务模式（AGENTLAB_DATABASE_DIRECT_URL、AGENTLAB_DB_PGBOUNCER）：地址怎么解析、哪些连接走直连、事务模式下
关掉了哪些东西。只建引擎、不连库，两种库都跑。"""
from __future__ import annotations

import pytest
from sqlalchemy.pool import NullPool

from app.core.config import Settings, settings

APP = "postgres://agentlab:secret@pgbouncer.internal:6432/agentlab"
DIRECT = "postgresql://agentlab:secret@db.internal:5432/agentlab"


def test_the_direct_url_is_normalized_and_used_for_session_bound_connections():
    s = Settings(database_url=APP, database_direct_url=DIRECT, _env_file=None)
    assert s.database_url.startswith("postgresql+asyncpg://") and "pgbouncer.internal:6432" in s.database_url
    assert s.database_direct_url.startswith("postgresql+asyncpg://") and "db.internal:5432" in s.database_direct_url
    assert s.session_db_url == s.database_direct_url
    assert "db.internal" in s.checkpoint_conninfo and s.checkpoint_conninfo.startswith("postgresql://"), "断点库缺省走直连"
    plain = Settings(database_url=APP, _env_file=None)
    assert plain.session_db_url == plain.database_url and "pgbouncer.internal" in plain.checkpoint_conninfo, "不设就全用同一个"
    with pytest.raises(ValueError, match="AGENTLAB_DATABASE_DIRECT_URL"):
        Settings(database_url=APP, database_direct_url="mysql://x@y/z", _env_file=None)


def test_transaction_mode_turns_off_client_pooling_and_statement_caches():
    from app.db.base import pg_engine_kwargs

    kw = pg_engine_kwargs(pgbouncer=True)
    assert kw["poolclass"] is NullPool
    args = kw["connect_args"]
    assert args["statement_cache_size"] == 0 and args["prepared_statement_cache_size"] == 0
    names = {args["prepared_statement_name_func"]() for _ in range(3)}
    assert len(names) == 3, "预备语句的名字每次都不一样，不和别的客户端撞名"
    direct = pg_engine_kwargs(pgbouncer=False, pool_size=5)
    assert direct["pool_size"] == 5 and "statement_cache_size" not in direct["connect_args"]


def test_listen_and_cluster_locks_go_direct(monkeypatch):
    from app.cluster.signal import listen_dsn
    from app.db import base

    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(base, "_session_engine", None)
    assert "db.internal:5432" in listen_dsn() and listen_dsn().startswith("postgresql://")
    made = base.session_engine()
    try:
        assert made is not base.engine and made.url.host == "db.internal" and made.url.port == 5432
        assert base.session_engine() is made, "直连那个池子只建一个"
    finally:
        monkeypatch.setattr(base, "_session_engine", None)
        made.sync_engine.dispose()
    monkeypatch.setattr(settings, "database_direct_url", "")
    assert base.session_engine() is base.engine and "pgbouncer.internal" in listen_dsn()


def test_transaction_mode_without_a_direct_url_refuses_to_start(monkeypatch):
    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)
    monkeypatch.setattr(settings, "database_direct_url", "")
    with pytest.raises(RuntimeError, match="AGENTLAB_DATABASE_DIRECT_URL"):
        base.check_pgbouncer()
    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, "目标是别的库：不套用应用库的直连地址"


def test_loopback_spellings_count_as_the_same_endpoint(monkeypatch):
    """R-b12 的尾巴：localhost、127.0.0.1、::1 是同一台机器——同一个端口时照样拒绝；IP 按地址比，解析不成 IP 的按字面比。"""
    from app.db import base

    def norm(url: str) -> str:
        return Settings(database_url=url, _env_file=None).database_url

    monkeypatch.setattr(settings, "db_pgbouncer", True)
    for app_url, direct in (("postgresql://u:p@localhost:6432/db", "postgresql://u:p@127.0.0.1:6432/db"),
                            ("postgresql://u:p@127.0.0.1:6432/db", "postgresql://u:p@[::1]:6432/db"),
                            ("postgresql://u:p@10.0.0.5:6432/db", "postgresql://u:p@[::ffff:10.0.0.5]:6432/db")):
        monkeypatch.setattr(settings, "database_url", norm(app_url))
        monkeypatch.setattr(settings, "database_direct_url", norm(direct))
        with pytest.raises(RuntimeError, match="指着同一个地址"):
            base.check_pgbouncer()
    monkeypatch.setattr(settings, "database_url", norm("postgresql://u:p@localhost:6432/db"))
    monkeypatch.setattr(settings, "database_direct_url", norm("postgresql://u:p@127.0.0.1:5432/db"))
    base.check_pgbouncer()   # 本机另一个端口：放行
    monkeypatch.setattr(settings, "database_direct_url", norm("postgresql://u:p@db.internal:6432/db"))
    base.check_pgbouncer()   # 名字不同、不查 DNS：按字面算不同