WealthySmart/backend/app/db.py

from pathlib import Path

from alembic import command
from alembic.config import Config
from sqlalchemy import inspect, text
from sqlmodel import Session, create_engine

from app.config import settings

engine = create_engine(settings.DATABASE_URL)

_MIGRATION_LOCK_KEY = 727274  # app-wide advisory lock id, arbitrary constant


def _alembic_config() -> Config:
    backend_dir = Path(__file__).resolve().parent.parent
    cfg = Config(str(backend_dir / "alembic.ini"))
    cfg.set_main_option("script_location", str(backend_dir / "alembic"))
    return cfg


def run_alembic_upgrade() -> None:
    """Bring the schema to head at startup.

    Databases that predate Alembic (prod, existing dev volumes) already match
    the baseline revision — the baseline was autogenerated from, and
    schema-diffed against, exactly that schema — so they are adopted by
    stamping instead of re-running DDL. Serialized with a Postgres advisory
    lock because prod uvicorn runs 2 workers, each executing the lifespan.
    """
    cfg = _alembic_config()
    is_postgres = engine.dialect.name == "postgresql"
    with engine.connect() as conn:
        if is_postgres:
            conn.execute(text(f"SELECT pg_advisory_lock({_MIGRATION_LOCK_KEY})"))
        try:
            insp = inspect(conn)
            if not insp.has_table("alembic_version") and insp.has_table("transaction"):
                command.stamp(cfg, "head")
            command.upgrade(cfg, "head")
        finally:
            if is_postgres:
                conn.execute(
                    text(f"SELECT pg_advisory_unlock({_MIGRATION_LOCK_KEY})")
                )


def get_session():
    with Session(engine) as session:
        yield session