def catch_signal(sender, **kwargs):
from fluff.pillow import get_fluff_pillow_configs
if settings.UNIT_TESTING or kwargs['using'] != DEFAULT_DB_ALIAS:
return
table_pillow_map = {}
for config in get_fluff_pillow_configs():
pillow = config.get_instance()
doc = pillow.indicator_class()
if doc.save_direct_to_sql:
table_pillow_map[doc._table.name] = {'doc': doc, 'pillow': pillow}
print '\tchecking fluff SQL tables for schema changes'
engine = sqlalchemy.create_engine(settings.SQL_REPORTING_DATABASE_URL)
with engine.begin() as connection:
migration_context = get_migration_context(connection,
table_pillow_map.keys())
raw_diffs = compare_metadata(migration_context, fluff_metadata)
diffs = reformat_alembic_diffs(raw_diffs)
tables_to_rebuild = get_tables_to_rebuild(diffs, table_pillow_map.keys())
for table in tables_to_rebuild:
info = table_pillow_map[table]
rebuild_table(engine, info['pillow'], info['doc'])
engine.dispose()
def catch_signal(sender, **kwargs):
from fluff.pillow import get_fluff_pillow_configs
if settings.UNIT_TESTING or kwargs['using'] != DEFAULT_DB_ALIAS:
return
table_pillow_map = {}
for config in get_fluff_pillow_configs():
pillow = config.get_instance()
for processor in pillow.processors:
doc = processor.indicator_class()
if doc.save_direct_to_sql:
table_pillow_map[doc._table.name] = {
'doc': doc,
'pillow': pillow
}
print('\tchecking fluff SQL tables for schema changes')
engine = connection_manager.get_engine('default')
with engine.begin() as connection:
migration_context = get_migration_context(connection,
list(table_pillow_map))
raw_diffs = compare_metadata(migration_context, fluff_metadata)
diffs = reformat_alembic_diffs(raw_diffs)
tables_to_rebuild = get_tables_to_rebuild(diffs, list(table_pillow_map))
for table in tables_to_rebuild:
info = table_pillow_map[table]
rebuild_table(engine, info['pillow'], info['doc'])
engine.dispose()
def detect_changed(self):
""" Detect the difference between the metadata and the database
:rtype: MigrationReport instance
"""
diff = compare_metadata(self.context, self.metadata)
return MigrationReport(self, diff)
def test_migration(self):
self.setup_base_db()
# we have no alembic base revision
self.assertTrue(self.current_db_revision() is None)
# run the migration, afterwards the DB is stamped
self.run_migration()
db_revision = self.current_db_revision()
self.assertTrue(db_revision is not None)
# db revision matches latest alembic revision
alembic_head = self.alembic_script().get_current_head()
self.assertEqual(db_revision, alembic_head)
# compare the db schema from a migrated database to
# one created fresh from the model definitions
opts = {
'compare_type': db_compare_type,
'compare_server_default': True,
}
with self.db.engine.connect() as conn:
context = MigrationContext.configure(connection=conn, opts=opts)
metadata_diff = compare_metadata(context, self.head_metadata)
self.assertEqual(metadata_diff, [])
def get_schema_diff(
metadata: MetaData,
database_url: str,
include_tables: Optional[Sequence[str]] = None,
exclude_tables: Optional[Sequence[str]] = None,
) -> List[tuple]:
"""
Parameters
-----
`metadata`: Metadata
Sqlalchemy Metadata
`database_url`:
Target databse url e.g. `sqlite://`
`include_tables`: Sequence[str] | None
List of table names to check against the database
`exclude_tables`: Sequence[str] | None
List of table names to be ignored
Return
------
List[tuple]
"""
engine = create_engine(database_url)
mc = migration.MigrationContext.configure(engine.connect())
diff = compare_metadata(mc, metadata)
if include_tables is not None and exclude_tables is not None:
raise Exception("`include_tables` and `exclude_tables` must not be used together") #TODO define custom error class
if exclude_tables is not None:
diff = list(filter(filter_out_excluded_tables(exclude_tables), diff))
if include_tables is not None:
diff = list(filter(filter_in_included_tables(include_tables), diff))
return list(diff)
def test_include_symbol(self):
diffs = []
def include_symbol(name, schema=None):
return name in ("address", "order")
context = MigrationContext.configure(
connection=self.bind.connect(),
opts={
"compare_type": True,
"compare_server_default": True,
"target_metadata": self.m2,
"include_symbol": include_symbol,
},
)
diffs = autogenerate.compare_metadata(context,
context.opts["target_metadata"])
alter_cols = set([
d[2] for d in self._flatten_diffs(diffs)
if d[0].startswith("modify")
])
eq_(alter_cols, set(["order"]))
def test_compare_metadata_schema(self):
metadata = self.m2
context = MigrationContext.configure(
connection=self.bind.connect(),
opts={
"include_schemas": True
}
)
diffs = autogenerate.compare_metadata(context, metadata)
eq_(
diffs[0],
('add_table', metadata.tables['test_schema.item'])
)
eq_(diffs[1][0], 'remove_table')
eq_(diffs[1][1].name, "extra")
eq_(diffs[2][0], "add_column")
eq_(diffs[2][1], "test_schema")
eq_(diffs[2][2], "address")
eq_(diffs[2][3], metadata.tables['test_schema.address'].c.street)
eq_(diffs[3][0], "add_column")
eq_(diffs[3][1], "test_schema")
eq_(diffs[3][2], "order")
eq_(diffs[3][3], metadata.tables['test_schema.order'].c.user_id)
eq_(diffs[4][0][0], 'modify_nullable')
eq_(diffs[4][0][5], False)
eq_(diffs[4][0][6], True)
def test_compare_metadata_schema(self):
metadata = self.m2
context = MigrationContext.configure(connection=self.bind.connect(),
opts={"include_schemas": True})
diffs = autogenerate.compare_metadata(context, metadata)
eq_(diffs[0], ('add_table', metadata.tables['test_schema.item']))
eq_(diffs[1][0], 'remove_table')
eq_(diffs[1][1].name, "extra")
eq_(diffs[2][0], "add_column")
eq_(diffs[2][1], "test_schema")
eq_(diffs[2][2], "address")
eq_(diffs[2][3], metadata.tables['test_schema.address'].c.street)
eq_(diffs[3][0], "add_constraint")
eq_(diffs[3][1].name, "uq_email")
eq_(diffs[4][0], "add_column")
eq_(diffs[4][1], "test_schema")
eq_(diffs[4][2], "order")
eq_(diffs[4][3], metadata.tables['test_schema.order'].c.user_id)
eq_(diffs[5][0][0], 'modify_nullable')
eq_(diffs[5][0][5], False)
eq_(diffs[5][0][6], True)
def test_include_symbol(self):
diffs = []
def include_symbol(name, schema=None):
return name in ('address', 'order')
context = MigrationContext.configure(
connection=self.bind.connect(),
opts={
'compare_type': True,
'compare_server_default': True,
'target_metadata': self.m2,
'include_symbol': include_symbol,
}
)
diffs = autogenerate.compare_metadata(
context, context.opts['target_metadata'])
alter_cols = set([
d[2] for d in self._flatten_diffs(diffs)
if d[0].startswith('modify')
])
eq_(alter_cols, set(['order']))
def test_compare_metadata_include_object(self):
metadata = self.m2
def include_object(obj, name, type_, reflected, compare_to):
if type_ == "table":
return name in ("extra", "order")
elif type_ == "column":
return name != "amount"
else:
return True
context = MigrationContext.configure(
connection=self.bind.connect(),
opts={
'compare_type': True,
'compare_server_default': True,
'include_object': include_object,
}
)
diffs = autogenerate.compare_metadata(context, metadata)
eq_(diffs[0][0], 'remove_table')
eq_(diffs[0][1].name, "extra")
eq_(diffs[1][0], "add_column")
eq_(diffs[1][1], None)
eq_(diffs[1][2], "order")
eq_(diffs[1][3], metadata.tables['order'].c.user_id)
def test_migration(self):
self.setup_base_db()
# we have no alembic base revision
self.assertTrue(self.current_db_revision() is None)
# run the migration, afterwards the DB is stamped
self.run_migration()
db_revision = self.current_db_revision()
self.assertTrue(db_revision is not None)
# db revision matches latest alembic revision
alembic_head = self.alembic_script().get_current_head()
self.assertEqual(db_revision, alembic_head)
# compare the db schema from a migrated database to
# one created fresh from the model definitions
opts = {
'compare_type': db_compare_type,
'compare_server_default': True,
}
with self.db.engine.connect() as conn:
context = MigrationContext.configure(connection=conn, opts=opts)
metadata_diff = compare_metadata(context, self.head_metadata)
# BBB until #353 is done, we have a minor expected difference
filtered_diff = []
for entry in metadata_diff:
if entry[0] == 'remove_column' and \
entry[2] in ('cell', 'cell_blacklist') and \
entry[3].name == 'id':
continue
filtered_diff.append(entry)
self.assertEqual(filtered_diff, [])
def catch_signal(sender, **kwargs):
if settings.UNIT_TESTING:
return
from fluff.pillow import FluffPillow
table_pillow_map = {}
pillow_configs = get_all_pillow_configs()
for pillow_config in pillow_configs:
pillow_class = pillow_config.get_class()
if issubclass(pillow_class, FluffPillow):
doc = pillow_class.indicator_class()
if doc.save_direct_to_sql:
table_pillow_map[doc._table.name] = {
'doc': doc,
'pillow': pillow_class
}
print '\tchecking fluff SQL tables for schema changes'
engine = sqlalchemy.create_engine(settings.SQL_REPORTING_DATABASE_URL)
with engine.begin() as connection:
migration_context = get_migration_context(connection, table_pillow_map.keys())
diffs = compare_metadata(migration_context, fluff_metadata)
tables_to_rebuild = get_tables_to_rebuild(diffs, table_pillow_map.keys())
for table in tables_to_rebuild:
info = table_pillow_map[table]
rebuild_table(engine, info['pillow'], info['doc'])
engine.dispose()
def test_include_symbol(self):
diffs = []
def include_symbol(name, schema=None):
return name in ("address", "order")
context = MigrationContext.configure(
connection=self.bind.connect(),
opts={
"compare_type": True,
"compare_server_default": True,
"target_metadata": self.m2,
"include_symbol": include_symbol,
},
)
diffs = autogenerate.compare_metadata(
context, context.opts["target_metadata"]
)
alter_cols = set(
[
d[2]
for d in self._flatten_diffs(diffs)
if d[0].startswith("modify")
]
)
eq_(alter_cols, set(["order"]))
def test_database_schema_and_sqlalchemy_model_are_in_sync(self):
all_meta_data = MetaData()
for (table_name, table) in airflow_base.metadata.tables.items():
all_meta_data._add_table(table_name, table.schema, table)
# create diff between database schema and SQLAlchemy model
mc = MigrationContext.configure(engine.connect())
diff = compare_metadata(mc, all_meta_data)
# known diffs to ignore
ignores = [
# users.password is not part of User model,
# otherwise it would show up in (old) UI
lambda t: (t[0] == 'remove_column' and t[2] == 'users' and t[3].
name == 'password'),
# ignore tables created by celery
lambda t:
(t[0] == 'remove_table' and t[1].name == 'celery_taskmeta'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'celery_tasksetmeta'),
# ignore indices created by celery
lambda t: (t[0] == 'remove_index' and t[1].name == 'task_id'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'taskset_id'),
# Ignore all the fab tables
lambda t:
(t[0] == 'remove_table' and t[1].name == 'ab_permission'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'ab_register_user'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_role'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'ab_permission_view'),
lambda t: (t[0] == 'remove_table' and t[1].name ==
'ab_permission_view_role'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_user_role'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_user'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_view_menu'),
# Ignore all the fab indices
lambda t:
(t[0] == 'remove_index' and t[1].name == 'permission_id'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'name'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'user_id'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'username'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'field_string'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'email'),
lambda t:
(t[0] == 'remove_index' and t[1].name == 'permission_view_id'),
# from test_security unit test
lambda t: (t[0] == 'remove_table' and t[1].name == 'some_model'),
]
for ignore in ignores:
diff = [d for d in diff if not ignore(d)]
self.assertFalse(
diff, 'Database schema and SQLAlchemy model are not in sync: ' +
str(diff))
def test_table_filter(self):
migration_context = get_migration_context(self.engine,
[self.table_name])
sqlalchemy.Table('new_table', self.metadata)
raw_diffs = compare_metadata(migration_context, self.metadata)
_, diffs = reformat_alembic_diffs(raw_diffs)
self.assertEqual(0, len(diffs))
def catch_signal(sender, **kwargs):
from fluff.pillow import get_fluff_pillow_configs
if settings.UNIT_TESTING or kwargs['using'] != DEFAULT_DB_ALIAS:
return
table_pillow_map = {}
for config in get_fluff_pillow_configs():
pillow = config.get_instance()
doc = pillow.indicator_class()
if doc.save_direct_to_sql:
table_pillow_map[doc._table.name] = {
'doc': doc,
'pillow': pillow
}
print '\tchecking fluff SQL tables for schema changes'
engine = sqlalchemy.create_engine(settings.SQL_REPORTING_DATABASE_URL)
with engine.begin() as connection:
migration_context = get_migration_context(connection, table_pillow_map.keys())
raw_diffs = compare_metadata(migration_context, fluff_metadata)
diffs = reformat_alembic_diffs(raw_diffs)
tables_to_rebuild = get_tables_to_rebuild(diffs, table_pillow_map.keys())
for table in tables_to_rebuild:
info = table_pillow_map[table]
rebuild_table(engine, info['pillow'], info['doc'])
engine.dispose()
def test_compare_metadata_include_object(self):
metadata = self.m2
def include_object(obj, name, type_, reflected, compare_to):
if type_ == "table":
return name in ("extra", "order")
elif type_ == "column":
return name != "amount"
else:
return True
context = MigrationContext.configure(connection=self.bind.connect(),
opts={
'compare_type': True,
'compare_server_default':
True,
'include_object':
include_object,
})
diffs = autogenerate.compare_metadata(context, metadata)
eq_(diffs[0][0], 'remove_table')
eq_(diffs[0][1].name, "extra")
eq_(diffs[1][0], "add_column")
eq_(diffs[1][1], None)
eq_(diffs[1][2], "order")
eq_(diffs[1][3], metadata.tables['order'].c.user_id)
def catch_signal(sender, **kwargs):
from fluff.pillow import get_fluff_pillow_configs
if settings.UNIT_TESTING or kwargs['using'] != DEFAULT_DB_ALIAS:
return
table_pillow_map = {}
for config in get_fluff_pillow_configs():
pillow = config.get_instance()
for processor in pillow.processors:
doc = processor.indicator_class()
if doc.save_direct_to_sql:
table_pillow_map[doc._table.name] = {
'doc': doc,
'pillow': pillow
}
print('\tchecking fluff SQL tables for schema changes')
engine = connection_manager.get_engine('default')
with engine.begin() as connection:
migration_context = get_migration_context(connection, list(table_pillow_map))
raw_diffs = compare_metadata(migration_context, fluff_metadata)
diffs = reformat_alembic_diffs(raw_diffs)
tables_to_rebuild = get_tables_to_rebuild(diffs, list(table_pillow_map))
for table in tables_to_rebuild:
info = table_pillow_map[table]
rebuild_table(engine, info['pillow'], info['doc'])
engine.dispose()
def test_include_symbol(self):
diffs = []
def include_symbol(name, schema=None):
return name in ('address', 'order')
context = MigrationContext.configure(connection=self.bind.connect(),
opts={
'compare_type': True,
'compare_server_default':
True,
'target_metadata': self.m2,
'include_symbol':
include_symbol,
})
diffs = autogenerate.compare_metadata(context,
context.opts['target_metadata'])
alter_cols = set([
d[2] for d in self._flatten_diffs(diffs)
if d[0].startswith('modify')
])
eq_(alter_cols, set(['order']))
def compare_metadata(self):
"""Generate a list of operations that would be present in a new
revision.
"""
db = current_app.extensions["sqlalchemy"].db
return autogenerate.compare_metadata(self.migration_context,
db.metadata)
def _test_diffs(self, metadata, expected_diffs, table_names=None):
migration_context = get_migration_context(
self.engine, table_names or [self.table_name, 'new_table'])
raw_diffs = compare_metadata(migration_context, metadata)
diffs = reformat_alembic_diffs(raw_diffs)
self.assertEqual(set(diffs), expected_diffs)
return diffs
def test_database_schema_and_sqlalchemy_model_are_in_sync(self):
all_meta_data = MetaData()
for (table_name, table) in airflow_base.metadata.tables.items():
all_meta_data._add_table(table_name, table.schema, table)
# create diff between database schema and SQLAlchemy model
mctx = MigrationContext.configure(engine.connect())
diff = compare_metadata(mctx, all_meta_data)
# known diffs to ignore
ignores = [
# ignore tables created by celery
lambda t:
(t[0] == 'remove_table' and t[1].name == 'celery_taskmeta'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'celery_tasksetmeta'),
# ignore indices created by celery
lambda t: (t[0] == 'remove_index' and t[1].name == 'task_id'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'taskset_id'),
# Ignore all the fab tables
lambda t:
(t[0] == 'remove_table' and t[1].name == 'ab_permission'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'ab_register_user'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_role'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'ab_permission_view'),
lambda t: (t[0] == 'remove_table' and t[1].name ==
'ab_permission_view_role'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_user_role'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_user'),
lambda t: (t[0] == 'remove_table' and t[1].name == 'ab_view_menu'),
# Ignore all the fab indices
lambda t:
(t[0] == 'remove_index' and t[1].name == 'permission_id'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'name'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'user_id'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'username'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'field_string'),
lambda t: (t[0] == 'remove_index' and t[1].name == 'email'),
lambda t:
(t[0] == 'remove_index' and t[1].name == 'permission_view_id'),
# from test_security unit test
lambda t: (t[0] == 'remove_table' and t[1].name == 'some_model'),
# MSSQL default tables
lambda t: (t[0] == 'remove_table' and t[1].name == 'spt_monitor'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'spt_fallback_db'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'spt_fallback_usg'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'MSreplication_options'),
lambda t:
(t[0] == 'remove_table' and t[1].name == 'spt_fallback_dev'),
]
for ignore in ignores:
diff = [d for d in diff if not ignore(d)]
assert not diff, 'Database schema and SQLAlchemy model are not in sync: ' + str(
diff)
def compare_schema(engine, metadata):
# compare the db schema from a migrated database to
# one created fresh from the model definitions
opts = {"compare_type": True, "compare_server_default": True}
with engine.connect() as conn:
context = MigrationContext.configure(connection=conn, opts=opts)
diff = compare_metadata(context, metadata)
return diff
def test_store_generated_schema_matches_base(tmpdir, db_url):
# Create a SQLAlchemyStore against tmpfile, directly verify that tmpfile contains a
# database with a valid schema
SqlAlchemyStore(db_url, tmpdir.join("ARTIFACTS").strpath)
engine = sqlalchemy.create_engine(db_url)
mc = MigrationContext.configure(engine.connect())
diff = compare_metadata(mc, Base.metadata)
assert len(diff) == 0
def get_table_diffs(engine, table_names, metadata):
with engine.begin() as connection:
migration_context = get_migration_context(connection, table_names)
raw_diffs = compare_metadata(migration_context, metadata)
return [
diff for diff in reformat_alembic_diffs(raw_diffs)
if diff.table_name in table_names
]
def all_diffs():
for diff in compare_metadata(mc, from_metadata):
if isinstance(diff[0], tuple):
op_name = diff[0][0]
else:
op_name = diff[0]
if op_name in supported_operations:
yield op_name, diff
def detect_changed(self):
""" Detect the difference between the metadata and the database
:rtype: MigrationReport instance
"""
diff = compare_metadata(self.context, self.metadata)
inspector = Inspector(self.conn)
diff.extend(self.detect_undetected_constraint_from_alembic(inspector))
return MigrationReport(self, diff)
def detect_changed(self):
""" Detect the difference between the metadata and the database
:rtype: MigrationReport instance
"""
diff = compare_metadata(self.context, self.metadata)
inspector = Inspector(self.conn)
diff.extend(self.detect_undetected_constraint_from_alembic(inspector))
return MigrationReport(self, diff)
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 compare_schema(engine, metadata):
# compare the db schema from a migrated database to
# one created fresh from the model definitions
opts = {
'compare_type': db_compare_type,
'compare_server_default': True,
}
with engine.connect() as conn:
context = MigrationContext.configure(connection=conn, opts=opts)
diff = compare_metadata(context, metadata)
return diff
def test_compare_metadata(self):
metadata = self.m2
diffs = autogenerate.compare_metadata(self.context, metadata)
eq_(
diffs[0],
('add_table', metadata.tables['item'])
)
eq_(diffs[1][0], 'remove_table')
eq_(diffs[1][1].name, "extra")
eq_(diffs[2][0], "add_column")
eq_(diffs[2][1], None)
eq_(diffs[2][2], "address")
eq_(diffs[2][3], metadata.tables['address'].c.street)
eq_(diffs[3][0], "add_constraint")
eq_(diffs[3][1].name, "uq_email")
eq_(diffs[4][0], "add_column")
eq_(diffs[4][1], None)
eq_(diffs[4][2], "order")
eq_(diffs[4][3], metadata.tables['order'].c.user_id)
eq_(diffs[5][0][0], "modify_type")
eq_(diffs[5][0][1], None)
eq_(diffs[5][0][2], "order")
eq_(diffs[5][0][3], "amount")
eq_(repr(diffs[5][0][5]), "NUMERIC(precision=8, scale=2)")
eq_(repr(diffs[5][0][6]), "Numeric(precision=10, scale=2)")
self._assert_fk_diff(
diffs[6], "add_fk",
"order", ["user_id"],
"user", ["id"]
)
eq_(diffs[7][0][0], "modify_default")
eq_(diffs[7][0][1], None)
eq_(diffs[7][0][2], "user")
eq_(diffs[7][0][3], "a1")
eq_(diffs[7][0][6].arg, "x")
eq_(diffs[8][0][0], 'modify_nullable')
eq_(diffs[8][0][5], True)
eq_(diffs[8][0][6], False)
eq_(diffs[9][0], 'remove_index')
eq_(diffs[9][1].name, 'pw_idx')
eq_(diffs[10][0], 'remove_column')
eq_(diffs[10][3].name, 'pw')
def test_compare_metadata(self):
metadata = self.m2
diffs = autogenerate.compare_metadata(self.context, metadata)
eq_(
diffs[0],
('add_table', metadata.tables['item'])
)
eq_(diffs[1][0], 'remove_table')
eq_(diffs[1][1].name, "extra")
eq_(diffs[2][0], "add_column")
eq_(diffs[2][1], None)
eq_(diffs[2][2], "address")
eq_(diffs[2][3], metadata.tables['address'].c.street)
eq_(diffs[3][0], "add_constraint")
eq_(diffs[3][1].name, "uq_email")
eq_(diffs[4][0], "add_column")
eq_(diffs[4][1], None)
eq_(diffs[4][2], "order")
eq_(diffs[4][3], metadata.tables['order'].c.user_id)
eq_(diffs[5][0][0], "modify_type")
eq_(diffs[5][0][1], None)
eq_(diffs[5][0][2], "order")
eq_(diffs[5][0][3], "amount")
eq_(repr(diffs[5][0][5]), "NUMERIC(precision=8, scale=2)")
eq_(repr(diffs[5][0][6]), "Numeric(precision=10, scale=2)")
self._assert_fk_diff(
diffs[6], "add_fk",
"order", ["user_id"],
"user", ["id"]
)
eq_(diffs[7][0][0], "modify_default")
eq_(diffs[7][0][1], None)
eq_(diffs[7][0][2], "user")
eq_(diffs[7][0][3], "a1")
eq_(diffs[7][0][6].arg, "x")
eq_(diffs[8][0][0], 'modify_nullable')
eq_(diffs[8][0][5], True)
eq_(diffs[8][0][6], False)
eq_(diffs[9][0], 'remove_index')
eq_(diffs[9][1].name, 'pw_idx')
eq_(diffs[10][0], 'remove_column')
eq_(diffs[10][3].name, 'pw')
def test_store_generated_schema_matches_base(tmpdir, db_url):
# Create a SQLAlchemyStore against tmpfile, directly verify that tmpfile contains a
# database with a valid schema
SqlAlchemyStore(db_url, tmpdir.join("ARTIFACTS").strpath)
engine = sqlalchemy.create_engine(db_url)
mc = MigrationContext.configure(engine.connect())
diff = compare_metadata(mc, Base.metadata)
# `diff` contains several `remove_index` operations because `Base.metadata` does not contain
# index metadata but `mc` does. Note this doesn't mean the MLflow database is missing indexes
# as tested in `test_create_index_on_run_uuid`.
diff = [d for d in diff if d[0] != "remove_index"]
assert len(diff) == 0
def assert_migrations_are_up_to_date():
with engine.begin() as conn:
context = migration.MigrationContext.configure(conn)
diff = compare_metadata(context, Base.metadata)
if len(diff) > 0:
logger.error(
"Migrations are not up to date. The following changes have been detected:\n"
+ "\n".join(str(d) for d in diff))
logger.warning("Create a new revision")
return True
else:
return False
def test_include_object(self):
def include_object(obj, name, type_, reflected, compare_to):
assert obj.name == name
if type_ == "table":
if reflected:
assert obj.metadata is not self.m2
else:
assert obj.metadata is self.m2
return name in ("address", "order", "user")
elif type_ == "column":
if reflected:
assert obj.table.metadata is not self.m2
else:
assert obj.table.metadata is self.m2
return name != "street"
else:
return True
context = MigrationContext.configure(
connection=self.bind.connect(),
opts={
"compare_type": True,
"compare_server_default": True,
"target_metadata": self.m2,
"include_object": include_object,
},
)
diffs = autogenerate.compare_metadata(
context, context.opts["target_metadata"]
)
alter_cols = (
set(
[
d[2]
for d in self._flatten_diffs(diffs)
if d[0].startswith("modify")
]
)
.union(
d[3].name
for d in self._flatten_diffs(diffs)
if d[0] == "add_column"
)
.union(
d[1].name
for d in self._flatten_diffs(diffs)
if d[0] == "add_table"
)
)
eq_(alter_cols, set(["user_id", "order", "user"]))
def test_database_schema_and_sqlalchemy_model_are_in_sync(self):
all_meta_data = MetaData()
for (table_name, table) in airflow_base.metadata.tables.items():
all_meta_data._add_table(table_name, table.schema, table)
# create diff between database schema and SQLAlchemy model
mc = MigrationContext.configure(engine.connect())
diff = compare_metadata(mc, all_meta_data)
# known diffs to ignore
ignores = [
# users.password is not part of User model,
# otherwise it would show up in (old) UI
lambda t: (t[0] == 'remove_column' and
t[2] == 'users' and
t[3].name == 'password'),
# ignore tables created by other tests
lambda t: (t[0] == 'remove_table' and
t[1].name == 't'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'test_airflow'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'test_postgres_to_postgres'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'test_mysql_to_mysql'),
# ignore tables created by celery
lambda t: (t[0] == 'remove_table' and
t[1].name == 'celery_taskmeta'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'celery_tasksetmeta'),
# Ignore all the fab tables
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_permission'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_register_user'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_role'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_permission_view'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_permission_view_role'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_user_role'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_user'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_view_menu'),
]
for ignore in ignores:
diff = [d for d in diff if not ignore(d)]
self.assertFalse(diff, 'Database schema and SQLAlchemy model are not in sync')
def _verify_schema(engine):
with engine.connect() as connection:
mc = MigrationContext.configure(connection)
diff = compare_metadata(mc, Base.metadata)
if len(diff) > 0:
_logger.error("Detected one or more differences between current database schema "
"and desired schema, exiting. Diff:\n %s",
pprint.pformat(diff, indent=2, width=20))
raise MlflowException(
"Detected out-of-date database schema. Take a backup of your database, then "
"run 'mlflow db upgrade %s' to migrate your database to the latest schema. "
"NOTE: schema migration may result in database downtime "
"- please consult your database's documentation for more detail." % engine.url)
def detect_changed(self, schema_only=False):
""" Detect the difference between the metadata and the database
:rtype: MigrationReport instance
"""
inspector = Inspector(self.conn)
if schema_only:
diff = self.detect_added_new_schema(inspector)
else:
diff = compare_metadata(self.context, self.metadata)
diff.extend(
self.detect_undetected_constraint_from_alembic(inspector))
return MigrationReport(self, diff)
def test_database_migration():
database_uri = conftest.create_temporary_database()
config = dict(database_uri=database_uri)
app = mock.Mock()
app = database.setup_database(app, config)
try:
database.upgrade_database(app, config, force_migration=True)
with app.engine.begin() as conn:
ctx = migration.MigrationContext.configure(conn)
diff = compare_metadata(ctx, Base.metadata)
assert diff == [], pprint.pformat(diff, indent=2, width=20)
finally:
database.SessionLocal.close_all()
app.engine.dispose()
def get_app():
app = FastAPI()
modules = tuple(Path().glob("modules/*"))
if modules:
for module in modules:
if module.is_dir():
module_path = ".".join(module.parts)
logger.info(f"loading module [{module_path}]")
try:
m = import_module(".".join([module_path, "main"]))
on_load_hook = getattr(m, "on_load", None)
if on_load_hook is not None:
# check if it's an awaitable
if asyncio.iscoroutinefunction(on_load_hook):
t = threading.Thread(target=asyncio.run,
args=[on_load_hook(app)])
t.start()
t.join()
# todo raise the exception from the thread
else:
on_load_hook(app)
# load "db.py" if it exists and it's a file
if (module / "db.py").is_file():
import_module(".".join([module_path, "db"]))
except ModuleNotFoundError:
logger.error(
f"Could not load module {module_path} missing main.py!"
)
else:
logger.warning("There is no modules folder !")
# ensure the database is up to date
engine = create_engine(get_setting().PG_DNS)
mc = MigrationContext.configure(engine.connect())
diff = compare_metadata(mc, db)
if diff:
logger.critical(
f"The database is not up to date ! use the Modular cli to update the database schema {pprint.pformat(diff)}"
)
exit(1)
db.init_app(app)
return app
def check_table(indicator_doc):
def check_diff(diff):
if diff[0] in ('add_table', 'remove_table'):
if diff[1].name == table_name:
raise RebuildTableException()
elif diff[2] == table_name:
raise RebuildTableException()
table_name = indicator_doc._table.name
diffs = compare_metadata(get_migration_context(), indicator_doc._table.metadata)
for diff in diffs:
if isinstance(diff, list):
for d in diff:
check_diff(d)
else:
check_diff(diff)
def heal():
# This is needed else the heal script will start spewing
# a lot of pointless warning messages from alembic.
LOG.setLevel(logging.INFO)
if context.is_offline_mode():
return
models_metadata = frozen_models.get_metadata()
# Compare metadata from models and metadata from migrations
# Diff example:
# [ ( 'add_table',
# Table('bat', MetaData(bind=None),
# Column('info', String(), table=<bat>), schema=None)),
# ( 'remove_table',
# Table(u'bar', MetaData(bind=None),
# Column(u'data', VARCHAR(), table=<bar>), schema=None)),
# ( 'add_column',
# None,
# 'foo',
# Column('data', Integer(), table=<foo>)),
# ( 'remove_column',
# None,
# 'foo',
# Column(u'old_data', VARCHAR(), table=None)),
# [ ( 'modify_nullable',
# None,
# 'foo',
# u'x',
# { 'existing_server_default': None,
# 'existing_type': INTEGER()},
# True,
# False)]]
opts = {
'compare_type': _compare_type,
'compare_server_default': _compare_server_default,
}
mc = alembic.migration.MigrationContext.configure(op.get_bind(), opts=opts)
set_storage_engine(op.get_bind(), "InnoDB")
diff1 = autogen.compare_metadata(mc, models_metadata)
# Alembic does not contain checks for foreign keys. Because of that it
# checks separately.
added_fks, dropped_fks = check_foreign_keys(models_metadata)
diff = dropped_fks + diff1 + added_fks
# For each difference run command
for el in diff:
execute_alembic_command(el)
def heal():
# This is needed else the heal script will start spewing
# a lot of pointless warning messages from alembic.
LOG.setLevel(logging.INFO)
if context.is_offline_mode():
return
models_metadata = frozen_models.get_metadata()
# Compare metadata from models and metadata from migrations
# Diff example:
# [ ( 'add_table',
# Table('bat', MetaData(bind=None),
# Column('info', String(), table=<bat>), schema=None)),
# ( 'remove_table',
# Table(u'bar', MetaData(bind=None),
# Column(u'data', VARCHAR(), table=<bar>), schema=None)),
# ( 'add_column',
# None,
# 'foo',
# Column('data', Integer(), table=<foo>)),
# ( 'remove_column',
# None,
# 'foo',
# Column(u'old_data', VARCHAR(), table=None)),
# [ ( 'modify_nullable',
# None,
# 'foo',
# u'x',
# { 'existing_server_default': None,
# 'existing_type': INTEGER()},
# True,
# False)]]
opts = {
'compare_type': _compare_type,
'compare_server_default': _compare_server_default,
}
mc = alembic.migration.MigrationContext.configure(op.get_bind(), opts=opts)
diff1 = autogen.compare_metadata(mc, models_metadata)
# Alembic does not contain checks for foreign keys. Because of that it
# checks separately.
diff2 = check_foreign_keys(models_metadata)
diff = diff1 + diff2
# For each difference run command
for el in diff:
execute_alembic_command(el)
def generate(self, message: str, allow_empty: bool) -> None:
"""
Generate upgrade scripts using alembic.
"""
if not self.__exists():
raise DBCreateException('Tables have not been created yet, use create to create them!')
# Verify that there are actual changes, and refuse to create empty migration scripts
context = MigrationContext.configure(self.__config['database']['engine'].connect(), opts={'compare_type': True})
diff = compare_metadata(context, metadata)
if (not allow_empty) and (len(diff) == 0):
raise DBCreateException('There is nothing different between code and the DB, refusing to create migration!')
self.__alembic_cmd(
'revision',
'--autogenerate',
'-m',
message,
)
def migrate():
from alembic.runtime.migration import MigrationContext
# use `db.session.connection()` instead of `db.engine.connect()`
# to avoid lock hang
context = MigrationContext.configure(
db.session.connection(),
opts={
"compare_type": True,
},
)
if request.method == "GET":
import pprint
from alembic.autogenerate import compare_metadata
diff = compare_metadata(context, db.metadata)
diff_str = pprint.pformat(diff, indent=2, width=20)
logger.info("Migrate steps: %s", diff_str)
return respond_success(migration=diff_str)
from alembic.autogenerate import produce_migrations
from alembic.operations import Operations
from alembic.operations.ops import OpContainer
migration = produce_migrations(context, db.metadata)
operation = Operations(context)
for outer_op in migration.upgrade_ops.ops:
logger.info("Invoking %s", outer_op)
if isinstance(outer_op, OpContainer):
for inner_op in outer_op.ops:
logger.info("Invoking %s", inner_op)
operation.invoke(inner_op)
else:
operation.invoke(outer_op)
db.session.commit()
db.session.close()
return respond_success()
def test_compare_metadata_include_symbol(self):
metadata = self.m2
def include_symbol(table_name, schema_name):
return table_name in ('extra', 'order')
context = MigrationContext.configure(
connection=self.bind.connect(),
opts={
'compare_type': True,
'compare_server_default': True,
'include_symbol': include_symbol,
}
)
diffs = autogenerate.compare_metadata(context, metadata)
eq_(diffs[0][0], 'remove_table')
eq_(diffs[0][1].name, "extra")
eq_(diffs[1][0], "add_column")
eq_(diffs[1][1], None)
eq_(diffs[1][2], "order")
eq_(diffs[1][3], metadata.tables['order'].c.user_id)
eq_(diffs[2][0][0], "modify_type")
eq_(diffs[2][0][1], None)
eq_(diffs[2][0][2], "order")
eq_(diffs[2][0][3], "amount")
eq_(repr(diffs[2][0][5]), "NUMERIC(precision=8, scale=2)")
eq_(repr(diffs[2][0][6]), "Numeric(precision=10, scale=2)")
eq_(diffs[2][1][0], 'modify_nullable')
eq_(diffs[2][1][2], 'order')
eq_(diffs[2][1][5], False)
eq_(diffs[2][1][6], True)
def get_table_diffs(engine, table_names, metadata):
with engine.begin() as connection:
migration_context = get_migration_context(connection, table_names)
raw_diffs = compare_metadata(migration_context, metadata)
diffs = reformat_alembic_diffs(raw_diffs)
return TableDiffs(raw=raw_diffs, formatted=diffs)
def _test_diffs(self, metadata, expected_diffs, table_names=None):
migration_context = get_migration_context(self.engine, table_names or [self.table_name, 'new_table'])
raw_diffs = compare_metadata(migration_context, metadata)
diffs = reformat_alembic_diffs(raw_diffs)
self.assertEqual(set(diffs), expected_diffs)
def test_table_filter(self):
migration_context = get_migration_context(self.engine, [self.table_name])
sqlalchemy.Table('new_table', self.metadata)
raw_diffs = compare_metadata(migration_context, self.metadata)
diffs = reformat_alembic_diffs(raw_diffs)
self.assertEqual(0, len(diffs))
if isinstance(op_, tuple):
op_name = op_[0]
if op_name.endswith('_table'):
op_obj = op_[1]
op_obj_name = op_obj.name
elif op_name.endswith('_column'):
op_obj = op_[3]
op_obj_name = op_[2] + '.' + op_obj.name
else:
op_obj = op_[1]
op_obj_name = op_obj.name
if op_obj_name is None:
op_obj_name = op_obj.table.name + '.' \
+ '_'.join(c.name for c in op_obj.columns)
else:
op_name = op_[0][0]
op_obj_name = op_[0][2] + '.' + op_[0][3]
return not op_obj_name in IGNORE_OPS.get(op_name, [])
# Set up a migration context.
alembic_cfg.set_main_option('script_location', 'thelma:db/schema/migrations')
script = ScriptDirectory.from_config(alembic_cfg)
env_ctxt = EnvironmentContext(alembic_cfg, script)
engine = create_engine(alembic_cfg)
env_ctxt.configure(engine.connect()) # , include_object=include_object)
mig_ctxt = env_ctxt.get_context()
ops = compare_metadata(mig_ctxt, metadata)
diff = [op for op in ops if include_op(op)]
pprint.pprint(diff, indent=2, width=20)
def test_database_schema_and_sqlalchemy_model_are_in_sync(self):
all_meta_data = MetaData()
for (table_name, table) in airflow_base.metadata.tables.items():
all_meta_data._add_table(table_name, table.schema, table)
# create diff between database schema and SQLAlchemy model
mc = MigrationContext.configure(engine.connect())
diff = compare_metadata(mc, all_meta_data)
# known diffs to ignore
ignores = [
# ignore tables created by celery
lambda t: (t[0] == 'remove_table' and
t[1].name == 'celery_taskmeta'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'celery_tasksetmeta'),
# ignore indices created by celery
lambda t: (t[0] == 'remove_index' and
t[1].name == 'task_id'),
lambda t: (t[0] == 'remove_index' and
t[1].name == 'taskset_id'),
# Ignore all the fab tables
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_permission'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_register_user'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_role'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_permission_view'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_permission_view_role'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_user_role'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_user'),
lambda t: (t[0] == 'remove_table' and
t[1].name == 'ab_view_menu'),
# Ignore all the fab indices
lambda t: (t[0] == 'remove_index' and
t[1].name == 'permission_id'),
lambda t: (t[0] == 'remove_index' and
t[1].name == 'name'),
lambda t: (t[0] == 'remove_index' and
t[1].name == 'user_id'),
lambda t: (t[0] == 'remove_index' and
t[1].name == 'username'),
lambda t: (t[0] == 'remove_index' and
t[1].name == 'field_string'),
lambda t: (t[0] == 'remove_index' and
t[1].name == 'email'),
lambda t: (t[0] == 'remove_index' and
t[1].name == 'permission_view_id'),
# from test_security unit test
lambda t: (t[0] == 'remove_table' and
t[1].name == 'some_model'),
]
for ignore in ignores:
diff = [d for d in diff if not ignore(d)]
self.assertFalse(
diff,
'Database schema and SQLAlchemy model are not in sync: ' + str(diff)
)
