From 45543983c2f3aa281e72d0b0e53c259909cbe491 Mon Sep 17 00:00:00 2001 From: 3252a8 <3252a8@proton.me> Date: Sun, 31 May 2026 14:57:32 +0300 Subject: [PATCH] fix: run migrations after database restore --- backend/bot/app/web/admin_api_impl/backups.py | 5 +- .../bot/services/backup_restore_service.py | 69 +++++++++++++++++++ tests/test_backup_restore_service.py | 40 +++++++++++ 3 files changed, 111 insertions(+), 3 deletions(-) diff --git a/backend/bot/app/web/admin_api_impl/backups.py b/backend/bot/app/web/admin_api_impl/backups.py index 63d4004..1504ac7 100644 --- a/backend/bot/app/web/admin_api_impl/backups.py +++ b/backend/bot/app/web/admin_api_impl/backups.py @@ -169,12 +169,11 @@ async def admin_backups_restore_route(request: web.Request) -> web.Response: except (OSError, subprocess.SubprocessError, TimeoutError) as exc: logger.exception("Backup restore failed") return _error(500, "backup_restore_failed", str(exc)) - - if result.database_restored: + finally: try: from db import database_setup - if database_setup.async_engine is not None: + if restore_database and database_setup.async_engine is not None: await database_setup.async_engine.dispose() except Exception: logger.exception("Failed to dispose DB engine after backup restore") diff --git a/backend/bot/services/backup_restore_service.py b/backend/bot/services/backup_restore_service.py index 4be5624..2bc52a7 100644 --- a/backend/bot/services/backup_restore_service.py +++ b/backend/bot/services/backup_restore_service.py @@ -9,11 +9,16 @@ import shutil import subprocess import tempfile import zipfile +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path, PurePosixPath from typing import Any, Optional +from sqlalchemy import inspect, text +from sqlalchemy.engine import Connection +from sqlalchemy.ext.asyncio import create_async_engine + from bot.services.backup_archive import ( BACKUP_APP_ID, BACKUP_FILENAME_PREFIX, @@ -29,6 +34,8 @@ from bot.services.backup_worker import ( DEFAULT_COMPOSE_EXCLUDED_DIRS, ) from config.settings import Settings +from db.migrator import MIGRATIONS, run_database_migrations +from db.models import Base logger = logging.getLogger(__name__) @@ -42,6 +49,23 @@ BACKUP_MAX_COMPRESSION_RATIO = 200 BACKUP_ZIP_BOMB_MIN_BYTES = 100 * 1024 * 1024 COMPOSE_PRE_RESTORE_PREFIX = "minishop-pre-restore-" SAFE_ARCHIVE_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.@+-]{0,220}\.zip$") +DB_RESTORE_MIGRATION_ADVISORY_LOCK_ID = 817512404897421337 + + +def _applied_migration_ids(connection: Connection) -> set[str]: + inspector = inspect(connection) + if "schema_migrations" not in inspector.get_table_names(): + return set() + return {row[0] for row in connection.execute(text("SELECT id FROM schema_migrations"))} + + +def _create_missing_tables_and_migrate(connection: Connection) -> list[str]: + before = _applied_migration_ids(connection) + Base.metadata.create_all(connection) + run_database_migrations(connection) + after = _applied_migration_ids(connection) + newly_applied = after - before + return [migration.id for migration in MIGRATIONS if migration.id in newly_applied] class BackupArchiveError(ValueError): @@ -92,6 +116,7 @@ class BackupRestoreResult: compose_files_restored: int = 0 compose_target_dir: Optional[str] = None compose_pre_restore_archive: Optional[str] = None + database_migrations_applied: list[str] = field(default_factory=list) warnings: list[str] = field(default_factory=list) def to_payload(self) -> dict[str, Any]: @@ -103,6 +128,7 @@ class BackupRestoreResult: "compose_files_restored": self.compose_files_restored, "compose_target_dir": self.compose_target_dir, "compose_pre_restore_archive": self.compose_pre_restore_archive, + "database_migrations_applied": self.database_migrations_applied, "warnings": self.warnings, } @@ -242,9 +268,11 @@ class BackupRestoreService: compose_pre_restore_archive = self._snapshot_current_compose(compose_target_dir) database_restored = False + database_migrations_applied: list[str] = [] if db_member is not None: dump_path = self._extract_database_dump(archive, db_member, temp_dir) self._run_pg_restore(dump_path) + database_migrations_applied = self._run_post_restore_migrations() database_restored = True compose_files_restored = 0 @@ -265,9 +293,50 @@ class BackupRestoreService: compose_pre_restore_archive=str(compose_pre_restore_archive) if compose_pre_restore_archive else None, + database_migrations_applied=database_migrations_applied, warnings=warnings, ) + def _run_post_restore_migrations(self) -> list[str]: + try: + asyncio.get_running_loop() + except RuntimeError: + run_migrations = lambda: asyncio.run(self._run_post_restore_migrations_async()) + else: + run_migrations = self._run_post_restore_migrations_in_thread + + try: + return run_migrations() + except BackupRestoreError: + raise + except Exception as exc: + raise BackupRestoreError( + f"Database restore completed, but post-restore migrations failed: {str(exc)[:500]}" + ) from exc + + def _run_post_restore_migrations_in_thread(self) -> list[str]: + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="backup-restore-migrate") as pool: + return pool.submit( + lambda: asyncio.run(self._run_post_restore_migrations_async()) + ).result() + + async def _run_post_restore_migrations_async(self) -> list[str]: + engine = create_async_engine( + self.settings.DATABASE_URL, + echo=False, + pool_pre_ping=True, + pool_size=1, + max_overflow=0, + ) + try: + async with engine.begin() as connection: + await connection.execute( + text(f"SELECT pg_advisory_xact_lock({DB_RESTORE_MIGRATION_ADVISORY_LOCK_ID})") + ) + return await connection.run_sync(_create_missing_tables_and_migrate) + finally: + await engine.dispose() + def _run_pg_restore(self, dump_path: Path) -> None: pg_restore_path = str(getattr(self.settings, "BACKUP_PG_RESTORE_PATH", "pg_restore") or "") pg_restore_path = pg_restore_path or "pg_restore" diff --git a/tests/test_backup_restore_service.py b/tests/test_backup_restore_service.py index f1afb70..eea141d 100644 --- a/tests/test_backup_restore_service.py +++ b/tests/test_backup_restore_service.py @@ -168,6 +168,7 @@ def test_backup_restore_service_runs_pg_restore_for_dump(tmp_path): restored_payloads.append(dump_path.read_bytes()) service._run_pg_restore = fake_pg_restore + service._run_post_restore_migrations = lambda: [] result = asyncio.run( service.restore_archive( @@ -179,6 +180,44 @@ def test_backup_restore_service_runs_pg_restore_for_dump(tmp_path): assert result.database_restored is True assert restored_payloads == [b"fake dump"] + assert result.database_migrations_applied == [] + + +def test_backup_restore_service_runs_migrations_after_database_restore(tmp_path): + compose_dir = tmp_path / "compose" + compose_dir.mkdir() + settings = _settings(tmp_path, compose_dir) + archive_path = Path(settings.BACKUP_DIR) / f"{BACKUP_FILENAME_PREFIX}20260527-12-00.zip" + _write_backup_archive(archive_path, include_compose=False) + service = BackupRestoreService(settings) + calls = [] + + def fake_pg_restore(dump_path: Path) -> None: + calls.append(("restore", dump_path.read_bytes())) + + def fake_migrations() -> list[str]: + calls.append(("migrate", None)) + return ["0031_add_subscription_notifications", "0032_add_telegram_notification_status"] + + service._run_pg_restore = fake_pg_restore + service._run_post_restore_migrations = fake_migrations + + result = service.restore_archive_sync( + archive_path.name, + restore_database=True, + restore_compose=False, + ) + + assert calls == [("restore", b"fake dump"), ("migrate", None)] + assert result.database_restored is True + assert result.database_migrations_applied == [ + "0031_add_subscription_notifications", + "0032_add_telegram_notification_status", + ] + assert result.to_payload()["database_migrations_applied"] == [ + "0031_add_subscription_notifications", + "0032_add_telegram_notification_status", + ] def test_backup_restore_service_accepts_archive_from_another_instance(tmp_path): @@ -197,6 +236,7 @@ def test_backup_restore_service_accepts_archive_from_another_instance(tmp_path): restored_payloads.append(dump_path.read_bytes()) service._run_pg_restore = fake_pg_restore + service._run_post_restore_migrations = lambda: [] result = service.restore_archive_sync( archive_path.name,