This repository has been archived on 2020-03-24. You can view files and clone it, but cannot push or open issues or pull requests.
quay/data/migrations/env.py
Joseph Schorr f6ff0d6ca0 Add support for populating test data during migration testing
This change ensures that the tables in the database during migration have at least one row of "real" data, which should help catch issues in the future where we forget to set column defaults and other such schema oversights that can only be truly tested with non-empty tables

Fixes https://jira.coreos.com/browse/QUAY-913
2018-06-18 14:14:32 -04:00

125 lines
3.9 KiB
Python

from __future__ import with_statement
import logging
import os
from alembic import context
from alembic.script.revision import ResolutionError
from alembic.util import CommandError
from sqlalchemy import engine_from_config, pool
from logging.config import fileConfig
from urllib import unquote, quote
from peewee import SqliteDatabase
from data.database import all_models, db
from data.migrations.tester import NoopTester, PopulateTestDataTester
from app import app
from data.model.sqlalchemybridge import gen_sqlalchemy_metadata
from release import GIT_HEAD, REGION, SERVICE
from util.morecollections import AttrDict
config = context.config
config.set_main_option('sqlalchemy.url', unquote(app.config['DB_URI']))
# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name:
fileConfig(config.config_file_name)
logger = logging.getLogger(__name__)
# add your model's MetaData object here
# for 'autogenerate' support
# from myapp import mymodel
# target_metadata = mymodel.Base.metadata
target_metadata = gen_sqlalchemy_metadata(all_models)
tables = AttrDict(target_metadata.tables)
# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.
def get_tester():
""" Returns the tester to use. We only return the tester that populates data
if the TEST_MIGRATE env var is set to `true` AND we make sure we're not
connecting to a production database.
"""
if os.environ.get('TEST_MIGRATE', '') == 'true':
url = unquote(app.config['DB_URI'])
if url.find('.quay.io') < 0:
return PopulateTestDataTester()
return NoopTester()
def run_migrations_offline():
"""Run migrations in 'offline' mode.
This configures the context with just a URL
and not an Engine, though an Engine is acceptable
here as well. By skipping the Engine creation
we don't even need a DBAPI to be available.
Calls to context.execute() here emit the given string to the
script output.
"""
url = unquote(app.config['DB_URI'])
context.configure(url=url, target_metadata=target_metadata, transactional_ddl=True)
with context.begin_transaction():
context.run_migrations(tables=tables, tester=get_tester())
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.
"""
if isinstance(db.obj, SqliteDatabase) and not 'GENMIGRATE' in os.environ and not 'DB_URI' in os.environ:
print ('Skipping Sqlite migration!')
return
engine = engine_from_config(
config.get_section(config.config_ini_section),
prefix='sqlalchemy.',
poolclass=pool.NullPool)
connection = engine.connect()
context.configure(
connection=connection,
target_metadata=target_metadata,
transactional_ddl=False,
)
try:
with context.begin_transaction():
try:
context.run_migrations(tables=tables, tester=get_tester())
except (CommandError, ResolutionError) as ex:
if 'No such revision' not in str(ex):
raise
if not REGION or not GIT_HEAD:
raise
from data.model.release import get_recent_releases
# ignore revision error if we're running the previous release
releases = list(get_recent_releases(SERVICE, REGION).offset(1).limit(1))
if releases and releases[0].version == GIT_HEAD:
logger.warn('Skipping database migration because revision not found')
else:
raise
finally:
connection.close()
if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()