AutoShop-Djimbo/migrations/env.py

# - *- coding: utf- 8 - *-
import asyncio
from logging.config import fileConfig
from typing import Any, List, Optional

from alembic import context
from alembic.operations import ops
from sqlalchemy import pool, text
from sqlalchemy.engine import Connection
from sqlalchemy.schema import DefaultClause
from sqlalchemy.ext.asyncio import async_engine_from_config

from tgbot.database.core import Base

import tgbot.database  # noqa: F401

config = context.config

if config.config_file_name is not None:
    fileConfig(config.config_file_name)

target_metadata = Base.metadata
INDEX_DIFF_NOISE = {"ix_storage_payments_id", "ix_storage_settings_id"}


# Отбрасывание SQLite-шума из автогенерации
def _filter_autogenerated_ops(container: Any) -> None:
    filtered_ops = []

    for operation in container.ops:
        if isinstance(operation, ops.AlterColumnOp):
            continue

        if _is_index_noise(operation):
            continue

        if isinstance(operation, ops.ModifyTableOps):
            _filter_autogenerated_ops(operation)

            if operation.ops:
                filtered_ops.append(operation)

            continue

        filtered_ops.append(operation)

    container.ops = filtered_ops


# Подготовка новых колонок к безопасному SQLite ADD COLUMN
def _prepare_sqlite_add_columns(container: Any) -> None:
    for operation in container.ops:
        if isinstance(operation, ops.AddColumnOp):
            _apply_server_default(operation.column)
            continue

        if isinstance(operation, ops.ModifyTableOps):
            _prepare_sqlite_add_columns(operation)


# Перенос default модели в server_default миграции
def _apply_server_default(column: Any) -> None:
    if column.nullable or column.server_default is not None or column.default is None:
        return

    default_value = _get_simple_default(column.default.arg)

    if default_value is None:
        return

    column.server_default = DefaultClause(text(default_value))


# Получение SQL-литерала для простых default-значений
def _get_simple_default(value: Any) -> Optional[str]:
    if callable(value):
        return None

    if isinstance(value, bool):
        return "1" if value else "0"

    if isinstance(value, (int, float)):
        return str(value)

    if isinstance(value, str):
        return repr(value)

    return None


# Игнорирование индексов, которые SQLite отражает иначе
def _is_index_noise(operation: Any) -> bool:
    if not isinstance(operation, (ops.CreateIndexOp, ops.DropIndexOp)):
        return False

    return getattr(operation, "index_name", "") in INDEX_DIFF_NOISE


# Ограничение сравнения SQLite только реальными добавлениями/удалениями
def _include_object(object_: Any, name: str, type_: str, reflected: bool, compare_to: Any) -> bool:
    if type_ == "column" and compare_to is not None:
        return False

    if type_ == "index" and name in INDEX_DIFF_NOISE:
        return False

    return True


# Очистка автосозданной миграции перед записью
def _process_revision_directives(_context: Any, _revision: Any, directives: List[Any]) -> None:
    if not directives:
        return

    script = directives[0]
    _prepare_sqlite_add_columns(script.upgrade_ops)
    _filter_autogenerated_ops(script.upgrade_ops)
    _filter_autogenerated_ops(script.downgrade_ops)

    if not script.upgrade_ops.ops and not script.downgrade_ops.ops:
        config.attributes["autogenerate_empty"] = True
        directives.clear()


# Запуск миграций без подключения
def run_migrations_offline() -> None:
    url = config.get_main_option("sqlalchemy.url")

    context.configure(
        url=url,
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
        render_as_batch=True,
        compare_type=False,
        include_object=_include_object,
        process_revision_directives=_process_revision_directives,
    )

    with context.begin_transaction():
        context.run_migrations()


# Запуск миграций на соединении
def do_run_migrations(connection: Connection) -> None:
    context.configure(
        connection=connection,
        target_metadata=target_metadata,
        render_as_batch=True,
        compare_type=False,
        include_object=_include_object,
        process_revision_directives=_process_revision_directives,
    )

    with context.begin_transaction():
        context.run_migrations()


# Запуск миграций через async-engine
async def run_async_migrations() -> None:
    connectable = async_engine_from_config(
        config.get_section(config.config_ini_section, {}),
        poolclass=pool.NullPool,
    )
    connection = connectable.connect()

    try:
        await connection.start()
        await connection.run_sync(do_run_migrations)
    finally:
        await connection.close()

    await connectable.dispose()


if context.is_offline_mode():
    run_migrations_offline()
else:
    external_connection = config.attributes.get("connection")

    if external_connection is not None:
        do_run_migrations(external_connection)
    else:
        asyncio.run(run_async_migrations())