from __future__ import annotations from collections.abc import Generator from typing import Any from sqlalchemy import create_engine, event from sqlalchemy.engine import Engine from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker from app.config import get_settings class Base(DeclarativeBase): pass def build_engine(database_url: str) -> Engine: is_sqlite = database_url.startswith("sqlite") connect_args: dict[str, Any] = {"check_same_thread": False} if is_sqlite else {} engine = create_engine( database_url, connect_args=connect_args, pool_pre_ping=not is_sqlite, ) if is_sqlite: @event.listens_for(engine, "connect") def _set_sqlite_pragma(dbapi_connection: Any, _connection_record: Any) -> None: cursor = dbapi_connection.cursor() cursor.execute("PRAGMA foreign_keys=ON") cursor.close() return engine engine = build_engine(get_settings().database_url) SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, expire_on_commit=False) def get_db() -> Generator[Session, None, None]: db = SessionLocal() try: yield db finally: db.close()