Python alembic.command.upgrade() Examples

def create_all_or_upgrade_db(engine, db_url):
    # alembic add a lot of import time, so we lazy import
    from alembic import command
    from alembic.config import Config
    from sqlalchemy import inspect

    alembic_config_file = os.path.join(os.path.dirname(__file__), 'alembic.ini')
    alembic_config = Config(alembic_config_file)
    alembic_config.set_main_option('sqlalchemy.url', db_url)

    inspector = inspect(engine)
    tables = inspector.get_table_names()

    if 'deployments' not in tables and 'bentos' not in tables:
        logger.debug('Creating tables')
        Base.metadata.create_all(engine)
        command.stamp(alembic_config, 'head')
    else:
        logger.debug('Upgrading tables to the latest revision')
        command.upgrade(alembic_config, 'heads') 

def init_db():
    """
    Initialize the database; call
    :py:meth:`sqlalchemy.schema.MetaData.create_all` on the metadata object.
    """
    logger.debug('Initializing database')
    # import all modules here that might define models so that
    # they will be registered properly on the metadata.  Otherwise
    # you will have to import them first before calling init_db()
    alembic_ini = pkg_resources.resource_filename(
        pkg_resources.Requirement.parse('biweeklybudget'),
        'biweeklybudget/alembic/alembic.ini'
    )
    topdir = os.path.abspath(
        os.path.join(os.path.dirname(alembic_ini), '..', '..')
    )
    logger.debug('Alembic configuration: %s', alembic_ini)
    with in_directory(topdir):
        alembic_config = Config(alembic_ini)
        script = ScriptDirectory.from_config(alembic_config)
        curr_rev = _alembic_get_current_rev(alembic_config, script)
        head_rev = script.get_revision("head").revision
        if curr_rev is None:
            # alembic not initialized at all; stamp with current version
            logger.warning(
                'Alembic not setup; creating all models and stamping'
            )
            logger.debug('Creating all models')
            Base.metadata.create_all(engine)
            command.stamp(alembic_config, "head")
            logger.debug("DB stamped at %s", head_rev)
        elif curr_rev != head_rev:
            logger.warning("Alembic head is %s but this DB is at %s; "
                           "running migrations", head_rev, curr_rev)
            command.upgrade(alembic_config, "head")
            logger.info("Migrations complete")
        else:
            logger.debug('Alembic is at the correct head version (%s)',
                         curr_rev)
    logger.debug('Done initializing DB')
    init_event_listeners(db_session, engine) 

def oneshot(cli_ctx, alembic_config):
    '''
    Set up your database with one-shot schema migration instead of
    iterating over multiple revisions if there is no existing database.
    It uses alembic.ini to configure database connection.

    Reference: http://alembic.zzzcomputing.com/en/latest/cookbook.html
               #building-an-up-to-date-database-from-scratch
    '''
    with cli_ctx.logger:
        alembic_cfg = Config(alembic_config)
        sa_url = alembic_cfg.get_main_option('sqlalchemy.url')

        engine = sa.create_engine(sa_url)
        engine.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp";')

        with engine.begin() as connection:
            context = MigrationContext.configure(connection)
            current_rev = context.get_current_revision()

        if current_rev is None:
            # For a fresh clean database, create all from scratch.
            # (it will raise error if tables already exist.)
            log.info('Detected a fresh new database.')
            log.info('Creating tables...')
            with engine.begin() as connection:
                alembic_cfg.attributes['connection'] = connection
                metadata.create_all(engine, checkfirst=False)
                log.info('Stamping alembic version to head...')
                command.stamp(alembic_cfg, 'head')
        else:
            # If alembic version info is already available, perform incremental upgrade.
            log.info('Detected an existing database.')
            log.info('Performing schema upgrade to head...')
            with engine.begin() as connection:
                alembic_cfg.attributes['connection'] = connection
                command.upgrade(alembic_cfg, 'head')

        log.info("If you don't need old migrations, delete them and set "
                 "\"down_revision\" value in the earliest migration to \"None\".") 

