feat: introduce universal database migration system

zyrarbot/storage/adapters/postgres/database.py

@@ -8,6 +8,9 @@ from typing import Any, Iterator
 
 import psycopg
 
+from zyrarbot.storage.adapters.postgres.migration_runner import MigrationRunner
+from zyrarbot.storage.adapters.postgres.migrations import MIGRATIONS
+
 logger = logging.getLogger(__name__)
 
 
@@ -30,33 +33,7 @@ class PostgresDatabase:
             "password": password,
         }
         with self.connect() as connection:
-            schema = """
-                CREATE TABLE IF NOT EXISTS guilds (
-                    id BIGINT PRIMARY KEY,
-                    notification_channel_id BIGINT,
-                    default_plays INTEGER
-                );
-                CREATE TABLE IF NOT EXISTS history_entries (
-                    guild_id BIGINT NOT NULL REFERENCES guilds(id) ON DELETE CASCADE,
-                    position INTEGER NOT NULL,
-                    id TEXT PRIMARY KEY,
-                    title TEXT NOT NULL,
-                    url TEXT,
-                    duration DOUBLE PRECISION,
-                    is_local BOOLEAN NOT NULL DEFAULT FALSE,
-                    CONSTRAINT history_entries_guild_position_key UNIQUE (guild_id, position)
-                );
-                CREATE TABLE IF NOT EXISTS queue_snapshots (
-                    guild_id BIGINT PRIMARY KEY REFERENCES guilds(id) ON DELETE CASCADE,
-                    tracks JSONB NOT NULL
-                );
-                CREATE TABLE IF NOT EXISTS migrations (
-                    name TEXT PRIMARY KEY
-                );
-                """
-            for statement in schema.split(";"):
-                if statement.strip():
-                    connection.execute(statement)
+            MigrationRunner(MIGRATIONS).apply_pending(connection)
 
     @contextmanager
     def connect(self) -> Iterator[psycopg.Connection[Any]]:
@@ -67,7 +44,7 @@ class PostgresDatabase:
         """Import legacy JSON files once, retaining the source files as a backup."""
         base = Path(config_dir)
         with self.connect() as connection:
-            if connection.execute("SELECT 1 FROM migrations WHERE name='json-v1'").fetchone():
+            if connection.execute("SELECT 1 FROM data_imports WHERE name='json-v1'").fetchone():
                 return
 
             config_path = base / "config.json"
@@ -107,7 +84,7 @@ class PostgresDatabase:
                             "ON CONFLICT(guild_id) DO NOTHING",
                             (int(path.stem), json.dumps(tracks, ensure_ascii=False)),
                         )
-            connection.execute("INSERT INTO migrations(name) VALUES ('json-v1')")
+            connection.execute("INSERT INTO data_imports(name) VALUES ('json-v1')")
 
 
 def _read_json(path: Path) -> dict[str, Any]:

