mirror of
https://github.com/ThearchyHelios/AgentLab.git
synced 2026-11-15 13:01:58 +00:00
加直连地址 AGENTLAB_DATABASE_DIRECT_URL 和 PgBouncer 事务模式开关 AGENTLAB_DB_PGBOUNCER（用户拍板）：依赖会话状态的连接——收发进程间通知的那一条（LISTEN）、跨进程的锁（会话级咨询锁，另建一个直连的小池子）、断点库（缺省）——走直连地址，普通查询走 AGENTLAB_DATABASE_URL，不设就全用同一个；开关打开时普通地址不用客户端连接池、关掉 SQLAlchemy 和 asyncpg 的预备语句缓存、预备语句用随机名字（照 SQLAlchemy 文档），事务级咨询锁照常成立；开关打开却没给直连地址时拒绝启动。README 写三种部署（直连、PgBouncer 会话模式、事务模式加直连地址）和要配的 server_reset_query，照实说没在真的 PgBouncer 上跑过。补测试：地址解析和规范化、断点库缺省走直连、事务模式关掉的几样、LISTEN 和跨进程锁走直连、事务模式没给直连地址拒绝启动
This commit is contained in:
ThearchyHelios
@@ -55,9 +55,9 @@ class ClusterLock:
|
||||
try:
|
||||
from sqlalchemy import text
|
||||
|
||||
from app.db.base import engine
|
||||
from app.db.base import session_engine
|
||||
|
||||
conn = await engine.connect()
|
||||
conn = await session_engine().connect()
|
||||
try:
|
||||
conn = await conn.execution_options(isolation_level="AUTOCOMMIT")
|
||||
got = bool(await conn.scalar(text("SELECT pg_try_advisory_lock(:k)"), {"k": self.key}))
|
||||
@@ -86,9 +86,11 @@ class ClusterLock:
|
||||
try:
|
||||
from sqlalchemy import text
|
||||
|
||||
from app.db.base import engine
|
||||
from app.db.base import session_engine
|
||||
|
||||
conn = await engine.connect()
|
||||
# 会话级的锁：走直连（AGENTLAB_DATABASE_DIRECT_URL，没设就是应用库）——PgBouncer 事务模式下取锁、放锁可能落在
|
||||
# 两条不同的服务端连接上
|
||||
conn = await session_engine().connect()
|
||||
try:
|
||||
conn = await conn.execution_options(isolation_level="AUTOCOMMIT")
|
||||
await conn.execute(text("SELECT pg_advisory_lock(:k)"), {"k": self.key})
|
||||
|
||||
@@ -37,6 +37,14 @@ RETRY_S = 2.0
|
||||
Handler = Callable[[list[str]], Any]
|
||||
|
||||
|
||||
def listen_dsn() -> str:
|
||||
"""收发通知那条专用连接的地址：LISTEN 要一条一直开着的会话，设了直连地址（AGENTLAB_DATABASE_DIRECT_URL）就走它——普通地址
|
||||
可能是 PgBouncer 的事务模式，那里收不到通知。"""
|
||||
from app.db.urls import libpq_conninfo
|
||||
|
||||
return libpq_conninfo(settings.session_db_url)
|
||||
|
||||
|
||||
def _tag() -> str:
|
||||
return identity.owner().rsplit(":", 1)[-1]
|
||||
|
||||
@@ -126,9 +134,8 @@ class Signals:
|
||||
import asyncpg
|
||||
|
||||
from app.db.base import APPLICATION_NAME
|
||||
from app.db.urls import libpq_conninfo
|
||||
|
||||
dsn = libpq_conninfo(settings.database_url)
|
||||
dsn = listen_dsn()
|
||||
first = True
|
||||
while not self._stopping:
|
||||
lost = asyncio.Event()
|
||||
|
||||
@@ -56,6 +56,13 @@ class Settings(BaseSettings):
|
||||
# LangGraph 的断点也放进同一个库；checkpoint_url 另给就放那里。数据目录照旧要有：工件、上传表格、沙箱等仍是文件
|
||||
database_url: str = ""
|
||||
checkpoint_url: str = ""
|
||||
# 直连地址（可选，写法同 database_url）：普通查询走 database_url（可以是 PgBouncer），依赖会话状态的连接走这里——进程间通知的
|
||||
# LISTEN 那一条、跨进程锁（会话级咨询锁）、断点库（预备语句、建表时的会话锁）。不设就全用 database_url
|
||||
database_direct_url: str = ""
|
||||
# database_url 是 PgBouncer 的事务模式（AGENTLAB_DB_PGBOUNCER=1）：应用库不用客户端连接池（交给 PgBouncer）、关掉 asyncpg 和
|
||||
# SQLAlchemy 的预备语句缓存、预备语句用不会撞名的名字。这时必须另给 database_direct_url（事务模式下 LISTEN、会话级咨询锁不成立），
|
||||
# 启动时检查
|
||||
db_pgbouncer: bool = False
|
||||
# PostgreSQL 的连接池（SQLite 不用这几项）。pool_size=0 表示不用连接池、每次新建连接（前面有 PgBouncer 时用，
|
||||
# 测试里也这么用：每个测试一个事件循环，asyncpg 的连接不能跨循环）
|
||||
db_pool_size: int = 10
|
||||
@@ -221,7 +228,7 @@ class Settings(BaseSettings):
|
||||
return True
|
||||
return value
|
||||
|
||||
@field_validator("database_url", "checkpoint_url", mode="before")
|
||||
@field_validator("database_url", "checkpoint_url", "database_direct_url", mode="before")
|
||||
@classmethod
|
||||
def _valid_postgres_url(cls, value: object, info: Any) -> object:
|
||||
# 写错了启动时就报出来，别等到第一次查库。只认 PostgreSQL：SQLite 就是不设
|
||||
@@ -232,7 +239,7 @@ class Settings(BaseSettings):
|
||||
from app.db.urls import normalize_database_url, normalize_checkpoint_url
|
||||
|
||||
name = f"AGENTLAB_{(info.field_name or '').upper()}"
|
||||
normalize = normalize_database_url if info.field_name == "database_url" else normalize_checkpoint_url
|
||||
normalize = normalize_checkpoint_url if info.field_name == "checkpoint_url" else normalize_database_url
|
||||
try:
|
||||
return normalize(value.strip())
|
||||
except ValueError as e:
|
||||
@@ -293,15 +300,21 @@ class Settings(BaseSettings):
|
||||
"""LangGraph checkpointer 用独立的库文件，避免和业务表抢写锁。只在 SQLite 时用。"""
|
||||
return self.data_dir / "checkpoints.db"
|
||||
|
||||
@property
|
||||
def session_db_url(self) -> str:
|
||||
"""依赖会话状态的连接（LISTEN、会话级咨询锁、断点库）用的应用库地址：设了直连地址用它，否则同 database_url。"""
|
||||
return self.database_direct_url or self.database_url
|
||||
|
||||
@property
|
||||
def checkpoint_conninfo(self) -> str | None:
|
||||
"""PostgreSQL 断点库的 libpq 连接串（psycopg 用）：checkpoint_url 优先，否则由 database_url 换算。SQLite 时 None。"""
|
||||
"""PostgreSQL 断点库的 libpq 连接串（psycopg 用）：checkpoint_url 优先，否则由直连地址（没设就是 database_url）换算。
|
||||
SQLite 时 None。断点库用预备语句、建表时取会话级的锁，所以缺省走直连。"""
|
||||
if self.checkpoint_url:
|
||||
return self.checkpoint_url
|
||||
if self.database_url:
|
||||
from app.db.urls import libpq_conninfo
|
||||
|
||||
return libpq_conninfo(self.database_url)
|
||||
return libpq_conninfo(self.session_db_url)
|
||||
return None
|
||||
|
||||
@property
|
||||
|
||||
+66
-10
@@ -188,12 +188,45 @@ def pg_json_dumps(value: Any) -> str:
|
||||
APPLICATION_NAME = "agentlab"
|
||||
|
||||
|
||||
def make_engine(url: str | None = None, *, application_name: str = APPLICATION_NAME) -> AsyncEngine:
|
||||
def pg_engine_kwargs(*, application_name: str = APPLICATION_NAME, pgbouncer: bool = False,
|
||||
pool_size: int | None = None) -> dict[str, Any]:
|
||||
"""PostgreSQL 引擎的连接池和连接参数。pgbouncer：这个地址是 PgBouncer 的事务模式（AGENTLAB_DB_PGBOUNCER）——
|
||||
|
||||
- 不用客户端的连接池（NullPool，连接池交给 PgBouncer；SQLAlchemy 的说明里同样这么建议）；
|
||||
- SQLAlchemy 的预备语句缓存（prepared_statement_cache_size）和 asyncpg 自己的（statement_cache_size）都关掉：事务模式下同一个
|
||||
客户端连接的下一条语句可能落在另一条服务端连接上，缓存着的预备语句在那边不存在；
|
||||
- asyncpg 每条语句仍要先预备再执行（SQLAlchemy 的 asyncpg 方言一律这么做，同一个事务里，落在同一条服务端连接上），名字用随机的
|
||||
（prepared_statement_name_func），不和别的客户端在同一条服务端连接上留下的撞名。PgBouncer 要配 server_reset_query=DISCARD ALL
|
||||
（且 server_reset_query_always=1）收掉留下的预备语句，或用 1.21 以上开 max_prepared_statements。
|
||||
事务级的咨询锁（pg_advisory_xact_lock，发名额、同源队、建表时用）在事务模式下照常成立：取锁和放锁都在同一个事务里。"""
|
||||
connect_args: dict[str, Any] = {
|
||||
# 会话时区定成 UTC：时间列是不带时区的 UTC（UTCDateTime），库里 now() 之类的也按 UTC 算
|
||||
"server_settings": {"application_name": application_name, "timezone": "UTC"}}
|
||||
size = settings.db_pool_size if pool_size is None else pool_size
|
||||
if pgbouncer:
|
||||
connect_args.update(statement_cache_size=0, prepared_statement_cache_size=0,
|
||||
prepared_statement_name_func=_unique_statement_name)
|
||||
return {"poolclass": NullPool, "connect_args": connect_args}
|
||||
if size <= 0:
|
||||
return {"poolclass": NullPool, "connect_args": connect_args}
|
||||
return {"pool_size": size, "max_overflow": max(settings.db_max_overflow, 0), "pool_timeout": settings.db_pool_timeout,
|
||||
"pool_recycle": settings.db_pool_recycle, "connect_args": connect_args}
|
||||
|
||||
|
||||
def _unique_statement_name() -> str:
|
||||
return f"__agentlab_{uuid.uuid4().hex}__"
|
||||
|
||||
|
||||
def make_engine(url: str | None = None, *, application_name: str = APPLICATION_NAME, pgbouncer: bool | None = None,
|
||||
pool_size: int | None = None) -> AsyncEngine:
|
||||
"""应用库的引擎。SQLite：WAL、外键、忙等，单个文件；PostgreSQL：连接池、方言调整（db/postgres.py）。
|
||||
|
||||
copy-to-postgres 建目标库的引擎也走这里，两边的列类型、排序规则是同一套。
|
||||
"""
|
||||
url = url or settings.db_url
|
||||
if pgbouncer is None:
|
||||
# 只有应用库的普通地址按 AGENTLAB_DB_PGBOUNCER 处理：直连地址、copy-to-postgres 给的目标库照常
|
||||
pgbouncer = bool(settings.db_pgbouncer) and url == settings.database_url
|
||||
if url.startswith("sqlite"):
|
||||
# pool_pre_ping：取连接前先探一下活。取消异常打进 SQLAlchemy 的清理流程时
|
||||
# （_close_connection 记完日志会把 CancelledError 原样抛出，跳过了把连接记录
|
||||
@@ -206,15 +239,9 @@ def make_engine(url: str | None = None, *, application_name: str = APPLICATION_N
|
||||
from app.db import postgres
|
||||
|
||||
_require_postgres_driver("asyncpg")
|
||||
pool: dict[str, Any] = (
|
||||
{"poolclass": NullPool} if settings.db_pool_size <= 0 else {
|
||||
"pool_size": settings.db_pool_size, "max_overflow": max(settings.db_max_overflow, 0),
|
||||
"pool_timeout": settings.db_pool_timeout, "pool_recycle": settings.db_pool_recycle})
|
||||
made = create_async_engine(
|
||||
url, echo=False, future=True, pool_pre_ping=True, json_serializer=pg_json_dumps,
|
||||
# 会话时区定成 UTC：时间列是不带时区的 UTC（UTCDateTime），库里 now() 之类的也按 UTC 算
|
||||
connect_args={"server_settings": {"application_name": application_name, "timezone": "UTC"}},
|
||||
**pool)
|
||||
made = create_async_engine(url, echo=False, future=True, pool_pre_ping=True, json_serializer=pg_json_dumps,
|
||||
**pg_engine_kwargs(application_name=application_name, pgbouncer=pgbouncer,
|
||||
pool_size=pool_size))
|
||||
postgres.install(made.sync_engine.dialect)
|
||||
event.listen(made.sync_engine, "before_cursor_execute", postgres.strip_nul, retval=True)
|
||||
return made
|
||||
@@ -250,6 +277,35 @@ engine = make_engine()
|
||||
|
||||
SessionLocal = async_sessionmaker(engine, expire_on_commit=False, class_=AsyncSession)
|
||||
|
||||
_session_engine: AsyncEngine | None = None
|
||||
|
||||
|
||||
def session_engine() -> AsyncEngine:
|
||||
"""依赖会话状态的连接用的引擎（跨进程锁的会话级咨询锁，app/cluster/mutex）：设了 AGENTLAB_DATABASE_DIRECT_URL 时是直连那个库的
|
||||
一个小池子，否则就是应用库的引擎。"""
|
||||
global _session_engine
|
||||
if not settings.database_direct_url or settings.database_direct_url == settings.database_url:
|
||||
return engine
|
||||
if _session_engine is None:
|
||||
_session_engine = make_engine(settings.database_direct_url, pgbouncer=False, pool_size=max(2, min(settings.db_pool_size or 5, 10)))
|
||||
return _session_engine
|
||||
|
||||
|
||||
async def dispose_session_engine() -> None:
|
||||
"""关停时关掉直连那个小池子（没另建就什么都不做）。"""
|
||||
global _session_engine
|
||||
if _session_engine is not None:
|
||||
await _session_engine.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:
|
||||
raise RuntimeError("AGENTLAB_DB_PGBOUNCER=1（应用库地址是 PgBouncer 的事务模式）时要另给 AGENTLAB_DATABASE_DIRECT_URL："
|
||||
"进程间通知（LISTEN）、跨进程的锁（会话级咨询锁）、断点库要直连 PostgreSQL（或 PgBouncer 的会话模式）。"
|
||||
"见 README「用 PostgreSQL」里的三种部署")
|
||||
|
||||
|
||||
async def get_session() -> AsyncIterator[AsyncSession]:
|
||||
"""FastAPI 依赖注入用的会话。"""
|
||||
|
||||
+6
-1
@@ -58,6 +58,10 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
||||
# SQLite 只支持单进程：同一个数据目录已经有一个进程在用（uvicorn --workers、两个服务指着同一个目录）就拒绝启动、照实说
|
||||
# （app/cluster/guard.py）。PostgreSQL 可以多进程，什么都不做
|
||||
guard.acquire()
|
||||
# PgBouncer 事务模式（AGENTLAB_DB_PGBOUNCER=1）要另给直连地址，没给就拒绝启动（db/base.check_pgbouncer）
|
||||
from app.db.base import check_pgbouncer
|
||||
|
||||
check_pgbouncer()
|
||||
await init_db()
|
||||
# 多进程协调（app/cluster）：登记这个工作进程、开始心跳和进程间通知、争调度当班。要在下面的启动整理之前：整理只碰属主不在了的
|
||||
# 行，要先知道别的进程谁还活着
|
||||
@@ -150,9 +154,10 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]:
|
||||
await cluster.stop()
|
||||
# 应用库的连接全关掉：SQLite 最后一个连接关掉时才把 -wal 并回、删掉（copy-to-postgres 按它认服务停没停），
|
||||
# PostgreSQL 的连接池也正常关
|
||||
from app.db.base import engine as db_engine
|
||||
from app.db.base import dispose_session_engine, engine as db_engine
|
||||
|
||||
await db_engine.dispose()
|
||||
await dispose_session_engine()
|
||||
logger.info("已停止")
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
"""直连地址和 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()
|
||||
Reference in New Issue
Block a user