Python flask_migrate.upgrade() Examples

def wait_on_postgres(retries=5, migrate=True):
    """Block until Postgres is ready (optionally, run any migrations)

    Shamelessly appropriated from https://github.com/agconti/wait-for-postgres
    """
    dsn = os.environ.get('DATABASE_URL')

    @retry(retries, exceptions=(psycopg2.OperationalError,))
    def wait():
        con = psycopg2.connect(**psycopg2.extensions.parse_dsn(dsn))
        con.close()
        log.info('Postgres is ready!')

    wait()
    if migrate:
        log.info('Running database migrations, if any')
        with application.app_context():
            flask_migrate.upgrade(directory=MIGRATION_DIR) 

def create_db():
    if not database_exists(DB_URI):
        LOG.debug("db not exist, create database")
        create_database(DB_URI)

        # Create all tables
        # Use flask-migrate instead of db.create_all()
        flask_migrate.upgrade() 

def setUp(self):
        app_ctx = manager.app.app_context()
        app_ctx.push()
        flask_migrate.init()
        flask_migrate.migrate()
        flask_migrate.upgrade()
        app_ctx.pop()
        manager.app.config['TESTING'] = True
        self.client = manager.app.test_client() 

def configure_db(app):
    """
    0.10 is the first version of ARA that ships with a stable database schema.
    We can identify a database that originates from before this by checking if
    there is an alembic revision available.
    If there is no alembic revision available, assume we are running the first
    revision which contains the latest state of the database prior to this.
    """
    db.init_app(app)
    log = logging.getLogger(app.logger_name)

    if app.config.get('ARA_AUTOCREATE_DATABASE'):
        with app.app_context():
            migrations = app.config['DB_MIGRATIONS']
            flask_migrate.Migrate(app, db, directory=migrations)
            config = app.extensions['migrate'].migrate.get_config(migrations)

            # Verify if the database tables have been created at all
            inspector = Inspector.from_engine(db.engine)
            if len(inspector.get_table_names()) == 0:
                log.info('Initializing new DB from scratch')
                flask_migrate.upgrade(directory=migrations)

            # Get current alembic head revision
            script = ScriptDirectory.from_config(config)
            head = script.get_current_head()

            # Get current revision, if available
            connection = db.engine.connect()
            context = MigrationContext.configure(connection)
            current = context.get_current_revision()

            if not current:
                log.info('Unstable DB schema, stamping original revision')
                flask_migrate.stamp(directory=migrations,
                                    revision='da9459a1f71c')

            if head != current:
                log.info('DB schema out of date, upgrading')
                flask_migrate.upgrade(directory=migrations) 

def init_db():
    "Creates datababase table (database must be created through PG client)."

    print("Creating database and tables...")
    from zou.app import app

    with app.app_context():
        import zou

        directory = os.path.join(os.path.dirname(zou.__file__), "migrations")
        flask_migrate.upgrade(directory=directory)
    print("Database and tables created.") 

def upgrade_db():
    "Upgrade database schema."

    from zou.app import app

    with app.app_context():
        import zou

        directory = os.path.join(os.path.dirname(zou.__file__), "migrations")
        flask_migrate.upgrade(directory=directory) 

def upgrade(args):
    """Upgrade database schema."""
    flask_migrate.upgrade(DIRECTORY, args['revision']) 

def init():
    "Initialize Puffin dependencies"
    wait()
    db_create()
    flask_migrate.upgrade()
    user_create("puffin")
    machine_network()
    machine_volume()
    machine_proxy()
    machine_mail() 

