from sqlmodel import SQLModel, create_engine, Session from typing import Generator from contextlib import contextmanager from config.settings import settings # The database_url from settings will be used here engine = create_engine(settings.DATABASE_URL, echo=True if settings.LOG_LEVEL == "DEBUG" else False) def create_db_and_tables(): SQLModel.metadata.create_all(engine) @contextmanager def get_session_context() -> Generator[Session, None, None]: session = Session(engine) try: yield session session.commit() except Exception: session.rollback() raise finally: session.close() def get_session() -> Generator[Session, None, None]: # This is a generator function that can be used with `Depends` in FastAPI # For Streamlit, we'll use the context manager directly or a simpler session management with Session(engine) as session: yield session