def check_alembic_version():
config = Config()
config.set_main_option("script_location", "migrations")
script = ScriptDirectory.from_config(config)
head_revision = script.get_current_head()
with get_app().app_context():
try:
conn = db.session.connection()
except ImportError:
if not faraday.server.config.database.connection_string:
print("\n\nNo database configuration found. Did you execute \"faraday-manage initdb\"? \n\n")
sys.exit(1)
except sqlalchemy.exc.OperationalError:
print("Bad Credentials, please check the .faraday/config/server.ini")
sys.exit(1)
context = MigrationContext.configure(conn)
current_revision = context.get_current_revision()
if head_revision != current_revision:
version_path = faraday.server.config.FARADAY_BASE / 'migrations'\
/ 'versions'
if list(version_path.glob(f'{current_revision}_*.py')):
print('--' * 20)
print('Missing migrations, please execute: \n\n')
print('faraday-manage migrate')
sys.exit(1)
else:
logger.warning(
"You are using an unknown schema version. If you are a "
"developer, this probably happened because you used branch "
"with a schema migration not merged yet. If you are a "
"normal user, consider reporting this bug back to us"
)
def drop_airflow_models(connection):
"""
Drops all airflow models.
@param connection:
@return: None
"""
from airflow.models.base import Base
# Drop connection and chart - those tables have been deleted and in case you
# run resetdb on schema with chart or users table will fail
chart = Table('chart', Base.metadata)
chart.drop(settings.engine, checkfirst=True)
user = Table('user', Base.metadata)
user.drop(settings.engine, checkfirst=True)
users = Table('users', Base.metadata)
users.drop(settings.engine, checkfirst=True)
dag_stats = Table('dag_stats', Base.metadata)
dag_stats.drop(settings.engine, checkfirst=True)
Base.metadata.drop_all(connection)
# we remove the Tables here so that if resetdb is run metadata does not keep the old tables.
Base.metadata.remove(dag_stats)
Base.metadata.remove(users)
Base.metadata.remove(user)
Base.metadata.remove(chart)
# alembic adds significant import time, so we import it lazily
from alembic.migration import MigrationContext # noqa
migration_ctx = MigrationContext.configure(connection)
version = migration_ctx._version # noqa pylint: disable=protected-access
if version.exists(connection):
version.drop(connection)
def initialize(settings, options):
engine = engine_from_config(settings, 'sqlalchemy.')
config = Configurator(settings=settings)
pyramid_dogpile_cache.includeme(config)
from gengine.metadata import (init_session, init_declarative_base, init_db)
init_caches()
init_session()
init_declarative_base()
init_db(engine)
from gengine.metadata import (Base, DBSession)
if options.get("reset_db", False):
Base.metadata.drop_all(engine)
engine.execute("DROP SCHEMA IF EXISTS public CASCADE")
engine.execute("CREATE SCHEMA IF NOT EXISTS public")
from alembic.config import Config
from alembic import command
from alembic.runtime.migration import MigrationContext
alembic_cfg = Config(attributes={'engine': engine, 'schema': 'public'})
script_location = os.path.join(
os.path.dirname(
os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
'app/alembic')
alembic_cfg.set_main_option("script_location", script_location)
context = MigrationContext.configure(engine.connect())
current_rev = context.get_current_revision()
if not current_rev:
#init
from gengine.app import model
tables = [
t for name, t in model.__dict__.items() if isinstance(t, Table)
]
Base.metadata.create_all(engine, tables=tables)
command.stamp(alembic_cfg, "head")
if options.get("populate_demo", False):
populate_demo(DBSession)
admin_user = options.get("admin_user", False)
admin_password = options.get("admin_password", False)
if admin_user and admin_password:
create_user(DBSession=DBSession,
user=admin_user,
password=admin_password)
else:
#upgrade
command.upgrade(alembic_cfg, 'head')
engine.dispose()
def spinner(timeout):
package_dir = os.path.dirname(importlib.util.find_spec('airflow').origin)
directory = os.path.join(package_dir, 'migrations')
config = Config(os.path.join(package_dir, 'alembic.ini'))
config.set_main_option('script_location', directory)
config.set_main_option('sqlalchemy.url', settings.SQL_ALCHEMY_CONN.replace('%', '%%'))
script_ = ScriptDirectory.from_config(config)
with settings.engine.connect() as connection:
context = MigrationContext.configure(connection)
ticker = 0
while True:
source_heads = set(script_.get_heads())
db_heads = set(context.get_current_heads())
if source_heads == db_heads:
logging.info('Airflow version: {}'.format(version.version))
logging.info('Current heads: {}'.format(db_heads))
break
elif ticker >= timeout:
raise TimeoutError(
"There are still unapplied migrations after {} seconds".format(timeout)
)
ticker += 1
time.sleep(1)
logging.info('Waiting for migrations... {} second(s)'.format(ticker))
def setUp(self):
self.conn = config.db.connect()
self.metadata = MetaData()
t1 = Table('foo',
self.metadata,
Column('id', Integer, primary_key=True),
Column('data', String(50)),
Column('x', Integer),
mysql_engine='InnoDB')
t1.create(self.conn)
self.conn.execute(t1.insert(), [{
"id": 1,
"data": "d1",
"x": 5
}, {
"id": 2,
"data": "22",
"x": 6
}, {
"id": 3,
"data": "8.5",
"x": 7
}, {
"id": 4,
"data": "9.46",
"x": 8
}, {
"id": 5,
"data": "d5",
"x": 9
}])
context = MigrationContext.configure(self.conn)
self.op = Operations(context)
def render_python_code(
up_or_down_op: "UpgradeOps",
sqlalchemy_module_prefix: str = "sa.",
alembic_module_prefix: str = "op.",
render_as_batch: bool = False,
imports: Tuple[str, ...] = (),
render_item: None = None,
migration_context: Optional["MigrationContext"] = None,
) -> str:
"""Render Python code given an :class:`.UpgradeOps` or
:class:`.DowngradeOps` object.
This is a convenience function that can be used to test the
autogenerate output of a user-defined :class:`.MigrationScript` structure.
"""
opts = {
"sqlalchemy_module_prefix": sqlalchemy_module_prefix,
"alembic_module_prefix": alembic_module_prefix,
"render_item": render_item,
"render_as_batch": render_as_batch,
}
if migration_context is None:
from ..runtime.migration import MigrationContext
from sqlalchemy.engine.default import DefaultDialect
migration_context = MigrationContext.configure(
dialect=DefaultDialect())
autogen_context = AutogenContext(migration_context, opts=opts)
autogen_context.imports = set(imports)
return render._indent(
render._render_cmd_body(up_or_down_op, autogen_context))
def add_new_columns(self):
"""
Adds columns present in df but not in the SQL table
for given instance of PandasSpecialEngine.
Notes
-----
Sadly, it seems that we cannot create JSON columns.
"""
# create deepcopies of the column because we are going to unbound
# them for the table model (otherwise alembic would think we add
# a column that already exists in the database)
cols_to_add = [deepcopy(col) for col in self.table.columns
if col.name not in self.get_db_columns_names()]
# check columns are not index levels
if any((c.name in self.df.index.names for c in cols_to_add)):
raise ValueError(('Cannot add any column that is part of the df index!\n'
"You'll have to update your table primary key or change your "
"df index"))
with self.engine.connect() as con:
ctx = MigrationContext.configure(con)
op = Operations(ctx)
for col in cols_to_add:
col.table = None # Important! unbound column from table
op.add_column(self.table.name, col, schema=self.schema)
log(f"Added column {col} (type: {col.type}) in table {self.table.name} "
f'(schema="{self.schema}")')
def check_migrations(timeout):
"""
Function to wait for all airflow migrations to complete.
@param timeout:
@return:
"""
from alembic.runtime.migration import MigrationContext
from alembic.script import ScriptDirectory
config = _get_alembic_config()
script_ = ScriptDirectory.from_config(config)
with settings.engine.connect() as connection:
context = MigrationContext.configure(connection)
ticker = 0
while True:
source_heads = set(script_.get_heads())
db_heads = set(context.get_current_heads())
if source_heads == db_heads:
break
if ticker >= timeout:
raise TimeoutError(
f"There are still unapplied migrations after {ticker} seconds."
)
ticker += 1
time.sleep(1)
log.info('Waiting for migrations... %s second(s)', ticker)
def alembic_tests():
print(f'Alembic tests')
conn = engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
try:
op.drop_table('waste')
except:
pass
t = op.create_table(
'waste',
Column('bools', sa.Boolean),
Column('ubytes', sa.Tinyint),
Column('shorts', sa.SmallInteger),
Column('ints', sa.Integer),
Column('bigints', sa.BigInteger),
Column('floats', sa.REAL),
Column('doubles', sa.Float),
Column('dates', sa.Date),
Column('datetimes', sa.DateTime),
Column('varchars', sa.String(10)),
Column('nvarchars', sa.UnicodeText),
Column('numerics', sa.Numeric(38, 10)),
)
data = [{
'bools': True,
'ubytes': 5,
'shorts': 55,
'ints': 555,
'bigints': 5555,
'floats': 5.0,
'doubles': 5.5555555,
'dates': date(2012, 11, 23),
'datetimes': datetime(2012, 11, 23, 16, 34, 56),
'varchars': 'bla',
'nvarchars': 'bla2',
'numerics': Decimal("1.1")
}, {
'bools': False,
'ubytes': 6,
'shorts': 66,
'ints': 666,
'bigints': 6666,
'floats': 6.0,
'doubles': 6.6666666,
'dates': date(2012, 11, 24),
'datetimes': datetime(2012, 11, 24, 16, 34, 57),
'varchars': 'bla',
'nvarchars': 'bla2',
'numerics': Decimal("-1.1")
}]
op.bulk_insert(t, data)
res = engine.execute('select * from waste').fetchall()
assert (res == [tuple(dikt.values()) for dikt in data])
def upgrade(self):
conn = self.engine.connect()
cfg = Config()
# this connection is used in `env.py` for the migrations
cfg.attributes["connection"] = conn
cfg.set_main_option("script_location", str(MIGRATION_DIR))
script = ScriptDirectory.from_config(cfg)
# let's make sure we actually need to migrate
context = MigrationContext.configure(conn)
try:
# try to find the locked version
HEAD = LOCK_PATH.open().read().strip()
self.ensure_migration_needed(script, context, HEAD)
click.secho(f"Upgrading database to {HEAD}")
command.upgrade(cfg, HEAD)
except FileNotFoundError:
raise click.ClickException(
"Cannot upgrade the system database, revision lock not found."
)
except MigrationUneededException:
click.secho(f"System database up-to-date.")
except Exception as err:
click.secho(
f"Cannot upgrade the system database. It might be corrupted or was created before database migrations where introduced (v0.34.0)",
fg="yellow",
err=True,
)
logging.exception(err)
finally:
conn.close()
def adapt_dtype_of_empty_db_columns(self):
"""
Changes the data types of empty columns in the SQL table defined
in given instance of a PandasSpecialEngine.
This should only happen in case of data type mismatches.
This means with columns for which the sqlalchemy table
model for df and the model for the SQL table have different data types.
"""
empty_db_columns = self.get_empty_columns()
db_table = self.get_db_table_schema()
# if column does not have value in db and there are values
# in the frame then change the column type if needed
for col in empty_db_columns:
# check if the column also exists in df
if col.name not in self.df.columns:
continue
# check same type
orig_type = db_table.columns[col.name].type.compile(self.engine.dialect)
dest_type = self.table.columns[col.name].type.compile(self.engine.dialect)
# remove character count e.g. "VARCHAR(50)" -> "VARCHAR"
orig_type = RE_CHARCOUNT_COL_TYPE.sub('', orig_type)
dest_type = RE_CHARCOUNT_COL_TYPE.sub('', dest_type)
# if same type or we want to insert TEXT instead of JSON continue
# (JSON is not supported on some DBs so it's normal to have TEXT instead)
if ((orig_type == dest_type) or
((orig_type == 'JSON') and (dest_type == 'TEXT'))):
continue
# grab the col/index from the df
# so we can check if there are any values
if col.name in self.df.index.names:
df_col = self.df.index.get_level_values(col.name)
else:
df_col = self.df[col.name]
if df_col.notna().any():
# raise error if we have to modify the dtype but we have a SQlite engine
# (SQLite does not support data type alteration)
if self._db_type == 'sqlite':
raise ValueError('SQlite does not support column data type alteration!')
with self.engine.connect() as con:
ctx = MigrationContext.configure(con)
op = Operations(ctx)
new_col = self.table.columns[col.name]
# check if postgres (in which case we have to use "using" syntax
# to alter columns data types)
if self._db_type == 'postgres':
escaped_col = str(new_col.compile(dialect=self.engine.dialect))
compiled_type = new_col.type.compile(dialect=self.engine.dialect)
alter_kwargs = {'postgresql_using':f'{escaped_col}::{compiled_type}'}
else:
alter_kwargs = {}
op.alter_column(table_name=self.table.name,
column_name=new_col.name,
type_=new_col.type,
schema=self.schema,
**alter_kwargs)
log(f"Changed type of column {new_col.name} "
f"from {col.type} to {new_col.type} "
f'in table {self.table.name} (schema="{self.schema}")')
def main():
with connectable.connect() as connection:
migration_context = MigrationContext.configure(
connection, opts={"include_object": include_object})
diff = compare_metadata(migration_context, db.metadata)
if len(diff) > 0:
raise Exception(
f"The database is not aligned with application model, create a migration or change the model to remove the following diff: {diff}"
)
def init_online(app: Flask):
db: SQLAlchemy = app.extensions['sqlalchemy'].db
conn = db.engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
try:
op.drop_table('seeds')
except:
pass
op.create_table('seeds', sa.Column('file', sa.String(255), primary_key=True))
def update(self):
code = self.migration_code
migrate = _get_migrate_function(code)
schema_name = self.context.content_metadata().schema
with self.engine.connect() as conn:
migration_context = MigrationContext.configure(
conn, opts={"include_schemas": False})
op = Operations(migration_context)
if not self.engine.dialect.has_schema(self.engine, schema_name):
conn.execute(CreateSchema(schema_name))
migrate(op)
def drop_models(self) -> None:
"""
Drop rds models
"""
LOGGER.info('Dropping rds models')
Base.metadata.drop_all(self.engine)
with self.engine.connect() as conn:
migration_ctx = MigrationContext.configure(conn)
version = migration_ctx._version
if version.exists(conn):
version.drop(conn)
def get_current_revision(db_url):
# create connecton
engine = create_engine(db_url)
connection = engine.connect()
durl = sqla_url.make_url(db_url)
mgrt_cxt = MigrationContext.configure(connection)
#mgrt_cxt = MigrationContext(environment_context=env_cxt,
#connection=connection,
#dialect=durl.get_dialect()(),
# opts=None
# )#
mgrt_cxt.get_current_revision()
def get_current_revision(db_url):
# create connecton
engine = create_engine(db_url)
connection = engine.connect()
durl = sqla_url.make_url(db_url)
mgrt_cxt = MigrationContext.configure(connection)
#mgrt_cxt = MigrationContext(environment_context=env_cxt,
#connection=connection,
#dialect=durl.get_dialect()(),
# opts=None
# )#
mgrt_cxt.get_current_revision()
def show(args):
'''Show the current schema information.'''
alembic_cfg = Config(args.config)
sa_url = alembic_cfg.get_main_option('sqlalchemy.url')
engine = sa.create_engine(sa_url)
with engine.begin() as connection:
context = MigrationContext.configure(connection)
current_rev = context.get_current_revision()
script = ScriptDirectory.from_config(alembic_cfg)
heads = script.get_heads()
head_rev = heads[0] if len(heads) > 0 else None
print(f'Current database revision: {current_rev}')
print(f'The head revision of available migrations: {head_rev}')
def get_current_revision(
database_url: str,
version_table: str = DEFAULT_ALEMBIC_VERSION_TABLE) -> str:
"""
Ask the database what its current revision is.
Arguments:
database_url: SQLAlchemy URL for the database
version_table: table name for Alembic versions
"""
engine = create_engine(database_url)
conn = engine.connect()
opts = {'version_table': version_table}
mig_context = MigrationContext.configure(conn, opts=opts)
return mig_context.get_current_revision()
def test_missing_column(self):
adapter = self._get_adapter()
adapter.build_table()
with adapter.engine.begin() as connection:
context = MigrationContext.configure(connection)
op = Operations(context)
op.drop_column(adapter.get_table().name, 'name')
doc = {
"_id": '123',
"domain": "domain",
"doc_type": "CommCareCase",
"name": 'bob'
}
with self.assertRaises(MissingColumnWarning):
adapter.best_effort_save(doc)
def upgrade_thd(engine):
with querylog.log_queries():
with sautils.withoutSqliteForeignKeys(engine):
with engine.connect() as conn:
def upgrade(rev, context):
log.msg(
f'Upgrading from {rev} to {target_revision}')
return alembic_scripts._upgrade_revs(
target_revision, rev)
context = MigrationContext.configure(
conn, opts={'fn': upgrade})
with context.begin_transaction():
context.run_migrations()
def test_missing_column(self):
adapter = get_indicator_adapter(self.config, raise_errors=True)
adapter.build_table()
with adapter.engine.begin() as connection:
context = MigrationContext.configure(connection)
op = Operations(context)
op.drop_column(adapter.get_table().name, 'name')
doc = {
"_id": '123',
"domain": "domain",
"doc_type": "CommCareCase",
"name": 'bob'
}
with self.assertRaises(MissingColumnWarning):
adapter.best_effort_save(doc)
def check_migration():
"""Check that current app has applied all migrations."""
config = Config()
config.set_main_option("script_location", "migrations")
script = ScriptDirectory.from_config(config)
head_revision = script.get_current_head()
engine = sqlalchemy.create_engine(
current_app.config.get('SQLALCHEMY_DATABASE_URI'))
conn = engine.connect()
context = MigrationContext.configure(conn)
current_rev = context.get_current_revision()
if head_revision != current_rev:
print(
f'Upgrade the database. head_revision: {head_revision}. current_rev: {current_rev}'
)
sys.exit(-1)
sys.exit(0)
def __init__(self):
"""
Load necessary alembic objects such ans Script and Context
"""
if not config.alembic_dir:
raise ValueError('Please setup config.alembic_dir param')
# loading alembic.ini config from installed path
self.init_config = Config(config.alembic_dir + 'alembic.ini')
# create database connection
self.engine = create_engine(self.__database_url__)
self.conn = self.engine.connect()
# and load Alembic Context
self.context = MigrationContext.configure(self.conn)
def create_user():
from walkoff.serverdb import add_user, User, Role, initialize_default_resources_admin, \
initialize_default_resources_guest
from sqlalchemy_utils import database_exists, create_database
if not database_exists(db.engine.url):
create_database(db.engine.url)
db.create_all()
alembic_cfg = Config(walkoff.config.Config.ALEMBIC_CONFIG,
ini_section="walkoff",
attributes={'configure_logger': False})
# This is necessary for a flask database
connection = db.engine.connect()
context = MigrationContext.configure(connection)
script = ScriptDirectory.from_config(alembic_cfg)
context.stamp(script, "head")
# Setup admin and guest roles
initialize_default_resources_admin()
initialize_default_resources_guest()
# Setup admin user
admin_role = Role.query.filter_by(id=1).first()
admin_user = User.query.filter_by(username="******").first()
if not admin_user:
add_user(username='******', password='******', roles=[1])
elif admin_role not in admin_user.roles:
admin_user.roles.append(admin_role)
db.session.commit()
apps = set(helpers.list_apps(walkoff.config.Config.APPS_PATH)) - set([
_app.name for _app in
current_app.running_context.execution_db.session.query(App).all()
])
current_app.logger.debug('Found new apps: {0}'.format(apps))
for app_name in apps:
current_app.running_context.execution_db.session.add(
App(name=app_name, devices=[]))
db.session.commit()
current_app.running_context.execution_db.session.commit()
reschedule_all_workflows()
current_app.logger.handlers = logging.getLogger('server').handlers
def __init__(self, context: ApplicationModel, request):
self.context = context
self.request = request
self.session = self.request.get_db_session("warehouse")
self.engine = create_engine(self.session.bind.url)
self.content_metadata = context.content_metadata()
with self.engine.connect() as conn:
# conn.dialect.default_schema_name = self.content_metadata.schema
migration_context = MigrationContext.configure(
conn, opts={"include_schemas": False})
self.upgrade_steps = self.get_upgrade_steps(migration_context)
if len(self.upgrade_steps.ops):
self.need_update = True
else:
self.need_update = False
self.migration_code = self.get_migration_code()
def execute_script(app, seed_file: str):
file_name = os.path.basename(seed_file)
file_name_without_extension = file_name.replace(".py", "")
working_directory = os.path.abspath(os.getcwd())
relative_path = os.path.relpath(os.path.dirname(seed_file), working_directory)
module_path = relative_path.replace("/", ".").replace("\\", ".")
db: SQLAlchemy = app.extensions['sqlalchemy'].db
exec(f"from {module_path}.{file_name_without_extension} import run")
eval(f"run(db)")
conn = db.engine.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)
op.bulk_insert(seed_table,
[
{'file': file_name_without_extension},
]
)
def upgrade(self, silent=False):
conn = self.engine.connect()
cfg = Config()
# this connection is used in `env.py` for the migrations
cfg.attributes["connection"] = conn
cfg.set_main_option("script_location", str(MIGRATION_DIR))
script = ScriptDirectory.from_config(cfg)
# let's make sure we actually need to migrate
migration_logger = logging.getLogger("alembic.runtime.migration")
original_log_level = migration_logger.getEffectiveLevel()
if silent:
migration_logger.setLevel(logging.ERROR)
context = MigrationContext.configure(conn)
if silent:
migration_logger.setLevel(original_log_level)
try:
# try to find the locked version
HEAD = LOCK_PATH.open().read().strip()
self.ensure_migration_needed(script, context, HEAD)
if not silent:
click.secho(f"Upgrading database to {HEAD}")
command.upgrade(cfg, HEAD)
except FileNotFoundError:
raise MigrationError(
"Cannot upgrade the system database, revision lock not found."
)
except MigrationUneededException:
if not silent:
click.secho(f"System database up-to-date.")
except Exception as err:
logging.exception(str(err))
click.secho(
f"Cannot upgrade the system database. It might be corrupted or was created before database migrations where introduced (v0.34.0)",
fg="yellow",
err=True,
)
finally:
conn.close()
def create_user():
from walkoff.serverdb import add_user, User, Role, initialize_default_resources_admin, \
initialize_default_resources_guest
from sqlalchemy_utils import database_exists, create_database
if not database_exists(db.engine.url):
create_database(db.engine.url)
db.create_all()
alembic_cfg = Config(walkoff.config.Config.ALEMBIC_CONFIG, ini_section="walkoff",
attributes={'configure_logger': False})
# This is necessary for a flask database
connection = db.engine.connect()
context = MigrationContext.configure(connection)
script = ScriptDirectory.from_config(alembic_cfg)
context.stamp(script, "head")
# Setup admin and guest roles
initialize_default_resources_admin()
initialize_default_resources_guest()
# Setup admin user
admin_role = Role.query.filter_by(id=1).first()
admin_user = User.query.filter_by(username="******").first()
if not admin_user:
add_user(username='******', password='******', roles=[1])
elif admin_role not in admin_user.roles:
admin_user.roles.append(admin_role)
db.session.commit()
apps = set(helpers.list_apps(walkoff.config.Config.APPS_PATH)) - set([_app.name
for _app in
current_app.running_context.execution_db.session.query(
App).all()])
current_app.logger.debug('Found new apps: {0}'.format(apps))
for app_name in apps:
current_app.running_context.execution_db.session.add(App(name=app_name, devices=[]))
db.session.commit()
current_app.running_context.execution_db.session.commit()
reschedule_all_workflows()
current_app.logger.handlers = logging.getLogger('server').handlers
def get_database_revision_heads(session: Session) -> Optional[List[Text]]:
"""Get the head revisions from a database. If all migration branches are merged,
then the result should be a single value.
Args:
session: Database session.
Returns:
List of head revisions for the connected database.
"""
from alembic.runtime.migration import MigrationContext
# noinspection PyBroadException
try:
context = MigrationContext.configure(session.connection())
return list(context.get_current_heads())
except Exception:
logger.warning("Unable to get database revision heads.")
return None
def add_measurement_columns(session):
"""Inspects all SampleMeasurements finding distinct values for
measurement_types, then creates new columns in the ExperimentMeasurement
table using an Alembic migration context."""
# TODO: Consider index on SampleMeasurement.measurement_types.
ExperimentMeasurement = get_ExperimentMeasurement()
measurement_types = session.query(
SampleMeasurement.measurement_type).distinct().all()
measurement_types = sorted([_ for (_, ) in measurement_types])
cols = set([f"measurement_{mt}" for mt in measurement_types])
new_cols = cols - set(ExperimentMeasurement.__table__.columns.keys())
if new_cols:
ctx = MigrationContext.configure(session.connection())
op = Operations(ctx)
for new_col in new_cols:
print(f"creating new_col: {new_col}")
op.add_column("experiment_measurements",
Column(new_col, DECIMAL(16, 6)))
def validate_schema_version(self) -> bool:
"""
Check if the schema from the given connection is the latest one
"""
LOGGER.info('Validating schema version.')
config = self._get_alembic_config()
script_directory = script.ScriptDirectory.from_config(config)
with self.engine.connect() as conn:
context = MigrationContext.configure(conn)
current_version = context.get_current_revision()
current_head = script_directory.get_current_head()
LOGGER.info(
f'current_version: {current_version}, current_head: {current_head}'
)
return current_version == current_head
def register_db_migrations(server, migrate, db):
with server.app_context():
conn = db.engine.connect()
context = MigrationContext.configure(conn)
db_revision = context.get_current_revision()
if server.config["DB_AUTO_CREATE"] and not database_exists(
server.config["SQLALCHEMY_DATABASE_URI"]):
with server.app_context():
print("creating db")
db.create_all()
flask_migrate.init(server.config["MIGRATIONS_DIR"])
command.stamp(migrate.get_config(server.config["MIGRATIONS_DIR"]),
"head")
add_admin(db)
if server.config["DB_AUTO_UPGRADE"]:
with server.app_context():
command.upgrade(
migrate.get_config(server.config["MIGRATIONS_DIR"]), "head")
def setUp(self):
self.conn = config.db.connect()
self.metadata = MetaData()
t1 = Table(
'foo', self.metadata,
Column('id', Integer, primary_key=True),
Column('data', String(50)),
Column('x', Integer),
mysql_engine='InnoDB'
)
t1.create(self.conn)
self.conn.execute(
t1.insert(),
[
{"id": 1, "data": "d1", "x": 5},
{"id": 2, "data": "22", "x": 6},
{"id": 3, "data": "8.5", "x": 7},
{"id": 4, "data": "9.46", "x": 8},
{"id": 5, "data": "d5", "x": 9}
]
)
context = MigrationContext.configure(self.conn)
self.op = Operations(context)
def _new_context(self, **kwargs):
return MigrationContext.configure(connection=self.connection, **kwargs)