mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 20:14:16 +00:00
fix: pg bouncer compatible
This commit is contained in:
1 parent
14c2f4890f
commit
ef19da9516
2 files changed
+24
-16
No files matched your search
@@ -16,7 +16,7 @@ don't need to know about SQLAlchemy dialect prefixes.
|
||||
|
||||
from typing import Optional
|
||||
|
||||
from sqlalchemy import Engine, create_engine
|
||||
from sqlalchemy import Engine, create_engine, event
|
||||
|
||||
from application.core.settings import settings
|
||||
|
||||
@@ -43,8 +43,7 @@ def _resolve_uri() -> str:
|
||||
#: Per-statement wall-clock cap applied to every connection handed out by
|
||||
#: the engine. 30s is generous for interactive hot paths (reads under a few
|
||||
#: hundred ms are normal) but still catches a runaway query before it
|
||||
#: stacks up on PgBouncer or holds locks indefinitely. Override by
|
||||
#: rebuilding the engine with a different ``connect_args`` in tests.
|
||||
#: stacks up on PgBouncer or holds locks indefinitely.
|
||||
STATEMENT_TIMEOUT_MS = 30_000
|
||||
|
||||
|
||||
@@ -52,8 +51,9 @@ def get_engine() -> Engine:
|
||||
"""Return the process-wide SQLAlchemy Engine, creating it if needed.
|
||||
|
||||
The engine applies a server-side ``statement_timeout`` to every
|
||||
connection it hands out, so both :func:`db_session` and
|
||||
:func:`db_readonly` inherit the same guardrail.
|
||||
connection it hands out via a ``connect`` event, so both
|
||||
:func:`db_session` and :func:`db_readonly` inherit the same
|
||||
guardrail.
|
||||
|
||||
Returns:
|
||||
A SQLAlchemy ``Engine`` configured with a pooled connection to
|
||||
@@ -68,13 +68,20 @@ def get_engine() -> Engine:
|
||||
pool_pre_ping=True, # survive PgBouncer / idle-disconnect recycles
|
||||
pool_recycle=1800,
|
||||
future=True,
|
||||
connect_args={
|
||||
# ``-c`` passes a GUC to the backend at connect time. This
|
||||
# covers *all* sessions — interactive, Celery, seeder — so
|
||||
# no route-handler can opt out by accident.
|
||||
"options": f"-c statement_timeout={STATEMENT_TIMEOUT_MS}",
|
||||
},
|
||||
)
|
||||
|
||||
@event.listens_for(_engine, "connect")
|
||||
def _apply_session_guardrails(dbapi_conn, _record):
|
||||
# Apply as a SQL ``SET`` (not a libpq ``options=-c ...``
|
||||
# startup parameter) so the engine works behind
|
||||
# PgBouncer-style poolers — notably Neon's ``-pooler``
|
||||
# endpoint, which rejects startup options. Explicit
|
||||
# ``commit()`` so the session-level SET survives SA's
|
||||
# transaction resets on pool return.
|
||||
with dbapi_conn.cursor() as cur:
|
||||
cur.execute(f"SET statement_timeout = {STATEMENT_TIMEOUT_MS}")
|
||||
dbapi_conn.commit()
|
||||
|
||||
return _engine
|
||||
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ Two invariants are covered here:
|
||||
|
||||
These tests run against a real ephemeral Postgres (via ``pg_engine``)
|
||||
rather than mocks. They rebuild the module-level engine cache so the
|
||||
``statement_timeout`` ``connect_args`` is actually exercised.
|
||||
engine factory's ``statement_timeout`` setup is actually exercised.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -35,9 +35,9 @@ def wired_engine(pg_engine, monkeypatch):
|
||||
"""Rebuild the module-level engine against the ephemeral DB.
|
||||
|
||||
``pg_engine`` already creates its own SQLAlchemy engine, but that
|
||||
engine does not carry our ``statement_timeout`` ``connect_args``. We
|
||||
reconstruct one via :func:`get_engine` so the real production factory
|
||||
code path is exercised.
|
||||
engine does not install our ``statement_timeout`` connect-event
|
||||
hook. We reconstruct one via :func:`get_engine` so the real
|
||||
production factory code path is exercised.
|
||||
"""
|
||||
# Reset the module-level cache so get_engine() re-reads the URL and
|
||||
# applies the production connect_args.
|
||||
@@ -129,7 +129,8 @@ class TestDbReadonlyEnforcement:
|
||||
|
||||
|
||||
class TestStatementTimeout:
|
||||
"""The engine factory installs ``statement_timeout`` via connect_args.
|
||||
"""The engine factory installs ``statement_timeout`` on every new
|
||||
connection.
|
||||
|
||||
We verify two things:
|
||||
|
||||
|
||||
Reference in new issue
Block a user