env.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. import asyncio
  2. from logging.config import fileConfig
  3. from sqlalchemy import engine_from_config, pool
  4. from sqlalchemy.ext.asyncio import AsyncEngine
  5. from alembic import context
  6. from app.core import config as app_config
  7. # this is the Alembic Config object, which provides
  8. # access to the values within the .ini file in use.
  9. config = context.config
  10. # Interpret the config file for Python logging.
  11. # This line sets up loggers basically.
  12. fileConfig(config.config_file_name) # type: ignore
  13. # add your model's MetaData object here
  14. # for 'autogenerate' support
  15. # from myapp import mymodel
  16. # target_metadata = mymodel.Base.metadata
  17. from app.models import Base # noqa
  18. target_metadata = Base.metadata
  19. # other values from the config, defined by the needs of env.py,
  20. # can be acquired:
  21. # my_important_option = config.get_main_option("my_important_option")
  22. # ... etc.
  23. def get_database_uri():
  24. return app_config.settings.DEFAULT_SQLALCHEMY_DATABASE_URI
  25. def run_migrations_offline():
  26. """Run migrations in 'offline' mode.
  27. This configures the context with just a URL
  28. and not an Engine, though an Engine is acceptable
  29. here as well. By skipping the Engine creation
  30. we don't even need a DBAPI to be available.
  31. Calls to context.execute() here emit the given string to the
  32. script output.
  33. """
  34. url = get_database_uri()
  35. context.configure(
  36. url=url,
  37. target_metadata=target_metadata,
  38. literal_binds=True,
  39. dialect_opts={"paramstyle": "named"},
  40. compare_type=True,
  41. compare_server_default=True,
  42. )
  43. with context.begin_transaction():
  44. context.run_migrations()
  45. def do_run_migrations(connection):
  46. context.configure(
  47. connection=connection, target_metadata=target_metadata, compare_type=True
  48. )
  49. with context.begin_transaction():
  50. context.run_migrations()
  51. async def run_migrations_online():
  52. """Run migrations in 'online' mode.
  53. In this scenario we need to create an Engine
  54. and associate a connection with the context.
  55. """
  56. configuration = config.get_section(config.config_ini_section)
  57. assert configuration
  58. configuration["sqlalchemy.url"] = get_database_uri()
  59. connectable = AsyncEngine(
  60. engine_from_config(
  61. configuration,
  62. prefix="sqlalchemy.",
  63. poolclass=pool.NullPool,
  64. future=True,
  65. ) # type: ignore
  66. )
  67. async with connectable.connect() as connection:
  68. await connection.run_sync(do_run_migrations)
  69. if context.is_offline_mode():
  70. run_migrations_offline()
  71. else:
  72. asyncio.run(run_migrations_online())