Skip to content

Commit

Permalink
feat: migrate charts on import (#24703)
Browse files Browse the repository at this point in the history
  • Loading branch information
betodealmeida authored Jul 21, 2023
1 parent e210da9 commit abb8e28
Show file tree
Hide file tree
Showing 3 changed files with 228 additions and 0 deletions.
51 changes: 51 additions & 0 deletions superset/charts/commands/importers/v1/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,14 +15,18 @@
# specific language governing permissions and limitations
# under the License.

import copy
import json
from inspect import isclass
from typing import Any

from flask import g
from sqlalchemy.orm import Session

from superset import security_manager
from superset.commands.exceptions import ImportFailedError
from superset.migrations.shared.migrate_viz import processors
from superset.migrations.shared.migrate_viz.base import MigrateViz
from superset.models.slice import Slice


Expand All @@ -46,6 +50,9 @@ def import_chart(
# TODO (betodealmeida): move this logic to import_from_dict
config["params"] = json.dumps(config["params"])

# migrate old viz types to new ones
config = migrate_chart(config)

chart = Slice.import_from_dict(session, config, recursive=False)
if chart.id is None:
session.flush()
Expand All @@ -54,3 +61,47 @@ def import_chart(
chart.owners.append(g.user)

return chart


def migrate_chart(config: dict[str, Any]) -> dict[str, Any]:
"""
Used to migrate old viz types to new ones.
"""
migrators = {
class_.source_viz_type: class_
for class_ in processors.__dict__.values()
if isclass(class_)
and issubclass(class_, MigrateViz)
and hasattr(class_, "source_viz_type")
and class_ != processors.MigrateAreaChart # incomplete
}

output = copy.deepcopy(config)
if config["viz_type"] not in migrators:
return output

migrator = migrators[config["viz_type"]](output["params"])
# pylint: disable=protected-access
migrator._pre_action()
migrator._migrate()
migrator._post_action()
params = migrator.data

params["viz_type"] = migrator.target_viz_type
output.update(
{
"params": json.dumps(params),
"viz_type": migrator.target_viz_type,
}
)

# also update `query_context`
try:
query_context = json.loads(output.get("query_context", "{}"))
except json.decoder.JSONDecodeError:
query_context = {}
if "form_data" in query_context:
query_context["form_data"] = output["params"]
output["query_context"] = json.dumps(query_context)

return output
12 changes: 12 additions & 0 deletions superset/migrations/shared/migrate_viz/processors.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,15 @@ def _pre_action(self) -> None:


class MigrateAreaChart(MigrateViz):
"""
Migrate area charts.
This migration is incomplete, see https://github.com/apache/superset/pull/24703#discussion_r1265222611
for more details. If you fix this migration, please update the ``migrate_chart``
function in ``superset/charts/commands/importers/v1/utils.py`` so that it gets
applied in chart imports.
"""

source_viz_type = "area"
target_viz_type = "echarts_area"
remove_keys = {"contribution", "stacked_style", "x_axis_label"}
Expand All @@ -51,6 +60,9 @@ def _pre_action(self) -> None:
self.data["show_extra_controls"] = True
self.data["stack"] = stacked_map.get(stacked)

if x_axis := self.data.get("granularity_sqla"):
self.data["x_axis"] = x_axis

if x_axis_label := self.data.get("x_axis_label"):
self.data["x_axis_title"] = x_axis_label
self.data["x_axis_title_margin"] = 30
Expand Down
165 changes: 165 additions & 0 deletions tests/unit_tests/charts/commands/importers/v1/utils_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,165 @@
# 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 json

from superset.charts.commands.importers.v1.utils import migrate_chart


def test_migrate_chart_area() -> None:
"""
Test the ``migrate_chart`` command when importing an area chart.
This is currently a no-op since the migration is not complete.
"""
chart_config = {
"slice_name": "Birth names by state",
"description": None,
"certified_by": None,
"certification_details": None,
"viz_type": "area",
"params": json.dumps(
{
"adhoc_filters": [],
"annotation_layers": [],
"bottom_margin": "auto",
"color_scheme": "supersetColors",
"comparison_type": "values",
"dashboards": [],
"datasource": "21__table",
"extra_form_data": {},
"granularity_sqla": "ds",
"groupby": ["state"],
"line_interpolation": "linear",
"metrics": ["count"],
"order_desc": True,
"rich_tooltip": True,
"rolling_type": "None",
"row_limit": 10000,
"show_brush": "auto",
"show_legend": True,
"stacked_style": "stack",
"time_grain_sqla": "P1D",
"time_range": "No filter",
"viz_type": "area",
"x_axis_format": "smart_date",
"x_ticks_layout": "auto",
"y_axis_bounds": [None, None],
"y_axis_format": "SMART_NUMBER",
}
),
"cache_timeout": None,
"uuid": "ffd15af2-2188-425c-b6b4-df28aac45872",
"version": "1.0.0",
"dataset_uuid": "a18b9cb0-b8d3-42ed-bd33-0f0fadbf0f6d",
}

new_config = migrate_chart(chart_config)
assert new_config == chart_config


def test_migrate_pivot_table() -> None:
"""
Test the ``migrate_chart`` command when importing an old pivot table.
"""
chart_config = {
"slice_name": "Pivot Table",
"description": None,
"certified_by": None,
"certification_details": None,
"viz_type": "pivot_table",
"params": json.dumps(
{
"columns": ["state"],
"compare_lag": "10",
"compare_suffix": "o10Y",
"granularity_sqla": "ds",
"groupby": ["name"],
"limit": "25",
"markup_type": "markdown",
"metrics": [
{
"aggregate": "SUM",
"column": {
"column_name": "num",
"type": "BIGINT",
},
"expressionType": "SIMPLE",
"label": "Births",
"optionName": "metric_11",
},
],
"row_limit": 50000,
"since": "100 years ago",
"time_range": "No filter",
"time_range_endpoints": ["inclusive", "exclusive"],
"until": "now",
"viz_type": "pivot_table",
},
),
"cache_timeout": None,
"uuid": "ffd15af2-2188-425c-b6b4-df28aac45872",
"version": "1.0.0",
"dataset_uuid": "a18b9cb0-b8d3-42ed-bd33-0f0fadbf0f6d",
}

new_config = migrate_chart(chart_config)
assert new_config == {
"slice_name": "Pivot Table",
"description": None,
"certified_by": None,
"certification_details": None,
"viz_type": "pivot_table_v2",
"params": json.dumps(
{
"groupbyColumns": ["state"],
"compare_lag": "10",
"compare_suffix": "o10Y",
"groupbyRows": ["name"],
"limit": "25",
"markup_type": "markdown",
"metrics": [
{
"aggregate": "SUM",
"column": {"column_name": "num", "type": "BIGINT"},
"expressionType": "SIMPLE",
"label": "Births",
"optionName": "metric_11",
}
],
"series_limit": 50000,
"since": "100 years ago",
"time_range_endpoints": ["inclusive", "exclusive"],
"until": "now",
"viz_type": "pivot_table_v2",
"rowOrder": "value_z_to_a",
"adhoc_filters": [
{
"clause": "WHERE",
"subject": "ds",
"operator": "TEMPORAL_RANGE",
"comparator": "No filter",
"expressionType": "SIMPLE",
}
],
}
),
"cache_timeout": None,
"uuid": "ffd15af2-2188-425c-b6b4-df28aac45872",
"version": "1.0.0",
"dataset_uuid": "a18b9cb0-b8d3-42ed-bd33-0f0fadbf0f6d",
}

0 comments on commit abb8e28

Please sign in to comment.