from __future__ import annotations import hashlib import sqlite3 from collections.abc import Callable, Iterable from dataclasses import dataclass from datetime import datetime, timezone MigrationAction = Callable[[sqlite3.Connection], None] class MigrationError(RuntimeError): pass @dataclass(frozen=True) class Migration: version: str name: str action: MigrationAction signature: str @property def checksum(self) -> str: return hashlib.sha256(self.signature.encode("utf-8")).hexdigest() class MigrationRunner: def apply( self, connection: sqlite3.Connection, migrations: Iterable[Migration], ) -> tuple[str, ...]: ordered = sorted(migrations, key=lambda item: item.version) versions = [item.version for item in ordered] if versions != sorted(set(versions)): raise MigrationError("Migration versions must be unique and ordered") self._ensure_ledger(connection) applied = { str(row["version"]): str(row["checksum"]) for row in connection.execute( "SELECT version, checksum FROM schema_migrations ORDER BY version" ) } known = set(versions) unknown = sorted(set(applied) - known) if unknown: raise MigrationError(f"Database contains unknown migrations: {', '.join(unknown)}") completed: list[str] = [] for migration in ordered: existing = applied.get(migration.version) if existing: if existing != migration.checksum: raise MigrationError( f"Migration checksum changed: {migration.version} {migration.name}" ) continue savepoint = f"migration_{migration.version.replace('-', '_')}" connection.execute(f"SAVEPOINT {savepoint}") try: migration.action(connection) connection.execute( """ INSERT INTO schema_migrations (version, name, checksum, applied_at) VALUES (?, ?, ?, ?) """, ( migration.version, migration.name, migration.checksum, datetime.now(timezone.utc).isoformat(), ), ) connection.execute(f"RELEASE SAVEPOINT {savepoint}") except Exception as exc: connection.execute(f"ROLLBACK TO SAVEPOINT {savepoint}") connection.execute(f"RELEASE SAVEPOINT {savepoint}") raise MigrationError( f"Migration failed: {migration.version} {migration.name}" ) from exc completed.append(migration.version) return tuple(completed) @staticmethod def _ensure_ledger(connection: sqlite3.Connection) -> None: connection.execute( """ CREATE TABLE IF NOT EXISTS schema_migrations ( version TEXT PRIMARY KEY, name TEXT NOT NULL, checksum TEXT NOT NULL, applied_at TEXT NOT NULL ) """ )