survey-extraction/etl/db/db.py
2025-03-26 17:11:31 +00:00

22 lines
No EOL
665 B
Python

from sqlmodel import SQLModel, create_engine, Session
from pydantic_settings import BaseSettings
from typing import Optional, List
class Settings(BaseSettings):
DATABASE_URL: Optional[str] = None # Default to None if not set
class Config:
env_file = ".env" # Load from an optional .env file
settings = Settings()
engine = create_engine(settings.DATABASE_URL) if settings.DATABASE_URL else None
def get_db_session():
if engine is None:
raise RuntimeError("Database is not configured. Set DATABASE_URL in environment variables.")
return Session(engine)
def init_db():
if engine:
SQLModel.metadata.create_all(engine)