feat: migrate charts on import

superset/charts/commands/importers/v1/utils.py

@@ -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 ChartMigratorStatus, MigrateViz
 from superset.models.slice import Slice
 
 
@@ -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()
@@ -54,3 +61,46 @@ 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 class_.status == ChartMigratorStatus.COMPLETE
+    }
+
+    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

superset/migrations/shared/migrate_viz/base.py

@@ -18,6 +18,7 @@
 
 import copy
 import json
+from enum import Enum
 from typing import Any
 
 from alembic import op
@@ -31,6 +32,12 @@
 Base = declarative_base()
 
 
+class ChartMigratorStatus(Enum):
+    NOT_IMPLEMENTED = "NOT_IMPLEMENTED"
+    INCOMPLETE = "INCOMPLETE"
+    COMPLETE = "COMPLETE"
+
+
 class Slice(Base):  # type: ignore
     __tablename__ = "slices"
 
@@ -51,6 +58,9 @@ class MigrateViz:
     target_viz_type: str
     has_x_axis_control: bool = False
 
+    # specific migrations need to override this
+    status = ChartMigratorStatus.NOT_IMPLEMENTED
+
     def __init__(self, form_data: str) -> None:
         self.data = try_load_json(form_data)
 

superset/migrations/shared/migrate_viz/processors.py

@@ -16,7 +16,7 @@
 # under the License.
 from typing import Any
 
-from .base import MigrateViz
+from .base import ChartMigratorStatus, MigrateViz
 
 
 class MigrateTreeMap(MigrateViz):
@@ -39,6 +39,8 @@ class MigrateAreaChart(MigrateViz):
     target_viz_type = "echarts_area"
     remove_keys = {"contribution", "stacked_style", "x_axis_label"}
 
+    status = ChartMigratorStatus.INCOMPLETE
+
     def _pre_action(self) -> None:
         if self.data.get("contribution"):
             self.data["contributionMode"] = "row"
@@ -51,6 +53,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
@@ -80,6 +85,8 @@ class MigratePivotTable(MigrateViz):
         "var": "Sample Variance",
     }
 
+    status = ChartMigratorStatus.COMPLETE
+
     def _pre_action(self) -> None:
         if pivot_margins := self.data.get("pivot_margins"):
             self.data["colTotals"] = pivot_margins
@@ -101,6 +108,8 @@ class MigrateDualLine(MigrateViz):
     }
     remove_keys = {"metric", "metric_2"}
 
+    status = ChartMigratorStatus.COMPLETE
+
     def _pre_action(self) -> None:
         self.data["yAxisIndex"] = 0
         self.data["yAxisIndexB"] = 1

tests/unit_tests/charts/commands/importers/v1/utils_test.py

@@ -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",
+    }