hale-api / app /core /database.py
Raunak211006's picture
Deploy HALE API v3.1
eaf56bf verified
Raw
History Blame Contribute Delete
2.09 kB
# app/core/database.py
"""
Database engine, session factory, and base model.
Supports SQLite (dev) and PostgreSQL (prod) via DATABASE_URL.
Production-hardened: connection pool tuning, pre-ping, recycle.
"""
import logging
from sqlalchemy import create_engine
from sqlalchemy.orm import declarative_base, sessionmaker
from app.core.config import settings
logger = logging.getLogger("hale.database")
# ── Connection args ────────────────────────────────────────────────────────
connect_args: dict = {}
pool_kwargs: dict = {}
if settings.is_sqlite:
# SQLite: single-thread safety, no connection pool
connect_args["check_same_thread"] = False
else:
# PostgreSQL: production-grade connection pooling
pool_kwargs = {
"pool_size": 20, # Base connections kept alive across workers
"max_overflow": 10, # Burst connections allowed above pool_size
"pool_timeout": 30, # Seconds to wait if pool is exhausted
"pool_recycle": 1800, # Recycle connections every 30 min (prevents stale TCP)
}
logger.info(
"PostgreSQL pool: pool_size=%d, max_overflow=%d, recycle=%ds",
pool_kwargs["pool_size"], pool_kwargs["max_overflow"], pool_kwargs["pool_recycle"],
)
engine = create_engine(
settings.DATABASE_URL,
connect_args=connect_args,
echo=False,
pool_pre_ping=True, # Validates connection before use β€” handles DB restarts
**pool_kwargs,
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
def init_db() -> None:
"""Create all tables (idempotent β€” safe for existing tables)."""
# Import all models so SQLAlchemy registers them before create_all
from app.models import user, goal, progress # noqa: F401
Base.metadata.create_all(bind=engine)
def get_db():
"""FastAPI dependency that yields a database session."""
db = SessionLocal()
try:
yield db
finally:
db.close()