def check_version(ev):
""" ApplicationCreated event handler """
if not Version.__table__.exists():
return
versions = dict((v.package, v.version_num)
for v in ptah.get_session().query(Version).all())
packages = ptah.get_cfg_storage(MIGRATION_ID).keys()
has_steps = False
log = logging.getLogger('ptah.alembic')
for pkg in packages:
version = versions.get(pkg)
script = ScriptDirectory(pkg)
for sc in script.walk_revisions():
if sc.is_head:
if sc.revision != version:
has_steps = True
log.error("Package '%s' current revision: '%s', head: '%s'",
pkg, version, sc.revision)
break
if has_steps:
config.shutdown()
log.error("Please run `ptah-migrate` script. Stopping...")
raise SystemExit(1)
def test_needs_flag(self):
a = util.rev_id()
script = ScriptDirectory.from_config(self.cfg)
script.generate_revision(a, None, refresh=True)
write_script(script, a, """
revision = '%s'
down_revision = None
from alembic import op
def upgrade():
op.execute("CREATE TABLE foo(id integer)")
def downgrade():
op.execute("DROP TABLE foo")
""" % a, sourceless=True)
script = ScriptDirectory.from_config(self.cfg)
eq_(script.get_heads(), [])
self.cfg.set_main_option("sourceless", "true")
script = ScriptDirectory.from_config(self.cfg)
eq_(script.get_heads(), [a])
def test_option(self):
self.cfg.set_main_option("file_template", "myfile_%%(slug)s")
script = ScriptDirectory.from_config(self.cfg)
a = util.rev_id()
script.generate_revision(a, "some message", refresh=True)
write_script(script, a, """
revision = '%s'
down_revision = None
from alembic import op
def upgrade():
op.execute("CREATE TABLE foo(id integer)")
def downgrade():
op.execute("DROP TABLE foo")
""" % a)
script = ScriptDirectory.from_config(self.cfg)
rev = script.get_revision(a)
eq_(rev.revision, a)
eq_(os.path.basename(rev.path), "myfile_some_message.py")
def test_lookup_legacy(self):
self.cfg.set_main_option("file_template", "%%(rev)s")
script = ScriptDirectory.from_config(self.cfg)
a = util.rev_id()
script.generate_revision(a, None, refresh=True)
write_script(
script,
a,
"""
down_revision = None
from alembic import op
def upgrade():
op.execute("CREATE TABLE foo(id integer)")
def downgrade():
op.execute("DROP TABLE foo")
""",
)
script = ScriptDirectory.from_config(self.cfg)
rev = script._get_rev(a)
eq_(rev.revision, a)
eq_(os.path.basename(rev.path), "%s.py" % a)
def test_args(self):
script = ScriptDirectory(
staging_directory,
file_template="%(rev)s_%(slug)s_" "%(year)s_%(month)s_" "%(day)s_%(hour)s_" "%(minute)s_%(second)s",
)
create_date = datetime.datetime(2012, 7, 25, 15, 8, 5)
eq_(
script._rev_path("12345", "this is a message", create_date),
"%s/versions/12345_this_is_a_" "message_2012_7_25_15_8_5.py" % staging_directory,
)
def revision(pkg, rev=None, message=None):
"""Create a new revision file."""
script = ScriptDirectory(pkg)
revs = [sc.revision for sc in script.walk_revisions()]
if not rev:
rev = alembic.util.rev_id()
if rev in revs:
raise KeyError('Revision already exists')
return script.generate_revision(rev, message, True).revision
def __new__(cls):
if cls.__instance is None:
i = object.__new__(cls)
i.SQLEngine = SQLEngine
i.DataBase = DataBase
i.Session = Session
i.connection = SQLEngine.connect()
i.context = MigrationContext.configure(i.connection)
i.current_revision = i.context.get_current_revision()
logger.boot('Database revision: %s', i.current_revision)
i.config = Config(ALEMBIC_CONFIG)
i.script = ScriptDirectory.from_config(i.config)
i.head_revision = i.script.get_current_head()
if i.current_revision is None or i.current_revision != i.head_revision:
logger.boot('Upgrading database to version %s.', i.head_revision)
command.upgrade(i.config, 'head')
from option import Option
from log import Log
session = Session()
options = session.query(Option).first()
if options is None:
options = Option()
session.add(options)
options.version = i.head_revision
session.commit()
i.current_revision = i.head_revision
cls.__instance = i
h = SQLAlchemyHandler()
logger.addHandler(h)
return cls.__instance
def test_args_propagate(self):
config = _no_sql_testing_config()
script = ScriptDirectory.from_config(config)
template_args = {"x": "x1", "y": "y1", "z": "z1"}
env = EnvironmentContext(config, script, template_args=template_args)
env.configure(dialect_name="sqlite", template_args={"y": "y2", "q": "q1"})
eq_(template_args, {"x": "x1", "y": "y2", "z": "z1", "q": "q1"})
def check_db():
""" Checks the database revision against the known alembic migrations. """
from inbox.ignition import db_uri
inbox_db_engine = sqlalchemy.create_engine(db_uri())
# top-level, with setup.sh
alembic_ini_filename = _absolute_path('../../alembic.ini')
assert os.path.isfile(alembic_ini_filename), \
'Must have alembic.ini file at {}'.format(alembic_ini_filename)
alembic_cfg = alembic_config(alembic_ini_filename)
try:
inbox_db_engine.dialect.has_table(inbox_db_engine, 'alembic_version')
except sqlalchemy.exc.OperationalError:
sys.exit("Databases don't exist! Run bin/create-db")
if inbox_db_engine.dialect.has_table(inbox_db_engine, 'alembic_version'):
res = inbox_db_engine.execute('SELECT version_num from alembic_version')
current_revision = [r for r in res][0][0]
assert current_revision, \
'Need current revision in alembic_version table...'
script = ScriptDirectory.from_config(alembic_cfg)
head_revision = script.get_current_head()
log.info('Head database revision: {0}'.format(head_revision))
log.info('Current database revision: {0}'.format(current_revision))
if current_revision != head_revision:
raise Exception(
'Outdated database! Migrate using `alembic upgrade head`')
else:
log.info('[OK] Database scheme matches latest')
else:
raise Exception(
'Un-stamped database! `bin/create-db` should have done this... bailing.')
def bootstrap_db(config_uri=None, engine=None, with_migration=True):
"""Bring a blank database to a functional state."""
if engine is None:
engine = create_engine(config_uri)
db.configure(bind=engine)
if with_migration:
context = MigrationContext.configure(engine.connect())
db_version = context.get_current_revision()
if db_version:
sys.stderr.write('Database already initialized. Bailing out.\n')
sys.exit(2)
config = Config(config_uri)
script_dir = ScriptDirectory.from_config(config)
heads = script_dir.get_heads()
if len(heads) > 1:
sys.stderr.write('Error: migration scripts have more than one '
'head.\nPlease resolve the situation before '
'attempting to bootstrap the database.\n')
sys.exit(2)
metadata.create_all(engine)
with transaction.manager:
model = MyModel(name='one', value=1)
db.add(model)
# Clean up the sccoped session to allow a later app instantiation.
db.remove()
if with_migration and heads:
command.stamp(config, 'head')
def check_migrations(timeout):
"""
Function to wait for all airflow migrations to complete.
:param timeout: Timeout for the migration in seconds
:return: None
"""
from alembic.runtime.migration import MigrationContext
from alembic.script import ScriptDirectory
config = _get_alembic_config()
script_ = ScriptDirectory.from_config(config)
with settings.engine.connect() as connection:
context = MigrationContext.configure(connection)
ticker = 0
while True:
source_heads = set(script_.get_heads())
db_heads = set(context.get_current_heads())
if source_heads == db_heads:
break
if ticker >= timeout:
raise TimeoutError(
f"There are still unapplied migrations after {ticker} seconds. "
f"Migration Head(s) in DB: {db_heads} | Migration Head(s) in Source Code: {source_heads}"
)
ticker += 1
time.sleep(1)
log.info('Waiting for migrations... %s second(s)', ticker)
def setUp(self):
global cfg, env, a
env = staging_env()
cfg = _no_sql_testing_config()
cfg.set_main_option('dialect_name', 'sqlite')
cfg.remove_main_option('url')
a = util.rev_id()
script = ScriptDirectory.from_config(cfg)
script.generate_revision(a, "revision a", refresh=True)
write_script(script,
a, ("""# coding: utf-8
from __future__ import unicode_literals
revision = '%s'
down_revision = None
from alembic import op
def upgrade():
op.execute("« S’il vous plaît…")
def downgrade():
op.execute("drôle de petite voix m’a réveillé")
""" % a),
encoding='utf-8')
def upgrade(self):
conn = self.engine.connect()
cfg = Config()
# this connection is used in `env.py` for the migrations
cfg.attributes["connection"] = conn
cfg.set_main_option("script_location", str(MIGRATION_DIR))
script = ScriptDirectory.from_config(cfg)
# let's make sure we actually need to migrate
context = MigrationContext.configure(conn)
try:
# try to find the locked version
HEAD = LOCK_PATH.open().read().strip()
self.ensure_migration_needed(script, context, HEAD)
click.secho(f"Upgrading database to {HEAD}")
command.upgrade(cfg, HEAD)
except FileNotFoundError:
raise click.ClickException(
"Cannot upgrade the system database, revision lock not found.")
except MigrationUneededException:
click.secho(f"System database up-to-date.")
except Exception as err:
click.secho(
f"Cannot upgrade the system database. It might be corrupted or was created before database migrations where introduced (v0.34.0)",
fg="yellow",
err=True,
)
logging.exception(err)
finally:
conn.close()
def test_001_revisions(self):
global a, b, c
a = util.rev_id()
b = util.rev_id()
c = util.rev_id()
script = ScriptDirectory.from_config(self.cfg)
script.generate_revision(a, None, refresh=True)
write_script(script,
a,
"""
revision = '%s'
down_revision = None
from alembic import op
def upgrade():
op.execute("CREATE TABLE foo(id integer)")
def downgrade():
op.execute("DROP TABLE foo")
""" % a,
sourceless=self.sourceless)
script.generate_revision(b, None, refresh=True)
write_script(script,
b,
"""
revision = '%s'
down_revision = '%s'
from alembic import op
def upgrade():
op.execute("CREATE TABLE bar(id integer)")
def downgrade():
op.execute("DROP TABLE bar")
""" % (b, a),
sourceless=self.sourceless)
script.generate_revision(c, None, refresh=True)
write_script(script,
c,
"""
revision = '%s'
down_revision = '%s'
from alembic import op
def upgrade():
op.execute("CREATE TABLE bat(id integer)")
def downgrade():
op.execute("DROP TABLE bat")
""" % (c, b),
sourceless=self.sourceless)
def check_alembic_revision(alembic_config, conn):
migration_context = MigrationContext.configure(conn)
db_revision = migration_context.get_current_revision()
script = ScriptDirectory.from_config(alembic_config)
head_revision = script.as_revision_number("head")
return (db_revision, head_revision)
def main(args): # pragma: no cover
"""local task to make changes from glottologcurator available to the production
site via an alembic migration script.
pulls the changelog from glottologcurator and creates a new alembic revision with it.
"""
kw = {}
if args.http_user and args.http_password:
kw['auth'] = (args.http_user, args.http_password)
changes = requests.get(args.log_url, **kw).json()
config = Config()
config.set_main_option("script_location", args.migrations_dir)
scriptdir = ScriptDirectory.from_config(config)
script = scriptdir.generate_revision(
rev_id(), "Glottolog Curator", refresh=True,
upgrades="""\
# from glottologcurator
conn = op.get_bind()
for sql, params in [
%s
]:
conn.execute(sql, params)
""" % '\n'.join(u' ("""{0}""", {1}),'.format(
event[0], parse_json_with_datetime(event[1])) for event in changes['events']))
args.log.info('new alembic migration script created:')
args.log.info(script.path)
args.log.info('run "alembic upgrade head" to merge changes')
def run_alembic_migration(self):
""" Migrate to latest Alembic revision if not up-to-date. """
def migrate_if_required(rev, context):
rev = script.get_revision(rev)
if not (rev and rev.is_head):
migration_required = True
return []
migration_required = False
config = Config(os.path.join(os.path.dirname(__file__), 'alembic.ini'))
config.set_section_option(
'alembic', 'script_location',
os.path.join(os.path.dirname(__file__), 'migrations'))
config.set_main_option('sqlalchemy.url', 'sqlite:///' + self.filepath)
script = ScriptDirectory.from_config(config)
with EnvironmentContext(config, script, fn=migrate_if_required):
script.run_env()
if migration_required:
logging.info('Migrating SQLite database to latest revision')
alembic.command.upgrade(config, 'head')
else:
logging.info('SQLite database is on the latest revision')
def alembic_revision(log_url):
"""local task to merge changes from glottologcurator available to the production
site via an alembic migration script.
pulls the changelog from glottologcurator and create a new alembic revision with it.
"""
user = raw_input('HTTP Basic auth user for glottologcurator: ')
password = getpass('HTTP Basic auth password for glottologcurator: ')
kw = {}
if user and password:
kw['auth'] = (user, password)
changes = requests.get(log_url, **kw).json()
config = Config()
config.set_main_option("script_location", path('.').joinpath('migrations'))
scriptdir = ScriptDirectory.from_config(config)
script = scriptdir.generate_revision(
rev_id(), "Glottolog Curator", refresh=True,
upgrades="""\
# from glottologcurator
conn = op.get_bind()
for sql, params in [
%s
]:
conn.execute(sql, params)
""" % '\n'.join(u' ("""{0}""", {1}),'.format(*event) for event in changes['events']))
print('new alembic migration script created:')
print(script.path)
def bootstrap_db(config_uri=None, with_migration=True):
"""Bring a blank database to a functional state."""
db = get_session_maker()
if with_migration:
context = MigrationContext.configure(db().connection())
db_version = context.get_current_revision()
if db_version:
sys.stderr.write('Database already initialized. Bailing out.\n')
sys.exit(0)
config = Config(config_uri)
script_dir = ScriptDirectory.from_config(config)
heads = script_dir.get_heads()
if len(heads) > 1:
sys.stderr.write('Error: migration scripts have more than one '
'head.\nPlease resolve the situation before '
'attempting to bootstrap the database.\n')
sys.exit(2)
import assembl.models
get_metadata().create_all(db().connection())
# Clean up the sccoped session to allow a later app instantiation.
if with_migration and heads:
context = MigrationContext.configure(db().connection())
context._ensure_version_table()
# The latter step seems optional?
# I am unclear as to why we'd migrate after creating tables
# on a clean database.
context.stamp(script_dir, heads[0])
return db
def upgrade_database(alembic_config_filename: str,
alembic_base_dir: str = None) -> None:
"""
Use Alembic to upgrade our database.
See http://alembic.readthedocs.org/en/latest/api/runtime.html
but also, in particular, site-packages/alembic/command.py
"""
if alembic_base_dir is None:
alembic_base_dir = os.path.dirname(alembic_config_filename)
os.chdir(alembic_base_dir) # so the directory in the config file works
config = Config(alembic_config_filename)
script = ScriptDirectory.from_config(config)
revision = 'head' # where we want to get to
# noinspection PyUnusedLocal,PyProtectedMember
def upgrade(rev, context):
return script._upgrade_revs(revision, rev)
log.info(
"Upgrading database to revision '{}' using Alembic".format(revision))
with EnvironmentContext(config,
script,
fn=upgrade,
as_sql=False,
starting_rev=None,
destination_rev=revision,
tag=None):
script.run_env()
log.info("Database upgrade completed")
def handle(self, *args, **options):
db_id = options.get("metadata", None)
metadata_module_name = get_module_from_db_name(db_id)
verbose = options.get("verbose")
use_metadata = options.get("use_metadata")
resolve_dependencies = options.get("resolve_dependencies")
try:
metadata_cls = import_string(metadata_module_name)
except ImportError:
raise CommandError("ImportError: Invalid Module Name '%s' " % metadata_module_name)
migrations_conf_dict = get_migrations_config(db_id)
if migrations_conf_dict is None:
raise CommandError("Migrations configuration not found.")
config = get_config_from_dict(metadata_cls, use_metadata=use_metadata, **migrations_conf_dict)
script = ScriptDirectory.from_config(config)
if resolve_dependencies:
heads = script.get_revisions("heads")
else:
heads = script.get_revisions(script.get_heads())
for rev in heads:
self.stdout.write(
rev.cmd_format(
verbose, include_branches=True, tree_indicators=False))
def create_session(db_url, debug=False, pool_recycle=3600):
""" Create the Session object to use to query the database.
:arg db_url: URL used to connect to the database. The URL contains
information with regards to the database engine, the host to connect
to, the user and password and the database name.
ie: <engine>://<user>:<password>@<host>/<dbname>
:kwarg debug: a boolean specifying wether we should have the verbose
output of sqlalchemy or not.
:return a Session that can be used to query the database.
"""
engine = sa.create_engine(
db_url, echo=debug,
pool_recycle=pool_recycle,
convert_unicode=True)
session = scoped_session(sessionmaker(
autocommit=False, autoflush=False, bind=engine))
# check that the database's schema is up-to-date
script_dir = ScriptDirectory.from_config(get_alembic_config(db_url))
head_rev = script_dir.get_current_head()
context = MigrationContext.configure(session.connection())
current_rev = context.get_current_revision()
if current_rev != head_rev:
raise DatabaseNeedsUpgrade
# everything looks good here
return session
def _expected_heads():
cfg_path = join(dirname(__file__), 'migrations', 'alembic.ini')
cfg = Config(cfg_path)
_configure_alembic(cfg)
cfg.set_main_option('script_location', dirname(cfg_path))
script_dir = ScriptDirectory.from_config(cfg)
return set(script_dir.get_heads())
def test_error_on_new_with_missing_revision(self):
self.cfg.set_main_option("file_template", "%%(slug)s_%%(rev)s")
script = ScriptDirectory.from_config(self.cfg)
a = util.rev_id()
script.generate_revision(a, "foobar", refresh=True)
path = script.get_revision(a).path
with open(path, 'w') as fp:
fp.write("""
down_revision = None
from alembic import op
def upgrade():
op.execute("CREATE TABLE foo(id integer)")
def downgrade():
op.execute("DROP TABLE foo")
""")
pyc_path = util.pyc_file_from_path(path)
if os.access(pyc_path, os.F_OK):
os.unlink(pyc_path)
assert_raises_message(
util.CommandError,
"Could not determine revision id from filename foobar_%s.py. "
"Be sure the 'revision' variable is declared "
"inside the script." % a,
Script._from_path, script, path)
def ensure_db_version(config_uri, session_maker):
"""Exit if database is not up-to-date."""
config = Config(config_uri)
script_dir = ScriptDirectory.from_config(config)
heads = script_dir.get_heads()
if len(heads) > 1:
sys.stderr.write('Error: migration scripts have more than one head.\n'
'Please resolve the situation before attempting to '
'start the application.\n')
sys.exit(2)
else:
repo_version = heads[0] if heads else None
context = MigrationContext.configure(session_maker()().connect())
db_version = context.get_current_revision()
if not db_version:
sys.stderr.write('Database not initialized.\n'
'Try this: "assembl-db-manage %s bootstrap".\n'
% config_uri)
sys.exit(2)
if db_version != repo_version:
sys.stderr.write('Stopping: DB version (%s) not up-to-date (%s).\n'
% (db_version, repo_version))
sys.stderr.write('Try this: "assembl-db-manage %s upgrade head".\n'
% config_uri)
sys.exit(2)
def _get_script_dir_and_config():
"""Get config and script directory"""
from alembic.script import ScriptDirectory
config = _get_alembic_config()
script_ = ScriptDirectory.from_config(config)
return script_, config
def init_db():
connection = SQLEngine.connect()
context = MigrationContext.configure(connection)
current_revision = context.get_current_revision()
logger.boot('Database revision: %s', current_revision)
if current_revision is None:
DataBase.metadata.create_all(SQLEngine)
config = Config(ALEMBIC_CONFIG)
script = ScriptDirectory.from_config(config)
head_revision = script.get_current_head()
if current_revision is None or current_revision != head_revision:
logger.boot('Upgrading database to version %s.', head_revision)
command.upgrade(config, 'head')
from option import Option
session = Session()
options = session.query(Option).first()
if options is None:
options = Option()
options.version = head_revision
session.add(options)
from pulse import Pulse
pulse = session.query(Pulse).first()
if pulse is None:
pulse = Pulse()
session.add(pulse)
session.commit()
def setUp(self):
self.env = staging_env()
self.cfg = cfg = _no_sql_testing_config()
cfg.set_main_option("dialect_name", "sqlite")
cfg.remove_main_option("url")
self.a = util.rev_id()
script = ScriptDirectory.from_config(cfg)
script.generate_revision(self.a, "revision a", refresh=True)
write_script(
script,
self.a,
(
compat.u(
"""# coding: utf-8
from __future__ import unicode_literals
revision = '%s'
down_revision = None
from alembic import op
def upgrade():
op.execute("« S’il vous plaît…")
def downgrade():
op.execute("drôle de petite voix m’a réveillé")
"""
)
% self.a
),
encoding="utf-8",
)
def current(config, head_only=False):
"""Display the current revision for each database."""
script = ScriptDirectory.from_config(config)
def display_version(rev, context):
rev = script.get_revision(rev)
if head_only:
config.print_stdout("%s%s" % (
rev.revision if rev else None,
" (head)" if rev and rev.is_head else ""))
else:
config.print_stdout("Current revision for %s: %s",
util.obfuscate_url_pw(
context.connection.engine.url),
rev)
return []
with EnvironmentContext(
config,
script,
fn=display_version
):
script.run_env()
def setUp(self):
self.bind = _sqlite_file_db()
self.env = staging_env()
self.cfg = _sqlite_testing_config()
self.a = a = util.rev_id()
self.b = b = util.rev_id()
script = ScriptDirectory.from_config(self.cfg)
script.generate_revision(a, None, refresh=True)
write_script(
script,
a,
"""
revision = '%s'
down_revision = None
"""
% a,
)
script.generate_revision(b, None, refresh=True)
write_script(
script,
b,
"""
revision = '%s'
down_revision = '%s'
"""
% (b, a),
)
def _make_script_dir(self, alembic_cfg):
"""
build and cast the script_directory
"""
script_dir = ScriptDirectory.from_config(alembic_cfg)
script_dir.__class__ = ScriptDirectoryWithDefaultEnvPy
return script_dir
def setUp(self):
self.env = staging_env()
self.cfg = _multi_dir_testing_config()
self.cfg.set_main_option("revision_environment", "true")
script = ScriptDirectory.from_config(self.cfg)
# MARKMARK
self.model1 = util.rev_id()
self.model2 = util.rev_id()
self.model3 = util.rev_id()
for model, name in [
(self.model1, "model1"),
(self.model2, "model2"),
(self.model3, "model3"),
]:
script.generate_revision(
model, name, refresh=True,
version_path=os.path.join(_get_staging_directory(), name),
head="base")
write_script(script, model, """\
"%s"
revision = '%s'
down_revision = None
branch_labels = ['%s']
from alembic import op
def upgrade():
pass
def downgrade():
pass
""" % (name, model, name))
def test_create_script_branches_old_template(self):
script = ScriptDirectory.from_config(self.cfg)
with open(os.path.join(script.dir, "script.py.mako"), "w") as file_:
file_.write(
"<%text>#</%text> ${message}\n"
"revision = ${repr(up_revision)}\n"
"down_revision = ${repr(down_revision)}\n"
"def upgrade():\n"
" ${upgrades if upgrades else 'pass'}\n\n"
"def downgrade():\n"
" ${downgrade if downgrades else 'pass'}\n\n"
)
# works OK if no branch names
command.revision(self.cfg, message="some message")
assert_raises_message(
util.CommandError,
r"Version \w+ specified branch_labels foobar, "
r"however the migration file .+?\b does not have them; have you "
"upgraded your script.py.mako to include the 'branch_labels' "
r"section\?",
command.revision,
self.cfg, message="some message", branch_label="foobar"
)
def revision(config, message=None, autogenerate=False, sql=False):
"""Create a new revision file."""
script = ScriptDirectory.from_config(config)
template_args = {}
imports = set()
environment = util.asbool(
config.get_main_option("revision_environment")
)
if autogenerate:
environment = True
util.requires_07("autogenerate")
def retrieve_migrations(rev, context):
if script.get_revision(rev) is not script.get_revision("head"):
raise util.CommandError("Target database is not up to date.")
autogen._produce_migration_diffs(context, template_args, imports)
return []
elif environment:
def retrieve_migrations(rev, context):
return []
if environment:
with EnvironmentContext(
config,
script,
fn=retrieve_migrations,
as_sql=sql,
template_args=template_args,
):
script.run_env()
script.generate_revision(util.rev_id(), message, **template_args)
def downgrade(config, revision, sql=False, tag=None):
"""Revert to a previous version."""
script = ScriptDirectory.from_config(config)
starting_rev = None
if ":" in revision:
if not sql:
raise util.CommandError("Range revision not allowed")
starting_rev, revision = revision.split(':', 2)
elif sql:
raise util.CommandError("downgrade with --sql requires <fromrev>:<torev>")
def downgrade(rev, context):
return script._downgrade_revs(revision, rev)
with EnvironmentContext(
config,
script,
fn=downgrade,
as_sql=sql,
starting_rev=starting_rev,
destination_rev=revision,
tag=tag
):
script.run_env()
def update(self):
"""
Performs the update
:returns: The update results
"""
if self._current_revision != self._newest_revision:
_log('DBUpdater: starting..')
try:
script_directory = ScriptDirectory.from_config(self._config)
revision_list = []
for script in script_directory.walk_revisions(self._current_revision, self._newest_revision):
if script.revision != self._current_revision:
revision_list.append(script.revision)
for rev in reversed(revision_list):
try:
_log('Applying database revision: {0}'.format(rev))
command.upgrade(self._config, rev)
except sqlalchemy.exc.OperationalError, err:
if 'already exists' in str(err):
_log('Table already exists.. stamping to revision.')
self._stamp_database(rev)
except sqlalchemy.exc.OperationalError, err:
_log('DBUpdater: failure - {0}'.format(err), logLevel=logging.ERROR)
return False
_log('DBUpdater: success')
def test_error_on_new_with_missing_revision(self):
self.cfg.set_main_option("file_template", "%%(slug)s_%%(rev)s")
script = ScriptDirectory.from_config(self.cfg)
a = util.rev_id()
script.generate_revision(a, "foobar", refresh=True)
assert_raises_message(
util.CommandError,
"Could not determine revision id from filename foobar_%s.py. "
"Be sure the 'revision' variable is declared "
"inside the script." % a,
write_script,
script,
a,
"""
down_revision = None
from alembic import op
def upgrade():
op.execute("CREATE TABLE foo(id integer)")
def downgrade():
op.execute("DROP TABLE foo")
""",
)
def _init_users():
"""Initialize Afterglow user datastore if AUTH_ENABLED = True"""
# noinspection PyUnresolvedReferences
from .. import oauth2 # register oauth token-related models
# All imports put here to avoid unnecessary loading of packages on startup
# if user auth is disabled
from alembic import (config as alembic_config, context as alembic_context)
from alembic.script import ScriptDirectory
from alembic.runtime.environment import EnvironmentContext
global user_datastore, security
user_datastore = SQLAlchemyUserDatastore(db, DbUser, DbRole)
security = Security(app, user_datastore, register_blueprint=False)
# Make sure that the database directory exists
try:
os.makedirs(os.path.abspath(app.config['DATA_ROOT']))
except OSError as e:
if e.errno != errno.EEXIST:
raise
# Create/upgrade tables via Alembic
cfg = alembic_config.Config()
cfg.set_main_option(
'script_location',
os.path.abspath(
os.path.join(__file__, '../..', 'db_migration', 'users')))
script = ScriptDirectory.from_config(cfg)
# noinspection PyProtectedMember
with EnvironmentContext(
cfg,
script,
fn=lambda rev, _: script._upgrade_revs('head', rev),
as_sql=False,
starting_rev=None,
destination_rev='head',
tag=None,
), db.engine.connect() as connection:
alembic_context.configure(connection=connection)
with alembic_context.begin_transaction():
alembic_context.run_migrations()
# Initialize user roles if missing
try:
roles_created = False
for name, descr in [('admin', 'Afterglow Administrator'),
('user', 'Afterglow User')]:
if not user_datastore.find_role(name):
user_datastore.create_role(name=name, description=descr)
roles_created = True
if roles_created:
user_datastore.commit()
except Exception:
db.session.rollback()
raise
def three_rev_fixture(cfg):
a = util.rev_id()
b = util.rev_id()
c = util.rev_id()
script = ScriptDirectory.from_config(cfg)
script.generate_revision(a, "revision a", refresh=True)
write_script(
script, a, """\
"Rev A"
revision = '%s'
down_revision = None
from alembic import op
def upgrade():
op.execute("CREATE STEP 1")
def downgrade():
op.execute("DROP STEP 1")
""" % a)
script.generate_revision(b, "revision b", refresh=True)
write_script(script,
b,
u("""# coding: utf-8
"Rev B, méil"
revision = '%s'
down_revision = '%s'
from alembic import op
def upgrade():
op.execute("CREATE STEP 2")
def downgrade():
op.execute("DROP STEP 2")
""") % (b, a),
encoding="utf-8")
script.generate_revision(c, "revision c", refresh=True)
write_script(
script, c, """\
"Rev C"
revision = '%s'
down_revision = '%s'
from alembic import op
def upgrade():
op.execute("CREATE STEP 3")
def downgrade():
op.execute("DROP STEP 3")
""" % (c, b))
return a, b, c
def upgrade(plugin_name=None, revision=None, lower="current"):
database_url = current_app.config.get("SQLALCHEMY_DATABASE_URI")
if database_url.startswith("sqlite"):
current_app.db.create_all()
return
if plugin_name is None:
# Get the directory name of the plugin if unspecified
# Doing it this way doesn't waste the rest of the inspect.stack call
frame = inspect.currentframe()
caller_info = inspect.getframeinfo(frame.f_back)
caller_path = caller_info[0]
plugin_name = os.path.basename(os.path.dirname(caller_path))
# Check if the plugin has migraitons
migrations_path = os.path.join(current_app.plugins_dir, plugin_name,
"migrations")
if os.path.isdir(migrations_path) is False:
return
engine = create_engine(database_url, poolclass=pool.NullPool)
conn = engine.connect()
context = MigrationContext.configure(conn)
op = Operations(context)
# Find the list of migrations to run
config = Config()
config.set_main_option("script_location", migrations_path)
config.set_main_option("version_locations", migrations_path)
script = ScriptDirectory.from_config(config)
# Choose base revision for plugin upgrade
# "current" points to the current plugin version stored in config
# None represents the absolute base layer (e.g. first installation)
if lower == "current":
lower = get_config(plugin_name + "_alembic_version")
# Do we upgrade to head or to a specific revision
if revision is None:
upper = script.get_current_head()
else:
upper = revision
# Apply from lower to upper
revs = list(script.iterate_revisions(lower=lower, upper=upper))
revs.reverse()
try:
for r in revs:
with context.begin_transaction():
r.module.upgrade(op=op)
# Set revision that succeeded so we don't need
# to start from the beginning on failure
set_config(plugin_name + "_alembic_version", r.revision)
finally:
conn.close()
# Set the new latest revision
set_config(plugin_name + "_alembic_version", upper)
def has_french_dem():
config = migrate.get_config()
script = ScriptDirectory.from_config(config)
migration_context = MigrationContext.configure(db.session.connection())
current_heads = migration_context.get_current_heads()
current_heads = set(
map(lambda rev: rev.revision, script.get_all_current(current_heads)))
return '1715cf31a75d' in current_heads # ign bd alti
def _get_current_rev(self, alembic_conf, engine):
script = ScriptDirectory.from_config(alembic_conf)
with engine.connect() as conn:
with EnvironmentContext(alembic_conf, script) as env_context:
env_context.configure(conn, version_table="alembic_version")
migration_context = env_context.get_context()
revision = migration_context.get_current_revision()
return revision
def history(config):
"""List changeset scripts in chronological order."""
script = ScriptDirectory.from_config(config)
for sc in script.walk_revisions():
if sc.is_head:
print
print sc
def test_only_single_head_revision_in_migrations(self):
config = Config()
config.set_main_option("script_location", "airflow:migrations")
script = ScriptDirectory.from_config(config)
# This will raise if there are multiple heads
# To resolve, use the command `alembic merge`
script.get_current_head()
def _fixture(self, **kw):
script = ScriptDirectory.from_config(self.cfg)
env = EnvironmentContext(
self.cfg,
script,
**kw
)
return env
def _script_and_context(env, db, script_location):
alembic_cfg = Config()
alembic_cfg.set_main_option("script_location", script_location)
script = ScriptDirectory.from_config(alembic_cfg)
connection = TracAlchemy(env).sqlalchemy_connection(db)
context = MigrationContext.configure(connection, opts={'script': script})
return (script, context)
def test_alembic_has_single_head(session):
"""
Avoid unintentional branches in the migration history.
"""
migrations_dir = current_app.extensions['migrate'].directory
heads = ScriptDirectory(migrations_dir).get_heads()
assert len(heads) == 1
def get_revision(self, app: str, index_: int) -> Optional[str]:
script = ScriptDirectory.from_config(self.config)
rev = script.walk_revisions(
f"{self.db_name}/{app}@base", f"{self.db_name}/{app}@head"
)
revision = next(filter(lambda x: x[0] == index_, enumerate(rev)), None)
if revision:
return revision[1]
def get_migration_head(config):
"""
This function returns the head of the migration scripts.
:param config: The alembic configuration.
:return: The version of the head.
"""
script = ScriptDirectory.from_config(config)
return script.get_current_head()
def _test_tz(self, timezone_arg, given, expected):
script = ScriptDirectory(
_get_staging_directory(),
file_template="%(rev)s_%(slug)s_"
"%(year)s_%(month)s_"
"%(day)s_%(hour)s_"
"%(minute)s_%(second)s",
timezone=timezone_arg,
)
with mock.patch(
"alembic.script.base.datetime",
mock.Mock(datetime=mock.Mock(utcnow=lambda: given,
now=lambda: given)),
):
create_date = script._generate_create_date()
eq_(create_date, expected)
def test_create_script_splice(self):
rev = command.revision(
self.cfg, message="some message", head=self.b, splice=True)
script = ScriptDirectory.from_config(self.cfg)
rev = script.get_revision(rev.revision)
eq_(rev.down_revision, self.b)
assert "some message" in rev.doc
eq_(set(script.get_heads()), set([rev.revision, self.c]))
def __init__(self, db, migration_dir, config_path=None):
self.db = db
self.migration_dir = migration_dir
self.config_path = config_path or os.path.join(migration_dir, 'alembic.ini')
self.logger = logging.getLogger(__name__)
self.config = Config(self.config_path)
self.config.set_main_option('script_location', self.migration_dir)
self.script = ScriptDirectory.from_config(self.config)
def setUp(self):
staging_env()
self.cfg = cfg = _no_sql_testing_config()
self.rid = rid = util.rev_id()
self.script = script = ScriptDirectory.from_config(cfg)
script.generate_revision(rid, None, refresh=True)
def list_migrations(cfg_path, head):
cfg = AlembicConfig(cfg_path)
script = ScriptDirectory.from_config(cfg)
migrations = [
x.revision for x in script.walk_revisions(base='base', head=head)
]
migrations.reverse()
return migrations
def _opened_transaction_fixture(self):
self.env = staging_env()
self.cfg = _sqlite_testing_config()
script = ScriptDirectory.from_config(self.cfg)
a = util.rev_id()
b = util.rev_id()
c = util.rev_id()
script.generate_revision(a, "revision a", refresh=True)
write_script(script, a, """
"rev a"
revision = '%s'
down_revision = None
def upgrade():
pass
def downgrade():
pass
""" % (a, ))
script.generate_revision(b, "revision b", refresh=True)
write_script(script, b, """
"rev b"
revision = '%s'
down_revision = '%s'
from alembic import op
def upgrade():
conn = op.get_bind()
trans = conn.begin()
def downgrade():
pass
""" % (b, a))
script.generate_revision(c, "revision c", refresh=True)
write_script(script, c, """
"rev c"
revision = '%s'
down_revision = '%s'
from alembic import op
def upgrade():
pass
def downgrade():
pass
""" % (c, b))
return a, b, c
def latest_schema_revision(self):
"""The current schema revision according to Alembic's migration scripts (a string)."""
logger.debug("Finding Alembic head revision ..")
migrations = ScriptDirectory.from_config(self.alembic_config)
revision = migrations.get_current_head()
logger.verbose(
"Current head (code base) database schema revision is %s.",
revision)
return revision
def _eq_cmd_output(self, buf, expected):
script = ScriptDirectory.from_config(self.cfg)
# test default encode/decode behavior as well,
# rev B has a non-ascii char in it + a coding header.
eq_(
buf.getvalue().decode("ascii", 'replace').strip(),
"\n".join([script.get_revision(rev).log_entry for rev in expected
]).encode("ascii", "replace").decode("ascii").strip())
def script_directory(self):
"""Get the Alembic :class:`~alembic.script.ScriptDirectory` for the current app."""
cache = self._get_cache()
if 'script' not in cache:
cache['script'] = ScriptDirectory.from_config(self.config)
return cache['script']
def test_args(self):
script = ScriptDirectory(
_get_staging_directory(),
file_template="%(rev)s_%(slug)s_"
"%(year)s_%(month)s_"
"%(day)s_%(hour)s_"
"%(minute)s_%(second)s",
)
create_date = datetime.datetime(2012, 7, 25, 15, 8, 5)
eq_(
script._rev_path(
script.versions, "12345", "this is a message", create_date
),
os.path.abspath(
"%s/versions/12345_this_is_a_"
"message_2012_7_25_15_8_5.py" % _get_staging_directory()
),
)
def test_args_propagate(self):
config = _no_sql_testing_config()
script = ScriptDirectory.from_config(config)
template_args = {"x": "x1", "y": "y1", "z": "z1"}
env = EnvironmentContext(config, script, template_args=template_args)
env.configure(
dialect_name="sqlite", template_args={"y": "y2", "q": "q1"}
)
eq_(template_args, {"x": "x1", "y": "y2", "z": "z1", "q": "q1"})
