add pool_pre_ping and dispose()

Author: rafizamankhan

.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/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,83 @@ def test_get_connection_url_missing_pool_vars(monkeypatch):
         get_connection_url()
 
 
+# --- 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(monkeypatch):
+    clear_engine_cache()
+    _direct_db_env(monkeypatch)
+    assert get_engine() is get_engine()
+    clear_engine_cache()
+
+
+def test_get_engine_different_urls_get_different_engines(monkeypatch):
+    clear_engine_cache()
+    _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
+    clear_engine_cache()
+
+
+def test_get_engine_not_keyed_by_masked_str_password(monkeypatch):
+    """str(URL) masks passwords; cache must still separate credentials."""
+    clear_engine_cache()
+    _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
+    clear_engine_cache()
+
+
+def test_clear_engine_cache_disposes_and_creates_new(monkeypatch):
+    clear_engine_cache()
+    _direct_db_env(monkeypatch)
+    first = get_engine()
+    clear_engine_cache()
+    second = get_engine()
+    assert first is not second
+    clear_engine_cache()
+
+
+def test_get_engine_applies_pool_settings(monkeypatch):
+    clear_engine_cache()
+    _direct_db_env(monkeypatch)
+    monkeypatch.setenv("DB_POOL_SIZE", "3")
+    monkeypatch.setenv("DB_MAX_OVERFLOW", "2")
+    engine = get_engine()
+    assert engine.pool.size() == 3
+    assert engine.pool._max_overflow == 2
+    assert engine.pool._pre_ping is True
+    clear_engine_cache()
+
+
 # --- Permission and Role Tests ---
 
 

utils/core/db.py

@@ -1,6 +1,5 @@
 import os
 import logging
-from functools import lru_cache
 from itertools import chain
 from typing import Union, Sequence
 from sqlalchemy.engine import Engine, URL
@@ -29,26 +28,36 @@
 # --- Database connection functions ---
 
 
-@lru_cache
-def _cached_engine(url: URL) -> Engine:
-    # URL objects hash/compare on their real fields (including password), so
-    # caching on the URL itself is safe. str(URL) masks the password as
-    # "***" for safe logging — passing that to create_engine() instead would
-    # make every connection attempt authenticate with the literal "***".
-    return create_engine(url)
+# Cache by URL object (not str(url)): str(URL) masks passwords as "***", so
+# different credentials would collide on one cache key.
+_engines: dict[URL, 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.
+    """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.
     """
-    return _cached_engine(get_connection_url())
+    url = get_connection_url()
+    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)."""
+    for engine in _engines.values():
+        engine.dispose()
+    _engines.clear()
 
 
 def ensure_database_exists(url: URL) -> None:
@@ -314,3 +323,4 @@ def tear_down_db() -> None:
             conn.commit()
     finally:
         engine.dispose()
+    clear_engine_cache()