env.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. from __future__ import with_statement
  2. from alembic import context
  3. from sqlalchemy import engine_from_config, pool, MetaData
  4. from logging.config import fileConfig
  5. import logging
  6. USE_TWOPHASE = False
  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)
  13. logger = logging.getLogger('alembic.env')
  14. # add your model's MetaData object here
  15. # for 'autogenerate' support
  16. # from myapp import mymodel
  17. # target_metadata = mymodel.Base.metadata
  18. from flask import current_app
  19. config.set_main_option('sqlalchemy.url',
  20. current_app.config.get('SQLALCHEMY_DATABASE_URI'))
  21. bind_names = []
  22. for name, url in current_app.config.get("SQLALCHEMY_BINDS").items():
  23. context.config.set_section_option(name, "sqlalchemy.url", url)
  24. bind_names.append(name)
  25. target_metadata = current_app.extensions['migrate'].db.metadata
  26. # other values from the config, defined by the needs of env.py,
  27. # can be acquired:
  28. # my_important_option = config.get_main_option("my_important_option")
  29. # ... etc.
  30. def get_metadata(bind):
  31. """Return the metadata for a bind."""
  32. if bind == '':
  33. bind = None
  34. m = MetaData()
  35. for t in target_metadata.tables.values():
  36. if t.info.get('bind_key') == bind:
  37. t.tometadata(m)
  38. return m
  39. def run_migrations_offline():
  40. """Run migrations in 'offline' mode.
  41. This configures the context with just a URL
  42. and not an Engine, though an Engine is acceptable
  43. here as well. By skipping the Engine creation
  44. we don't even need a DBAPI to be available.
  45. Calls to context.execute() here emit the given string to the
  46. script output.
  47. """
  48. # for the --sql use case, run migrations for each URL into
  49. # individual files.
  50. engines = {'': {'url': context.config.get_main_option('sqlalchemy.url')}}
  51. for name in bind_names:
  52. engines[name] = rec = {}
  53. rec['url'] = context.config.get_section_option(name,
  54. "sqlalchemy.url")
  55. for name, rec in engines.items():
  56. logger.info("Migrating database %s" % (name or '<default>'))
  57. file_ = "%s.sql" % name
  58. logger.info("Writing output to %s" % file_)
  59. with open(file_, 'w') as buffer:
  60. context.configure(url=rec['url'], output_buffer=buffer,
  61. target_metadata=get_metadata(name),
  62. literal_binds=True)
  63. with context.begin_transaction():
  64. context.run_migrations(engine_name=name)
  65. def run_migrations_online():
  66. """Run migrations in 'online' mode.
  67. In this scenario we need to create an Engine
  68. and associate a connection with the context.
  69. """
  70. # this callback is used to prevent an auto-migration from being generated
  71. # when there are no changes to the schema
  72. # reference: http://alembic.readthedocs.org/en/latest/cookbook.html
  73. def process_revision_directives(context, revision, directives):
  74. if getattr(config.cmd_opts, 'autogenerate', False):
  75. script = directives[0]
  76. if len(script.upgrade_ops_list) >= len(bind_names) + 1:
  77. empty = True
  78. for upgrade_ops in script.upgrade_ops_list:
  79. if not upgrade_ops.is_empty():
  80. empty = False
  81. if empty:
  82. directives[:] = []
  83. logger.info('No changes in schema detected.')
  84. # for the direct-to-DB use case, start a transaction on all
  85. # engines, then run all migrations, then commit all transactions.
  86. engines = {'': {'engine': engine_from_config(
  87. config.get_section(config.config_ini_section),
  88. prefix='sqlalchemy.',
  89. poolclass=pool.NullPool)}}
  90. for name in bind_names:
  91. engines[name] = rec = {}
  92. rec['engine'] = engine_from_config(
  93. context.config.get_section(name),
  94. prefix='sqlalchemy.',
  95. poolclass=pool.NullPool)
  96. for name, rec in engines.items():
  97. engine = rec['engine']
  98. rec['connection'] = conn = engine.connect()
  99. if USE_TWOPHASE:
  100. rec['transaction'] = conn.begin_twophase()
  101. else:
  102. rec['transaction'] = conn.begin()
  103. try:
  104. for name, rec in engines.items():
  105. logger.info("Migrating database %s" % (name or '<default>'))
  106. context.configure(
  107. connection=rec['connection'],
  108. upgrade_token="%s_upgrades" % name,
  109. downgrade_token="%s_downgrades" % name,
  110. target_metadata=get_metadata(name),
  111. process_revision_directives=process_revision_directives,
  112. **current_app.extensions['migrate'].configure_args
  113. )
  114. context.run_migrations(engine_name=name)
  115. if USE_TWOPHASE:
  116. for rec in engines.values():
  117. rec['transaction'].prepare()
  118. for rec in engines.values():
  119. rec['transaction'].commit()
  120. except:
  121. for rec in engines.values():
  122. rec['transaction'].rollback()
  123. raise
  124. finally:
  125. for rec in engines.values():
  126. rec['connection'].close()
  127. if context.is_offline_mode():
  128. run_migrations_offline()
  129. else:
  130. run_migrations_online()