backend/core/db_migrate.py

"""Idempotent Alembic upgrade on app startup (Postgres / HF Spaces).

Gunicorn runs multiple workers — only the leader runs migrations (fcntl lock).
SQLite and pytest are skipped; disable with RUN_ALEMBIC_ON_STARTUP=false.
"""
from __future__ import annotations

import logging
import os
import sys
from typing import Any

logger = logging.getLogger(__name__)


def _env_flag(name: str, default: bool = True) -> bool:
    val = os.environ.get(name)
    if val is None:
        return default
    return val.lower() not in ("0", "false", "no", "off")


def _running_under_pytest() -> bool:
    return (
        "PYTEST_CURRENT_TEST" in os.environ
        or os.environ.get("ENV", "").lower() == "test"
        or "pytest" in sys.modules
    )


def _is_postgres_url(url: str | None) -> bool:
    if not url:
        return False
    return url.startswith("postgresql://") or url.startswith("postgres://")


def _acquire_migrate_leader() -> bool:
    """Only one worker runs Alembic (Gunicorn multi-worker)."""
    lock_path = os.environ.get(
        "ALEMBIC_LOCK_PATH",
        os.path.join(os.environ.get("TMPDIR", "/tmp"), "grantforge_alembic.lock"),
    )
    try:
        import fcntl

        lock_file = open(lock_path, "w")
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
        lock_file.write(str(os.getpid()))
        lock_file.flush()
        # Keep FD open for process lifetime so the lock is held.
        globals()["_alembic_lock_fd"] = lock_file
        return True
    except OSError:
        logger.info("[db_migrate] Inny worker już uruchomił Alembic — pomijam.")
        return False
    except ImportError:
        return True


def _backend_root() -> str:
    return os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))


def run_alembic_upgrade_head() -> dict[str, Any]:
    """Run ``alembic upgrade head`` when DATABASE_URL is Postgres.

    Returns a small status dict (no credentials). Idempotent when already at head.
    """
    status: dict[str, Any] = {"ran": False, "skipped": True, "reason": None}

    if not _env_flag("RUN_ALEMBIC_ON_STARTUP", default=True):
        status["reason"] = "RUN_ALEMBIC_ON_STARTUP=false"
        logger.info("[db_migrate] Pomijam Alembic (%s).", status["reason"])
        return status

    db_url = os.environ.get("DATABASE_URL", "")
    if not _is_postgres_url(db_url):
        status["reason"] = "not_postgres"
        logger.info("[db_migrate] DATABASE_URL nie jest Postgres — pomijam Alembic.")
        return status

    if _running_under_pytest():
        status["reason"] = "pytest"
        logger.info("[db_migrate] Środowisko testowe — pomijam Alembic.")
        return status

    if not _acquire_migrate_leader():
        status["reason"] = "not_leader"
        return status

    backend_root = _backend_root()
    ini_path = os.path.join(backend_root, "alembic.ini")
    if not os.path.isfile(ini_path):
        status["reason"] = "alembic_ini_missing"
        status["skipped"] = False
        logger.error("[db_migrate] Brak pliku %s — nie mogę uruchomić migracji.", ini_path)
        return status

    try:
        from alembic import command
        from alembic.config import Config

        cfg = Config(ini_path)
        cfg.set_main_option("script_location", os.path.join(backend_root, "alembic"))
        # env.py reads DATABASE_URL from the environment; do not log the URL.
        logger.info("[db_migrate] Uruchamiam alembic upgrade head (Postgres)...")
        command.upgrade(cfg, "head")
        status["ran"] = True
        status["skipped"] = False
        status["reason"] = "ok"
        logger.info("[db_migrate] Alembic upgrade head zakończony pomyślnie.")
    except Exception as e:
        status["ran"] = False
        status["skipped"] = False
        status["reason"] = "error"
        status["error"] = type(e).__name__
        logger.error(
            "[db_migrate] Alembic upgrade head nieudany: %s: %s. "
            "Serwer kontynuuje z create_all — sprawdź logi / schema.",
            type(e).__name__,
            e,
        )
    return status