File size: 928 Bytes
2e88094 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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 |