diff --git a/scripts/benchmark_migration.py b/scripts/benchmark_migration.py index 72cbc68d2..163fc2983 100644 --- a/scripts/benchmark_migration.py +++ b/scripts/benchmark_migration.py @@ -30,7 +30,7 @@ from flask_appbuilder import Model from flask_migrate import downgrade, upgrade from graphlib import TopologicalSorter # pylint: disable=wrong-import-order from progress.bar import ChargingBar -from sqlalchemy import create_engine, inspect, Table +from sqlalchemy import create_engine, inspect from sqlalchemy.ext.automap import automap_base from superset import db @@ -172,7 +172,6 @@ def main( rows = session.query(model).count() print(f"- {model.__name__} ({rows} rows in table {model.__tablename__})") model_rows[model] = rows - session.close() print("Benchmarking migration") results: Dict[str, float] = {} diff --git a/superset/migrations/versions/49b5a32daba5_add_report_schedules.py b/superset/migrations/versions/49b5a32daba5_add_report_schedules.py index f20745471..edf65728b 100644 --- a/superset/migrations/versions/49b5a32daba5_add_report_schedules.py +++ b/superset/migrations/versions/49b5a32daba5_add_report_schedules.py @@ -28,6 +28,7 @@ down_revision = "96e99fb176a0" import sqlalchemy as sa from alembic import op +from sqlalchemy.engine.reflection import Inspector from sqlalchemy.exc import OperationalError @@ -119,13 +120,17 @@ def upgrade(): ) +def has_unique_constraint(constraint_name: str, table_name: str) -> bool: + bind = op.get_bind() + inspector = Inspector.from_engine(bind) + unique_constraints = inspector.get_unique_constraints(table_name) + return constraint_name in {constraint["name"] for constraint in unique_constraints} + + def downgrade(): op.drop_index(op.f("ix_report_schedule_active"), table_name="report_schedule") - try: + if has_unique_constraint("uq_report_schedule_name", "report_schedule"): op.drop_constraint("uq_report_schedule_name", "report_schedule", type_="unique") - except Exception: - # Expected to fail on SQLite - pass op.drop_table("report_execution_log") op.drop_table("report_recipient")