scripts/benchmark_migration.py
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import importlib.util
import logging
import re
import time
from collections import defaultdict
from graphlib import TopologicalSorter
from inspect import getsource
from pathlib import Path
from types import ModuleType
from typing import Any
import click
from flask import current_app
from flask_appbuilder import Model
from flask_migrate import downgrade, upgrade
from progress.bar import ChargingBar
from sqlalchemy import create_engine, inspect
from sqlalchemy.ext.automap import automap_base
from superset import db
from superset.utils.mock_data import add_sample_rows
logger = logging.getLogger(__name__)
def import_migration_script(filepath: Path) -> ModuleType:
"""
Import migration script as if it were a module.
"""
spec = importlib.util.spec_from_file_location(filepath.stem, filepath)
if spec:
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module) # type: ignore
return module
raise Exception(f"No module spec found in location: `{str(filepath)}`")
def extract_modified_tables(module: ModuleType) -> set[str]:
"""
Extract the tables being modified by a migration script.
This function uses a simple approach of looking at the source code of
the migration script looking for patterns. It could be improved by
actually traversing the AST.
"""
tables: set[str] = set()
for function in {"upgrade", "downgrade"}:
source = getsource(getattr(module, function))
tables.update(re.findall(r'alter_table\(\s*"(\w+?)"\s*\)', source, re.DOTALL))
tables.update(re.findall(r'add_column\(\s*"(\w+?)"\s*,', source, re.DOTALL))
tables.update(re.findall(r'drop_column\(\s*"(\w+?)"\s*,', source, re.DOTALL))
return tables
def find_models(module: ModuleType) -> list[type[Model]]:
"""
Find all models in a migration script.
"""
models: list[type[Model]] = []
tables = extract_modified_tables(module)
# add models defined explicitly in the migration script
queue = list(module.__dict__.values())
while queue:
obj = queue.pop()
if hasattr(obj, "__tablename__"):
tables.add(obj.__tablename__)
elif isinstance(obj, list):
queue.extend(obj)
elif isinstance(obj, dict):
queue.extend(obj.values())
# build models by automapping the existing tables, instead of using current
# code; this is needed for migrations that modify schemas (eg, add a column),
# where the current model is out-of-sync with the existing table after a
# downgrade
sqlalchemy_uri = current_app.config["SQLALCHEMY_DATABASE_URI"]
engine = create_engine(sqlalchemy_uri)
Base = automap_base()
Base.prepare(engine, reflect=True)
seen = set()
while tables:
table = tables.pop()
seen.add(table)
try:
model = getattr(Base.classes, table)
except AttributeError:
continue
model.__tablename__ = table
models.append(model)
# add other models referenced in foreign keys
inspector = inspect(model)
for column in inspector.columns.values():
for foreign_key in column.foreign_keys:
table = foreign_key.column.table.name
if table not in seen:
tables.add(table)
# sort topologically so we can create entities in order and
# maintain relationships (eg, create a database before creating
# a slice)
sorter: TopologicalSorter[Any] = TopologicalSorter()
for model in models:
inspector = inspect(model)
dependent_tables: list[str] = []
for column in inspector.columns.values():
for foreign_key in column.foreign_keys:
if foreign_key.column.table.name != model.__tablename__:
dependent_tables.append(foreign_key.column.table.name)
sorter.add(model.__tablename__, *dependent_tables)
order = list(sorter.static_order())
models.sort(key=lambda model: order.index(model.__tablename__))
return models
@click.command()
@click.argument("filepath")
@click.option("--limit", default=1000, help="Maximum number of entities.")
@click.option("--force", is_flag=True, help="Do not prompt for confirmation.")
@click.option("--no-auto-cleanup", is_flag=True, help="Do not remove created models.")
def main(
filepath: str, limit: int = 1000, force: bool = False, no_auto_cleanup: bool = False
) -> None:
auto_cleanup = not no_auto_cleanup
print(f"Importing migration script: {filepath}")
module = import_migration_script(Path(filepath))
revision: str = getattr(module, "revision", "")
down_revision: str = getattr(module, "down_revision", "")
if not revision or not down_revision:
raise Exception(
"Not a valid migration script, couldn't find down_revision/revision"
)
print(f"Migration goes from {down_revision} to {revision}")
current_revision = db.engine.execute(
"SELECT version_num FROM alembic_version"
).scalar()
print(f"Current version of the DB is {current_revision}")
if current_revision != down_revision:
if not force:
click.confirm(
"\nRunning benchmark will downgrade the Superset DB to "
f"{down_revision} and upgrade to {revision} again. There may "
"be data loss in downgrades. Continue?",
abort=True,
)
downgrade(revision=down_revision)
print("\nIdentifying models used in the migration:")
models = find_models(module)
model_rows: dict[type[Model], int] = {}
for model in models:
rows = db.session.query(model).count()
print(f"- {model.__name__} ({rows} rows in table {model.__tablename__})")
model_rows[model] = rows
print("Benchmarking migration")
results: dict[str, float] = {}
start = time.time()
upgrade(revision=revision)
duration = time.time() - start
results["Current"] = duration
print(f"Migration on current DB took: {duration:.2f} seconds")
min_entities = 10
new_models: dict[type[Model], list[Model]] = defaultdict(list)
while min_entities <= limit:
downgrade(revision=down_revision)
print(f"Running with at least {min_entities} entities of each model")
for model in models:
missing = min_entities - model_rows[model]
if missing > 0:
entities: list[Model] = []
print(f"- Adding {missing} entities to the {model.__name__} model")
bar = ChargingBar("Processing", max=missing)
try:
for entity in add_sample_rows(model, missing):
entities.append(entity)
bar.next()
except Exception:
db.session.rollback()
raise
bar.finish()
model_rows[model] = min_entities
db.session.add_all(entities)
db.session.commit()
if auto_cleanup:
new_models[model].extend(entities)
start = time.time()
upgrade(revision=revision)
duration = time.time() - start
print(f"Migration for {min_entities}+ entities took: {duration:.2f} seconds")
results[f"{min_entities}+"] = duration
min_entities *= 10
print("\nResults:\n")
for label, duration in results.items():
print(f"{label}: {duration:.2f} s")
if auto_cleanup:
print("Cleaning up DB")
# delete in reverse order of creation to handle relationships
for model, entities in list(new_models.items())[::-1]:
db.session.query(model).filter(
model.id.in_(entity.id for entity in entities)
).delete(synchronize_session=False)
db.session.commit()
if current_revision != revision and not force:
click.confirm(f"\nRevert DB to {revision}?", abort=True)
upgrade(revision=revision)
print("Reverted")
if __name__ == "__main__":
from superset.app import create_app
app = create_app()
with app.app_context():
# pylint: disable=no-value-for-parameter
main()