env.py 1.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263
  1. import asyncio
  2. from logging.config import fileConfig
  3. from alembic import context
  4. from sqlalchemy import pool
  5. from sqlalchemy.ext.asyncio import async_engine_from_config
  6. from app.config import settings
  7. from app.core.database import Base
  8. # Alembic Config 对象
  9. config = context.config
  10. # 从 settings 注入数据库 URL(覆盖 alembic.ini 中的占位符)
  11. config.set_main_option("sqlalchemy.url", settings.database_url)
  12. if config.config_file_name is not None:
  13. fileConfig(config.config_file_name)
  14. # 导入所有 model,让 Base.metadata 能感知到所有表
  15. import app.models.document # noqa: F401
  16. target_metadata = Base.metadata
  17. def run_migrations_offline() -> None:
  18. url = config.get_main_option("sqlalchemy.url")
  19. context.configure(
  20. url=url,
  21. target_metadata=target_metadata,
  22. literal_binds=True,
  23. dialect_opts={"paramstyle": "named"},
  24. )
  25. with context.begin_transaction():
  26. context.run_migrations()
  27. def do_run_migrations(connection) -> None:
  28. context.configure(connection=connection, target_metadata=target_metadata)
  29. with context.begin_transaction():
  30. context.run_migrations()
  31. async def run_async_migrations() -> None:
  32. connectable = async_engine_from_config(
  33. config.get_section(config.config_ini_section, {}),
  34. prefix="sqlalchemy.",
  35. poolclass=pool.NullPool,
  36. )
  37. async with connectable.connect() as connection:
  38. await connection.run_sync(do_run_migrations)
  39. await connectable.dispose()
  40. def run_migrations_online() -> None:
  41. asyncio.run(run_async_migrations())
  42. if context.is_offline_mode():
  43. run_migrations_offline()
  44. else:
  45. run_migrations_online()