fix: engine too many connections

Fix create_engine() connection leaks per request/call

Author: bhtabor

tests/browser/db_helpers.py

@@ -8,10 +8,10 @@
 from datetime import timedelta
 
 from dotenv import load_dotenv
-from sqlmodel import Session, create_engine, select
+from sqlmodel import Session, select
 
 from utils.core.auth import get_password_hash
-from utils.core.db import create_default_roles, get_connection_url
+from utils.core.db import create_default_roles, get_engine
 from utils.core.models import (
     Account,
     Invitation,
@@ -39,8 +39,7 @@ def browser_db_session():
     saved = {key: os.environ.get(key) for key in env}
     os.environ.update(env)
     try:
-        engine = create_engine(get_connection_url())
-        with Session(engine) as session:
+        with Session(get_engine()) as session:
             yield session
     finally:
         for key, value in saved.items():
@@ -168,8 +167,7 @@ def browser_csrf_db_session():
     saved = {key: os.environ.get(key) for key in env}
     os.environ.update(env)
     try:
-        engine = create_engine(get_connection_url())
-        with Session(engine) as session:
+        with Session(get_engine()) as session:
             yield session
     finally:
         for key, value in saved.items():

tests/conftest.py

@@ -3,11 +3,12 @@
 from typing import Generator, cast
 
 pytest_plugins = ["tests.frontend.fixtures"]
-from sqlmodel import create_engine, Session, select
+from sqlmodel import Session, select
 from fastapi.testclient import TestClient
 from dotenv import load_dotenv
 from utils.core.db import (
     get_connection_url,
+    get_engine,
     tear_down_db,
     set_up_db,
     create_default_roles,
@@ -74,9 +75,12 @@ def engine(env_vars):
     Create a new SQLModel engine for the test database.
     Use PostgreSQL for testing to match production environment.
     """
-    # Use PostgreSQL for testing to match production environment
+    # Use PostgreSQL for testing to match production environment. get_engine()
+    # returns the same cached pool the app itself uses via get_session(), so
+    # TestClient requests and direct test queries share one pool instead of
+    # opening a second one for the same database.
     ensure_database_exists(get_connection_url())
-    engine = create_engine(get_connection_url())
+    engine = get_engine()
     set_up_db(drop=True)
 
     yield engine

utils/core/auth.py

@@ -13,7 +13,7 @@
 from fastapi.templating import Jinja2Templates
 from fastapi import Cookie
 from starlette.responses import Response
-from utils.core.db import create_engine, get_connection_url
+from utils.core.db import get_engine
 from utils.core.models import (
     AccountRecoveryToken,
     EmailVerificationToken,
@@ -339,8 +339,7 @@ def send_reset_email_task(email: str) -> None:
     FastAPI background tasks should not reuse request-scoped resources from
     `yield` dependencies, because cleanup may run before the task executes.
     """
-    engine = create_engine(get_connection_url())
-    with Session(engine) as session:
+    with Session(get_engine()) as session:
         send_reset_email(email, session)
 
 

utils/core/db.py

@@ -2,7 +2,7 @@
 import logging
 from itertools import chain
 from typing import Union, Sequence
-from sqlalchemy.engine import URL
+from sqlalchemy.engine import Engine, URL
 from sqlmodel import create_engine, Session, SQLModel, select, text
 from utils.core.models import (
     Account,
@@ -28,17 +28,42 @@
 # --- Database connection functions ---
 
 
+_engine_cache: dict[str, Engine] = {}
+
+
+def get_engine() -> Engine:
+    """Returns a process-wide cached engine for the current connection URL.
+
+    Calling create_engine() per request/call leaks pooled connections until
+    Python's GC reclaims them (engines aren't freed by simple refcounting),
+    which can exhaust Postgres's max_connections under sustained load. Cache
+    by URL rather than as a single global so tests that swap DB_NAME within
+    the same process (e.g. browser vs. browser-csrf DBs) still get one
+    engine per database instead of sharing a mismatched pool.
+    """
+    url = get_connection_url()
+    key = str(url)
+    engine = _engine_cache.get(key)
+    if engine is None:
+        engine = create_engine(url)
+        _engine_cache[key] = engine
+    return engine
+
+
 def ensure_database_exists(url: URL) -> None:
     dbname = url.database
     server_url = url.set(database="postgres")
     engine = create_engine(server_url, isolation_level="AUTOCOMMIT")
-    with engine.connect() as conn:
-        exists = conn.execute(
-            text("SELECT 1 FROM pg_database WHERE datname = :n"),
-            {"n": dbname},
-        ).scalar()
-        if not exists:
-            conn.execute(text(f'CREATE DATABASE "{dbname}"'))
+    try:
+        with engine.connect() as conn:
+            exists = conn.execute(
+                text("SELECT 1 FROM pg_database WHERE datname = :n"),
+                {"n": dbname},
+            ).scalar()
+            if not exists:
+                conn.execute(text(f'CREATE DATABASE "{dbname}"'))
+    finally:
+        engine.dispose()
 
 
 def get_connection_url() -> URL:
@@ -259,28 +284,32 @@ def set_up_db(drop: bool = False) -> None:
         drop (bool): If True, drops all existing tables before creating new ones.
     """
     engine = create_engine(get_connection_url())
-    if drop:
-        SQLModel.metadata.drop_all(engine)
-    # Ensure the private schema exists before creating tables
-    with engine.connect() as conn:
-        conn.execute(text("CREATE SCHEMA IF NOT EXISTS private"))
-        conn.commit()
-    SQLModel.metadata.create_all(engine)
-    # Create default permissions and seed account emails
-    with Session(engine) as session:
-        create_permissions(session)
-        session.commit()
-        seed_account_emails(session)
-    engine.dispose()
+    try:
+        if drop:
+            SQLModel.metadata.drop_all(engine)
+        # Ensure the private schema exists before creating tables
+        with engine.connect() as conn:
+            conn.execute(text("CREATE SCHEMA IF NOT EXISTS private"))
+            conn.commit()
+        SQLModel.metadata.create_all(engine)
+        # Create default permissions and seed account emails
+        with Session(engine) as session:
+            create_permissions(session)
+            session.commit()
+            seed_account_emails(session)
+    finally:
+        engine.dispose()
 
 
 def tear_down_db() -> None:
     """
     Tears down the database by dropping all tables and the private schema.
     """
     engine = create_engine(get_connection_url())
-    SQLModel.metadata.drop_all(engine)
-    with engine.connect() as conn:
-        conn.execute(text("DROP SCHEMA IF EXISTS private CASCADE"))
-        conn.commit()
-    engine.dispose()
+    try:
+        SQLModel.metadata.drop_all(engine)
+        with engine.connect() as conn:
+            conn.execute(text("DROP SCHEMA IF EXISTS private CASCADE"))
+            conn.commit()
+    finally:
+        engine.dispose()

utils/core/dependencies.py

@@ -13,7 +13,7 @@
     oauth2_scheme_cookie,
     verify_password,
 )
-from utils.core.db import create_engine, get_connection_url
+from utils.core.db import get_engine
 from utils.core.models import (
     User,
     Role,
@@ -43,8 +43,7 @@ def get_session() -> Generator[Session, None, None]:
     Yields:
         Session: A SQLModel session object for database operations.
     """
-    engine = create_engine(get_connection_url())
-    with Session(engine) as session:
+    with Session(get_engine()) as session:
         yield session
 
 
@@ -434,8 +433,7 @@ async def get_user_from_request(request: Request) -> Optional[User]:
     tokens = (access_token, refresh_token)
 
     # Get a database session
-    engine = create_engine(get_connection_url())
-    with Session(engine) as session:
+    with Session(get_engine()) as session:
         user, new_access_token, new_refresh_token = get_user_from_tokens(
             tokens, session
         )