SupplyAgent/scripts/mysql_schema_migration_smoke.py

from __future__ import annotations

import argparse
import os
import re

from alembic import command
from alembic.config import Config
from sqlalchemy import create_engine, inspect, text
from sqlalchemy.engine import make_url

from supply_infra.config import InfraSettings
from supply_infra.pipeline.preflight import REQUIRED_PIPELINE_TABLES

_SAFE_DATABASE = re.compile(r"^supply_agent_preflight_[a-z0-9_]+$")


def _quoted(identifier: str) -> str:
    if not re.fullmatch(r"[A-Za-z0-9_]+", identifier):
        raise ValueError(f"unsafe SQL identifier: {identifier!r}")
    return f"`{identifier}`"


def run_smoke(
    target_database: str,
    *,
    downgrade_to: str | None = None,
) -> dict[str, object]:
    settings = InfraSettings()
    source_database = settings.mysql_database
    if not _SAFE_DATABASE.fullmatch(target_database):
        raise ValueError(
            "temporary database name must start with "
            "supply_agent_preflight_ and contain only lowercase letters, "
            "digits and underscores"
        )
    if target_database == source_database:
        raise ValueError("temporary database must differ from source database")

    source_url = make_url(settings.mysql_url)
    server_url = source_url.set(database=None)
    target_url = source_url.set(database=target_database)
    server_engine = create_engine(server_url)
    target_engine = None
    created = False
    try:
        with server_engine.begin() as connection:
            exists = connection.scalar(
                text(
                    "SELECT COUNT(*) FROM INFORMATION_SCHEMA.SCHEMATA "
                    "WHERE SCHEMA_NAME = :database"
                ),
                {"database": target_database},
            )
            if int(exists or 0) > 0:
                raise RuntimeError(
                    f"temporary database already exists: {target_database}"
                )
            connection.execute(
                text(
                    f"CREATE DATABASE {_quoted(target_database)} "
                    "CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci"
                )
            )
            created = True

        source_engine = create_engine(source_url)
        try:
            source_tables = inspect(source_engine).get_table_names()
            with source_engine.connect() as connection:
                source_revision = connection.scalar(
                    text("SELECT version_num FROM alembic_version")
                )
        finally:
            source_engine.dispose()
        with server_engine.begin() as connection:
            for table_name in source_tables:
                connection.execute(
                    text(
                        f"CREATE TABLE {_quoted(target_database)}."
                        f"{_quoted(table_name)} LIKE "
                        f"{_quoted(source_database)}.{_quoted(table_name)}"
                    )
                )
            connection.execute(
                text(
                    f"INSERT INTO {_quoted(target_database)}.`alembic_version` "
                    f"SELECT * FROM {_quoted(source_database)}.`alembic_version`"
                )
            )

        alembic_config = Config("alembic.ini")
        target_url_text = target_url.render_as_string(hide_password=False)
        previous_alembic_url = os.environ.get("ALEMBIC_DATABASE_URL")
        os.environ["ALEMBIC_DATABASE_URL"] = target_url_text
        try:
            if downgrade_to is not None:
                command.downgrade(alembic_config, downgrade_to)
            command.upgrade(alembic_config, "head")
        finally:
            if previous_alembic_url is None:
                os.environ.pop("ALEMBIC_DATABASE_URL", None)
            else:
                os.environ["ALEMBIC_DATABASE_URL"] = previous_alembic_url

        target_engine = create_engine(target_url)
        with target_engine.connect() as connection:
            target_tables = set(inspect(connection).get_table_names())
            revision = connection.scalar(
                text("SELECT version_num FROM alembic_version")
            )
        missing = sorted(REQUIRED_PIPELINE_TABLES - target_tables)
        if missing:
            raise RuntimeError(
                f"migrated schema is missing required tables: {missing}"
            )
        source_engine = create_engine(source_url)
        try:
            current_source_tables = inspect(source_engine).get_table_names()
            with source_engine.connect() as connection:
                current_source_revision = connection.scalar(
                    text("SELECT version_num FROM alembic_version")
                )
        finally:
            source_engine.dispose()
        if (
            current_source_tables != source_tables
            or current_source_revision != source_revision
        ):
            raise RuntimeError(
                "source schema changed during isolated migration smoke test"
            )
        return {
            "success": True,
            "source_database": source_database,
            "temporary_database": target_database,
            "downgrade_to": downgrade_to,
            "source_table_count": len(source_tables),
            "target_table_count": len(target_tables),
            "revision": revision,
            "missing_required_tables": missing,
        }
    finally:
        if target_engine is not None:
            target_engine.dispose()
        if created:
            with server_engine.begin() as connection:
                connection.execute(
                    text(f"DROP DATABASE {_quoted(target_database)}")
                )
        server_engine.dispose()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Clone the current MySQL schema and test Alembic upgrade"
    )
    parser.add_argument("--database", required=True)
    parser.add_argument(
        "--downgrade-to",
        help="First downgrade the isolated clone to this revision",
    )
    args = parser.parse_args()
    result = run_smoke(args.database, downgrade_to=args.downgrade_to)
    for key, value in result.items():
        print(f"{key}={value}")


if __name__ == "__main__":
    main()