from __future__ import annotations import asyncio import os from logging.config import fileConfig from alembic import context from app.db import Base, create_postgres_engine config = context.config if config.config_file_name: fileConfig(config.config_file_name) database_url = os.environ["SMS_DATABASE_URL"] if database_url.startswith("postgresql://"): database_url = database_url.replace("postgresql://", "postgresql+asyncpg://", 1) config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%")) target_metadata = Base.metadata def run_migrations_offline() -> None: context.configure( url=config.get_main_option("sqlalchemy.url"), target_metadata=target_metadata, literal_binds=True, dialect_opts={"paramstyle": "named"}, version_table_schema="sms", include_schemas=True, ) with context.begin_transaction(): context.run_migrations() def do_run_migrations(connection) -> None: context.configure( connection=connection, target_metadata=target_metadata, version_table_schema="sms", include_schemas=True, compare_type=True, ) with context.begin_transaction(): context.run_migrations() async def run_async_migrations() -> None: connectable = create_postgres_engine(database_url) async with connectable.connect() as connection: await connection.run_sync(do_run_migrations) await connectable.dispose() if context.is_offline_mode(): run_migrations_offline() else: asyncio.run(run_async_migrations())