feat(env): migrate database setup to alembic

2026-02-08 23:27:04 +03:00
parent 8ba0e465b1
commit f870e4d8b7
9 changed files with 730 additions and 289 deletions
@@ -0,0 +1,277 @@
+import logging
+from pathlib import Path
+from typing import Tuple
+
+from alembic import command
+from alembic.config import Config
+from sqlalchemy import inspect, text
+from sqlalchemy.engine import Connection
+from sqlalchemy.ext.asyncio import AsyncEngine
+
+from config.settings import Settings
+
+
+_BASELINE_REVISION = "0001_initial_schema"
+
+
+def _build_alembic_config(settings: Settings) -> Config:
+    project_root = Path(__file__).resolve().parents[1]
+    config = Config(str(project_root / "alembic.ini"))
+    config.set_main_option("script_location", str(project_root / "alembic"))
+    config.set_main_option("sqlalchemy.url", settings.DATABASE_URL)
+    return config
+
+
+def _inspect_database_state(connection: Connection) -> Tuple[bool, bool, bool]:
+    db_inspector = inspect(connection)
+    has_alembic_version = db_inspector.has_table("alembic_version")
+    has_users_table = db_inspector.has_table("users")
+    has_legacy_migrator_table = db_inspector.has_table("schema_migrations")
+    return has_alembic_version, has_users_table, has_legacy_migrator_table
+
+
+def _run_legacy_migrator_compatibility(connection: Connection) -> None:
+    db_inspector = inspect(connection)
+    if not db_inspector.has_table("users"):
+        return
+
+    users_columns = {
+        column["name"]
+        for column in db_inspector.get_columns("users")
+    }
+    user_alter_statements = []
+
+    if "channel_subscription_verified" not in users_columns:
+        user_alter_statements.append(
+            "ALTER TABLE users ADD COLUMN channel_subscription_verified BOOLEAN"
+        )
+    if "channel_subscription_checked_at" not in users_columns:
+        user_alter_statements.append(
+            "ALTER TABLE users ADD COLUMN channel_subscription_checked_at TIMESTAMPTZ"
+        )
+    if "channel_subscription_verified_for" not in users_columns:
+        user_alter_statements.append(
+            "ALTER TABLE users ADD COLUMN channel_subscription_verified_for BIGINT"
+        )
+    if "referral_code" not in users_columns:
+        user_alter_statements.append(
+            "ALTER TABLE users ADD COLUMN referral_code VARCHAR(16)"
+        )
+
+    for statement in user_alter_statements:
+        connection.execute(text(statement))
+
+    users_columns = {
+        column["name"]
+        for column in inspect(connection).get_columns("users")
+    }
+    if "referral_code" in users_columns:
+        connection.execute(
+            text(
+                """
+                WITH generated_codes AS (
+                    SELECT
+                        user_id,
+                        UPPER(
+                            SUBSTRING(
+                                md5(
+                                    user_id::text
+                                    || clock_timestamp()::text
+                                    || random()::text
+                                )
+                                FROM 1 FOR 9
+                            )
+                        ) AS referral_code
+                    FROM users
+                    WHERE referral_code IS NULL OR referral_code = ''
+                )
+                UPDATE users AS u
+                SET referral_code = g.referral_code
+                FROM generated_codes AS g
+                WHERE u.user_id = g.user_id
+                """
+            )
+        )
+        connection.execute(
+            text(
+                """
+                CREATE UNIQUE INDEX IF NOT EXISTS uq_users_referral_code
+                ON users (referral_code)
+                WHERE referral_code IS NOT NULL
+                """
+            )
+        )
+        connection.execute(
+            text(
+                """
+                UPDATE users
+                SET referral_code = UPPER(referral_code)
+                WHERE referral_code IS NOT NULL
+                  AND referral_code <> UPPER(referral_code)
+                """
+            )
+        )
+
+    db_inspector = inspect(connection)
+    if db_inspector.has_table("payments"):
+        payments_columns = {
+            column["name"]
+            for column in db_inspector.get_columns("payments")
+        }
+        if "original_amount" not in payments_columns:
+            connection.execute(text("ALTER TABLE payments ADD COLUMN original_amount FLOAT"))
+        if "discount_applied" not in payments_columns:
+            connection.execute(text("ALTER TABLE payments ADD COLUMN discount_applied FLOAT"))
+
+    db_inspector = inspect(connection)
+    has_promo_codes = db_inspector.has_table("promo_codes")
+    if has_promo_codes:
+        promo_columns = {
+            column["name"]
+            for column in db_inspector.get_columns("promo_codes")
+        }
+        if "promo_type" not in promo_columns:
+            connection.execute(
+                text(
+                    "ALTER TABLE promo_codes ADD COLUMN promo_type VARCHAR NOT NULL DEFAULT 'bonus_days'"
+                )
+            )
+        if "discount_percentage" not in promo_columns:
+            connection.execute(
+                text("ALTER TABLE promo_codes ADD COLUMN discount_percentage INTEGER")
+            )
+        if "bonus_days" in promo_columns:
+            connection.execute(
+                text("ALTER TABLE promo_codes ALTER COLUMN bonus_days DROP NOT NULL")
+            )
+        connection.execute(
+            text(
+                "CREATE INDEX IF NOT EXISTS idx_promo_codes_promo_type ON promo_codes (promo_type)"
+            )
+        )
+
+    db_inspector = inspect(connection)
+    has_active_discounts = db_inspector.has_table("active_discounts")
+    if not has_active_discounts and has_promo_codes:
+        connection.execute(
+            text(
+                """
+                CREATE TABLE IF NOT EXISTS active_discounts (
+                    user_id BIGINT PRIMARY KEY,
+                    promo_code_id INTEGER NOT NULL,
+                    discount_percentage INTEGER NOT NULL,
+                    activated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                    CONSTRAINT fk_active_discounts_user
+                        FOREIGN KEY (user_id) REFERENCES users (user_id) ON DELETE CASCADE,
+                    CONSTRAINT fk_active_discounts_promo_code
+                        FOREIGN KEY (promo_code_id) REFERENCES promo_codes (promo_code_id) ON DELETE CASCADE
+                )
+                """
+            )
+        )
+        has_active_discounts = True
+
+    if has_active_discounts and has_promo_codes:
+        connection.execute(
+            text(
+                "DELETE FROM active_discounts ad "
+                "WHERE NOT EXISTS (SELECT 1 FROM users u WHERE u.user_id = ad.user_id) "
+                "OR NOT EXISTS (SELECT 1 FROM promo_codes p WHERE p.promo_code_id = ad.promo_code_id)"
+            )
+        )
+        connection.execute(
+            text(
+                "ALTER TABLE active_discounts "
+                "DROP CONSTRAINT IF EXISTS active_discounts_user_id_fkey"
+            )
+        )
+        connection.execute(
+            text(
+                "ALTER TABLE active_discounts "
+                "DROP CONSTRAINT IF EXISTS fk_active_discounts_user"
+            )
+        )
+        connection.execute(
+            text(
+                "ALTER TABLE active_discounts "
+                "DROP CONSTRAINT IF EXISTS active_discounts_promo_code_id_fkey"
+            )
+        )
+        connection.execute(
+            text(
+                "ALTER TABLE active_discounts "
+                "DROP CONSTRAINT IF EXISTS fk_active_discounts_promo_code"
+            )
+        )
+        connection.execute(
+            text(
+                "ALTER TABLE active_discounts "
+                "ADD CONSTRAINT fk_active_discounts_user "
+                "FOREIGN KEY (user_id) REFERENCES users (user_id) ON DELETE CASCADE"
+            )
+        )
+        connection.execute(
+            text(
+                "ALTER TABLE active_discounts "
+                "ADD CONSTRAINT fk_active_discounts_promo_code "
+                "FOREIGN KEY (promo_code_id) REFERENCES promo_codes (promo_code_id) ON DELETE CASCADE"
+            )
+        )
+    elif has_active_discounts and not has_promo_codes:
+        logging.warning(
+            "Alembic legacy compatibility: skipped active_discounts FK repair "
+            "because promo_codes table is missing."
+        )
+
+
+def _run_stamp(connection: Connection, alembic_config: Config, revision: str) -> None:
+    alembic_config.attributes["connection"] = connection
+    command.stamp(alembic_config, revision)
+
+
+def _run_upgrade(connection: Connection, alembic_config: Config) -> None:
+    alembic_config.attributes["connection"] = connection
+    command.upgrade(alembic_config, "head")
+
+
+async def run_alembic_migrations(settings: Settings, async_engine: AsyncEngine) -> None:
+    """Apply Alembic migrations with bootstrap for existing installations."""
+
+    alembic_config = _build_alembic_config(settings)
+
+    async with async_engine.begin() as async_connection:
+        (
+            has_alembic_version,
+            has_users_table,
+            has_legacy_migrator_table,
+        ) = await async_connection.run_sync(
+            _inspect_database_state
+        )
+
+        if not has_alembic_version and has_users_table:
+            if not has_legacy_migrator_table:
+                raise RuntimeError(
+                    "Alembic bootstrap refused: found existing users table without "
+                    "alembic_version and without legacy schema_migrations marker. "
+                    "Cannot safely determine migration baseline."
+                )
+
+            logging.info(
+                "Alembic: applying legacy migrator compatibility fixes before stamp."
+            )
+            await async_connection.run_sync(_run_legacy_migrator_compatibility)
+
+            logging.info(
+                "Alembic: existing schema detected without alembic_version; stamping %s.",
+                _BASELINE_REVISION,
+            )
+            await async_connection.run_sync(
+                _run_stamp,
+                alembic_config,
+                _BASELINE_REVISION,
+            )
+
+        logging.info("Alembic: running upgrade to head...")
+        await async_connection.run_sync(_run_upgrade, alembic_config)
+
+    logging.info("Alembic: migrations applied successfully.")
@@ -4,8 +4,7 @@ from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sess
 from sqlalchemy.orm import sessionmaker

 from config.settings import Settings
