From 8a3bba139eaf62889587f10c0ed7b75ab6f50657 Mon Sep 17 00:00:00 2001 From: busya Date: Tue, 25 Aug 2026 16:06:14 +0300 Subject: [PATCH] fix(database): auto-reset orphaned Alembic revisions --- backend/src/scripts/prepare_database.py | 29 ++++++++++++++----- .../test_alembic_user_id_migration.py | 3 +- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/backend/src/scripts/prepare_database.py b/backend/src/scripts/prepare_database.py index f190ce637..1af01348e 100644 --- a/backend/src/scripts/prepare_database.py +++ b/backend/src/scripts/prepare_database.py @@ -7,6 +7,9 @@ from pathlib import Path import sys from alembic.config import Config +from alembic.script import ScriptDirectory +from alembic.script.revision import ResolutionError +from alembic.util.exc import CommandError from sqlalchemy import create_engine, inspect, text from alembic import command @@ -20,16 +23,22 @@ def _enabled(value: str | None) -> bool: return str(value or "").strip().lower() in {"1", "true", "yes", "y"} -def _maybe_reset(connection) -> None: +def _revision_is_known(config: Config, revision: str) -> bool: + try: + return ScriptDirectory.from_config(config).get_revision(revision) is not None + except (CommandError, ResolutionError): + return False + + +def _maybe_reset(connection, config: Config) -> None: reset_enabled = _enabled(os.getenv("RESET_DATABASE_SCHEMA")) sys.stderr.write( f"[database] RESET_DATABASE_SCHEMA={'enabled' if reset_enabled else 'disabled'}\n" ) - if not reset_enabled: - return - if connection.dialect.name != "postgresql": - raise RuntimeError("RESET_DATABASE_SCHEMA is supported only for PostgreSQL") + if reset_enabled: + raise RuntimeError("RESET_DATABASE_SCHEMA is supported only for PostgreSQL") + return tables = set(inspect(connection).get_table_names()) revision = None @@ -42,6 +51,9 @@ def _maybe_reset(connection) -> None: sys.stderr.write("[database] Baseline already applied; one-shot reset skipped\n") connection.commit() return + orphaned_revision = bool(revision) and not _revision_is_known(config, revision) + if not reset_enabled and not orphaned_revision: + return if tables and revision is None: raise RuntimeError( "Database schema is non-empty but has no Alembic revision; " @@ -52,7 +64,8 @@ def _maybe_reset(connection) -> None: connection.execute(text("CREATE SCHEMA public")) connection.commit() sys.stderr.write( - f"[database] Legacy Alembic revision {revision or ''} reset successfully\n" + f"[database] {'Orphaned' if orphaned_revision else 'Legacy'} Alembic revision " + f"{revision or ''} reset successfully\n" ) @@ -60,14 +73,14 @@ def prepare_database() -> None: engine = create_engine(database_url(), pool_pre_ping=True) try: with engine.connect() as connection: + config = Config(str(Path(__file__).resolve().parents[2] / "alembic.ini")) locked = connection.dialect.name == "postgresql" if locked: connection.execute(text("SELECT pg_advisory_lock(:lock_id)"), {"lock_id": _MIGRATION_LOCK_ID}) connection.commit() try: - _maybe_reset(connection) + _maybe_reset(connection, config) connection.commit() - config = Config(str(Path(__file__).resolve().parents[2] / "alembic.ini")) config.attributes["connection"] = connection command.upgrade(config, "head") finally: diff --git a/backend/tests/integration/test_alembic_user_id_migration.py b/backend/tests/integration/test_alembic_user_id_migration.py index d329ad9cd..281b04748 100644 --- a/backend/tests/integration/test_alembic_user_id_migration.py +++ b/backend/tests/integration/test_alembic_user_id_migration.py @@ -74,7 +74,8 @@ def test_prepare_database_resets_any_legacy_revision( connection.execute(text("CREATE TABLE legacy_probe (id INTEGER PRIMARY KEY)")) monkeypatch.setenv("DATABASE_URL", migration_database_url) - monkeypatch.setenv("RESET_DATABASE_SCHEMA", "true") + # Orphaned revisions are reset independently of stale Compose defaults. + monkeypatch.setenv("RESET_DATABASE_SCHEMA", "false") monkeypatch.delenv("RESET_DATABASE_FROM_REVISION", raising=False) monkeypatch.delenv("RESET_DATABASE_NAME", raising=False)