diff --git a/etl/db/db.py b/etl/db/db.py index 0bc1e6b..d4f5ab7 100644 --- a/etl/db/db.py +++ b/etl/db/db.py @@ -1,6 +1,7 @@ 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 @@ -10,8 +11,19 @@ class Settings(BaseSettings): settings = Settings() # engine to the dabatase, currently set up to connect via settings. database -engine = create_engine(settings.DATABASE_URL) if settings.DATABASE_URL else None - +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: