def create_migration_ctx(**kwargs):
"""Create an alembic migration context."""
env = EnvironmentContext(Config(), None)
env.configure(connection=db.engine.connect(),
sqlalchemy_module_prefix='db.',
**kwargs)
return env.get_context()
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 create_migration_ctx(**kwargs):
"""Create an alembic migration context."""
env = EnvironmentContext(Config(), None)
env.configure(
connection=db.engine.connect(),
sqlalchemy_module_prefix='db.',
**kwargs
)
return env.get_context()
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 upgrade(pkg, sql=False):
"""Upgrade to a later version."""
from alembic.config import Config
from alembic.environment import EnvironmentContext
if ':' in pkg:
pkg, rev = pkg.split(':', 1)
else:
rev = 'head'
script = ScriptDirectory(pkg)
env = EnvironmentContext(Config(''), script)
conn = ptah.get_base().metadata.bind.connect()
def upgrade(revision, context):
return script._upgrade_revs(rev, revision)
env.configure(
connection=conn,
fn=upgrade,
as_sql=sql,
starting_rev=None,
destination_rev=rev,
)
mc = env._migration_context
env._migration_context = MigrationContext(pkg, conn.dialect, conn, mc.opts)
with env:
try:
with env.begin_transaction():
env.run_migrations()
finally:
conn.close()
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 is_head(db_url, upgrade=False):
"""Check whether or not db_url is head.
Args:
db_url (str): database url
upgrade (Bool): if True, upgrade db to head
"""
config = Config()
config.set_main_option("script_location", "spinedb_api:alembic")
script = ScriptDirectory.from_config(config)
head = script.get_current_head()
engine = create_engine(db_url)
with engine.connect() as connection:
migration_context = MigrationContext.configure(connection)
current_rev = migration_context.get_current_revision()
if current_rev == head:
return True
if not upgrade:
return False
# Upgrade function
def fn(rev, context):
return script._upgrade_revs("head", rev)
with EnvironmentContext(config,
script,
fn=fn,
as_sql=False,
starting_rev=None,
destination_rev="head",
tag=None) as environment_context:
environment_context.configure(connection=connection,
target_metadata=None)
with environment_context.begin_transaction():
environment_context.run_migrations()
return True
def upgrade_database(engine: Engine, revision: str = 'head'):
config = get_alembic_config(engine)
script = ScriptDirectory.from_config(config)
def upgrade(rev, context):
def update_current_rev(old, new):
if old == new:
return
if new is None:
context.impl._exec(context._version.delete())
elif old is None:
context.impl._exec(context._version.insert().values(
version_num=literal_column("'%s'" % new)))
else:
context.impl._exec(context._version.update().values(
version_num=literal_column("'%s'" % new)))
if not rev and revision == 'head':
import_all_modules()
Base.metadata.create_all(engine)
dest = script.get_revision(revision)
update_current_rev(None, dest and dest.revision)
return []
return script._upgrade_revs(revision, rev)
with EnvironmentContext(config,
script,
fn=upgrade,
as_sql=False,
destination_rev=revision,
tag=None):
script.run_env()
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 move(revision, after):
""" Reorder revisions """
# P(ver) -> ver -> C(ver) ... P(after) -> after -> C(after)
# becomes
# P(ver) -> C(ver) ... P(after) -> after -> ver -> C(after)
if revision == after:
print(
"Cannot move revision {0} after itself".format(revision),
file=sys.stderr
)
sys.exit(1)
def _move(r, ctx):
what_rev, after_rev = rev(revision), rev(after)
pwhat_id, cwhat_ids = what_rev.down_revision, list(what_rev.nextrev)
cafter_ids = list(after_rev.nextrev)
for rev_id in cwhat_ids:
replace(rev(rev_id).path, revision, pwhat_id)
replace(what_rev.path, pwhat_id, after)
for rev_id in cafter_ids:
replace(rev(rev_id).path, after, revision)
return []
with EnvironmentContext(config, script, fn=_move):
script.run_env()
def run_env(self, fn: Callable, **kw) -> None:
with EnvironmentContext(self.config,
self.script_dir,
fn=fn,
version_table=self.version_table,
**kw):
self.script_dir.run_env()
def get_database_revision(engine):
"""Gets the current revision of the database.
:param engine: the database engine to get the current revision
:type engine: :class:`sqlalchemy.engine.base.Engine`
:returns: the script of the current revision
:rtype: :class:`alembic.script.Script`
"""
config = get_alembic_config(engine)
script = ScriptDirectory.from_config(config)
result = [None]
def get_revision(rev, context):
result[0] = rev and script.get_revision(rev)
return []
with EnvironmentContext(config,
script,
fn=get_revision,
as_sql=False,
destination_rev=None,
tag=None):
script.run_env()
return result[0]
def _fixture(self, **kw):
script = ScriptDirectory.from_config(self.cfg)
env = EnvironmentContext(
self.cfg,
script,
**kw
)
return env
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 upgrade(pkg, sql=False):
"""Upgrade to a later version."""
from alembic.config import Config
from alembic.environment import EnvironmentContext
if ':' in pkg:
pkg, rev = pkg.split(':',1)
else:
rev = 'head'
script = ScriptDirectory(pkg)
env = EnvironmentContext(Config(''), script)
conn = ptah.get_base().metadata.bind.connect()
def upgrade(revision, context):
return script._upgrade_revs(rev, revision)
env.configure(
connection = conn,
fn = upgrade,
as_sql = sql,
starting_rev = None,
destination_rev = rev,
)
mc = env._migration_context
env._migration_context = MigrationContext(pkg, conn.dialect, conn, mc.opts)
with env:
try:
with env.begin_transaction():
env.run_migrations()
finally:
conn.close()
def test_db_metadata_differences(models, settings):
db_name = settings['db_name']
# first we drop anything there might be
dropdb(db_name)
# then we create a clean DB from the metadata
createdb(db_name)
metadata = models.metadata
engine = engine_from_config(settings)
metadata.bind = engine
metadata.create_all(engine, tables=[table for name, table
in metadata.tables.items() if not name.startswith('test_')])
# and store the results
create_all_result = dumpdb(db_name)
engine.dispose()
# now we do it again, but this time using migrations
dropdb(db_name)
createdb(db_name)
config = Config()
config.set_main_option('script_location', 'backrest:migrations')
script = ScriptDirectory.from_config(config)
connection = engine.connect()
environment = EnvironmentContext(config, script,
starting_rev='base', destination_rev='head')
context = MigrationContext.configure(connection)
def upgrade(rev, context):
return script._upgrade_revs('head', rev)
context._migrations_fn = upgrade
environment._migration_context = context
with environment.begin_transaction():
environment.run_migrations()
# we drop alembic_version to avoid it showing up in the diff
engine.execute('DROP TABLE alembic_version;')
# we store these results
alembic_result = dumpdb(db_name)
del context
del environment
connection.close()
del connection
engine.dispose()
# now we check whether there are differences and output them if there are
diff = unified_diff(alembic_result, create_all_result)
assert alembic_result == create_all_result, \
'Migration output differs:\n' + ''.join(diff)
def run_env(self, fn, **kw):
"""
run alembic's context
"""
with EnvironmentContext(self.config,
self.script_dir,
fn=fn,
version_table=self.version_table,
**kw):
self.script_dir.run_env()
def test_sys_path_prepend(self, config_value, expected):
self.cfg.set_main_option("prepend_sys_path", config_value)
script = ScriptDirectory.from_config(self.cfg)
env = EnvironmentContext(self.cfg, script)
target = os.path.abspath(_get_staging_directory())
def assert_(heads, context):
eq_(
[os.path.abspath(p) for p in sys.path[0:len(expected)]],
[os.path.abspath(p) for p in expected],
)
return []
p = [p for p in sys.path if os.path.abspath(p) != target]
with mock.patch.object(sys, "path", p):
env.configure(url="sqlite://", fn=assert_)
with env:
script.run_env()
def alembic_process(self, config):
from alembic.command import history
from alembic.environment import EnvironmentContext
from alembic.script import ScriptDirectory
from cliche.web.app import app
from cliche.web.db import get_database_engine
with app.app_context():
engine = get_database_engine()
script = ScriptDirectory.from_config(config)
with EnvironmentContext(config, script, as_sql=False,
destination_rev=None, tag=None) as context:
context.configure(engine.contextual_connect(), engine.url)
history(config)
def bootstrap_db(config_uri, with_migration=True):
"""Bring a blank database to a functional state."""
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)
elif len(heads) == 0:
sys.stderr.write('Error: migration scripts have no head.\n')
sys.exit(2)
head = heads[0]
db = get_session_maker()
db.flush()
if not has_tables(db()):
with locked_transaction(db, 1234) as session:
context = MigrationContext.configure(session.connection())
if not has_tables(session):
import assembl.models
get_metadata().create_all(session.connection())
assert has_tables(session)
context._ensure_version_table()
context.stamp(script_dir, head)
elif with_migration:
context = MigrationContext.configure(db().connection())
db_version = context.get_current_revision()
# artefact: in tests, db_version may be none.
if db_version and db_version != head:
with locked_transaction(db, 1235) as session:
context = MigrationContext.configure(session.connection())
db_version = context.get_current_revision()
if db_version != head:
with EnvironmentContext(
config,
script_dir,
as_sql=False,
fn=lambda heads, context: script_dir._upgrade_revs(
head, db_version),
destination_rev=head):
script_dir.run_env()
context = MigrationContext.configure(db().connection())
db_version = context.get_current_revision()
assert db_version == head
return db
def _upgrade_database(db, revision="head"):
"""Upgrade ThreediDatabase instance"""
with db.get_engine().begin() as connection:
config = get_alembic_config(connection)
script = ScriptDirectory.from_config(config)
def upgrade(rev, context):
return script._upgrade_revs(revision, rev)
with EnvironmentContext(
config,
script,
fn=upgrade,
destination_rev=revision,
version_table=constants.VERSION_TABLE_NAME,
):
script.run_env()
def init_environment_context(cls, conf):
file_template = conf.get('alembic.file_template', '%%(day).3d-%%(rev)s-%%(slug)s')
script_location = conf.get('alembic.script_location', 'mediadrop:migrations')
version_table = conf.get('alembic.version_table', 'alembic_migrations')
alembic_cfg = Config(ini_section='main')
alembic_cfg.set_main_option('script_location', script_location)
alembic_cfg.set_main_option('sqlalchemy.url', conf['sqlalchemy.url'])
# TODO: add other sqlalchemy options
alembic_cfg.set_main_option('file_template', file_template)
script = ScriptDirectory.from_config(alembic_cfg)
def upgrade(current_db_revision, context):
return script._upgrade_revs('head', current_db_revision)
table_name = prefix_table_name(conf, table_name=version_table)
return EnvironmentContext(alembic_cfg, script, fn=upgrade, version_table=table_name)
def upgrade(config, engine, revision):
"""Upgrade to the given revision."""
script = from_config(config)
def upgrade(rev, context):
return script._upgrade_revs(revision, rev)
# Always run full upgrade
starting_rev = None
with EnvironmentContext(config,
script,
fn=upgrade,
as_sql=False,
starting_rev=starting_rev,
destination_rev=revision):
run_migration(engine)
def current(config):
"""Display the current revision for each database."""
script = ScriptDirectory.from_config(config)
def display_version(rev, context):
print "Current revision for %s: %s" % (
util.obfuscate_url_pw(
context.connection.engine.url),
script.get_revision(rev))
return []
with EnvironmentContext(
config,
script,
fn = display_version
):
script.run_env()
def get_database_revision(engine):
config = get_alembic_config(engine)
script = ScriptDirectory.from_config(config)
result = [None]
def get_revision(rev, context):
result[0] = rev and script.get_revision(rev)
return []
with EnvironmentContext(config,
script,
fn=get_revision,
as_sql=False,
destination_rev=None,
tag=None):
script.run_env()
return None if result[0] == () else result[0]
def __init__(self):
# Defer importing the current app since other utilities might not
# require being inside a flask application context
from flask import current_app
# Grab the current alembic configuration set by flask_migrate
config = current_app.extensions['migrate'].migrate.get_config()
script = ScriptDirectory.from_config(config)
self.head_revision = script.get_current_head()
# Current revision will be set via a callback once the alembic
# environment is entered
self.current_revision = None
# Enter the alembic EnvironmentContext so that we can pull the latest
# revision from the DB.
with EnvironmentContext(config, script, fn=self._set_current_revision):
script.run_env()
def apply_migrations(engine):
"""
Apply Alembic migrations to the connected database
"""
target = 'head'
config = Config()
config.set_main_option("script_location", "metadb:migrations")
script = ScriptDirectory.from_config(config)
def upgrade_fn(rev, context):
return script._upgrade_revs(target, rev)
with EnvironmentContext(config, script) as env:
with engine.connect() as conn:
env.configure(conn, fn=upgrade_fn, destination_rev=target)
with env.begin_transaction():
env.run_migrations()
def upgrade_database(
config: Config,
engine: Engine,
metadata: MetaData,
*,
revision: str = 'head',
module_name: typing.Optional[str] = None,
) -> None:
"""Upgrades the database schema to the chosen ``revision`` (default is
head).
"""
script = ScriptDirectory.from_config(config)
def upgrade(rev, context):
def update_current_rev(old, new):
if old == new:
return
if new is None:
context.impl._exec(context._version.delete())
elif old is None:
context.impl._exec(context._version.insert().values(
version_num=literal_column("'%s'" % new)))
else:
context.impl._exec(context._version.update().values(
version_num=literal_column("'%s'" % new)))
if not rev and revision == 'head':
if module_name:
import_all_modules(module_name)
metadata.create_all(engine)
dest = script.get_revision(revision)
update_current_rev(None, dest and dest.revision)
return []
return script._upgrade_revs(revision, rev)
with EnvironmentContext(config,
script,
fn=upgrade,
as_sql=False,
destination_rev=revision,
tag=None):
script.run_env()
def _check_db_version(self, upgrade=False):
"""Check if database is the latest version and raise a `SpineDBVersionError` if not.
If upgrade is `True`, then don't raise the error and upgrade the database instead.
"""
config = Config()
config.set_main_option("script_location", "spinedb_api:alembic")
script = ScriptDirectory.from_config(config)
head = script.get_current_head()
with self.engine.connect() as connection:
migration_context = MigrationContext.configure(connection)
current = migration_context.get_current_revision()
if current is None:
# No revision information. Check if the schema of the given url corresponds to
# a non-upgraded new Spine db --otherwise we can't go on.
ref_engine = create_new_spine_database("sqlite://",
upgrade=False)
if not schemas_are_equal(self.engine, ref_engine):
raise SpineDBAPIError(
"The db at '{0}' doesn't seem like a valid Spine db.".
format(self.db_url))
if current == head:
return
if not upgrade:
raise SpineDBVersionError(url=self.db_url,
current=current,
expected=head)
# Upgrade function
def upgrade_to_head(rev, context):
return script._upgrade_revs("head", rev)
with EnvironmentContext(config,
script,
fn=upgrade_to_head,
as_sql=False,
starting_rev=None,
destination_rev="head",
tag=None) as environment_context:
environment_context.configure(connection=connection,
target_metadata=None)
with environment_context.begin_transaction():
environment_context.run_migrations()
def run_alembic_migration(db_uri, log_handler=None, setup_app=True):
if log_handler:
logging.getLogger(migration_name).addHandler(log_handler)
config = Config()
config.set_main_option("script_location", "data:migrations")
config.set_main_option("db_uri", db_uri)
if setup_app:
config.set_main_option('alembic_setup_app', 'True')
else:
config.set_main_option('alembic_setup_app', '')
script = ScriptDirectory.from_config(config)
def fn(rev, context):
return script._upgrade_revs('head', rev)
with EnvironmentContext(config, script, fn=fn, destination_rev='head'):
script.run_env()
def current(pkg):
"""Display the current revision."""
def display_version(rev, context):
rev = script._get_rev(rev)
if rev is None:
print("Package '{0}' rev: None".format(pkg))
else:
print("Package '{0}' rev: {1}{2} {3}".format(
pkg, rev.revision, '(head)' if rev.is_head else "", rev.doc))
return []
conn = ptah.get_base().metadata.bind.connect()
script = ptah.migrate.ScriptDirectory(pkg)
env = EnvironmentContext(Config(''), script)
env.configure(connection=conn, fn=display_version)
mc = env._migration_context
env._migration_context = MigrationContext(pkg, conn.dialect, conn, mc.opts)
with env.begin_transaction():
env.run_migrations()
def upgrade_database(engine, revision='head'):
"""Upgrades the database schema to the chosen ``revision`` (default is
head).
:param engine: the database engine to upgrade
:type engine: :class:`sqlalchemy.engine.base.Engine`
:param revision: the revision to upgrade to. default is ``'head'``
:type revision: :class:`str`
"""
config = get_alembic_config(engine)
script = ScriptDirectory.from_config(config)
def upgrade(rev, context):
def update_current_rev(old, new):
if old == new:
return
if new is None:
context.impl._exec(context._version.delete())
elif old is None:
context.impl._exec(context._version.insert().values(
version_num=literal_column("'%s'" % new)))
else:
context.impl._exec(context._version.update().values(
version_num=literal_column("'%s'" % new)))
if not rev and revision == 'head':
import_all_modules()
Base.metadata.create_all(engine)
dest = script.get_revision(revision)
update_current_rev(None, dest and dest.revision)
return []
return script._upgrade_revs(revision, rev)
with EnvironmentContext(config,
script,
fn=upgrade,
as_sql=False,
destination_rev=revision,
tag=None):
script.run_env()
def rename(**kwargs):
""" Change a revision's ID """
from_, to = kwargs.get('from'), kwargs.get('to')
def _rename(rev, context):
from_rev = script.get_revision(from_)
from_path = from_rev.path
to_path = from_path.replace(from_, to)
for rev_id in from_rev.nextrev:
replace(script.get_revision(rev_id).path, from_, to)
replace(from_path, from_, to)
replace(
script.get_revision(from_rev.down_revision).path, from_, to
)
os.rename(from_path, to_path)
return []
with EnvironmentContext(config, script, fn=_rename):
script.run_env()
def current(pkg):
"""Display the current revision."""
def display_version(rev, context):
rev = script._get_rev(rev)
if rev is None:
print("Package '{0}' rev: None".format(pkg))
else:
print("Package '{0}' rev: {1}{2} {3}".format(pkg, rev.revision, "(head)" if rev.is_head else "", rev.doc))
return []
conn = ptah.get_base().metadata.bind.connect()
script = ptah.migrate.ScriptDirectory(pkg)
env = EnvironmentContext(Config(""), script)
env.configure(connection=conn, fn=display_version)
mc = env._migration_context
env._migration_context = MigrationContext(pkg, conn.dialect, conn, mc.opts)
with env.begin_transaction():
env.run_migrations()
if isinstance(op_, tuple):
op_name = op_[0]
if op_name.endswith('_table'):
op_obj = op_[1]
op_obj_name = op_obj.name
elif op_name.endswith('_column'):
op_obj = op_[3]
op_obj_name = op_[2] + '.' + op_obj.name
else:
op_obj = op_[1]
op_obj_name = op_obj.name
if op_obj_name is None:
op_obj_name = op_obj.table.name + '.' \
+ '_'.join(c.name for c in op_obj.columns)
else:
op_name = op_[0][0]
op_obj_name = op_[0][2] + '.' + op_[0][3]
return not op_obj_name in IGNORE_OPS.get(op_name, [])
# Set up a migration context.
alembic_cfg.set_main_option('script_location', 'thelma:db/schema/migrations')
script = ScriptDirectory.from_config(alembic_cfg)
env_ctxt = EnvironmentContext(alembic_cfg, script)
engine = create_engine(alembic_cfg)
env_ctxt.configure(engine.connect()) # , include_object=include_object)
mig_ctxt = env_ctxt.get_context()
ops = compare_metadata(mig_ctxt, metadata)
diff = [op for op in ops if include_op(op)]
pprint.pprint(diff, indent=2, width=20)
