diff --git a/quantara/web_app/db/crud/base.py b/quantara/web_app/db/crud/base.py index a1726380..491dfde1 100644 --- a/quantara/web_app/db/crud/base.py +++ b/quantara/web_app/db/crud/base.py @@ -6,10 +6,11 @@ import uuid from typing import Type, TypeVar -from sqlalchemy import create_engine from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import scoped_session, sessionmaker +from web_app.db.database import get_database_url, init_engine + from web_app.db.database import get_database_url from web_app.db.models import AirDrop, Base from web_app.utils.logger import get_logger @@ -32,11 +33,20 @@ class DBConnector: def __init__(self, db_url: str = None): """ Initialize the database connection and session factory. + + The engine is built via :func:`web_app.db.database.init_engine` + so every ``DBConnector`` instance inherits the same pool policy as + the module-level engine configured by ``web_app.db.database.init_db`` + (configurable via ``DB_POOL_SIZE``, ``DB_MAX_OVERFLOW``, + ``DB_POOL_RECYCLE`` env vars with reasonable defaults and an + always-on ``pool_pre_ping``). When ``db_url`` is omitted, the + URL is derived from environment variables. + :param db_url: Optional database URL. If not provided, fetches from environment. """ if db_url is None: db_url = get_database_url() - self.engine = create_engine(db_url) + self.engine = init_engine(db_url) self.session_factory = sessionmaker(bind=self.engine) self.Session = scoped_session(self.session_factory) diff --git a/quantara/web_app/db/database.py b/quantara/web_app/db/database.py index 82f5e73b..f3633077 100644 --- a/quantara/web_app/db/database.py +++ b/quantara/web_app/db/database.py @@ -31,13 +31,39 @@ def get_database_url() -> str: f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_SERVER}:{DB_PORT}/{DB_NAME}" ) +def _engine_pool_kwargs() -> dict: + """Return the standard pool-related kwargs for every SQLAlchemy engine. + + Pool sizing is read from the ``DB_POOL_SIZE``, ``DB_MAX_OVERFLOW`` and + ``DB_POOL_RECYCLE`` environment variables with sensible defaults so + deployment configs can tune them per environment without code + changes. ``pool_pre_ping`` is always enabled because it is universally + safe for a web service and prevents hard-to-debug failures after + database restarts or network hiccups. Centralising the kwargs keeps + the two engine constructions in this codebase consistent. + """ + return { + "pool_size": int(os.environ.get("DB_POOL_SIZE", "5")), + "max_overflow": int(os.environ.get("DB_MAX_OVERFLOW", "10")), + "pool_recycle": int(os.environ.get("DB_POOL_RECYCLE", "1800")), + "pool_pre_ping": True, + } + + +def init_engine(db_url: str = None): + """Construct a SQLAlchemy engine that uses the project's pool policy.""" + if db_url is None: + db_url = get_database_url() + return create_engine(db_url, **_engine_pool_kwargs()) + + def init_db() -> None: - """Initialize database connection and session factory.""" + """Initialize the module-level database connection and session factory.""" global engine, SessionLocal if engine is not None: return - engine = create_engine(get_database_url()) + engine = create_engine(get_database_url(), **_engine_pool_kwargs()) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) def get_database() -> Generator[Session, None, None]: diff --git a/quantara/web_app/tests/test_db_pool.py b/quantara/web_app/tests/test_db_pool.py new file mode 100644 index 00000000..ed78d509 --- /dev/null +++ b/quantara/web_app/tests/test_db_pool.py @@ -0,0 +1,111 @@ +""" +Tests for SQLAlchemy connection pool configuration. + +Both engine construction sites (``web_app.db.database.init_db`` and the +``DBConnector`` constructor) route through +``web_app.db.database.create_engine`` so that the pool policy defined in +:func:`web_app.db.database._engine_pool_kwargs` is applied uniformly. + +The actual pool behaviour is exercised in the integration suite against a +real PostgreSQL container. This file focuses narrowly on configuration by +patching ``create_engine`` so no database is required. +""" + +from unittest.mock import patch + +import pytest + + +# Default values must stay in lock-step with the implementation in +# ``web_app/db/database.py``. +DEFAULT_POOL_SIZE = 5 +DEFAULT_MAX_OVERFLOW = 10 +DEFAULT_POOL_RECYCLE = 1800 + + +@pytest.fixture(autouse=True) +def _scrub_database_module_state(monkeypatch): + """Reset module-level engine state so ``init_db`` re-runs in tests.""" + from web_app.db import database + + monkeypatch.setattr(database, "engine", None) + monkeypatch.setattr(database, "SessionLocal", None) + yield + + +def test_database_init_db_uses_env_pool_settings(monkeypatch): + """``init_db`` forwards environment-driven pool kwargs to ``create_engine``.""" + monkeypatch.setenv("DB_POOL_SIZE", "7") + monkeypatch.setenv("DB_MAX_OVERFLOW", "11") + monkeypatch.setenv("DB_POOL_RECYCLE", "900") + + from web_app.db import database + + with patch("web_app.db.database.create_engine") as create_engine: + database.init_db() + assert create_engine.call_count == 1 + kwargs = create_engine.call_args.kwargs + assert kwargs["pool_size"] == 7 + assert kwargs["max_overflow"] == 11 + assert kwargs["pool_recycle"] == 900 + assert kwargs["pool_pre_ping"] is True + + +def test_database_init_db_falls_back_to_defaults(monkeypatch): + """Defaults apply when no environment variables are set.""" + for var in ("DB_POOL_SIZE", "DB_MAX_OVERFLOW", "DB_POOL_RECYCLE"): + monkeypatch.delenv(var, raising=False) + + from web_app.db import database + + with patch("web_app.db.database.create_engine") as create_engine: + database.init_db() + kwargs = create_engine.call_args.kwargs + assert kwargs["pool_size"] == DEFAULT_POOL_SIZE + assert kwargs["max_overflow"] == DEFAULT_MAX_OVERFLOW + assert kwargs["pool_recycle"] == DEFAULT_POOL_RECYCLE + assert kwargs["pool_pre_ping"] is True + + +def test_database_init_db_rejects_non_integer_env_values(monkeypatch): + """Non-integer env vars surface as ``ValueError`` rather than silently falling back.""" + monkeypatch.setenv("DB_POOL_SIZE", "not-an-int") + + from web_app.db import database + + with pytest.raises(ValueError): + database.init_db() + + +def test_db_connector_routes_through_init_engine(monkeypatch): + """``DBConnector`` constructs its engine via ``init_engine``.""" + monkeypatch.setenv("DB_POOL_SIZE", "3") + monkeypatch.setenv("DB_MAX_OVERFLOW", "8") + monkeypatch.setenv("DB_POOL_RECYCLE", "600") + + from web_app.db.crud.base import DBConnector + + with patch("web_app.db.database.create_engine") as create_engine: + DBConnector(db_url="postgresql://user:pwd@host:5432/dbname") + assert create_engine.call_count == 1 + kwargs = create_engine.call_args.kwargs + assert kwargs["pool_size"] == 3 + assert kwargs["max_overflow"] == 8 + assert kwargs["pool_recycle"] == 600 + assert kwargs["pool_pre_ping"] is True + + +def test_db_connector_falls_back_to_defaults(monkeypatch): + """Defaults apply in ``DBConnector`` as well when env vars are unset.""" + for var in ("DB_POOL_SIZE", "DB_MAX_OVERFLOW", "DB_POOL_RECYCLE"): + monkeypatch.delenv(var, raising=False) + + from web_app.db.crud.base import DBConnector + + with patch("web_app.db.database.create_engine") as create_engine: + DBConnector(db_url="postgresql://user:pwd@host:5432/dbname") + kwargs = create_engine.call_args.kwargs + assert kwargs["pool_size"] == DEFAULT_POOL_SIZE + assert kwargs["max_overflow"] == DEFAULT_MAX_OVERFLOW + assert kwargs["pool_recycle"] == DEFAULT_POOL_RECYCLE + assert kwargs["pool_pre_ping"] is True