def init_development_data(context, upgrade_db=True, skip_on_failure=False):
    """
    Fill a database with development data like default users.
    """
    if upgrade_db:
        context.invoke_execute(context, 'app.db.upgrade')

    log.info("Initializing development data...")

    from migrations import initial_development_data
    try:
        initial_development_data.init()
    except AssertionError as exception:
        if not skip_on_failure:
            log.error("%s", exception)
        else:
            log.debug(
                "The following error was ignored due to the `skip_on_failure` flag: %s",
                exception
            )
            log.info("Initializing development data step is skipped.")
    else:
        log.info("Fixtures have been successfully applied.") 

def _test_005_nextrev(self):
        script = env.generate_revision(
            def_, "this is the next rev", refresh=True
        )
        assert os.access(
            os.path.join(
                env.dir, "versions", "%s_this_is_the_next_rev.py" % def_
            ),
            os.F_OK,
        )
        eq_(script.revision, def_)
        eq_(script.down_revision, abc)
        eq_(env.get_revision(abc).nextrev, set([def_]))
        assert script.module.down_revision == abc
        assert callable(script.module.upgrade)
        assert callable(script.module.downgrade)
        eq_(env.get_heads(), [def_])
        eq_(env.get_base(), abc) 

def test_renders_added_directives_no_autogen(self):
        m = sa.MetaData()

        def process_revision_directives(context, rev, generate_revisions):
            generate_revisions[0].upgrade_ops.ops.append(
                ops.CreateIndexOp("some_index", "some_table", ["a", "b"])
            )

        with self._env_fixture(process_revision_directives, m):
            rev = command.revision(
                self.cfg, message="some message", head="model1@head", sql=True
            )

        with mock.patch.object(rev.module, "op") as op_mock:
            rev.module.upgrade()
        eq_(
            op_mock.mock_calls,
            [
                mock.call.create_index(
                    "some_index", "some_table", ["a", "b"], unique=False
                )
            ],
        ) 

def test_create_rev_autogen_db_not_up_to_date_multi_heads(self):
        self._env_fixture()
        command.revision(self.cfg)
        rev2 = command.revision(self.cfg)
        rev3a = command.revision(self.cfg)
        command.revision(self.cfg, head=rev2.revision, splice=True)
        command.upgrade(self.cfg, "heads")
        command.revision(self.cfg, head=rev3a.revision)

        assert_raises_message(
            util.CommandError,
            "Target database is not up to date.",
            command.revision,
            self.cfg,
            autogenerate=True,
        ) 

def test_create_rev_plain_db_not_up_to_date_multi_heads(self):
        self._env_fixture()
        command.revision(self.cfg)
        rev2 = command.revision(self.cfg)
        rev3a = command.revision(self.cfg)
        command.revision(self.cfg, head=rev2.revision, splice=True)
        command.upgrade(self.cfg, "heads")
        command.revision(self.cfg, head=rev3a.revision)

        assert_raises_message(
            util.CommandError,
            "Multiple heads are present; please specify the head revision "
            "on which the new revision should be based, or perform a merge.",
            command.revision,
            self.cfg,
        ) 

def test_create_rev_autogen_need_to_select_head(self):
        self._env_fixture()
        command.revision(self.cfg)
        rev2 = command.revision(self.cfg)
        command.revision(self.cfg)
        command.revision(self.cfg, head=rev2.revision, splice=True)
        command.upgrade(self.cfg, "heads")
        # there's multiple heads present
        assert_raises_message(
            util.CommandError,
            "Multiple heads are present; please specify the head revision "
            "on which the new revision should be based, or perform a merge.",
            command.revision,
            self.cfg,
            autogenerate=True,
        ) 

def test_err_correctly_raised_on_dupe_rows_no_pk(self):
        self._env_fixture(version_table_pk=False)
        command.revision(self.cfg)
        r2 = command.revision(self.cfg)
        db = _sqlite_file_db()
        command.upgrade(self.cfg, "head")
        with db.connect() as conn:
            conn.execute(
                text("insert into alembic_version values ('%s')" % r2.revision)
            )
        assert_raises_message(
            util.CommandError,
            "Online migration expected to match one row when "
            "updating .* in 'alembic_version'; 2 found",
            command.downgrade,
            self.cfg,
            "-1",
        ) 