-from .models import Base
-from .migrator import run_database_migrations
+from .alembic_runner import run_alembic_migrations

 async_engine = None

@@ -46,7 +45,7 @@ def init_db_connection(settings: Settings) -> sessionmaker:
        autoflush=False,
    )
    logging.info(
-        f"SQLAlchemy Async Engine and SessionFactory configured for PostgreSQL."
+        "SQLAlchemy Async Engine and SessionFactory configured for PostgreSQL."
    )
    return local_async_session_factory

@@ -77,12 +76,8 @@ async def init_db(settings: Settings, session_factory: sessionmaker):
            "async_engine is not initialized. Call init_db_connection and get session_factory first."
        )

-    async with async_engine.begin() as conn:
-        await conn.run_sync(Base.metadata.create_all)
-        await conn.run_sync(run_database_migrations)
-    logging.info(
-        "PostgreSQL database initialized/checked successfully using SQLAlchemy."
-    )
+    await run_alembic_migrations(settings, async_engine)
+    logging.info("PostgreSQL database migrations checked/applied via Alembic.")

    async with session_factory() as session:
        from .dal.panel_sync_dal import get_panel_sync_status, update_panel_sync_status
@@ -1,280 +0,0 @@
-import logging
-from dataclasses import dataclass
-from typing import Callable, List, Set
-
-from sqlalchemy import inspect, text
-from sqlalchemy.engine import Connection
-
-
-@dataclass(frozen=True)
-class Migration:
-    id: str
-    description: str
-    upgrade: Callable[[Connection], None]
-
-
-def _ensure_migrations_table(connection: Connection) -> None:
-    connection.execute(
-        text(
-            """
-            CREATE TABLE IF NOT EXISTS schema_migrations (
-                id VARCHAR(255) PRIMARY KEY,
-                applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
-            )
-            """
-        )
-    )
-
-
-def _migration_0001_add_channel_subscription_fields(connection: Connection) -> None:
-    inspector = inspect(connection)
-    columns: Set[str] = {col["name"] for col in inspector.get_columns("users")}
-    statements: List[str] = []
-
-    if "channel_subscription_verified" not in columns:
-        statements.append(
-            "ALTER TABLE users ADD COLUMN channel_subscription_verified BOOLEAN"
-        )
-    if "channel_subscription_checked_at" not in columns:
-        statements.append(
-            "ALTER TABLE users ADD COLUMN channel_subscription_checked_at TIMESTAMPTZ"
-        )
-    if "channel_subscription_verified_for" not in columns:
-        statements.append(
-            "ALTER TABLE users ADD COLUMN channel_subscription_verified_for BIGINT"
-        )
-
-    for stmt in statements:
-        connection.execute(text(stmt))
-
-
-def _migration_0002_add_referral_code(connection: Connection) -> None:
-    inspector = inspect(connection)
-    columns: Set[str] = {col["name"] for col in inspector.get_columns("users")}
-
-    if "referral_code" not in columns:
-        connection.execute(
-            text("ALTER TABLE users ADD COLUMN referral_code VARCHAR(16)")
-        )
-
-    connection.execute(
-        text(
-            """
-            WITH generated_codes AS (
-                SELECT
-                    user_id,
-                    UPPER(
-                        SUBSTRING(
-                            md5(
-                                user_id::text
-                                || clock_timestamp()::text
-                                || random()::text
-                            )
-                            FROM 1 FOR 9
-                        )
-                    ) AS referral_code
-                FROM users
-                WHERE referral_code IS NULL OR referral_code = ''
-            )
-            UPDATE users AS u
-            SET referral_code = g.referral_code
-            FROM generated_codes AS g
-            WHERE u.user_id = g.user_id
-            """
-        )
-    )
-
-    connection.execute(
-        text(
-            """
-            CREATE UNIQUE INDEX IF NOT EXISTS uq_users_referral_code
-            ON users (referral_code)
-            WHERE referral_code IS NOT NULL
-            """
-        )
-    )
-
-
-def _migration_0003_normalize_referral_codes(connection: Connection) -> None:
-    inspector = inspect(connection)
-    columns: Set[str] = {col["name"] for col in inspector.get_columns("users")}
-    if "referral_code" not in columns:
-        return
-
-    connection.execute(
-        text(
-            """
-            UPDATE users
-            SET referral_code = UPPER(referral_code)
-            WHERE referral_code IS NOT NULL
-              AND referral_code <> UPPER(referral_code)
-            """
-        )
-    )
-
-
-def _migration_0004_add_discount_promo_codes(connection: Connection) -> None:
-    inspector = inspect(connection)
-
-    # 1. Добавить поля в payments
-    payment_columns: Set[str] = {col["name"] for col in inspector.get_columns("payments")}
-    if "original_amount" not in payment_columns:
-        connection.execute(text("ALTER TABLE payments ADD COLUMN original_amount FLOAT"))
-    if "discount_applied" not in payment_columns:
-        connection.execute(text("ALTER TABLE payments ADD COLUMN discount_applied FLOAT"))
-
-    # 2. Модифицировать promo_codes
-    promo_columns: Set[str] = {col["name"] for col in inspector.get_columns("promo_codes")}
-
-    if "promo_type" not in promo_columns:
-        connection.execute(
-            text(
-                "ALTER TABLE promo_codes ADD COLUMN promo_type VARCHAR NOT NULL DEFAULT 'bonus_days'"
-            )
-        )
-
-    if "discount_percentage" not in promo_columns:
-        connection.execute(
-            text("ALTER TABLE promo_codes ADD COLUMN discount_percentage INTEGER")
-        )
-
-    # Изменить bonus_days на nullable (если еще не nullable)
-    connection.execute(
-        text("ALTER TABLE promo_codes ALTER COLUMN bonus_days DROP NOT NULL")
-    )
-
-    # Создать индекс на promo_type
-    connection.execute(
-        text(
-            "CREATE INDEX IF NOT EXISTS idx_promo_codes_promo_type ON promo_codes (promo_type)"
-        )
-    )
-
-    # 3. Создать таблицу active_discounts
-    connection.execute(
-        text(
-            """
-            CREATE TABLE IF NOT EXISTS active_discounts (
-                user_id BIGINT PRIMARY KEY,
-                promo_code_id INTEGER NOT NULL,
-                discount_percentage INTEGER NOT NULL,
-                activated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
-                CONSTRAINT fk_active_discounts_user
-                    FOREIGN KEY (user_id) REFERENCES users (user_id) ON DELETE CASCADE,
-                CONSTRAINT fk_active_discounts_promo_code
-                    FOREIGN KEY (promo_code_id) REFERENCES promo_codes (promo_code_id) ON DELETE CASCADE
-            )
-            """
-        )
-    )
-
-
-def _migration_0005_fix_active_discounts_fk_cascade(connection: Connection) -> None:
-    inspector = inspect(connection)
-    if not inspector.has_table("active_discounts"):
-        return
-
-    connection.execute(
-        text(
-            "DELETE FROM active_discounts ad "
-            "WHERE NOT EXISTS (SELECT 1 FROM users u WHERE u.user_id = ad.user_id) "
-            "OR NOT EXISTS (SELECT 1 FROM promo_codes p WHERE p.promo_code_id = ad.promo_code_id)"
-        )
-    )
-
-    connection.execute(
-        text("ALTER TABLE active_discounts DROP CONSTRAINT IF EXISTS active_discounts_user_id_fkey")
-    )
-    connection.execute(
-        text("ALTER TABLE active_discounts DROP CONSTRAINT IF EXISTS fk_active_discounts_user")
-    )
-    connection.execute(
-        text("ALTER TABLE active_discounts DROP CONSTRAINT IF EXISTS active_discounts_promo_code_id_fkey")
-    )
-    connection.execute(
-        text("ALTER TABLE active_discounts DROP CONSTRAINT IF EXISTS fk_active_discounts_promo_code")
-    )
-
-    connection.execute(
-        text(
-            "ALTER TABLE active_discounts "
-            "ADD CONSTRAINT fk_active_discounts_user "
-            "FOREIGN KEY (user_id) REFERENCES users (user_id) ON DELETE CASCADE"
-        )
-    )
-    connection.execute(
-        text(
-            "ALTER TABLE active_discounts "
-            "ADD CONSTRAINT fk_active_discounts_promo_code "
-            "FOREIGN KEY (promo_code_id) REFERENCES promo_codes (promo_code_id) ON DELETE CASCADE"
-        )
-    )
-
-MIGRATIONS: List[Migration] = [
-    Migration(
-        id="0001_add_channel_subscription_fields",
-        description="Add columns to track required channel subscription verification",
-        upgrade=_migration_0001_add_channel_subscription_fields,
-    ),
-    Migration(
-        id="0002_add_referral_code",
-        description="Store short referral codes for users and backfill existing rows",
-        upgrade=_migration_0002_add_referral_code,
-    ),
-    Migration(
-        id="0003_normalize_referral_codes",
-        description="Normalize referral codes to uppercase for consistent lookups",
-        upgrade=_migration_0003_normalize_referral_codes,
-    ),
-    Migration(
-        id="0004_add_discount_promo_codes",
-        description="Add support for percentage discount promo codes",
-        upgrade=_migration_0004_add_discount_promo_codes,
-    ),
-    Migration(
-        id="0005_fix_active_discounts_fk_cascade",
-        description="Ensure active_discounts FKs cascade on user/promo delete",
-        upgrade=_migration_0005_fix_active_discounts_fk_cascade,
-    ),
-]
-
-
-def run_database_migrations(connection: Connection) -> None:
-    """
-    Apply pending migrations sequentially. Already applied revisions are skipped.
-    """
-    _ensure_migrations_table(connection)
-
-    applied_revisions: Set[str] = {
-        row[0]
-        for row in connection.execute(
-            text("SELECT id FROM schema_migrations")
-        )
-    }
-
-    for migration in MIGRATIONS:
-        if migration.id in applied_revisions:
-            continue
-
-        logging.info(
-            "Migrator: applying %s – %s", migration.id, migration.description
-        )
-        try:
-            with connection.begin_nested():
-                migration.upgrade(connection)
-                connection.execute(
-                    text(
-                        "INSERT INTO schema_migrations (id) VALUES (:revision)"
-                    ),
-                    {"revision": migration.id},
-                )
-        except Exception as exc:
-            logging.error(
-                "Migrator: failed to apply %s (%s)",
-                migration.id,
-                migration.description,
-                exc_info=True,
-            )
-            raise exc
-        else:
-            logging.info("Migrator: migration %s applied successfully", migration.id)