File size: 4,110 Bytes
ce8f04a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
"""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