import asyncio

from alembic import context

from zhifu.oms import models  # noqa: F401
from zhifu.oms.database import OmsBase, oms_engine


def migrate(connection):
    context.configure(connection=connection, target_metadata=OmsBase.metadata, compare_type=True)
    with context.begin_transaction():
        context.run_migrations()


async def main():
    engine = oms_engine()
    try:
        async with engine.connect() as connection:
            await connection.run_sync(migrate)
    finally:
        await engine.dispose()


asyncio.run(main())