def get_alembic_version_table(config): script_dir = alembic_script.ScriptDirectory.from_config(config) alembic_version_table = [None] def alembic_version_table_from_env(rev, context): alembic_version_table[0] = context.version_table return [] with environment.EnvironmentContext(config, script_dir, fn=alembic_version_table_from_env): script_dir.run_env() return alembic_version_table[0]
def run_sanity_checks(config, revision): script_dir = alembic_script.ScriptDirectory.from_config(config) def check_sanity(rev, context): for script in script_dir.iterate_revisions(revision, rev): if hasattr(script.module, 'check_sanity'): script.module.check_sanity(context.connection) return [] with environment.EnvironmentContext(config, script_dir, fn=check_sanity, starting_rev=None, destination_rev=revision): script_dir.run_env()
def db_sync(connection, revision='head'): script = a_script.ScriptDirectory.from_config(_alembic_config()) def upgrade(rev, context): return script._upgrade_revs(revision, rev) config = _alembic_config() with a_env.EnvironmentContext(config, script, fn=upgrade, as_sql=False, starting_rev=None, destination_rev=revision, tag=None) as context: context.configure(connection=connection) context.run_migrations()
def run_sanity_checks(config, revision): script_dir = alembic_script.ScriptDirectory.from_config(config) def check_sanity(rev, context): # TODO(ihrachyshka): here we use internal API for alembic; we may need # alembic to expose implicit_base= argument into public # iterate_revisions() call for script in script_dir.revision_map.iterate_revisions( revision, rev, implicit_base=True): if hasattr(script.module, 'check_sanity'): script.module.check_sanity(context.connection) return [] with environment.EnvironmentContext(config, script_dir, fn=check_sanity, starting_rev=None, destination_rev=revision): script_dir.run_env()
class SchemaManager: """ Database manager, based on alembic. """ SCRIPT_LOCATION = os.path.join( os.path.dirname(niamoto.__file__), 'migrations', ) CONFIG = Config(os.path.join(SCRIPT_LOCATION, 'alembic.ini')) CONFIG.set_main_option("script_location", SCRIPT_LOCATION) SCRIPT = ScriptDirectory.from_config(CONFIG) ENV = environment.EnvironmentContext( CONFIG, SCRIPT, ) @classmethod def get_current_revision(cls): LOGGER.debug("Getting the current database revision using alembic...") with Connector.get_connection() as connection: context = migration.MigrationContext.configure(connection) rev = context.get_current_revision() LOGGER.debug("Current revision is '{}'".format(rev)) return rev @classmethod def get_head_revision(cls): LOGGER.debug("Getting head revision using alembic...") rev = cls.ENV.get_head_revision() LOGGER.debug("Head revision is '{}'".format(rev)) return rev @classmethod def upgrade_db(cls, revision): LOGGER.debug("Upgrading database to revision '{}'...".format(revision)) wd = cls.SCRIPT_LOCATION result = subprocess.Popen( ["alembic", "upgrade", revision], cwd=wd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) output = result.communicate() LOGGER.debug(output) LOGGER.debug("Database is now at revision '{}'".format(revision)) @classmethod def downgrade_db(cls, revision): LOGGER.debug( "Downgrading database to revision '{}'...".format(revision)) wd = cls.SCRIPT_LOCATION result = subprocess.Popen( ["alembic", "downgrade", revision], cwd=wd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, ) output = result.communicate() LOGGER.debug(output) LOGGER.debug("Database is now at revision '{}'".format(revision))