def parse_arguments(argv):
    """Parse command line arguments.

    :param argv: Command line arguments
    :type argv: list(str)
    :returns: Parsed arguments
    :rtype: argparse.Namespace

    """
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--postgresql-host',
                        dest='postgresql_host',
                        help='Address the database is listening on',
                        default='localhost')
    parser.add_argument('--postgresql-username',
                        dest='postgresql_username',
                        help='Username for the database connection')
    parser.add_argument('--postgresql-password',
                        dest='postgresql_password',
                        help='Password for the database connection')
    parser.add_argument('--postgresql-db-name',
                        dest='postgresql_db_name',
                        help='Database name')

    subparsers = parser.add_subparsers(help='Migration subcommands')

    downgrade_parser = subparsers.add_parser(
        'downgrade', help='Downgrade schema to target revision')
    downgrade_parser.add_argument('revision', help='Target schema revision')
    downgrade_parser.set_defaults(func=downgrade)

    upgrade_parser = subparsers.add_parser(
        'upgrade', help='Upgrade schema to target revision')
    upgrade_parser.add_argument('revision', help='Target schema revision')
    upgrade_parser.set_defaults(func=upgrade)

    current_parser = subparsers.add_parser(
        'current', help='Get current database schema revision')
    current_parser.set_defaults(func=current)

    log_levels = ['debug', 'info', 'warning', 'error', 'critical']
    parser.add_argument(
        '-l', '--log-level',
        dest='log_level',
        choices=log_levels,
        default='debug',
        help=('Log level. One of {0} or {1} '
              '(%(default)s by default)'
              .format(', '.join(log_levels[:-1]), log_levels[-1])))

    args = vars(parser.parse_args(argv))
    args['log_level'] = getattr(logging, args['log_level'].upper())
    return args 

def initial_data():
    pdns_proto = os.environ['PDNS_PROTO']
    pdns_host = os.environ['PDNS_HOST']
    pdns_port = os.environ['PDNS_PORT']
    pdns_api_url = '{0}://{1}:{2}'.format(pdns_proto, pdns_host, pdns_port)

    api_url_setting = Setting('pdns_api_url', pdns_api_url)
    api_key_setting = Setting('pdns_api_key', os.environ['PDNS_API_KEY'])
    allow_create_domain_setting = Setting('allow_user_create_domain', True)

    try:
        flask_migrate.upgrade()

        db.session.add(api_url_setting)
        db.session.add(api_key_setting)
        db.session.add(allow_create_domain_setting)

        test_user = app.config.get('TEST_USER')
        test_user_pass = app.config.get('TEST_USER_PASSWORD')
        test_admin_user = app.config.get('TEST_ADMIN_USER')
        test_admin_pass = app.config.get('TEST_ADMIN_PASSWORD')

        admin_user = User(username=test_admin_user,
                          plain_text_password=test_admin_pass,
                          email="admin@admin.com")
        msg = admin_user.create_local_user()

        if not msg:
            raise Exception("Error occurred creating user {0}".format(msg))

        ordinary_user = User(username=test_user,
                             plain_text_password=test_user_pass,
                             email="test@test.com")
        msg = ordinary_user.create_local_user()

        if not msg:
            raise Exception("Error occurred creating user {0}".format(msg))

    except Exception as e:
        print("Unexpected ERROR: {0}".format(e))
        raise e

    yield

    db.session.close()
    os.unlink(app.config['TEST_DB_LOCATION']) 

def initial_apikey_data():
    pdns_proto = os.environ['PDNS_PROTO']
    pdns_host = os.environ['PDNS_HOST']
    pdns_port = os.environ['PDNS_PORT']
    pdns_api_url = '{0}://{1}:{2}'.format(pdns_proto, pdns_host, pdns_port)

    api_url_setting = Setting('pdns_api_url', pdns_api_url)
    api_key_setting = Setting('pdns_api_key', os.environ['PDNS_API_KEY'])
    allow_create_domain_setting = Setting('allow_user_create_domain', True)

    try:
        flask_migrate.upgrade()

        db.session.add(api_url_setting)
        db.session.add(api_key_setting)
        db.session.add(allow_create_domain_setting)

        test_user_apikey = app.config.get('TEST_USER_APIKEY')
        test_admin_apikey = app.config.get('TEST_ADMIN_APIKEY')

        dummy_apikey = ApiKey(desc="dummy", role_name="Administrator")

        admin_key = dummy_apikey.get_hashed_password(
            plain_text_password=test_admin_apikey).decode('utf-8')

        admin_apikey = ApiKey(key=admin_key,
                              desc="test admin apikey",
                              role_name="Administrator")
        admin_apikey.create()

        user_key = dummy_apikey.get_hashed_password(
            plain_text_password=test_user_apikey).decode('utf-8')

        user_apikey = ApiKey(key=user_key,
                             desc="test user apikey",
                             role_name="User")
        user_apikey.create()

    except Exception as e:
        print("Unexpected ERROR: {0}".format(e))
        raise e

    yield

    db.session.close()
    os.unlink(app.config['TEST_DB_LOCATION'])