44 lines
1.0 KiB
Python
44 lines
1.0 KiB
Python
import os
|
|
from collections.abc import Generator
|
|
from contextlib import contextmanager
|
|
|
|
from sqlalchemy import Engine, create_engine
|
|
from sqlalchemy.orm import Session, sessionmaker
|
|
|
|
DATABASE_URL_DEFAULT = "sqlite:///finanz.db"
|
|
|
|
# Optional: allow an external engine (e.g. for testing with in-memory SQLite)
|
|
_engine: Engine | None = None
|
|
|
|
|
|
def set_engine(engine: Engine) -> None:
|
|
global _engine
|
|
_engine = engine
|
|
|
|
|
|
def get_set_engine() -> Engine | None:
|
|
return _engine
|
|
|
|
|
|
def get_engine() -> Engine:
|
|
url = os.environ.get("DATABASE_URL", DATABASE_URL_DEFAULT)
|
|
return create_engine(url)
|
|
|
|
|
|
def get_sessionmaker(engine: Engine | None = None) -> sessionmaker[Session]:
|
|
engine = engine or get_engine()
|
|
return sessionmaker(bind=engine)
|
|
|
|
|
|
@contextmanager
|
|
def get_session(engine: Engine | None = None) -> Generator[Session]:
|
|
session_factory = get_sessionmaker(engine)
|
|
session = session_factory()
|
|
try:
|
|
yield session
|
|
except Exception:
|
|
session.rollback()
|
|
raise
|
|
finally:
|
|
session.close()
|