env.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133
  1. from __future__ import with_statement
  2. from alembic import context
  3. from sqlalchemy import engine_from_config, pool
  4. from logging.config import fileConfig
  5. import logging
  6. import re
  7. USE_TWOPHASE = False
  8. # this is the Alembic Config object, which provides
  9. # access to the values within the .ini file in use.
  10. config = context.config
  11. # Interpret the config file for Python logging.
  12. # This line sets up loggers basically.
  13. fileConfig(config.config_file_name)
  14. logger = logging.getLogger('alembic.env')
  15. # gather section names referring to different
  16. # databases. These are named "engine1", "engine2"
  17. # in the sample .ini file.
  18. db_names = config.get_main_option('databases')
  19. # add your model's MetaData objects here
  20. # for 'autogenerate' support. These must be set
  21. # up to hold just those tables targeting a
  22. # particular database. table.tometadata() may be
  23. # helpful here in case a "copy" of
  24. # a MetaData is needed.
  25. # from myapp import mymodel
  26. # target_metadata = {
  27. # 'engine1':mymodel.metadata1,
  28. # 'engine2':mymodel.metadata2
  29. #}
  30. target_metadata = {}
  31. # other values from the config, defined by the needs of env.py,
  32. # can be acquired:
  33. # my_important_option = config.get_main_option("my_important_option")
  34. # ... etc.
  35. def run_migrations_offline():
  36. """Run migrations in 'offline' mode.
  37. This configures the context with just a URL
  38. and not an Engine, though an Engine is acceptable
  39. here as well. By skipping the Engine creation
  40. we don't even need a DBAPI to be available.
  41. Calls to context.execute() here emit the given string to the
  42. script output.
  43. """
  44. # for the --sql use case, run migrations for each URL into
  45. # individual files.
  46. engines = {}
  47. for name in re.split(r',\s*', db_names):
  48. engines[name] = rec = {}
  49. rec['url'] = context.config.get_section_option(name,
  50. "sqlalchemy.url")
  51. for name, rec in engines.items():
  52. logger.info("Migrating database %s" % name)
  53. file_ = "%s.sql" % name
  54. logger.info("Writing output to %s" % file_)
  55. with open(file_, 'w') as buffer:
  56. context.configure(url=rec['url'], output_buffer=buffer,
  57. target_metadata=target_metadata.get(name),
  58. literal_binds=True)
  59. with context.begin_transaction():
  60. context.run_migrations(engine_name=name)
  61. def run_migrations_online():
  62. """Run migrations in 'online' mode.
  63. In this scenario we need to create an Engine
  64. and associate a connection with the context.
  65. """
  66. # for the direct-to-DB use case, start a transaction on all
  67. # engines, then run all migrations, then commit all transactions.
  68. engines = {}
  69. for name in re.split(r',\s*', db_names):
  70. engines[name] = rec = {}
  71. rec['engine'] = engine_from_config(
  72. context.config.get_section(name),
  73. prefix='sqlalchemy.',
  74. poolclass=pool.NullPool)
  75. for name, rec in engines.items():
  76. engine = rec['engine']
  77. rec['connection'] = conn = engine.connect()
  78. if USE_TWOPHASE:
  79. rec['transaction'] = conn.begin_twophase()
  80. else:
  81. rec['transaction'] = conn.begin()
  82. try:
  83. for name, rec in engines.items():
  84. logger.info("Migrating database %s" % name)
  85. context.configure(
  86. connection=rec['connection'],
  87. upgrade_token="%s_upgrades" % name,
  88. downgrade_token="%s_downgrades" % name,
  89. target_metadata=target_metadata.get(name)
  90. )
  91. context.run_migrations(engine_name=name)
  92. if USE_TWOPHASE:
  93. for rec in engines.values():
  94. rec['transaction'].prepare()
  95. for rec in engines.values():
  96. rec['transaction'].commit()
  97. except:
  98. for rec in engines.values():
  99. rec['transaction'].rollback()
  100. raise
  101. finally:
  102. for rec in engines.values():
  103. rec['connection'].close()
  104. if context.is_offline_mode():
  105. run_migrations_offline()
  106. else:
  107. run_migrations_online()