env.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849
  1. import os
  2. from logging.config import fileConfig
  3. from alembic import context
  4. from app.db.models import Base
  5. from sqlalchemy import engine_from_config, pool
  6. SERVICE_VERSION_TABLE = "scheduler_alembic_version"
  7. config = context.config
  8. database_url = os.getenv("AGENT_PLATFORM_DATABASE_URL")
  9. if database_url:
  10. config.set_main_option("sqlalchemy.url", database_url.replace("%", "%%"))
  11. if config.config_file_name is not None:
  12. fileConfig(config.config_file_name)
  13. target_metadata = Base.metadata
  14. def run_migrations_offline() -> None:
  15. url = config.get_main_option("sqlalchemy.url")
  16. context.configure(
  17. url=url,
  18. target_metadata=target_metadata,
  19. literal_binds=True,
  20. version_table=SERVICE_VERSION_TABLE)
  21. with context.begin_transaction():
  22. context.run_migrations()
  23. def run_migrations_online() -> None:
  24. connectable = engine_from_config(
  25. config.get_section(config.config_ini_section, {}),
  26. prefix="sqlalchemy.",
  27. poolclass=pool.NullPool)
  28. with connectable.connect() as connection:
  29. context.configure(
  30. connection=connection,
  31. target_metadata=target_metadata,
  32. version_table=SERVICE_VERSION_TABLE)
  33. with context.begin_transaction():
  34. context.run_migrations()
  35. if context.is_offline_mode():
  36. run_migrations_offline()
  37. else:
  38. run_migrations_online()