53 lines
1.3 KiB
Python
53 lines
1.3 KiB
Python
"""SQLAlchemy engine/session plumbing and ORM base class."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Generator
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import DeclarativeBase, Session, sessionmaker
|
|
|
|
from app.config import get_settings
|
|
|
|
_engine = None
|
|
_SessionFactory: sessionmaker[Session] | None = None
|
|
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
|
|
def _url_with_driver() -> str:
|
|
url = get_settings().database_url
|
|
# Accept plain postgresql:// URLs by normalising to the psycopg driver.
|
|
if url.startswith("postgresql://"):
|
|
url = url.replace("postgresql://", "postgresql+psycopg://", 1)
|
|
return url
|
|
|
|
|
|
def get_engine():
|
|
global _engine
|
|
if _engine is None:
|
|
_engine = create_engine(_url_with_driver(), pool_pre_ping=True, future=True)
|
|
return _engine
|
|
|
|
|
|
def get_session_factory() -> sessionmaker[Session]:
|
|
global _SessionFactory
|
|
if _SessionFactory is None:
|
|
_SessionFactory = sessionmaker(bind=get_engine(), future=True)
|
|
return _SessionFactory
|
|
|
|
|
|
def get_db() -> Generator[Session, None, None]:
|
|
"""FastAPI dependency yielding a database session per request."""
|
|
factory = get_session_factory()
|
|
session = factory()
|
|
try:
|
|
yield session
|
|
session.commit()
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close() |