contrib/ast-db-manage/env.py

   1 from __future__ import with_statement
   2 from alembic import context
   3 from alembic.script import ScriptDirectory
   4 from alembic.operations import Operations
   5 from sqlalchemy import engine_from_config, pool, MetaData
   6 from sqlalchemy.ext.declarative import declarative_base
   7 from logging.config import fileConfig
   8 import logging
   9
  10 # this is the Alembic Config object, which provides
  11 # access to the values within the .ini file in use.
  12 config = context.config
  13
  14 # Interpret the config file for Python logging.
  15 # This line sets up loggers basically.
  16 try:
  17     fileConfig(config.config_file_name)
  18 except:
  19     pass
  20
  21 ## below block is needed for mssql
  22 meta = MetaData(naming_convention = {
  23         "ix": 'ix_%(column_0_label)s',
  24         "uq": "uq_%(table_name)s_%(column_0_name)s",
  25         "ck": "ck_%(table_name)s_%(column_0_name)s_%(constraint_name)s",
  26         "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
  27         "pk": "pk_%(table_name)s"
  28 })
  29 Base = declarative_base(metadata=meta)
  30
  31 logger = logging.getLogger('alembic.runtime.setup')
  32 # add your model's MetaData object here
  33 # for 'autogenerate' support
  34 # from myapp import mymodel
  35 # target_metadata = mymodel.Base.metadata
  36 target_metadata = None
  37 #Comment above line and uncomment below line for mssql
  38 #target_metadata = Base.metadata
  39
  40 # other values from the config, defined by the needs of env.py,
  41 # can be acquired:
  42 # my_important_option = config.get_main_option("my_important_option")
  43 # ... etc.
  44
  45 def run_migrations_offline():
  46     """Run migrations in 'offline' mode.
  47
  48     This configures the context with just a URL
  49     and not an Engine, though an Engine is acceptable
  50     here as well.  By skipping the Engine creation
  51     we don't even need a DBAPI to be available.
  52
  53     Calls to context.execute() here emit the given string to the
  54     script output.
  55
  56     """
  57     url = config.get_main_option("sqlalchemy.url")
  58     context.configure(url=url,target_metadata=target_metadata)
  59
  60     with context.begin_transaction():
  61         context.run_migrations()
  62
  63 def run_migrations_online():
  64     """Run migrations in 'online' mode.
  65
  66     In this scenario we need to create an Engine
  67     and associate a connection with the context.
  68
  69     """
  70     engine = engine_from_config(
  71                 config.get_section(config.config_ini_section),
  72                 prefix='sqlalchemy.',
  73                 poolclass=pool.NullPool)
  74
  75     logger.info('Testing for an old alembic_version table.')
  76
  77     connection = engine.connect()
  78     context.configure(
  79                 connection=connection,
  80                 target_metadata=target_metadata,
  81                 version_table='alembic_version'
  82                 )
  83
  84     script_location = config.get_main_option('script_location')
  85     found = False
  86     mc = context.get_context()
  87     current_db_revision = mc.get_current_revision()
  88     script = ScriptDirectory.from_config(config)
  89     """ If there was an existing alembic_version table, we need to
  90     check that it's current revision is in the history for the tree
  91     we're working with.
  92     """
  93     for x in script.iterate_revisions('head', 'base'):
  94         if x.revision == current_db_revision:
  95             """ An alembic_versions table was found and it belongs to
  96             this alembic tree
  97             """
  98             logger.info(
  99                 ('An old alembic_version table at revision %s was '
 100                  'found for %s.  Renaming to alembic_version_%s.'),
 101                         current_db_revision, script_location,
 102                         script_location)
 103             op = Operations(mc)
 104             try:
 105                 with context.begin_transaction():
 106                     op.rename_table(
 107                         'alembic_version', 'alembic_version_%s'
 108                         % script_location)
 109                 found = True
 110             except:
 111                 logger.error(('Unable to rename alembic_version to '
 112                              'alembic_version_%s.'),
 113                              script_location)
 114                 connection.close()
 115                 return
 116
 117             break
 118
 119     if not found:
 120         logger.info('Didn\'t find an old alembic_version table.')
 121     logger.info('Trying alembic_version_%s.' % script_location)
 122
 123     """ We MAY have an alembic_version table that doesn't belong to
 124     this tree but if we still don't have an alembic_version_<tree>
 125     table, alembic will create it.
 126     """
 127     context.configure(
 128                 connection=connection,
 129                 target_metadata=target_metadata,
 130                 version_table='alembic_version_' + script_location
 131                 )
 132     mc = context.get_context()
 133     current_db_revision = mc.get_current_revision()
 134     if current_db_revision:
 135         logger.info(
 136             'Using the alembic_version_%s table at revision %s.',
 137             script_location, current_db_revision)
 138     else:
 139         logger.info('Creating new alembic_version_%s table.',
 140                     script_location)
 141
 142     try:
 143         with context.begin_transaction():
 144             context.run_migrations()
 145     finally:
 146         connection.close()
 147
 148
 149 if context.is_offline_mode():
 150     run_migrations_offline()
 151 else:
 152     run_migrations_online()