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 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 get_metadata(app: str) -> Optional[MetaData]:
meta = []
control = []
if callable(MetaData):
metadata = MetaData()
for name in resources.contents(app):
if name == "models.py":
_module = import_module(f"{app}.{name[:-3]}")
for model in dir(_module):
table = getattr(_module, model)
if hasattr(table, "metadata") and hasattr(table, "database"):
table_module = getattr(table, "__module__", None)
if (
table_module == _module.__name__
and table.database == os.environ["CRAX_DB_NAME"]
):
meta.append(table)
control.append(table.table.name)
for base in meta:
for (table_name, table) in base.metadata.tables.items():
if table_name in control:
metadata._add_table(table_name, table.schema, table)
return metadata
def _mergeMetadata(metas: Iterable[MetaData]) -> MetaData:
globalMetaData = MetaData()
for metaData in metas:
for (tableName, table) in metaData.tables.items():
globalMetaData._add_table(table.name, table.schema, table)
return globalMetaData
def merge_base(*args):
from sqlalchemy import MetaData
combined_meta_data = MetaData()
for declarative_base in args:
for (table_name, table) in declarative_base.metadata.tables.items():
combined_meta_data._add_table(table_name, table.schema, table)
return combined_meta_data
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')
class MetadataCommands(DataBaseCommands):
def __init__(self, opts: list, **kwargs) -> None:
super(MetadataCommands, self).__init__(opts, **kwargs)
engine = create_engine(self.default_connection, echo=False)
self.metadata = MetaData(bind=engine)
self.dependency_map = self.create_dependency_map()
self.collect_metadata()
def collect_metadata(self) -> None:
meta = []
sorted_applications = sort_applications(self.dependency_map,
revers=self.args.down)
for app in sorted_applications:
m = get_metadata(app)
meta.append(m)
for m in meta:
for (table_name, table) in m.tables.items():
self.metadata._add_table(table_name, table.schema, table)
def test_detect_foreign_key_options_changed(self, registry):
with cnx(registry) as conn:
registry.TestFK2.__table__.drop(bind=conn)
meta = MetaData()
meta._add_table('testfktarget', None,
registry.TestFKTarget.__table__)
registry.TestFK2.__table__ = Table(
'testfk2',
meta,
Column('integer', Integer, primary_key=True),
Column('other', Integer, ForeignKey('testfktarget.integer')),
)
registry.TestFK2.__table__.create(bind=conn)
report = registry.migration.detect_changed()
assert report.log_has(
"Drop Foreign keys on testfk2.other => testfktarget.integer")
assert report.log_has(
"Add Foreign keys on (testfk2.other) => (testfktarget.integer)")
def test_detect_drop_foreign_key_with_protected_column(self, registry):
with cnx(registry) as conn:
registry.Test.__table__.drop(bind=conn)
meta = MetaData(naming_convention=naming_convention)
meta._add_table('system_blok', None,
registry.System.Blok.__table__)
registry.Test.__table__ = Table(
'test',
meta,
Column('integer', Integer, primary_key=True),
Column('other', String(64), ForeignKey('system_blok.name')),
)
registry.Test.__table__.create(bind=conn)
# anyblok_fk_test__other_on_system_blok__name
registry.migration.ignore_migration_for = {'test': ['other']}
report = registry.migration.detect_changed()
assert not report.log_has(
"Drop Foreign keys on test.other => system_blok.name")
def test_detect_m2m_primary_key(self, registry):
with cnx(registry) as conn:
Table('reltable',
registry.declarativebase.metadata,
autoload_with=conn).drop(bind=conn)
meta = MetaData()
meta._add_table('testm2m1', None, registry.TestM2M1.__table__)
meta._add_table('testm2m2', None, registry.TestM2M2.__table__)
Table(
'reltable',
meta,
Column('idmodel1', Integer, ForeignKey('testm2m1.idmodel1')),
Column('idmodel2', Integer, ForeignKey('testm2m2.idmodel2')),
).create(bind=conn)
with pytest.raises(MigrationException):
registry.migration.detect_changed()
def test_detect_drop_foreign_key(self, registry):
with cnx(registry) as conn:
registry.Test.__table__.drop(bind=conn)
meta = MetaData()
meta._add_table('system_blok', None,
registry.System.Blok.__table__)
registry.Test.__table__ = Table(
'test',
meta,
Column('integer', Integer, primary_key=True),
Column('other', String(64), ForeignKey('system_blok.name')),
)
registry.Test.__table__.create(bind=conn)
report = registry.migration.detect_changed()
assert report.log_has(
"Drop Foreign keys on test.other => system_blok.name")
report.apply_change()
report = registry.migration.detect_changed()
assert report.log_has(
"Drop Foreign keys on test.other => system_blok.name")
def test_detect_drop_column_with_foreign_key(self, registry):
with cnx(registry) as conn:
registry.Test.__table__.drop(bind=conn)
meta = MetaData(naming_convention=naming_convention)
meta._add_table('system_blok', None,
registry.System.Blok.__table__)
registry.Test.__table__ = Table(
'test',
meta,
Column('integer', Integer, primary_key=True),
Column('other', String(64)),
Column('other2', String(64), ForeignKey('system_blok.name')),
schema='test_db_schema')
registry.Test.__table__.create(bind=conn)
report = registry.migration.detect_changed()
message = ("Drop Foreign keys on test.other2 => %ssystem_blok.name"
) % ('dbo.' if sgdb_in(['MsSQL']) else '')
assert report.log_has(message)
report.apply_change()
report = registry.migration.detect_changed()
assert not report.log_has(message)
class DBInitializer:
def __init__(self):
self.engine = engine
self.BaseTable = MetaData()
for (table_name, table) in BaseTable.metadata.tables.items():
self.BaseTable._add_table(table_name, table.schema, table)
def create_database(self):
"""
Create database if it does not exist
:return success: True if it succeeded or if database already exists
"""
# Create database if it does not exist
try:
if database_exists(self.engine.url):
logger.info("Database already exists. Skipping")
create_database(self.engine.url)
except Exception as err:
logger.info(f"Failed to create database: {err}")
def create_schema(self):
"""
Create database schema
"""
try:
schema_name = os.environ.get('SCHEMA_NAME')
if not schema_name:
raise Exception(f"schema name not given")
# Create schema
self.engine.execute(CreateSchema(schema_name))
except Exception as err:
logger.info(f"Failed to create database: {err}")
def create_tables(self):
"""
Create database tables based on ORM models
"""
self.BaseTable.create_all(self.engine)
def test_database_schema_and_sqlalchemy_model_are_in_sync(self):
# combine Airflow and Flask-AppBuilder (if rbac enabled) models
all_meta_data = MetaData()
for (table_name, table) in airflow_base.metadata.tables.items():
all_meta_data._add_table(table_name, table.schema, table)
if RBAC:
for (table_name, table) in fab_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'),
]
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')
# target_metadata = mymodel.Base.metadata
import os
import sys
from os.path import abspath, dirname
sys.path.insert(0, dirname(dirname(abspath(__file__))))
from api.db import models
from api.settings import Settings
from api.auth import Base
# Combine metadata from auth and containers/templates
combined_meta_data = MetaData()
for declarative_base in [models.Base, Base]:
for (table_name, table) in declarative_base.metadata.tables.items():
combined_meta_data._add_table(table_name, table.schema, table)
target_metadata = combined_meta_data
config.set_main_option(
"sqlalchemy.url",
os.environ.get("DATABASE_URL", "sqlite:///config/data.sqlite"))
# 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 run_migrations_offline():
"""Run migrations in 'offline' mode.
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)
)