zyrarbot/storage/adapters/postgres/migration_runner.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+import argparse
+import os
+
+import psycopg
+
+from zyrarbot.storage.adapters.postgres.migrations.base import Migration
+
+
+class MigrationRunner:
+    """Apply registered PostgreSQL migrations once, in version order."""
+
+    def __init__(self, migrations: tuple[Migration, ...]):
+        versions = [migration.version for migration in migrations]
+        if len(versions) != len(set(versions)):
+            raise ValueError("Migration versions must be unique")
+        self._migrations = tuple(sorted(migrations, key=lambda migration: migration.version))
+
+    def apply_pending(self, connection: psycopg.Connection) -> None:
+        connection.execute(
+            """
+            CREATE TABLE IF NOT EXISTS schema_migrations (
+                version INTEGER PRIMARY KEY,
+                name TEXT NOT NULL,
+                applied_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+            )
+            """
+        )
+        rows = connection.execute("SELECT version FROM schema_migrations").fetchall()
+        applied_versions = {row[0] for row in rows}
+        for migration in self._migrations:
+            if migration.version in applied_versions:
+                continue
+            with connection.transaction():
+                migration.apply(connection)
+                connection.execute(
+                    "INSERT INTO schema_migrations(version, name) VALUES (%s, %s)",
+                    (migration.version, migration.name),
+                )
+
+    def rollback_to(self, connection: psycopg.Connection, target_version: int = 0) -> None:
+        """Roll back applied migrations newer than target_version."""
+        rows = connection.execute("SELECT version FROM schema_migrations").fetchall()
+        applied_versions = {row[0] for row in rows}
+        known_versions = {migration.version for migration in self._migrations}
+        unknown_versions = applied_versions - known_versions
+        if unknown_versions:
+            raise ValueError(f"No registered migrations for versions: {sorted(unknown_versions)}")
+
+        for migration in reversed(self._migrations):
+            if migration.version <= target_version or migration.version not in applied_versions:
+                continue
+            with connection.transaction():
+                migration.rollback(connection)
+                connection.execute(
+                    "DELETE FROM schema_migrations WHERE version = %s", (migration.version,)
+                )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Manage PostgreSQL schema migrations")
+    parser.add_argument("action", choices=("rollback",))
+    parser.add_argument("--to", type=int, default=0, dest="target_version")
+    args = parser.parse_args()
+
+    with psycopg.connect(
+        host=os.environ.get("POSTGRES_HOST", "localhost"),
+        port=int(os.environ.get("POSTGRES_PORT", "5432")),
+        dbname=os.environ.get("POSTGRES_DB", "zyrarbot"),
+        user=os.environ.get("POSTGRES_USER", "zyrarbot"),
+        password=os.environ.get("POSTGRES_PASSWORD", "zyrarbot"),
+    ) as connection:
+        MigrationRunner(MIGRATIONS).rollback_to(connection, args.target_version)
+
+
+if __name__ == "__main__":
+    from zyrarbot.storage.adapters.postgres.migrations import MIGRATIONS
+
+    main()

zyrarbot/storage/adapters/postgres/migrations/0001_initial_schema.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import psycopg
+
+
+class InitialSchemaMigration:
+    """Create the initial application schema."""
+
+    version = 1
+    name = "initial_schema"
+
+    def apply(self, connection: psycopg.Connection) -> None:
+        connection.execute(
+            """
+            CREATE TABLE guilds (
+                id BIGINT PRIMARY KEY,
+                notification_channel_id BIGINT,
+                default_plays INTEGER
+            )
+            """
+        )
+        connection.execute(
+            """
+            CREATE TABLE history_entries (
+                guild_id BIGINT NOT NULL REFERENCES guilds(id) ON DELETE CASCADE,
+                position INTEGER NOT NULL,
+                id TEXT PRIMARY KEY,
+                title TEXT NOT NULL,
+                url TEXT,
+                duration DOUBLE PRECISION,
+                is_local BOOLEAN NOT NULL DEFAULT FALSE,
+                CONSTRAINT history_entries_guild_position_key UNIQUE (guild_id, position)
+            )
+            """
+        )
+        connection.execute(
+            """
+            CREATE TABLE queue_snapshots (
+                guild_id BIGINT PRIMARY KEY REFERENCES guilds(id) ON DELETE CASCADE,
+                tracks JSONB NOT NULL
+            )
+            """
+        )
+        connection.execute(
+            """
+            CREATE TABLE data_imports (
+                name TEXT PRIMARY KEY,
+                imported_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP
+            )
+            """
+        )
+
+    def rollback(self, connection: psycopg.Connection) -> None:
+        connection.execute("DROP TABLE data_imports")
+        connection.execute("DROP TABLE queue_snapshots")
+        connection.execute("DROP TABLE history_entries")
+        connection.execute("DROP TABLE guilds")

zyrarbot/storage/adapters/postgres/migrations/__init__.py

@@ -0,0 +1,9 @@
+"""Ordered PostgreSQL schema migrations."""
+
+from importlib import import_module
+
+InitialSchemaMigration = import_module(
+    ".0001_initial_schema", __name__
+).InitialSchemaMigration
+
+MIGRATIONS = (InitialSchemaMigration(),)

zyrarbot/storage/adapters/postgres/migrations/base.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from typing import Protocol
+
+import psycopg
+
+
+class Migration(Protocol):
+    version: int
+    name: str
+
+    def apply(self, connection: psycopg.Connection) -> None: ...
+
+    def rollback(self, connection: psycopg.Connection) -> None: ...