zhifu-platform/migrations/env.py

import asyncio

from alembic import context

from zhifu.core.config import Settings
from zhifu.core.database import Base, create_engine
from zhifu.modules.after_sales import models  # noqa: F401
from zhifu.modules.agent_runtime import models as agent_models  # noqa: F401
from zhifu.modules.billing import models as billing_models  # noqa: F401
from zhifu.modules.channels import models as channel_models  # noqa: F401
from zhifu.modules.conversations import models as conversation_models  # noqa: F401
from zhifu.modules.identity import access_models  # noqa: F401
from zhifu.modules.identity.models import Store, Tenant  # noqa: F401
from zhifu.modules.integrations import models as integration_models  # noqa: F401
from zhifu.modules.knowledge import models as knowledge_models  # noqa: F401
from zhifu.modules.memory import models as memory_models  # noqa: F401
from zhifu.modules.merchant import models as merchant_models  # noqa: F401
from zhifu.modules.notifications import models as notification_models  # noqa: F401
from zhifu.modules.platform import models as platform_models  # noqa: F401
from zhifu.modules.storefront import models as storefront_models  # noqa: F401

target_metadata = Base.metadata


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


async def run_online():
    engine = create_engine(Settings())
    try:
        async with engine.connect() as connection:
            await connection.run_sync(run_migrations)
    finally:
        await engine.dispose()


if context.is_offline_mode():
    context.configure(
        url=Settings().database_url.get_secret_value(),
        target_metadata=target_metadata,
        literal_binds=True,
        dialect_opts={"paramstyle": "named"},
    )
    with context.begin_transaction():
        context.run_migrations()
else:
    asyncio.run(run_online())