54 lines
1.7 KiB
Python
54 lines
1.7 KiB
Python
"""Database engine, sessions, and FastAPI dependency."""
|
|
|
|
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import Engine
|
|
from sqlmodel import Session, SQLModel, create_engine
|
|
|
|
from .settings import settings
|
|
|
|
_engine: Engine | None = None
|
|
|
|
|
|
def _engine_kwargs(url: str) -> dict:
|
|
if url.startswith("sqlite"):
|
|
# FastAPI runs sync endpoints in a threadpool; connections are shared
|
|
# across threads, so defer the same-thread check. File-based SQLite
|
|
# needs an existing parent dir (SQLite won't create it).
|
|
return {"connect_args": {"check_same_thread": False}}
|
|
return {}
|
|
|
|
|
|
def get_engine(url: str | None = None) -> Engine:
|
|
"""Process-wide engine singleton (override URL in tests via settings)."""
|
|
global _engine
|
|
if _engine is None:
|
|
url = url or settings.database_url
|
|
if url.startswith("sqlite:///") and not url.startswith("sqlite:///:memory:"):
|
|
Path(url.removeprefix("sqlite:///")).parent.mkdir(parents=True, exist_ok=True)
|
|
_engine = create_engine(url, **_engine_kwargs(url))
|
|
return _engine
|
|
|
|
|
|
def reset_engine() -> None:
|
|
"""Drop the cached engine (used by the test suite between DB targets)."""
|
|
global _engine
|
|
if _engine is not None:
|
|
_engine.dispose()
|
|
_engine = None
|
|
|
|
|
|
def get_session() -> Iterator[Session]:
|
|
with Session(get_engine()) as session:
|
|
yield session
|
|
|
|
|
|
def get_db() -> Iterator[Session]:
|
|
"""Single FastAPI DB dependency (all routers import this)."""
|
|
yield from get_session()
|
|
|
|
|
|
def create_all(url: str | None = None) -> None:
|
|
"""Create tables directly (dev convenience; migrations own prod schema)."""
|
|
SQLModel.metadata.create_all(get_engine(url))
|