def upgrade()

in superset/migrations/versions/2020-08-12_00-24_978245563a02_migrate_iframe_to_dash_markdown.py [0:0]

• 78 lines of code
• 22 McCabe index (conditional complexity)

def upgrade():
    bind = op.get_bind()
    session = db.Session(bind=bind)

    dash_to_migrate = defaultdict(list)
    iframe_urls = defaultdict(list)

    try:
        # find iframe viz_type and its url
        iframes = session.query(Slice).filter_by(viz_type="iframe").all()
        iframe_ids = [slc.id for slc in iframes]

        for iframe in iframes:
            iframe_params = json.loads(iframe.params or "{}")
            url = iframe_params.get("url")
            iframe_urls[iframe.id] = url

        # find iframe viz_type that used in dashboard
        dash_slc = (
            session.query(dashboard_slices)
            .filter(dashboard_slices.c.slice_id.in_(iframe_ids))
            .all()
        )
        for entry in dash_slc:
            dash_to_migrate[entry.dashboard_id].append(entry.slice_id)
        dashboard_ids = list(dash_to_migrate.keys())

        # replace iframe in dashboard metadata
        dashboards = (
            session.query(Dashboard).filter(Dashboard.id.in_(dashboard_ids)).all()
        )
        for i, dashboard in enumerate(dashboards):
            print(
                f"scanning dashboard ({i + 1}/{len(dashboards)}) dashboard: {dashboard.id} >>>>"  # noqa: E501
            )

            # remove iframe slices from dashboard
            dashboard.slices = [
                slc for slc in dashboard.slices if slc.id not in iframe_ids
            ]

            # find iframe chart position in metadata
            # and replace it with markdown component
            position_dict = json.loads(dashboard.position_json or "{}")
            keys_to_remove = []
            for key, chart_position in position_dict.items():
                if (
                    chart_position
                    and isinstance(chart_position, dict)
                    and chart_position["type"] == "CHART"
                    and chart_position["meta"]
                    and chart_position["meta"]["chartId"] in iframe_ids
                ):
                    iframe_id = chart_position["meta"]["chartId"]
                    # make new markdown component
                    markdown = create_new_markdown_component(
                        chart_position, iframe_urls[iframe_id]
                    )
                    keys_to_remove.append(key)
                    position_dict[markdown["id"]] = markdown

                    # add markdown to layout tree
                    parent_id = markdown["parents"][-1]
                    children = position_dict[parent_id]["children"]
                    children.remove(key)
                    children.append(markdown["id"])

            if keys_to_remove:
                for key_to_remove in keys_to_remove:
                    del position_dict[key_to_remove]
                dashboard.position_json = json.dumps(
                    position_dict,
                    indent=None,
                    separators=(",", ":"),
                    sort_keys=True,
                )

        # remove iframe, separator and markup charts
        slices_to_remove = (
            session.query(Slice)
            .filter(Slice.viz_type.in_(["iframe", "separator", "markup"]))
            .all()
        )
        slices_ids = [slc.id for slc in slices_to_remove]

        # remove dependencies first
        session.query(dashboard_slices).filter(
            dashboard_slices.c.slice_id.in_(slices_ids)
        ).delete(synchronize_session=False)

        session.query(slice_user).filter(slice_user.c.slice_id.in_(slices_ids)).delete(
            synchronize_session=False
        )

        # remove slices
        session.query(Slice).filter(Slice.id.in_(slices_ids)).delete(
            synchronize_session=False
        )

    except Exception as ex:
        logging.exception(f"dashboard {dashboard.id} has error: {ex}")

    session.commit()
    session.close()