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_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 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 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)
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 check_db_revision(engine):
"""Check the JupyterHub database revision
After calling this function, an alembic tag is guaranteed to be stored in the db.
- Checks the alembic tag and raises a ValueError if it's not the current revision
- If no tag is stored (Bug in Hub prior to 0.8),
guess revision based on db contents and tag the revision.
- Empty databases are tagged with the current revision
"""
# Check database schema version
current_table_names = set(engine.table_names())
my_table_names = set(Base.metadata.tables.keys())
with _temp_alembic_ini(engine.url) as ini:
cfg = alembic.config.Config(ini)
scripts = ScriptDirectory.from_config(cfg)
head = scripts.get_heads()[0]
base = scripts.get_base()
if not my_table_names.intersection(current_table_names):
# no tables have been created, stamp with current revision
app_log.debug("Stamping empty database with alembic revision %s", head)
alembic.command.stamp(cfg, head)
return
if 'alembic_version' not in current_table_names:
# Has not been tagged or upgraded before.
# we didn't start tagging revisions correctly except during `upgrade-db`
# until 0.8
# This should only occur for databases created prior to JupyterHub 0.8
msg_t = "Database schema version not found, guessing that JupyterHub %s created this database."
if 'spawners' in current_table_names:
# 0.8
app_log.warning(msg_t, '0.8.dev')
rev = head
elif 'services' in current_table_names:
# services is present, tag for 0.7
app_log.warning(msg_t, '0.7.x')
rev = 'af4cbdb2d13c'
else:
# it's old, mark as first revision
app_log.warning(msg_t, '0.6 or earlier')
rev = base
app_log.debug("Stamping database schema version %s", rev)
alembic.command.stamp(cfg, rev)
# check database schema version
# it should always be defined at this point
alembic_revision = engine.execute('SELECT version_num FROM alembic_version').first()[0]
if alembic_revision == head:
app_log.debug("database schema version found: %s", alembic_revision)
pass
else:
raise DatabaseSchemaMismatch("Found database schema version {found} != {head}. "
"Backup your database and run `jupyterhub upgrade-db`"
" to upgrade to the latest schema.".format(
found=alembic_revision,
head=head,
))
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 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 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 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 _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 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 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 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 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 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 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 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 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 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 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 __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 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")
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)
for sc in script.walk_revisions():
if sc.is_branch_point:
self.stdout.write(
"%s\n%s\n",
sc.cmd_format(verbose, include_branches=True),
"\n".join(
"%s -> %s"
% (
" " * len(str(sc.revision)),
rev_obj.cmd_format(False, include_branches=True, include_doc=verbose),
)
for rev_obj in (script.get_revision(rev) for rev in sc.nextrev)
),
)
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 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():
# Check database revision
from inbox.server.models.ignition import db_uri # needs to be after load_config()
inbox_db_engine = sqlalchemy.create_engine(db_uri())
# top-level, with setup.sh
alembic_ini_filename = config.get('ALEMBIC_INI', None)
assert alembic_ini_filename, 'Must set ALEMBIC_INI config var'
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 create_db.py")
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! `create_db.py` 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 _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())
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
'dirs',
metavar='DIRS',
help='path to user data file director(ies) to migrate; '
'may include wildcards')
args = parser.parse_args()
cfg = alembic_config.Config()
cfg.set_main_option(
'script_location',
os.path.abspath(
os.path.join(os.path.dirname(__file__), '..', 'afterglow_core',
'db_migration', 'data_files')))
script = ScriptDirectory.from_config(cfg)
for root in glob(args.dirs):
if not os.listdir(root):
continue
db_path = os.path.join(root, 'data_files.db')
if not os.path.isfile(db_path):
continue
# Make sure that the db files are writable
os.system('sudo chmod g+w {}*'.format(db_path))
print('Migrating', root, '...', end=' ', file=sys.stderr)
db_url = 'sqlite:///{}'.format(db_path)
engine = create_engine(
db_url,
def has_revisions(self):
try:
return bool(ScriptDirectory.from_config(self.config).get_heads())
except OSError:
return False
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 check_db_revision(engine):
"""Check the JupyterHub database revision
After calling this function, an alembic tag is guaranteed to be stored in the db.
- Checks the alembic tag and raises a ValueError if it's not the current revision
- If no tag is stored (Bug in Hub prior to 0.8),
guess revision based on db contents and tag the revision.
- Empty databases are tagged with the current revision
"""
# Check database schema version
current_table_names = set(engine.table_names())
from .dbutil import _temp_alembic_ini
with _temp_alembic_ini(engine.url) as ini:
cfg = alembic.config.Config(ini)
scripts = ScriptDirectory.from_config(cfg)
head = scripts.get_heads()[0]
base = scripts.get_base()
if not my_table_names.intersection(current_table_names):
# no tables have been created, stamp with current revision
app_log.debug(
"Stamping empty dashboards database with alembic revision %s",
head)
alembic.command.stamp(cfg, head)
return
if 'cds_alembic_version' not in current_table_names:
# Has not been tagged or upgraded before.
# This should only occur for databases created on cdsdashboards 0.0.11 or earlier
# Need to identify if this is really an old version or not
rev = head
if 'dashboards' in current_table_names:
inspector = inspect(engine)
cols = inspector.get_columns('dashboards')
colnames = [c.get('name', '') for c in cols]
if not 'presentation_type' in colnames:
rev = base
# presentation_type was added in v0.0.13, so the reason we don't have cds_alembic_version
# is because the db was created before we had db versioning
# If we DO have dashboards.presentation_type but no cds_alembic_version then this is just because
# it's a new installation, Dashboards has been created before this first check.
app_log.debug("Stamping dashboards database schema version %s",
rev)
alembic.command.stamp(cfg, rev)
else:
alembic_revision = engine.execute(
'SELECT version_num FROM cds_alembic_version').first()[0]
if alembic_revision == base:
if 'dashboards' in current_table_names:
inspector = inspect(engine)
cols = inspector.get_columns('dashboards')
colnames = [c.get('name', '') for c in cols]
if 'presentation_type' in colnames:
# For people who got stuck in the broken upgrade before - actually they are NOT on base...
rev = '260ac5c1a9e0'
app_log.debug(
"Stamping dashboards database schema version %s",
rev)
alembic.command.stamp(cfg, rev)
# check database schema version
# it should always be defined at this point
alembic_revision = engine.execute(
'SELECT version_num FROM cds_alembic_version').first()[0]
if alembic_revision == head:
app_log.debug("database dashboards schema version found: %s",
alembic_revision)
pass
else:
raise DatabaseSchemaMismatch(
"Found database schema version {found} != {head}. "
"Backup your database and run `jupyterhub upgrade-db`"
" to upgrade to the latest schema.".format(found=alembic_revision,
head=head))
def _opened_transaction_fixture(self, future=False):
self.env = staging_env()
if future:
self.cfg = _sqlite_testing_config(future=future)
else:
self.cfg = _sqlite_testing_config()
if self.branched_connection:
self._branched_connection_env()
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()
# this should fail for a SQLAlchemy 2.0 connection b.c. there is
# already a transaction.
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 _make_script_dir(alembic_cfg):
script_dir = ScriptDirectory.from_config(alembic_cfg)
script_dir.__class__ = ScriptDirectoryWithDefaultEnvPy # O_o
return script_dir
def test_create_script_basic(self):
rev = command.revision(self.cfg, message="some message")
script = ScriptDirectory.from_config(self.cfg)
rev = script.get_revision(rev.revision)
eq_(rev.down_revision, self.c)
assert "some message" in rev.doc
def alembic_script(self):
return ScriptDirectory.from_config(self.alembic_config())
def script_directory(self) -> ScriptDirectory:
if isinstance(self.__script_directory, ScriptDirectory):
return self.__script_directory
self.__script_directory = ScriptDirectory.from_config(self.config)
return self.__script_directory
def get_head_revision():
alembic_cfg = AlembicConfig(ALEMBIC_CONFIG)
script = AlembicScriptDirectory.from_config(alembic_cfg)
head_revision = script.get_current_head()
return head_revision
def check_db_version(install_path, db_file):
"""
Check if db exists and is up to date.
If it doesn't exist create it. If it's out of date update it.
"""
from alembic.script import ScriptDirectory
from alembic.config import Config
import sqlite3
import flask_migrate
# db_file = job.config.DBFILE
mig_dir = os.path.join(install_path, "arm/migrations")
config = Config()
config.set_main_option("script_location", mig_dir)
script = ScriptDirectory.from_config(config)
# create db file if it doesn't exist
if not os.path.isfile(db_file):
logging.info("No database found. Initializing arm.db...")
make_dir(os.path.dirname(db_file))
with app.app_context():
flask_migrate.upgrade(mig_dir)
if not os.path.isfile(db_file):
logging.error(
"Can't create database file. This could be a permissions issue. Exiting..."
)
sys.exit()
# check to see if db is at current revision
head_revision = script.get_current_head()
logging.debug("Head is: " + head_revision)
conn = sqlite3.connect(db_file)
c = conn.cursor()
c.execute("SELECT {cn} FROM {tn}".format(cn="version_num",
tn="alembic_version"))
db_version = c.fetchone()[0]
logging.debug("Database version is: " + db_version)
if head_revision == db_version:
logging.info("Database is up to date")
else:
logging.info("Database out of date. Head is " + head_revision +
" and database is " + db_version +
". Upgrading database...")
with app.app_context():
ts = round(time.time() * 100)
logging.info("Backuping up database '" + db_file + "' to '" +
db_file + str(ts) + "'.")
shutil.copy(db_file, db_file + "_" + str(ts))
flask_migrate.upgrade(mig_dir)
logging.info("Upgrade complete. Validating version level...")
c.execute("SELECT {cn} FROM {tn}".format(tn="alembic_version",
cn="version_num"))
db_version = c.fetchone()[0]
logging.debug("Database version is: " + db_version)
if head_revision == db_version:
logging.info("Database is now up to date")
else:
logging.error("Database is still out of date. Head is " +
head_revision + " and database is " + db_version +
". Exiting arm.")
sys.exit()
def _fixture(self, **kw):
script = ScriptDirectory.from_config(self.cfg)
env = EnvironmentContext(self.cfg, script, **kw)
return env
def alembic_script(self):
"""Context manager to return an instance of an Alembic `ScriptDirectory`."""
from alembic.script import ScriptDirectory
with self.alembic_config() as config:
yield ScriptDirectory.from_config(config)
def head(self):
script = ScriptDirectory.from_config(self.config)
return script.get_current_head()
def get_revisions(rootdir, db_url):
revisions_dir = ScriptDirectory.from_config(
get_alembic_config(rootdir, db_url))
for revision in revisions_dir.walk_revisions("base", "heads"):
yield revision
def __init__(self, database):
self._database = database
self._script = ScriptDirectory.from_config(alembic_cfg)
def test_setting(self):
self.cfg.set_main_option('output_encoding', 'latin-1')
script = ScriptDirectory.from_config(self.cfg)
eq_(script.output_encoding, 'latin-1')
def run_migrations_online():
"""Run migrations in 'online' mode.
In this scenario we need to create an Engine
and associate a connection with the context.
"""
engine = engine_from_config(config.get_section(config.config_ini_section),
prefix='sqlalchemy.',
poolclass=pool.NullPool)
logger.info('Testing for an old alembic_version table.')
connection = engine.connect()
context.configure(connection=connection,
target_metadata=target_metadata,
version_table='alembic_version')
script_location = config.get_main_option('script_location')
found = False
mc = context.get_context()
current_db_revision = mc.get_current_revision()
script = ScriptDirectory.from_config(config)
""" If there was an existing alembic_version table, we need to
check that it's current revision is in the history for the tree
we're working with.
"""
for x in script.iterate_revisions('head', 'base'):
if x.revision == current_db_revision:
""" An alembic_versions table was found and it belongs to
this alembic tree
"""
logger.info(('An old alembic_version table at revision %s was '
'found for %s. Renaming to alembic_version_%s.'),
current_db_revision, script_location, script_location)
op = Operations(mc)
try:
with context.begin_transaction():
op.rename_table('alembic_version',
'alembic_version_%s' % script_location)
found = True
except:
logger.error(('Unable to rename alembic_version to '
'alembic_version_%s.'), script_location)
connection.close()
return
break
if not found:
logger.info('Didn\'t find an old alembic_version table.')
logger.info('Trying alembic_version_%s.' % script_location)
""" We MAY have an alembic_version table that doesn't belong to
this tree but if we still don't have an alembic_version_<tree>
table, alembic will create it.
"""
context.configure(connection=connection,
target_metadata=target_metadata,
version_table='alembic_version_' + script_location)
mc = context.get_context()
current_db_revision = mc.get_current_revision()
if current_db_revision:
logger.info('Using the alembic_version_%s table at revision %s.',
script_location, current_db_revision)
else:
logger.info('Creating new alembic_version_%s table.', script_location)
try:
with context.begin_transaction():
context.run_migrations()
finally:
connection.close()
def get_schema_version():
"""Returns the version of the schema in this library"""
config = get_alembic_config()
script = ScriptDirectory.from_config(config)
with EnvironmentContext(config=config, script=script) as env:
return int(env.get_head_revision())
def test_default(self):
script = ScriptDirectory.from_config(self.cfg)
eq_(script.output_encoding, 'utf-8')
def check_branch_exists(self, branch: str) -> bool:
script = ScriptDirectory.from_config(self.config)
bases = script.get_bases()
branches = list(
chain(*[script.get_revision(x).branch_labels for x in bases]))
return branch in branches
def get_revisions() -> Iterable["Script"]:
config = _get_alembic_config()
script = ScriptDirectory.from_config(config)
yield from script.walk_revisions()
def test_multi_head(config, config_dir, entry_points: Path, alembic_config,
engine, refresh_db):
quetz_migrations_path = Path(config_dir) / "migrations"
quetz_versions_path = quetz_migrations_path / "versions"
alembic_config.config_file_name = quetz_migrations_path / "alembic.ini"
os.makedirs(quetz_versions_path)
plugin_versions_path = Path(entry_points) / "versions"
with open(quetz_migrations_path / "env.py", "w") as fid:
fid.write(alembic_env)
with open(quetz_migrations_path / "script.py.mako", "w") as fid:
fid.write(script_mako)
with open(quetz_versions_path / "0000_initial.py", 'w') as fid:
fid.write(quetz_rev)
alembic_config.set_main_option(
"version_locations",
" ".join(map(str, [plugin_versions_path, quetz_versions_path])),
)
alembic_config.set_main_option("script_location",
str(quetz_migrations_path))
cli._run_migrations(alembic_config=alembic_config)
# initialize a plugin
cli._make_migrations(
None,
message="test revision",
plugin_name="quetz-plugin",
initialize=True,
alembic_config=alembic_config,
)
cli._run_migrations(alembic_config=alembic_config)
rev_file = next((plugin_versions_path).glob("*test_revision.py"))
with open(rev_file) as fid:
content = fid.read()
assert 'down_revision = None' in content
assert "depends_on = 'quetz'" in content
import re
m = re.search("revision = '(.*)'", content)
assert m
plugin_rev_1 = m.groups()[0]
# second revision quetz
cli._make_migrations(
None,
message="test revision",
alembic_config=alembic_config,
)
cli._run_migrations(alembic_config=alembic_config)
rev_file = next((quetz_versions_path).glob("*test_revision.py"))
with open(rev_file) as fid:
content = fid.read()
assert "down_revision = '0000'" in content
# second revision plugin
cli._make_migrations(
None,
message="plugin rev 2",
plugin_name="quetz-plugin",
alembic_config=alembic_config,
)
rev_file = next((plugin_versions_path).glob("*plugin_rev_2.py"))
with open(rev_file) as fid:
content = fid.read()
assert f"down_revision = '{plugin_rev_1}'" in content
cli._run_migrations(alembic_config=alembic_config)
# check heads
script_directory = ScriptDirectory.from_config(alembic_config)
heads = script_directory.get_revisions("heads")
assert len(heads) == 2
for p in (plugin_versions_path).glob("*.py"):
os.remove(p)
try:
engine.execute("DROP TABLE alembic_version")
except sa.exc.DatabaseError:
pass
def alembic_scriptdir(self):
from alembic.script import ScriptDirectory
return ScriptDirectory.from_config(self.alembic_config)
def alembic_db_revision(self):
return ScriptDirectory.from_config(
self.alembic_config).as_revision_number("head")
def connect(db_url):
"""Connect to the database using an environment variable.
"""
logger.info("Connecting to SQL database %r", db_url)
kwargs = {}
if db_url.startswith('sqlite:'):
kwargs['connect_args'] = {'check_same_thread': False}
engine = create_engine(db_url, **kwargs)
# logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)
if db_url.startswith('sqlite:'):
@sqlalchemy.event.listens_for(sqlalchemy.engine.Engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
alembic_cfg = alembic.config.Config()
alembic_cfg.set_main_option('script_location', 'taguette:migrations')
alembic_cfg.set_main_option('sqlalchemy.url', db_url)
conn = engine.connect()
if not engine.dialect.has_table(conn, Project.__tablename__):
logger.warning("The tables don't seem to exist; creating")
Base.metadata.create_all(bind=engine)
# Mark this as the most recent Alembic version
alembic.command.stamp(alembic_cfg, "head")
# Set SQLite's "application ID"
if db_url.startswith('sqlite:'):
conn.execute("PRAGMA application_id=0x54677474;") # 'Tgtt'
else:
# Perform Alembic migrations if needed
context = MigrationContext.configure(conn)
current_rev = context.get_current_revision()
scripts = ScriptDirectory.from_config(alembic_cfg)
if [current_rev] != scripts.get_heads():
if db_url.startswith('sqlite:'):
logger.warning("The database schema used by Taguette has "
"changed! We will try to update your workspace "
"automatically.")
assert db_url.startswith('sqlite:///')
assert os.path.exists(db_url[10:])
backup = db_url[10:] + '.bak'
shutil.copy2(db_url[10:], backup)
logger.warning(
"A backup copy of your database file has been "
"created. If the update goes horribly wrong, "
"make sure to keep that file, and let us know: "
"%s", backup)
alembic.command.upgrade(alembic_cfg, 'head')
else:
logger.critical(
"The database schema used by Taguette has "
"changed! Because you are not using SQLite, "
"we will not attempt a migration "
"automatically; back up your data and use "
"`taguette --database=%s migrate` if you want "
"to proceed.", db_url)
sys.exit(3)
else:
logger.info("Database is up to date: %s", current_rev)
# Record to Prometheus
conn.close()
conn = engine.connect()
revision = MigrationContext.configure(conn).get_current_revision()
PROM_DATABASE_VERSION.labels(revision).set(1)
DBSession = sessionmaker(bind=engine)
return DBSession
def _test_001_revisions(self):
self.a = a = util.rev_id()
self.b = b = util.rev_id()
self.c = 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 get_migration_heads(config: AlembicConfig) -> Tuple[Script]:
script = ScriptDirectory.from_config(config)
return script.get_revisions("heads")
def check_db_revision(engine):
"""Check the JupyterHub database revision
After calling this function, an alembic tag is guaranteed to be stored in the db.
- Checks the alembic tag and raises a ValueError if it's not the current revision
- If no tag is stored (Bug in Hub prior to 0.8),
guess revision based on db contents and tag the revision.
- Empty databases are tagged with the current revision
"""
# Check database schema version
current_table_names = set(engine.table_names())
my_table_names = set(Base.metadata.tables.keys())
from .dbutil import _temp_alembic_ini
with _temp_alembic_ini(engine.url) as ini:
cfg = alembic.config.Config(ini)
scripts = ScriptDirectory.from_config(cfg)
head = scripts.get_heads()[0]
base = scripts.get_base()
if not my_table_names.intersection(current_table_names):
# no tables have been created, stamp with current revision
app_log.debug("Stamping empty database with alembic revision %s",
head)
alembic.command.stamp(cfg, head)
return
if 'alembic_version' not in current_table_names:
# Has not been tagged or upgraded before.
# we didn't start tagging revisions correctly except during `upgrade-db`
# until 0.8
# This should only occur for databases created prior to JupyterHub 0.8
msg_t = "Database schema version not found, guessing that JupyterHub %s created this database."
if 'spawners' in current_table_names:
# 0.8
app_log.warning(msg_t, '0.8.dev')
rev = head
elif 'services' in current_table_names:
# services is present, tag for 0.7
app_log.warning(msg_t, '0.7.x')
rev = 'af4cbdb2d13c'
else:
# it's old, mark as first revision
app_log.warning(msg_t, '0.6 or earlier')
rev = base
app_log.debug("Stamping database schema version %s", rev)
alembic.command.stamp(cfg, rev)
# check database schema version
# it should always be defined at this point
alembic_revision = engine.execute(
'SELECT version_num FROM alembic_version').first()[0]
if alembic_revision == head:
app_log.debug("database schema version found: %s", alembic_revision)
pass
else:
raise DatabaseSchemaMismatch(
"Found database schema version {found} != {head}. "
"Backup your database and run `jupyterhub upgrade-db`"
" to upgrade to the latest schema.".format(
found=alembic_revision,
head=head,
))
def init_model(sqlalchemy_url: str, drop: bool = False, check_version: bool = True):
"""
Initializes the tables and classes of the model using configured engine.
:param sqlalchemy_url: The database URI.
:param drop: Whether to drop the tables first.
:param check_version: Whether to ensure that the database version is up-to-date.
"""
engine = create_engine(
sqlalchemy_url, pool_recycle=3600
) # pool_recycle is for mysql
sm = sessionmaker(autoflush=True, autocommit=False, bind=engine)
meta.engine = engine
meta.scoped_session = scoped_session(sm)
alembic_cfg = migrationsutil.create_config(sqlalchemy_url=sqlalchemy_url)
alembic_script = ScriptDirectory.from_config(alembic_cfg)
# Check to see whether the database has already been created or not.
# Based on this, we know whether we need to upgrade the database or mark the database
# as the latest version.
inspector = Inspector.from_engine(engine)
db_objects_created = len(inspector.get_table_names()) > 1
fresh_db = False
if not db_objects_created:
log.info("Database appears uninitialized, creating database tables")
meta.metadata.create_all(engine, tables=MANAGED_TABLES, checkfirst=True)
create_supplemental_db_objects(engine)
fresh_db = True
elif drop:
log.info("Dropping database tables and re-creating.")
drop_supplemental_db_objects(engine)
meta.metadata.drop_all(engine, tables=MANAGED_TABLES, checkfirst=True)
meta.metadata.create_all(engine, tables=MANAGED_TABLES)
fresh_db = True
if fresh_db:
command.stamp(alembic_cfg, "head")
else:
if check_version:
latest = alembic_script.get_current_head()
installed = migrationsutil.get_database_version()
try:
alembic_script.get_revisions(installed)
except CommandError:
warnings.warn(
"Unknown db revision {} installed, ignoring db upgrade.".format(
installed
)
)
else:
if latest != installed:
log.info(
"Installed database ({0}) does not match latest available ({1}). (UPGRADING)".format(
installed, latest
),
UserWarning,
)
command.upgrade(alembic_cfg, "head")
else:
log.info("Skipping database upgrade.")
