fix: engine too many connections (#231)

* fix: engine too many connections

Fix create_engine() connection leaks per request/call

* refactor: use lru cached engine

Use lru cached engine

* add pool_pre_ping and dispose()

* fix: address review feedback on engine cache

- Route rate_limit.py through the shared get_engine() cache instead of
  its own module-level singleton, so no second unmanaged pool survives
  clear_engine_cache() or goes stale when the connection URL changes.
- Guard the engine cache with a threading.Lock: sync dependencies run in
  FastAPI's threadpool, so concurrent first requests could race in
  get_engine() and orphan an undisposed engine, and clear_engine_cache()
  could hit a dict-mutated-during-iteration error.
- Remove the clear_engine_cache() side effect from tear_down_db(); test
  fixtures that relied on it now call it explicitly in their teardown.
- Make engine-cache tests isolation-safe via an engine_cache fixture so
  cleanup runs even when an assertion fails mid-test.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Rafiuzzaman Khan <rafiuzzamank@gmail.com>
Co-authored-by: chriscarrollsmith <chriscarrollsmith@gmail.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: 4 people

.env.example

@@ -7,6 +7,10 @@ BASE_URL=http://localhost:8000
 # Database connection mode (0 = direct, 1 = pooled)
 USE_POOL=0
 
+# SQLAlchemy pool (per process; keep small when using PgBouncer)
+# DB_POOL_SIZE=5
+# DB_MAX_OVERFLOW=5
+
 # Shared database settings
 DB_HOST=127.0.0.1
 DB_SSLMODE=prefer

main.py

@@ -46,7 +46,7 @@
     RateLimitError,
 )
 from exceptions.exceptions import NeedsNewTokens
-from utils.core.db import set_up_db
+from utils.core.db import set_up_db, clear_engine_cache
 
 logger = logging.getLogger("uvicorn.error")
 logger.setLevel(logging.DEBUG)
@@ -58,7 +58,7 @@ async def lifespan(app: FastAPI):
     load_dotenv()
     set_up_db()
     yield
-    # Optional shutdown logic
+    clear_engine_cache()
 
 
 # Initialize the FastAPI app

tests/browser/conftest.py

