Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .env.example
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand Down
3 changes: 3 additions & 0 deletions tests/browser/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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")
Expand All @@ -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()
10 changes: 4 additions & 6 deletions tests/browser/db_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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():
Expand Down Expand Up @@ -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():
Expand Down
12 changes: 9 additions & 3 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
80 changes: 80 additions & 0 deletions tests/utils/test_db.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 ---


Expand Down
5 changes: 2 additions & 3 deletions utils/core/auth.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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)


Expand Down
100 changes: 74 additions & 26 deletions utils/core/db.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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:
Expand Down Expand Up @@ -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()
8 changes: 3 additions & 5 deletions utils/core/dependencies.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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
)
Expand Down
Loading
Loading