env.py 1.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. """Alembic environment.
  2. Runs migrations against the async engine. The database URL comes from app
  3. settings (``DATABASE_URL`` env var with a local default). The asyncpg async URL
  4. is used directly via SQLAlchemy's async engine + ``connection.run_sync``.
  5. """
  6. from __future__ import annotations
  7. import asyncio
  8. from logging.config import fileConfig
  9. from alembic import context
  10. from sqlalchemy.ext.asyncio import async_engine_from_config
  11. from sqlalchemy import pool
  12. from app.config import get_settings
  13. from app.db.models import Base
  14. config = context.config
  15. if config.config_file_name is not None:
  16. fileConfig(config.config_file_name)
  17. # Inject the runtime database URL.
  18. config.set_main_option("sqlalchemy.url", get_settings().database_url)
  19. target_metadata = Base.metadata
  20. def run_migrations_offline() -> None:
  21. """Run migrations in 'offline' mode (emit SQL, no DBAPI)."""
  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. compare_type=True,
  29. )
  30. with context.begin_transaction():
  31. context.run_migrations()
  32. def do_run_migrations(connection) -> None:
  33. context.configure(
  34. connection=connection,
  35. target_metadata=target_metadata,
  36. compare_type=True,
  37. )
  38. with context.begin_transaction():
  39. context.run_migrations()
  40. async def run_migrations_online() -> None:
  41. """Run migrations in 'online' mode using the async engine."""
  42. connectable = async_engine_from_config(
  43. config.get_section(config.config_ini_section, {}),
  44. prefix="sqlalchemy.",
  45. poolclass=pool.NullPool,
  46. )
  47. async with connectable.connect() as connection:
  48. await connection.run_sync(do_run_migrations)
  49. await connectable.dispose()
  50. if context.is_offline_mode():
  51. run_migrations_offline()
  52. else:
  53. asyncio.run(run_migrations_online())