Python alembic.command.stamp() Examples

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 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 main(argv=sys.argv):
    if len(argv) < 2:
        usage(argv)
    config_uri = argv[1]
    options = parse_vars(argv[2:])
    setup_logging(config_uri)
    settings = load_full_settings(config_uri, options=options)

    engine = engine_from_config(settings, "sqlalchemy.")
    wait_for_db(engine)

    with engine.begin() as connection:
        initdb(connection, drop_all="--force" in options)

    app_name = urllib.parse.urlparse(config_uri).fragment

    # generate the Alembic version table and stamp it with the latest revision
    alembic_cfg = Config("alembic.ini", ini_section=app_name)
    command.stamp(alembic_cfg, "head") 

def create_schema(config=None, engine=None):
    """Create database schema from models description.

    Can be used for initial installation instead of upgrade('head').
    """
    if engine is None:
        engine = enginefacade.writer.get_engine()

    # NOTE(viktors): If we will use metadata.create_all() for non empty db
    #                schema, it will only add the new tables, but leave
    #                existing as is. So we should avoid of this situation.
    if version(engine=engine) is not None:
        raise db_exc.DBMigrationError("DB schema is already under version"
                                      " control. Use upgrade() instead")

    models.Base.metadata.create_all(engine)
    stamp('head', config=config) 

def init():
    from sqlalchemy import create_engine
    from moxie.models import Base
    from moxie.core import DATABASE_URL
    engine = create_engine(DATABASE_URL)
    for table in Base.metadata.tables:
        engine.execute("DROP TABLE IF EXISTS \"{}\" CASCADE;".format(table))
    Base.metadata.create_all(engine)
    import os
    from alembic.config import Config
    from alembic import command
    alembic_cfg = Config(os.path.join(
        os.path.abspath(os.path.dirname(__file__)),
        "..",
        "alembic.ini"
    ))
    command.stamp(alembic_cfg, "head") 

def _upgrade_db_initialized_before_mlflow_1(engine):
    """
    Upgrades the schema of an MLflow tracking database created prior to MLflow 1.0, removing
    duplicate constraint names. This method performs a one-time update for pre-1.0 users that we
    plan to make available in MLflow 1.0 but remove in successive versions (e.g. MLflow 1.1),
    after which we will assume that effectively all databases have been initialized using the schema
    in mlflow.store.dbmodels.initial_models (with a small number of special-case databases
    initialized pre-1.0 and migrated to have the same schema as mlflow.store.dbmodels.initial_models
    via this method).
    TODO: remove this method in MLflow 1.1.
    """
    # alembic adds significant import time, so we import it lazily
    from alembic import command
    _logger.info("Updating database tables in preparation for MLflow 1.0 schema migrations ")
    alembic_dir = os.path.join(_get_package_dir(), 'temporary_db_migrations_for_pre_1_users')
    config = _get_alembic_config(str(engine.url), alembic_dir)
    command.upgrade(config, 'heads')
    # Reset the alembic version to "base" (the 'first' version) so that a) the versioning system
    # is unaware that this migration occurred and b) subsequent migrations, like the migration to
    # add metric steps, do not need to depend on this one. This allows us to eventually remove this
    # method and the associated migration e.g. in MLflow 1.1.
    command.stamp(config, "base") 

def run(self, parser, args):
        # alembic.ini uses relative paths, so set the working directory
        os.chdir(os.path.dirname(os.path.dirname(relengapi.__file__)))
        for dbname in current_app.db.database_names:
            logger.info("creating tables for database %s", dbname)
            meta = current_app.db.metadata[dbname]
            engine = current_app.db.engine(dbname)
            meta.create_all(bind=engine)

            # load the Alembic config and stamp it with the most recent rev
            config_path = os.path.join(os.path.dirname(relengapi.__file__),
                                       'alembic', dbname, 'alembic.ini')
            if os.path.isfile(config_path):
                logger.info("stamping database %s with head", dbname)
                alembic_cfg = Config(config_path)
                command.stamp(alembic_cfg, "head") 

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 build_database(echo=True, tests=False):
    """ Applies the schema to the database. Run this command once to build the database. """
    engine = session.get_engine(echo=echo)

    schema = config_get('database', 'schema', raise_exception=False)
    if schema:
        print('Schema set in config, trying to create schema:', schema)
        try:
            engine.execute(CreateSchema(schema))
        except Exception as e:
            print('Cannot create schema, please validate manually if schema creation is needed, continuing:', e)

    models.register_models(engine)

    # Put the database under version control
    alembic_cfg = Config(config_get('alembic', 'cfg'))
    command.stamp(alembic_cfg, "head") 

