"""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\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