99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
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
|
|
)
|
|
"""
|
|
)
|