fix: pg bouncer compatible

This commit is contained in:
Alex committed 2026-04-19 17:53:22 +01:00
1 parent 14c2f4890f
commit ef19da9516
2 files changed
+24 -16

No files matched your search

+18 -11
View File
@@ -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
+6 -5
View File
@@ -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: