survey-extraction/etl/db/db.py
2025-11-04 21:27:47 +00:00

36 lines
No EOL
1.3 KiB
Python

from sqlmodel import SQLModel, create_engine, Session
from pydantic_settings import BaseSettings
from typing import Optional, List
from sqlalchemy.pool import QueuePool
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 to the dabatase, currently set up to connect via settings. database
engine = (
create_engine(
settings.DATABASE_URL,
poolclass=QueuePool, # use standard connection pool
pool_pre_ping=True, # test connection before use
pool_recycle=300, # reconnect every 5 minutes
pool_size=5, # limit pool size for CI/CD or serverless
max_overflow=2, # allow brief overuse
connect_args={"sslmode": "require"}, # enforce SSL for cloud DBs
)
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:
# Links SQLModel and metadata defined in sqlmodel instance
SQLModel.metadata.create_all(engine)