models/db.py

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