superset 2020-09-28_17-57_b56500de1855_add_uuid_column_to_import_mixin 源码

superset 2020-09-28_17-57_b56500de1855_add_uuid_column_to_import_mixin 代码

文件路径：/superset/migrations/versions/2020-09-28_17-57_b56500de1855_add_uuid_column_to_import_mixin.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.
"""add_uuid_column_to_import_mixin

Revision ID: b56500de1855
Revises: 18532d70ab98
Create Date: 2020-09-28 17:57:23.128142

"""
import json
import os
from json.decoder import JSONDecodeError
from uuid import uuid4

import sqlalchemy as sa
from alembic import op
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import load_only
from sqlalchemy_utils import UUIDType

from superset import db
from superset.migrations.shared.utils import assign_uuids
from superset.utils import core as utils

# revision identifiers, used by Alembic.
revision = "b56500de1855"
down_revision = "18532d70ab98"


Base = declarative_base()


class ImportMixin:
    id = sa.Column(sa.Integer, primary_key=True)
    uuid = sa.Column(UUIDType(binary=True), primary_key=False, default=uuid4)


table_names = [
    # Core models
    "dbs",
    "dashboards",
    "slices",
    # SQLAlchemy connectors
    "tables",
    "table_columns",
    "sql_metrics",
    # Druid connector
    "clusters",
    "datasources",
    "columns",
    "metrics",
    # Dashboard email schedules
    "dashboard_email_schedules",
    "slice_email_schedules",
]
models = {
    table_name: type(table_name, (Base, ImportMixin), {"__tablename__": table_name})
    for table_name in table_names
}

models["dashboards"].position_json = sa.Column(utils.MediumText())

default_batch_size = int(os.environ.get("BATCH_SIZE", 200))


def update_position_json(dashboard, session, uuid_map):
    try:
        layout = json.loads(dashboard.position_json or "{}")
    except JSONDecodeError:
        layout = {}

    for object_ in layout.values():
        if (
            isinstance(object_, dict)
            and object_["type"] == "CHART"
            and object_["meta"]["chartId"]
        ):
            chart_id = object_["meta"]["chartId"]
            if chart_id in uuid_map:
                object_["meta"]["uuid"] = str(uuid_map[chart_id])
            elif object_["meta"].get("uuid"):
                del object_["meta"]["uuid"]

    dashboard.position_json = json.dumps(layout, indent=4)
    session.merge(dashboard)


def update_dashboards(session, uuid_map):
    message = (
        "Updating dasboard position json with slice uuid.."
        if uuid_map
        else "Cleaning up slice uuid from dashboard position json.."
    )
    print(f"\n{message}\r", end="")

    query = session.query(models["dashboards"])
    dashboard_count = query.count()
    for i, dashboard in enumerate(query.all()):
        update_position_json(dashboard, session, uuid_map)
        if i and i % default_batch_size == 0:
            session.commit()
        print(f"{message} {i+1}/{dashboard_count}\r", end="")

    session.commit()
    # Extra whitespace to override very long numbers, e.g. 99999/99999.
    print(f"{message} Done.      \n")


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

    for table_name, model in models.items():
        with op.batch_alter_table(table_name) as batch_op:
            batch_op.add_column(
                sa.Column(
                    "uuid",
                    UUIDType(binary=True),
                    primary_key=False,
                    default=uuid4,
                ),
            )

        assign_uuids(model, session)

        # add uniqueness constraint
        with op.batch_alter_table(table_name) as batch_op:
            # batch mode is required for sqlite
            batch_op.create_unique_constraint(f"uq_{table_name}_uuid", ["uuid"])

    # add UUID to Dashboard.position_json
    slice_uuid_map = {
        slc.id: slc.uuid
        for slc in session.query(models["slices"])
        .options(load_only("id", "uuid"))
        .all()
    }
    update_dashboards(session, slice_uuid_map)


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

    # remove uuid from position_json
    update_dashboards(session, {})

    # remove uuid column
    for table_name in models:
        with op.batch_alter_table(table_name) as batch_op:
            batch_op.drop_constraint(f"uq_{table_name}_uuid", type_="unique")
            batch_op.drop_column("uuid")

相关信息

superset 源码目录

相关文章

superset 2015-09-21_17-30_4e6a06bad7a8_init 源码

superset 2015-10-05_10-325a7bad26f2a7 源码

superset 2015-10-05_22-111e2841a4128 源码

superset 2015-10-19_20-54_2929af7925ed_tz_offsets_in_data_sources 源码

superset 2015-11-21_11-18_289ce07647b_add_encrypted_password_field 源码

superset 2015-12-04_09-42_1a48a5411020_adding_slug_to_dash 源码

superset 2015-12-04_11-16_315b3f4da9b0_adding_log_model 源码

superset 2015-12-13_08-38_55179c7f25c7_sqla_descr 源码

superset 2015-12-14_13-37_12d55656cbca_is_featured 源码

superset 2015-12-15_17-02_2591d77e9831_user_id 源码