feat: refactor database session management and introduce global state for engine

Author: umbertogriffo

src/api/deps.py

@@ -1,7 +1,7 @@
 from typing import Annotated, Generator
 
+import state
 from config import settings
-from database import engine
 from fastapi import Depends, HTTPException
 from fastapi.security import APIKeyHeader
 from sqlmodel import Session
@@ -16,7 +16,7 @@ def get_db_session() -> Generator[Session, None, None]:
     """
     Create a new database session and close the session after the operation has ended.
     """
-    with Session(engine) as session:
+    with Session(state.engine) as session:
         yield session
 
 

src/api/endpoints/health.py

@@ -1,6 +1,7 @@
-from database import check_health_safe, engine
+from database import check_health_safe
 from fastapi import APIRouter
 from starlette import status
+from state import engine
 
 router = APIRouter()
 

src/database.py

@@ -78,6 +78,3 @@ def create_db_and_tables(engine: Engine):
     Create database tables based on the defined SQLModel models.
     """
     SQLModel.metadata.create_all(engine)
-
-
-engine = create_db_engine()

src/main.py

@@ -1,16 +1,20 @@
 import os
 from contextlib import asynccontextmanager
 
+import state
 import uvicorn
 from api.exception_handlers import validation_exception_handler
 from api.router import router
 from config import init_whatever, settings
-from database import engine
+from database import create_db_engine
 from fastapi import FastAPI
 from fastapi.exceptions import RequestValidationError
 from fastapi.middleware.cors import CORSMiddleware
 from middlewares.logging import LogMiddleware
 from pydantic import ValidationError
+from utils.log import get_logger
+
+logger = get_logger()
 
 
 @asynccontextmanager
@@ -20,9 +24,14 @@ async def lifespan(app: FastAPI):
     # It may be stored in app state object:
     # app.state.some_resource = SomeResource()
     # https://github.com/fastapi/fastapi/discussions/13029
+
+    state.engine = create_db_engine()
     yield
+
     # Clean up
-    engine.dispose()
+    if state.engine:
+        state.engine.dispose()
+        logger.info("Database engine disposed")
 
 
 app = FastAPI(title=settings.SERVICE_NAME, version="0.1.0", lifespan=lifespan)

src/state.py

@@ -0,0 +1,9 @@
+"""
+Global application state.
+Holds singleton instances that are initialized during app startup.
+"""
+
+from sqlalchemy import Engine
+
+# Global singleton instances
+engine: Engine | None = None

tests/conftest.py

@@ -1,5 +1,3 @@
-import os
-import tempfile
 from pathlib import Path
 
 import pytest
@@ -28,25 +26,26 @@ def data_folder_path():
     return Path(__file__).parent.parent / "data"
 
 
-@pytest.fixture(name="session")
-def session_fixture(request, monkeypatch) -> Session:
-    """Create a new database session for a test."""
+@pytest.fixture(scope="session")
+def db_engine(request, tmp_path_factory, session_mocker):
+    """
+    Create a session-scoped database engine.
+    Database is created once and migrations run once for all tests.
+    """
     # TODO: Use an in-memory SQLite database for faster tests if possible.
     #       https://sqlmodel.tiangolo.com/tutorial/fastapi/tests/#memory-database
 
     db_type = request.config.getoption("--db")
-
     if db_type == "postgres":
         db_url = "postgresql://develop:develop_secret@localhost:5432/develop"
-        path = None
     else:
         # Create a temporary database file for SQLite
-        fd, path = tempfile.mkstemp(suffix=".db")
-        os.close(fd)
-        db_url = f"sqlite:///{path}"
+        temp_dir = tmp_path_factory.mktemp("db")
+        db_path = temp_dir / "test.db"
+        db_url = f"sqlite:///{db_path}"
 
     # Use monkeypatch to set DATABASE_URL environment variable
-    monkeypatch.setattr("config.settings.DATABASE_URL", db_url)
+    session_mocker.patch("config.settings.DATABASE_URL", db_url)
 
     # Get path to alembic.ini
     src_dir = Path(__file__).parents[1] / "src"
@@ -61,12 +60,22 @@ def session_fixture(request, monkeypatch) -> Session:
     # TODO: Alternatively, you can create tables directly without migrations for simpler setups.
     # create_db_and_tables(engine)
 
-    # Ensure that changes made during tests do not persist and affect other tests using a nested transaction
-    # This is needed for PostgreSQL since the SQLite is erased after each test by deleting the temp file
-    connection = engine.connect()
+    yield engine
+
+    # Clean up at the end of the test session
+    engine.dispose()
+
+
+@pytest.fixture(name="session")
+def session_fixture(db_engine) -> Session:
+    """
+    Create a new database session for a test, wrapped in a transaction that is rolled back after the test.
+    """
+
+    connection = db_engine.connect()
     transaction = connection.begin()
     session = Session(bind=connection)
-
+    # Ensure that changes made during tests do not persist and affect other tests using a nested transaction
     nested = connection.begin_nested()
 
     @sa.event.listens_for(session, "after_transaction_end")
@@ -77,15 +86,11 @@ def end_savepoint(session, transaction):
 
     yield session
 
+    # Rollback the transaction (this undoes all changes made during the test)
     session.close()
     transaction.rollback()
     connection.close()
 
-    # Clean up
-    engine.dispose()
-    if path:
-        os.unlink(path)
-
 
 @pytest.fixture(name="client_with_db")
 def client_fixture(session: Session):