env.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061
  1. """Alembic运行环境。"""
  2. from __future__ import annotations
  3. from logging.config import fileConfig
  4. from alembic import context
  5. from sqlalchemy import engine_from_config, pool
  6. from dms.config import load_dms_config
  7. from dms.extensions import db
  8. import dms.models # noqa: F401
  9. config = context.config
  10. if config.config_file_name is not None:
  11. fileConfig(config.config_file_name)
  12. config.set_main_option(
  13. "sqlalchemy.url",
  14. load_dms_config()["SQLALCHEMY_DATABASE_URI"].replace("%", "%%"),
  15. )
  16. target_metadata = db.metadata
  17. def run_migrations_offline() -> None:
  18. context.configure(
  19. url=config.get_main_option("sqlalchemy.url"),
  20. target_metadata=target_metadata,
  21. literal_binds=True,
  22. dialect_opts={"paramstyle": "named"},
  23. compare_type=True,
  24. compare_server_default=True,
  25. )
  26. with context.begin_transaction():
  27. context.run_migrations()
  28. def run_migrations_online() -> None:
  29. connectable = engine_from_config(
  30. config.get_section(config.config_ini_section, {}),
  31. prefix="sqlalchemy.",
  32. poolclass=pool.NullPool,
  33. isolation_level="READ COMMITTED",
  34. connect_args={"init_command": "SET time_zone = '+00:00'"},
  35. )
  36. with connectable.connect() as connection:
  37. context.configure(
  38. connection=connection,
  39. target_metadata=target_metadata,
  40. compare_type=True,
  41. compare_server_default=True,
  42. )
  43. with context.begin_transaction():
  44. context.run_migrations()
  45. if context.is_offline_mode():
  46. run_migrations_offline()
  47. else:
  48. run_migrations_online()