From 6947983d5faadc78758bc05fb1e35dccb9c69855 Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Fri, 21 Jul 2023 09:33:34 -0700 Subject: [PATCH] feat: migrate charts on import (#24703) (cherry picked from commit abb8e28e4914ad46ef50e33934ec97c1e8fcf5b4) --- .../charts/commands/importers/v1/utils.py | 51 ++++++ .../shared/migrate_viz/processors.py | 12 ++ .../commands/importers/v1/utils_test.py | 165 ++++++++++++++++++ 3 files changed, 228 insertions(+) create mode 100644 tests/unit_tests/charts/commands/importers/v1/utils_test.py diff --git a/superset/charts/commands/importers/v1/utils.py b/superset/charts/commands/importers/v1/utils.py index 399e6c2243fa2..589ae76a310a0 100644 --- a/superset/charts/commands/importers/v1/utils.py +++ b/superset/charts/commands/importers/v1/utils.py @@ -15,7 +15,9 @@ # 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 @@ -23,6 +25,8 @@ 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 @@ -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,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 diff --git a/superset/migrations/shared/migrate_viz/processors.py b/superset/migrations/shared/migrate_viz/processors.py index fa7043fc38dd1..70c3c2705542c 100644 --- a/superset/migrations/shared/migrate_viz/processors.py +++ b/superset/migrations/shared/migrate_viz/processors.py @@ -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"} @@ -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 diff --git a/tests/unit_tests/charts/commands/importers/v1/utils_test.py b/tests/unit_tests/charts/commands/importers/v1/utils_test.py new file mode 100644 index 0000000000000..c99ecaf6d6da7 --- /dev/null +++ b/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", + }