mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 06:21:58 +00:00
ThearchyHelios
134 lines
7.8 KiB
Python
134 lines
7.8 KiB
Python
"""直连地址和 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：按字面算不同
|