Lamont/core/migrations/__init__.py

"""Versioned PostgreSQL migrations for the framework and feature modules."""

from dataclasses import dataclass
from hashlib import sha256
from pathlib import Path
import re

import psycopg2.extras

from core import postgres


MIGRATION_NAME = re.compile(r"^(?P<version>\d+)(?:[_-].*)?\.sql$")
LOCK_NAME = "llm-bot-framework:schema-migrations"


class MigrationError(RuntimeError):
    """Raised when migration history is inconsistent or cannot be applied."""


@dataclass(frozen=True)
class Migration:
    namespace: str
    version: int
    name: str
    path: Path
    checksum: str
    sql: str


def _project_root(project_root=None):
    if project_root is not None:
        return Path(project_root).resolve()
    return Path(__file__).resolve().parents[2]


def _read_namespace(namespace, directory):
    migrations = []
    if not directory.is_dir():
        return migrations

    for path in sorted(directory.glob("*.sql")):
        match = MIGRATION_NAME.match(path.name)
        if not match:
            raise MigrationError(f"Invalid migration filename: {path}")
        sql = path.read_text(encoding="utf-8")
        migrations.append(
            Migration(
                namespace=namespace,
                version=int(match.group("version")),
                name=path.name,
                path=path,
                checksum=sha256(sql.encode("utf-8")).hexdigest(),
                sql=sql,
            )
        )
    return migrations


def discover_migrations(project_root=None):
    """Return core and feature migrations in deterministic application order."""
    root = _project_root(project_root)
    migrations = _read_namespace("core", root / "config" / "migrations")

    modules_root = root / "modules"
    if modules_root.is_dir():
        for module_path in sorted(modules_root.iterdir(), key=lambda path: path.name):
            if module_path.is_dir() and not module_path.name.startswith("_"):
                migrations.extend(
                    _read_namespace(module_path.name, module_path / "migrations")
                )

    seen = set()
    for migration in migrations:
        key = (migration.namespace, migration.version)
        if key in seen:
            raise MigrationError(
                f"Duplicate migration {migration.namespace}:{migration.version}"
            )
        seen.add(key)

    return sorted(
        migrations,
        key=lambda item: (
            item.namespace != "core",
            item.namespace,
            item.version,
            item.name,
        ),
    )


def _ensure_history_table(cursor):
    cursor.execute(
        """
        CREATE TABLE IF NOT EXISTS schema_migrations (
            namespace VARCHAR(255) NOT NULL,
            version INTEGER NOT NULL,
            checksum CHAR(64) NOT NULL,
            applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            PRIMARY KEY (namespace, version)
        )
        """
    )


def _applied_migrations(cursor):
    cursor.execute(
        """
        SELECT namespace, version, checksum, applied_at
        FROM schema_migrations
        ORDER BY namespace, version
        """
    )
    return {
        (row["namespace"], row["version"]): dict(row) for row in cursor.fetchall()
    }


def upgrade(project_root=None):
    """Apply pending migrations transactionally under a PostgreSQL advisory lock."""
    migrations = discover_migrations(project_root)
    applied_now = []

    with postgres.get_connection() as connection:
        with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
            cursor.execute("SELECT pg_advisory_xact_lock(hashtext(%s))", (LOCK_NAME,))
            _ensure_history_table(cursor)
            applied = _applied_migrations(cursor)

            for migration in migrations:
                key = (migration.namespace, migration.version)
                previous = applied.get(key)
                if previous:
                    if previous["checksum"] != migration.checksum:
                        raise MigrationError(
                            "Applied migration checksum changed: "
                            f"{migration.namespace}:{migration.version}"
                        )
                    continue

                cursor.execute(migration.sql)
                cursor.execute(
                    """
                    INSERT INTO schema_migrations (namespace, version, checksum)
                    VALUES (%s, %s, %s)
                    """,
                    (migration.namespace, migration.version, migration.checksum),
                )
                applied_now.append(migration)

    return applied_now


def migration_status(project_root=None):
    """Return applied, pending, changed, and source-missing migration records."""
    migrations = discover_migrations(project_root)
    with postgres.get_connection() as connection:
        with connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cursor:
            cursor.execute(
                "SELECT to_regclass('public.schema_migrations') AS table_name"
            )
            if cursor.fetchone()["table_name"] is None:
                applied = {}
            else:
                applied = _applied_migrations(cursor)

    status = []
    source_keys = set()
    for migration in migrations:
        key = (migration.namespace, migration.version)
        source_keys.add(key)
        record = applied.get(key)
        state = "pending"
        applied_at = None
        if record:
            state = "applied" if record["checksum"] == migration.checksum else "changed"
            applied_at = record["applied_at"]
        status.append(
            {
                "namespace": migration.namespace,
                "version": migration.version,
                "name": migration.name,
                "checksum": migration.checksum,
                "applied_at": applied_at,
                "state": state,
            }
        )

    for key, record in sorted(applied.items()):
        if key not in source_keys:
            status.append(
                {
                    "namespace": key[0],
                    "version": key[1],
                    "name": None,
                    "checksum": record["checksum"],
                    "applied_at": record["applied_at"],
                    "state": "missing",
                }
            )
    return status