Update models/db.py

models/db.py

@@ -0,0 +1,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