rex/migrations.py

Rama master · réplica de lectura del repositorio original.

"""Database migration state validation.

Tracks which schema migrations have been applied and validates that all
expected migrations are present before the application accepts traffic.

Environment variables
---------------------
SKIP_MIGRATION_CHECK    Set to "1", "true", or "yes" to bypass the check
                        (emergency use only).

Usage
-----
Call :func:`validate_migration_state` early in application startup, before
any request handler is registered::

    from rex.migrations import validate_migration_state
    validate_migration_state()

If unapplied migrations exist the function logs each pending migration name
at WARNING level and exits the process with code 1.
"""

from __future__ import annotations

import logging
import os
import sqlite3
import sys
from datetime import UTC, datetime
from pathlib import Path

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Default database path (matches DashboardStore default)
# ---------------------------------------------------------------------------

_DEFAULT_DB_PATH = Path("data/dashboard_notifications.db")

# ---------------------------------------------------------------------------
# Migration registry
# Each entry is a (name, sql) tuple.  Names must be unique and are stored in
# the schema_migrations table when applied.
# ---------------------------------------------------------------------------

MIGRATIONS: list[tuple[str, str]] = [
    (
        "001_create_notifications",
        """
        CREATE TABLE IF NOT EXISTS notifications (
            id TEXT PRIMARY KEY,
            priority TEXT NOT NULL DEFAULT 'normal',
            title TEXT NOT NULL,
            body TEXT NOT NULL DEFAULT '',
            channel TEXT NOT NULL DEFAULT 'dashboard',
            timestamp TEXT NOT NULL,
            read INTEGER NOT NULL DEFAULT 0,
            user_id TEXT,
            metadata_json TEXT NOT NULL DEFAULT '{}'
        )
        """,
    ),
    (
        "002_idx_notifications_timestamp",
        """
        CREATE INDEX IF NOT EXISTS idx_notifications_timestamp
        ON notifications (timestamp DESC)
        """,
    ),
    (
        "003_idx_notifications_user_id",
        """
        CREATE INDEX IF NOT EXISTS idx_notifications_user_id
        ON notifications (user_id)
        """,
    ),
]

# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

_SCHEMA_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS schema_migrations (
    name TEXT PRIMARY KEY,
    applied_at TEXT NOT NULL
)
"""

_EXPECTED_NAMES: tuple[str, ...] = tuple(name for name, _ in MIGRATIONS)


def _open_db(db_path: Path) -> sqlite3.Connection:
    """Open a direct (non-pooled) connection for migration checks."""
    db_path.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(str(db_path), check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def get_pending_migrations(db_path: Path | None = None) -> list[str]:
    """Return names of migrations that have not yet been applied.

    Creates the ``schema_migrations`` tracking table if it does not exist so
    that bootstrapping a fresh database works without a separate step.

    Args:
        db_path: Path to the SQLite database file.  Defaults to the
            dashboard store default path.

    Returns:
        A list of migration names (in order) that are not recorded in
        ``schema_migrations``.
    """
    resolved = db_path or _DEFAULT_DB_PATH
    conn = _open_db(resolved)
    try:
        conn.execute(_SCHEMA_TABLE_SQL)
        conn.commit()
        rows = conn.execute("SELECT name FROM schema_migrations").fetchall()
        applied: set[str] = {row["name"] for row in rows}
    finally:
        conn.close()

    return [name for name in _EXPECTED_NAMES if name not in applied]


def mark_migration_applied(name: str, db_path: Path | None = None) -> None:
    """Record a migration as applied in the tracking table.

    Args:
        name: Migration name (must be in :data:`MIGRATIONS`).
        db_path: Path to the SQLite database file.

    Raises:
        ValueError: If *name* is not a known migration.
    """
    if name not in _EXPECTED_NAMES:
        raise ValueError(f"Unknown migration: {name!r}")
    resolved = db_path or _DEFAULT_DB_PATH
    conn = _open_db(resolved)
    try:
        conn.execute(_SCHEMA_TABLE_SQL)
        applied_at = datetime.now(UTC).isoformat()
        conn.execute(
            "INSERT OR IGNORE INTO schema_migrations (name, applied_at) VALUES (?, ?)",
            (name, applied_at),
        )
        conn.commit()
    finally:
        conn.close()


def validate_migration_state(db_path: Path | None = None) -> None:
    """Validate that all migrations have been applied.

    Should be called on application startup, before any request handler is
    registered.  Exits with code 1 if unapplied migrations are found (unless
    ``SKIP_MIGRATION_CHECK`` is set).

    Args:
        db_path: Path to the SQLite database file.  Defaults to the
            dashboard store default path.
    """
    skip = os.environ.get("SKIP_MIGRATION_CHECK", "").lower() in {"1", "true", "yes"}
    if skip:
        logger.warning("Migration check skipped (SKIP_MIGRATION_CHECK is set)")
        return

    pending = get_pending_migrations(db_path)
    if not pending:
        logger.debug("Migration state OK — all %d migrations applied", len(_EXPECTED_NAMES))
        return

    logger.warning(
        "Unapplied database migrations detected (%d pending). "
        "Run migrations before starting the application.",
        len(pending),
    )
    for name in pending:
        logger.warning("  Pending migration: %s", name)

    sys.exit(1)


__all__ = [
    "MIGRATIONS",
    "get_pending_migrations",
    "mark_migration_applied",
    "validate_migration_state",
]