Sergey Antropoff
71 lines
2.2 KiB
Python
71 lines
2.2 KiB
Python
"""Checksummed asynchronous PostgreSQL migration runner."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import asyncpg # type: ignore[import-untyped]
|
|
from asyncpg import Connection
|
|
|
|
|
|
class MigrationError(RuntimeError):
|
|
pass
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class Migration:
|
|
version: int
|
|
name: str
|
|
sql: str
|
|
checksum: str
|
|
|
|
|
|
def load_migrations(root: Path | None = None) -> tuple[Migration, ...]:
|
|
directory = root or Path(__file__).with_name("migrations")
|
|
migrations = []
|
|
for path in sorted(directory.glob("[0-9][0-9][0-9]_*.sql")):
|
|
version = int(path.name.split("_", 1)[0])
|
|
sql = path.read_text()
|
|
migrations.append(
|
|
Migration(version, path.stem, sql, hashlib.sha256(sql.encode()).hexdigest())
|
|
)
|
|
return tuple(migrations)
|
|
|
|
|
|
async def migrate(connection: Connection, migrations: tuple[Migration, ...] | None = None) -> int:
|
|
await connection.execute(
|
|
"""CREATE TABLE IF NOT EXISTS schema_migrations (
|
|
version integer PRIMARY KEY, name text NOT NULL, checksum text NOT NULL,
|
|
applied_at timestamptz NOT NULL DEFAULT now())"""
|
|
)
|
|
applied = {
|
|
int(row["version"]): str(row["checksum"])
|
|
for row in await connection.fetch("SELECT version, checksum FROM schema_migrations")
|
|
}
|
|
count = 0
|
|
for migration in migrations or load_migrations():
|
|
if migration.version in applied:
|
|
if applied[migration.version] != migration.checksum:
|
|
raise MigrationError(f"migration {migration.version} checksum mismatch")
|
|
continue
|
|
async with connection.transaction():
|
|
await connection.execute(migration.sql)
|
|
await connection.execute(
|
|
"INSERT INTO schema_migrations(version, name, checksum) VALUES($1, $2, $3)",
|
|
migration.version,
|
|
migration.name,
|
|
migration.checksum,
|
|
)
|
|
count += 1
|
|
return count
|
|
|
|
|
|
async def migrate_url(database_url: str) -> int:
|
|
connection = await asyncpg.connect(database_url)
|
|
try:
|
|
return await migrate(connection)
|
|
finally:
|
|
await connection.close()
|