@@ -8,6 +8,7 @@
 from playwright.sync_api import Browser, Page
 from tests.browser.db_helpers import browser_db_env
 from utils.core.db import (
+    clear_engine_cache,
     ensure_database_exists,
     get_connection_url,
     set_up_db,
@@ -149,6 +150,7 @@ def live_server(browser_env):
         proc.wait(timeout=5)
         with _temporary_env(browser_env):
             tear_down_db()
+        clear_engine_cache()
 
 
 @pytest.fixture(scope="session")
@@ -162,3 +164,4 @@ def live_server_csrf(browser_csrf_env):
         proc.wait(timeout=5)
         with _temporary_env(browser_csrf_env):
             tear_down_db()
+        clear_engine_cache()

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,13 @@
 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 (
+    clear_engine_cache,
     get_connection_url,
+    get_engine,
     tear_down_db,
     set_up_db,
     create_default_roles,
@@ -74,15 +76,19 @@ 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
 
     # Clean up after tests
     tear_down_db()
+    clear_engine_cache()
 
 
 @pytest.fixture

tests/utils/test_db.py

@@ -3,6 +3,8 @@
 from sqlalchemy import Engine
 from utils.core.db import (
     get_connection_url,
+    get_engine,
+    clear_engine_cache,
     assign_permissions_to_role,
     create_default_roles,
     create_permissions,
@@ -146,6 +148,84 @@ def test_get_connection_url_missing_pool_vars(monkeypatch):
         get_connection_url()
 
 
+# --- Engine cache ---
+
+
+@pytest.fixture
+def engine_cache():
+    """Ensure engine-cache tests start clean and never leak cached engines,
+    even when an assertion fails mid-test."""
+    clear_engine_cache()
+    yield
+    clear_engine_cache()
+
+
+def _direct_db_env(monkeypatch, *, name: str = "testdb", password: str = "testpass"):
+    for var in (
+        "USE_POOL",
+        "DB_HOST",
+        "DB_PORT",
+        "DB_NAME",
+        "DB_USER",
+        "DB_PASSWORD",
+        "DB_POOL_PORT",
+        "DB_POOL_NAME",
+        "DB_APPUSER",
+        "DB_APPUSER_PASSWORD",
+        "DB_POOL_SIZE",
+        "DB_MAX_OVERFLOW",
+    ):
+        monkeypatch.delenv(var, raising=False)
+    monkeypatch.setenv("DB_HOST", "localhost")
+    monkeypatch.setenv("DB_PORT", "5432")
+    monkeypatch.setenv("DB_NAME", name)
+    monkeypatch.setenv("DB_USER", "testuser")
+    monkeypatch.setenv("DB_PASSWORD", password)
+
+
+def test_get_engine_reuses_same_instance(engine_cache, monkeypatch):
+    _direct_db_env(monkeypatch)
+    assert get_engine() is get_engine()
+
+
+def test_get_engine_different_urls_get_different_engines(engine_cache, monkeypatch):
+    _direct_db_env(monkeypatch, name="db_a")
+    engine_a = get_engine()
+    _direct_db_env(monkeypatch, name="db_b")
+    engine_b = get_engine()
+    assert engine_a is not engine_b
+
+
+def test_get_engine_not_keyed_by_masked_str_password(engine_cache, monkeypatch):
+    """str(URL) masks passwords; cache must still separate credentials."""
+    _direct_db_env(monkeypatch, password="secretA")
+    engine_a = get_engine()
+    assert "***" in str(get_connection_url())
+    _direct_db_env(monkeypatch, password="secretB")
+    engine_b = get_engine()
+    assert engine_a is not engine_b
+
+
+def test_clear_engine_cache_disposes_and_creates_new(engine_cache, monkeypatch):
+    _direct_db_env(monkeypatch)
+    first = get_engine()
+    clear_engine_cache()
+    second = get_engine()
+    assert first is not second
+
+
+def test_get_engine_applies_pool_settings(engine_cache, monkeypatch):
+    _direct_db_env(monkeypatch)
+    monkeypatch.setenv("DB_POOL_SIZE", "3")
+    monkeypatch.setenv("DB_MAX_OVERFLOW", "2")
+    engine = get_engine()
+    assert engine.pool.size() == 3
+    # No public accessors for these; private attrs are stable in practice but
+    # may need updating on a SQLAlchemy major upgrade.
+    assert engine.pool._max_overflow == 2
+    assert engine.pool._pre_ping is True
+
+
 # --- Permission and Role Tests ---
 
 

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

@@ -1,8 +1,9 @@
 import os
 import logging
+import threading
 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 +29,60 @@
 # --- Database connection functions ---
 
 
+# Cache by URL object (not str(url)): str(URL) masks passwords as "***", so
+# different credentials would collide on one cache key. Guarded by a lock:
+# sync dependencies run in FastAPI's threadpool, so concurrent first requests
+# would otherwise race and orphan an undisposed engine.
+_engines: dict[URL, Engine] = {}
+_engines_lock = threading.Lock()
+
+
+def get_engine() -> Engine:
+    """Return a process-wide cached engine for the current connection URL.
+
+    Calling create_engine() per request leaks pooled connections until GC
+    reclaims them. Cache one engine per URL so tests that swap DB_NAME still
+    get a separate pool per database. Pool settings (DB_POOL_SIZE,
+    DB_MAX_OVERFLOW) are read only when the engine for a URL is first
+    created; later environment changes have no effect on a cached engine.
+    """
+    url = get_connection_url()
+    with _engines_lock:
+        engine = _engines.get(url)
+        if engine is None:
+            engine = create_engine(
+                url,
+                pool_pre_ping=True,
+                pool_size=int(os.getenv("DB_POOL_SIZE", "5")),
+                max_overflow=int(os.getenv("DB_MAX_OVERFLOW", "5")),
+            )
+            _engines[url] = engine
+    return engine
+
+
+def clear_engine_cache() -> None:
+    """Dispose and drop all cached engines (shutdown / tests)."""
+    with _engines_lock:
+        engines = list(_engines.values())
+        _engines.clear()
+    for engine in engines:
+        engine.dispose()
+
+
 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 +303,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
         )