def upgrade_db(url, py_file, engine, name=None):
    if name is None:
        name = 'alembic'
    packagedir = os.path.dirname(os.path.abspath(py_file))
    alembic_path = os.path.join(packagedir, name)
    if not os.path.isdir(alembic_path):
        logmessage(name + " directory not found in package directory " + packagedir)
        return
    ini_file = os.path.join(packagedir, name + '.ini')
    if not os.path.isfile(ini_file):
        logmessage(name + ".ini file not found at " + ini_file)
        return
    versions_path = os.path.join(alembic_path, 'versions')
    if not os.path.isdir(versions_path):
        os.makedirs(versions_path)
    from alembic.config import Config
    from alembic import command
    alembic_cfg = Config(ini_file)
    alembic_cfg.set_main_option("sqlalchemy.url", url)
    alembic_cfg.set_main_option("script_location", alembic_path)
    if not engine.has_table('alembic_version'):
        command.stamp(alembic_cfg, "head")
    command.upgrade(alembic_cfg, "head") 

def _test_007_stamp_upgrade(self):
        command.stamp(self.cfg, self.c)
        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 upgrade(self, nocreate=False):
        from alembic import command
        from alembic import migration

        cfg = self._get_alembic_config()
        cfg.conf = self.conf
        if nocreate:
            command.upgrade(cfg, "head")
        else:
            with self.facade.writer_connection() as connection:
                ctxt = migration.MigrationContext.configure(connection)
                current_version = ctxt.get_current_revision()
                if current_version is None:
                    Base.metadata.create_all(connection)
                    command.stamp(cfg, "head")
                else:
                    command.upgrade(cfg, "head")

        try:
            with self.facade.writer() as session:
                session.add(
                    ResourceType(
                        name="generic",
                        tablename="generic",
                        state="active",
                        attributes=resource_type.ResourceTypeAttributes()))
        except exception.DBDuplicateEntry:
            pass

    # NOTE(jd) We can have deadlock errors either here or later in
    # map_and_create_tables(). We can't decorate create_resource_type()
    # directly or each part might retry later on its own and cause a
    # duplicate. And it seems there's no way to use the same session for
    # both adding the resource_type in our table and calling
    # map_and_create_tables() :-( 

def stamp(self, directory=None, revision='head', sql=False, tag=None,
              **kwargs):  # pragma: no cover
        """'stamp' the revision table with the given revision; don't run any
        migrations"""
        config = _get_config(directory)
        command.stamp(config, revision, sql=sql, tag=tag) 

def stamp(to_version='head', sql_url=None):
    """Stamp the specified version, with no migration performed."""
    alembic_cfg = init_config(sql_url)
    alembic_command.stamp(alembic_cfg, to_version) 

def make_parser(self, subparsers):
        parser = subparsers.add_parser('stamp', help=self.stamp.__doc__)
        parser.add_argument('--tag', dest='tag', default=None,
                            help=("Arbitrary 'tag' name - can be used by custom "
                                  "env.py scripts"))
        parser.add_argument('--sql', dest='sql', action='store_true', default=False,
                            help=("Don't emit SQL to database - dump to standard "
                                  "output instead"))
        parser.add_argument('revision', default=None, help="revision identifier")
        parser.add_argument('-d', '--directory', dest='directory', default=None,
                            help=("migration script directory (default is "
                                  "'migrations')"))
        return parser 

def execute(self):
        from alembic.config import Config
        from alembic import command
        create_all()
        alembic_cfg = Config(get_config('alembic.alembic_ini'))
        command.stamp(alembic_cfg, "head") 