def _inline_enum_script(self):
        write_script(
            self.script,
            self.rid,
            """
revision = '%s'
down_revision = None

from alembic import op
from sqlalchemy.dialects.postgresql import ENUM
from sqlalchemy import Column


def upgrade():
    op.create_table("sometable",
        Column("data", ENUM("one", "two", "three", name="pgenum"))
    )


def downgrade():
    op.drop_table("sometable")
"""
            % self.rid,
        ) 

def _reset_database(request, app):
    db_name = "test_freight"

    session.close_all_sessions()

    # 9.1 does not support --if-exists
    if subprocess.call(f"psql -l | grep '{db_name}'", shell=True) == 0:
        engine = db.engine
        engine.connect().close()
        engine.dispose()
        subprocess.check_call(f"dropdb {db_name}", shell=True)

    subprocess.check_call(f"createdb -E utf-8 {db_name}", shell=True)

    command.upgrade(ALEMBIC_CONFIG, "head")
    return lambda: reset_database(request, app) 

def _test_005_upgrade(self):
        command.upgrade(self.cfg, self.b)
        db = self.bind
        assert db.dialect.has_table(db.connect(), "foo")
        assert db.dialect.has_table(db.connect(), "bar")
        assert not db.dialect.has_table(db.connect(), "bat") 

def _test_002_upgrade(self):
        command.upgrade(self.cfg, self.c)
        db = self.bind
        assert db.dialect.has_table(db.connect(), "foo")
        assert db.dialect.has_table(db.connect(), "bar")
        assert db.dialect.has_table(db.connect(), "bat") 

def test_offline_distinct_enum_create(self):
        self._distinct_enum_script()
        with capture_context_buffer() as buf:
            command.upgrade(self.cfg, self.rid, sql=True)
        assert (
            "CREATE TYPE pgenum AS ENUM "
            "('one', 'two', 'three')" in buf.getvalue()
        )
        assert "CREATE TABLE sometable (\n    data pgenum\n)" in buf.getvalue() 

def _distinct_enum_script(self):
        write_script(
            self.script,
            self.rid,
            """
revision = '%s'
down_revision = None

from alembic import op
from sqlalchemy.dialects.postgresql import ENUM
from sqlalchemy import Column


def upgrade():
    enum = ENUM("one", "two", "three", name="pgenum", create_type=False)
    enum.create(op.get_bind(), checkfirst=False)
    op.create_table("sometable",
        Column("data", enum)
    )


def downgrade():
    op.drop_table("sometable")
    ENUM(name="pgenum").drop(op.get_bind(), checkfirst=False)

"""
            % self.rid,
        ) 

def upgrade_tables(ctx):
    """ Alias for 'alembic upgrade head'.
    Upgrade to latest model version
    """
    current_session = get_current_session()
    command.upgrade(setup_alembic_config(url=current_session.url),
                    revision='head', sql=False, tag=None) 

def test_help_text(self):
        commands = {
            fn.__name__
            for fn in [getattr(command, n) for n in dir(command)]
            if inspect.isfunction(fn)
            and fn.__name__[0] != "_"
            and fn.__module__ == "alembic.command"
        }
        # make sure we found them
        assert commands.intersection(
            {"upgrade", "downgrade", "merge", "revision"}
        )

        # catch help text coming intersection
        with mock.patch("alembic.config.ArgumentParser") as argparse:
            config.CommandLine()
            for kall in argparse().add_subparsers().mock_calls:
                for sub_kall in kall.call_list():
                    if sub_kall[0] == "add_parser":
                        cmdname = sub_kall[1][0]
                        help_text = sub_kall[2]["help"]
                        if help_text:
                            commands.remove(cmdname)
                            # more than two spaces
                            assert not re.search(r"   ", help_text)

                            # no markup stuff
                            assert ":" not in help_text

                            # no newlines
                            assert "\n" not in help_text

                            # ends with a period
                            assert help_text.endswith(".")

                            # not too long
                            assert len(help_text) < 80
        assert not commands, "Commands without help text: %s" % commands 

def test_base_to_head_sql(self):
        with capture_context_buffer() as buf:
            command.upgrade(self.cfg, "base:head", sql=True)
        assert "CREATE TABLE alembic_version" in buf.getvalue()
        assert "UPDATE alembic_version" in buf.getvalue()
        assert "CREATE STEP 1" in buf.getvalue()
        assert "CREATE STEP 2" in buf.getvalue()
        assert "CREATE STEP 3" in buf.getvalue() 

