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.")