def upgrade(self, nocreate=False):
        cfg = self._get_alembic_config()
        cfg.conf = self.conf
        if nocreate:
            command.upgrade(cfg, "head")
        else:
            engine = self._engine_facade.get_engine()
            ctxt = migration.MigrationContext.configure(engine.connect())
            current_version = ctxt.get_current_revision()
            if current_version is None:
                models.Base.metadata.create_all(engine, checkfirst=False)
                command.stamp(cfg, "head")
            else:
                command.upgrade(cfg, "head") 

def run(self, parser, args):
        self.stamp(**vars(args)) 

def test_sql_stamp_from_rev(self):
        with capture_context_buffer() as buf:
            command.stamp(self.cfg, "%s:head" % self.a, sql=True)
        assert (
            "UPDATE alembic_version "
            "SET version_num='%s' "
            "WHERE alembic_version.version_num = '%s';" % (self.c, self.a)
        ) in buf.getvalue() 

def _createtables():
    db.init_app(app)

    db.engine.execute('''CREATE EXTENSION IF NOT EXISTS "uuid-ossp";''')
    db.create_all()
    db.session.commit()

    from alembic.config import Config
    from alembic import command
    alembic_cfg = Config('alembic.ini')
    command.stamp(alembic_cfg, 'head') 

def stamp(self, revision):
        alembic_cmd.stamp(self.cfg, revision) 

def test_noerr_transaction_opened_externally(self):
        a, b, c = self._opened_transaction_fixture()

        env_file_fixture(
            """
from sqlalchemy import engine_from_config, pool

def run_migrations_online():
    connectable = engine_from_config(
        config.get_section(config.config_ini_section),
        prefix='sqlalchemy.',
        poolclass=pool.NullPool)

    with connectable.connect() as connection:
        with connection.begin() as real_trans:
            context.configure(
                connection=connection,
                transactional_ddl=False,
                transaction_per_migration=False
            )

            with context.begin_transaction():
                context.run_migrations()

run_migrations_online()

"""
        )

        command.stamp(self.cfg, c) 

def test_sql_stamp_revision_as_kw(self):
        with capture_context_buffer() as buf:
            command.stamp(self.cfg, revision="head", sql=True)
        assert (
            "INSERT INTO alembic_version (version_num) VALUES ('%s')" % self.c
            in buf.getvalue()
        ) 

def setUp(self):
        command.stamp(self.cfg, "base") 

def test_stamp_version_already_there(self):
        command.stamp(self.cfg, self.b)
        command.stamp(self.cfg, self.b)

        with self.bind.connect() as conn:
            eq_(
                conn.scalar(text("select version_num from alembic_version")),
                self.b,
            ) 

def test_stamp_existing_downgrade(self):
        command.stamp(self.cfg, self.b)
        command.stamp(self.cfg, self.a)
        with self.bind.connect() as conn:
            eq_(
                conn.scalar(text("select version_num from alembic_version")),
                self.a,
            ) 

def test_sql_stamp_from_partial_rev(self):
        with capture_context_buffer() as buf:
            command.stamp(self.cfg, "%s:head" % self.a[0:7], sql=True)
        assert (
            "UPDATE alembic_version "
            "SET version_num='%s' "
            "WHERE alembic_version.version_num = '%s';" % (self.c, self.a)
        ) in buf.getvalue() 

def test_stamp_existing_upgrade(self):
        command.stamp(self.cfg, self.a)
        command.stamp(self.cfg, self.b)
        with self.bind.connect() as conn:
            eq_(
                conn.scalar(text("select version_num from alembic_version")),
                self.b,
            ) 

def test_stamp_creates_table(self):
        command.stamp(self.cfg, "head")
        with self.bind.connect() as conn:
            eq_(
                conn.scalar(text("select version_num from alembic_version")),
                self.b,
            ) 

def test_stamp_argparser_single_rev(self):
        cmd = config.CommandLine()
        options = cmd.parser.parse_args(["stamp", self.c, "--sql"])
        with capture_context_buffer() as buf:
            cmd.run_cmd(self.cfg, options)
        assert (
            "INSERT INTO alembic_version (version_num) VALUES ('%s')" % self.c
            in buf.getvalue()
        )