def test_needs_flag(self):
        a = util.rev_id()

        script = ScriptDirectory.from_config(self.cfg)
        script.generate_revision(a, None, refresh=True)
        write_script(
            script,
            a,
            """
    revision = '%s'
    down_revision = None

    from alembic import op


    def upgrade():
        op.execute("CREATE TABLE foo(id integer)")


    def downgrade():
        op.execute("DROP TABLE foo")

    """
            % a,
            sourceless=True,
        )

        script = ScriptDirectory.from_config(self.cfg)
        eq_(script.get_heads(), [])

        self.cfg.set_main_option("sourceless", "true")
        script = ScriptDirectory.from_config(self.cfg)
        eq_(script.get_heads(), [a]) 

def test_none_to_head_sql(self):
        with capture_context_buffer() as buf:
            command.upgrade(self.cfg, "head", sql=True)
        assert "CREATE TABLE alembic_version" in buf.getvalue()
        assert "UPDATE alembic_version" in buf.getvalue()
        assert "CREATE STEP 1" in buf.getvalue()
        assert "CREATE STEP 2" in buf.getvalue()
        assert "CREATE STEP 3" in buf.getvalue() 

def test_version_from_none_insert(self):
        with capture_context_buffer() as buf:
            command.upgrade(self.cfg, self.a, sql=True)
        assert "CREATE TABLE alembic_version" in buf.getvalue()
        assert "INSERT INTO alembic_version" in buf.getvalue()
        assert "CREATE STEP 1" in buf.getvalue()
        assert "CREATE STEP 2" not in buf.getvalue()
        assert "CREATE STEP 3" not in buf.getvalue() 

def test_create_rev_autogenerate_db_not_up_to_date_post_merge(self):
        self._env_fixture()
        command.revision(self.cfg)
        rev2 = command.revision(self.cfg)
        command.revision(self.cfg)
        command.revision(self.cfg, head=rev2.revision, splice=True)
        command.upgrade(self.cfg, "heads")
        command.merge(self.cfg, "heads")
        assert_raises_message(
            util.CommandError,
            "Target database is not up to date.",
            command.revision,
            self.cfg,
            autogenerate=True,
        ) 

def test_create_rev_autogenerate_post_merge(self):
        self._env_fixture()
        command.revision(self.cfg)
        rev2 = command.revision(self.cfg)
        command.revision(self.cfg)
        command.revision(self.cfg, head=rev2.revision, splice=True)
        command.merge(self.cfg, "heads")
        command.upgrade(self.cfg, "heads")
        command.revision(self.cfg, autogenerate=True) 

def test_create_rev_plain_need_to_select_head(self):
        self._env_fixture()
        command.revision(self.cfg)
        rev2 = command.revision(self.cfg)
        command.revision(self.cfg)
        command.revision(self.cfg, head=rev2.revision, splice=True)
        command.upgrade(self.cfg, "heads")
        # there's multiple heads present
        assert_raises_message(
            util.CommandError,
            "Multiple heads are present; please specify the head revision "
            "on which the new revision should be based, or perform a merge.",
            command.revision,
            self.cfg,
        ) 

def upgrade(config, revision: str, sql: bool, tag: Optional[str]):
    """Upgrade to a later version."""

    bot = Bot(config)

    directory = os.path.join('yui', 'migrations')
    c = Config(os.path.join(directory, 'alembic.ini'))
    c.set_main_option('script_location', directory)
    c.set_main_option('sqlalchemy.url', bot.config.DATABASE_URL)
    c.attributes['Base'] = bot.orm_base

    command.upgrade(c, revision, sql=sql, tag=tag) 

def upgrade(self, revision):
        alembic_command.upgrade(self.config(), revision) 

def test_install(self):
        self.upgrade('heads') 

def test_upgrade_heads_downgrade_base(self):
        self.upgrade('heads')
        self.downgrade('base') 

def _test_006_upgrade_again(self):
        command.upgrade(self.cfg, self.b)
        db = self.bind
        assert db.dialect.has_table(db.connect(), "foo")
        assert db.dialect.has_table(db.connect(), "bar")
        assert not db.dialect.has_table(db.connect(), "bat")