59 lines
1.4 KiB
Python
59 lines
1.4 KiB
Python
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker, declarative_base
|
|
|
|
Base = declarative_base()
|
|
|
|
# Engine and session factory are created lazily so that tests can
|
|
# override DATABASE_URL via environment *before* any import triggers
|
|
# the real PostgreSQL engine creation (which requires psycopg2).
|
|
_engine = None
|
|
_SessionLocal = None
|
|
|
|
|
|
def _get_engine():
|
|
global _engine
|
|
if _engine is None:
|
|
from app.config import settings
|
|
_engine = create_engine(settings.DATABASE_URL)
|
|
return _engine
|
|
|
|
|
|
def _get_session_factory():
|
|
global _SessionLocal
|
|
if _SessionLocal is None:
|
|
_SessionLocal = sessionmaker(
|
|
autocommit=False, autoflush=False, bind=_get_engine()
|
|
)
|
|
return _SessionLocal
|
|
|
|
|
|
class _LazySessionLocal:
|
|
"""Proxy so ``SessionLocal()`` keeps working as before but the real
|
|
sessionmaker is only created on first call."""
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return _get_session_factory()(*args, **kwargs)
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(_get_session_factory(), name)
|
|
|
|
|
|
SessionLocal = _LazySessionLocal()
|
|
|
|
|
|
class _EngineProxy:
|
|
"""Thin proxy so ``from app.database import engine`` still works."""
|
|
def __getattr__(self, name):
|
|
return getattr(_get_engine(), name)
|
|
|
|
|
|
engine = _EngineProxy() # type: ignore[assignment]
|
|
|
|
|
|
def get_db():
|
|
db = SessionLocal()
|
|
try:
|
|
yield db
|
|
finally:
|
|
db.close()
|