env.py 1.8 KB

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