From ef19da951620d935f37b6941825d9b9a14f45927 Mon Sep 17 00:00:00 2001
From: Alex <a@tushynski.me>
Date: Sun, 19 Apr 2026 17:53:22 +0100
Subject: [PATCH] fix: pg bouncer compatible

---
 application/storage/db/engine.py            | 29 +++++++++++++--------
 tests/storage/db/test_session_guardrails.py | 11 ++++----
 2 files changed, 24 insertions(+), 16 deletions(-)

diff --git a/application/storage/db/engine.py b/application/storage/db/engine.py
index a1c81330..fefceda5 100644
--- a/application/storage/db/engine.py
+++ b/application/storage/db/engine.py
@@ -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
 
 
diff --git a/tests/storage/db/test_session_guardrails.py b/tests/storage/db/test_session_guardrails.py
index e24dd895..c3d3a77b 100644
--- a/tests/storage/db/test_session_guardrails.py
+++ b/tests/storage/db/test_session_guardrails.py
@@ -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: