48 lines
1.2 KiB
Python
48 lines
1.2 KiB
Python
"""Engine + sessionmaker. Lazily initialised so test code can swap the URL."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Iterator
|
|
from contextlib import contextmanager
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.engine import Engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
from app.core.config import settings
|
|
|
|
_engine: Engine | None = None
|
|
_SessionLocal: sessionmaker[Session] | None = None
|
|
|
|
|
|
def get_engine() -> Engine:
|
|
global _engine
|
|
if _engine is None:
|
|
_engine = create_engine(
|
|
settings.database_url,
|
|
pool_pre_ping=True,
|
|
future=True,
|
|
)
|
|
return _engine
|
|
|
|
|
|
def get_sessionmaker() -> sessionmaker[Session]:
|
|
global _SessionLocal
|
|
if _SessionLocal is None:
|
|
_SessionLocal = sessionmaker(bind=get_engine(), expire_on_commit=False, future=True)
|
|
return _SessionLocal
|
|
|
|
|
|
@contextmanager
|
|
def session_scope() -> Iterator[Session]:
|
|
"""Context manager that commits on success, rolls back on error."""
|
|
s = get_sessionmaker()()
|
|
try:
|
|
yield s
|
|
s.commit()
|
|
except Exception:
|
|
s.rollback()
|
|
raise
|
|
finally:
|
|
s.close()
|