Files
superset2/superset/cli/viz_migrations.py
2024-06-26 08:36:07 -03:00

116 lines
3.3 KiB
Python

# 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.
from enum import Enum
import click
from click_option_group import optgroup, RequiredMutuallyExclusiveOptionGroup
from flask.cli import with_appcontext
from superset import db
class VizType(str, Enum):
AREA = "area"
BAR = "bar"
BUBBLE = "bubble"
DIST_BAR = "dist_bar"
DUAL_LINE = "dual_line"
HEATMAP = "heatmap"
HISTOGRAM = "histogram"
LINE = "line"
PIVOT_TABLE = "pivot_table"
SANKEY = "sankey"
SUNBURST = "sunburst"
TREEMAP = "treemap"
@click.group()
def migrate_viz() -> None:
"""
Migrate a viz from one type to another.
"""
@migrate_viz.command()
@with_appcontext
@optgroup.group(
"Grouped options",
cls=RequiredMutuallyExclusiveOptionGroup,
)
@optgroup.option(
"--viz_type",
"-t",
help=f"The viz type to migrate: {', '.join(list(VizType))}",
)
def upgrade(viz_type: str) -> None:
"""Upgrade a viz to the latest version."""
migrate(VizType(viz_type))
@migrate_viz.command()
@with_appcontext
@optgroup.group(
"Grouped options",
cls=RequiredMutuallyExclusiveOptionGroup,
)
@optgroup.option(
"--viz_type",
"-t",
help=f"The viz type to migrate: {', '.join(list(VizType))}",
)
def downgrade(viz_type: str) -> None:
"""Downgrade a viz to the previous version."""
migrate(VizType(viz_type), is_downgrade=True)
def migrate(viz_type: VizType, is_downgrade: bool = False) -> None:
"""Migrate a viz from one type to another."""
# pylint: disable=import-outside-toplevel
from superset.migrations.shared.migrate_viz.processors import (
MigrateAreaChart,
MigrateBarChart,
MigrateBubbleChart,
MigrateDistBarChart,
MigrateDualLine,
MigrateHeatmapChart,
MigrateHistogramChart,
MigrateLineChart,
MigratePivotTable,
MigrateSankey,
MigrateSunburst,
MigrateTreeMap,
)
migrations = {
VizType.AREA: MigrateAreaChart,
VizType.BAR: MigrateBarChart,
VizType.BUBBLE: MigrateBubbleChart,
VizType.DIST_BAR: MigrateDistBarChart,
VizType.DUAL_LINE: MigrateDualLine,
VizType.HEATMAP: MigrateHeatmapChart,
VizType.HISTOGRAM: MigrateHistogramChart,
VizType.LINE: MigrateLineChart,
VizType.PIVOT_TABLE: MigratePivotTable,
VizType.SANKEY: MigrateSankey,
VizType.SUNBURST: MigrateSunburst,
VizType.TREEMAP: MigrateTreeMap,
}
if is_downgrade:
migrations[viz_type].downgrade(db.session)
else:
migrations[viz_type].upgrade(db.session)