56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
# - *- coding: utf- 8 - *-
|
|
import asyncio
|
|
from pathlib import Path
|
|
from typing import Optional
|
|
|
|
from alembic import command
|
|
from alembic.config import Config
|
|
from sqlalchemy.engine import Connection
|
|
from sqlalchemy.ext.asyncio import AsyncEngine
|
|
|
|
from tgbot.database.core import database_url
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
ALEMBIC_INI = PROJECT_ROOT / "alembic.ini"
|
|
MIGRATIONS_DIR = PROJECT_ROOT / "migrations"
|
|
|
|
|
|
# Сбор Alembic-конфига от корня проекта
|
|
def get_alembic_config(url: str = database_url) -> Config:
|
|
config = Config(str(ALEMBIC_INI))
|
|
config.set_main_option("script_location", str(MIGRATIONS_DIR))
|
|
config.set_main_option("sqlalchemy.url", url)
|
|
|
|
return config
|
|
|
|
|
|
# Применение миграций до последней версии
|
|
async def run_migrations(engine: Optional[AsyncEngine] = None) -> None:
|
|
if engine is None:
|
|
config = get_alembic_config()
|
|
await asyncio.to_thread(command.upgrade, config, "head")
|
|
return
|
|
|
|
config = get_alembic_config(str(engine.url))
|
|
|
|
connection = engine.connect()
|
|
await connection.start()
|
|
transaction = connection.begin()
|
|
await transaction.start()
|
|
|
|
try:
|
|
await connection.run_sync(_upgrade_with_connection, config)
|
|
except Exception:
|
|
await transaction.rollback()
|
|
raise
|
|
else:
|
|
await transaction.commit()
|
|
finally:
|
|
await connection.close()
|
|
|
|
|
|
# Запуск Alembic на готовом соединении
|
|
def _upgrade_with_connection(connection: Connection, config: Config) -> None:
|
|
config.attributes["connection"] = connection
|
|
command.upgrade(config, "head")
|