149 lines
5.7 KiB
Python
149 lines
5.7 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import sqlite3
|
|
from collections.abc import Callable, Iterable
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime
|
|
|
|
from backend.database.connection import Database
|
|
|
|
MigrationAction = Callable[[sqlite3.Connection], None]
|
|
|
|
|
|
class MigrationError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Migration:
|
|
version: int
|
|
name: str
|
|
signature: str
|
|
upgrade: MigrationAction
|
|
downgrade: MigrationAction
|
|
|
|
@property
|
|
def checksum(self) -> str:
|
|
source = f"{self.version}:{self.name}:{self.signature}"
|
|
return hashlib.sha256(source.encode("utf-8")).hexdigest()
|
|
|
|
|
|
class MigrationRunner:
|
|
def __init__(self, database: Database) -> None:
|
|
self._database = database
|
|
|
|
def upgrade(self, migrations: Iterable[Migration]) -> tuple[int, ...]:
|
|
ordered = self._validate(migrations)
|
|
with self._database.transaction() as connection:
|
|
self._ensure_ledger(connection)
|
|
applied = self._applied(connection)
|
|
self._verify_history(ordered, applied)
|
|
completed: list[int] = []
|
|
for migration in ordered:
|
|
if migration.version in applied:
|
|
continue
|
|
self._run_action(connection, migration, migration.upgrade, "upgrade")
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO schema_migrations (version, name, checksum, applied_at)
|
|
VALUES (?, ?, ?, ?)
|
|
""",
|
|
(
|
|
migration.version,
|
|
migration.name,
|
|
migration.checksum,
|
|
datetime.now(UTC).isoformat(timespec="seconds"),
|
|
),
|
|
)
|
|
completed.append(migration.version)
|
|
return tuple(completed)
|
|
|
|
def downgrade(self, migrations: Iterable[Migration], target_version: int) -> tuple[int, ...]:
|
|
if target_version < 0:
|
|
raise MigrationError("Target version cannot be negative")
|
|
ordered = self._validate(migrations)
|
|
by_version = {migration.version: migration for migration in ordered}
|
|
with self._database.transaction() as connection:
|
|
self._ensure_ledger(connection)
|
|
applied = self._applied(connection)
|
|
self._verify_history(ordered, applied)
|
|
pending = sorted(
|
|
(version for version in applied if version > target_version), reverse=True
|
|
)
|
|
rolled_back: list[int] = []
|
|
for version in pending:
|
|
migration = by_version[version]
|
|
self._run_action(connection, migration, migration.downgrade, "downgrade")
|
|
connection.execute("DELETE FROM schema_migrations WHERE version = ?", (version,))
|
|
rolled_back.append(version)
|
|
return tuple(rolled_back)
|
|
|
|
@staticmethod
|
|
def _validate(migrations: Iterable[Migration]) -> tuple[Migration, ...]:
|
|
ordered = tuple(sorted(migrations, key=lambda migration: migration.version))
|
|
versions = [migration.version for migration in ordered]
|
|
if any(version <= 0 for version in versions):
|
|
raise MigrationError("Migration versions must be positive integers")
|
|
if len(versions) != len(set(versions)):
|
|
raise MigrationError("Migration versions must be unique")
|
|
return ordered
|
|
|
|
@staticmethod
|
|
def _ensure_ledger(connection: sqlite3.Connection) -> None:
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS schema_migrations (
|
|
version INTEGER PRIMARY KEY,
|
|
name TEXT NOT NULL,
|
|
checksum TEXT NOT NULL,
|
|
applied_at TEXT NOT NULL
|
|
)
|
|
"""
|
|
)
|
|
|
|
@staticmethod
|
|
def _applied(connection: sqlite3.Connection) -> dict[int, str]:
|
|
return {
|
|
int(row["version"]): str(row["checksum"])
|
|
for row in connection.execute(
|
|
"SELECT version, checksum FROM schema_migrations ORDER BY version"
|
|
)
|
|
}
|
|
|
|
@staticmethod
|
|
def _verify_history(migrations: tuple[Migration, ...], applied: dict[int, str]) -> None:
|
|
by_version = {migration.version: migration for migration in migrations}
|
|
known_versions = sorted(by_version)
|
|
applied_versions = sorted(applied)
|
|
unknown = sorted(set(applied) - set(known_versions))
|
|
if unknown:
|
|
values = ", ".join(f"{version:04d}" for version in unknown)
|
|
raise MigrationError(f"Database contains unknown migrations: {values}")
|
|
if applied_versions != known_versions[: len(applied_versions)]:
|
|
raise MigrationError("Database migration history is not contiguous")
|
|
for version in applied_versions:
|
|
if applied[version] != by_version[version].checksum:
|
|
raise MigrationError(
|
|
f"Migration checksum changed: {version:04d} {by_version[version].name}"
|
|
)
|
|
|
|
@staticmethod
|
|
def _run_action(
|
|
connection: sqlite3.Connection,
|
|
migration: Migration,
|
|
action: MigrationAction,
|
|
direction: str,
|
|
) -> None:
|
|
savepoint = f"migration_{migration.version:04d}_{direction}"
|
|
connection.execute(f"SAVEPOINT {savepoint}")
|
|
try:
|
|
action(connection)
|
|
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 {direction} failed: {migration.version:04d} {migration.name}"
|
|
) from exc
|