From 0f5d96f86072d245161b048fe1df4c5b3de154c4 Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Fri, 14 Jul 2023 16:01:59 -0700 Subject: [PATCH 1/3] feat: migrate charts on import --- .../charts/commands/importers/v1/utils.py | 48 +++++++ .../shared/migrate_viz/processors.py | 3 + .../commands/importers/v1/utils_test.py | 118 ++++++++++++++++++ 3 files changed, 169 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..079376fab67f5 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,44 @@ 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_ != MigrateViz + } + + 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..4db308e04c093 100644 --- a/superset/migrations/shared/migrate_viz/processors.py +++ b/superset/migrations/shared/migrate_viz/processors.py @@ -51,6 +51,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..ab8fd0d1f7ded --- /dev/null +++ b/tests/unit_tests/charts/commands/importers/v1/utils_test.py @@ -0,0 +1,118 @@ +# 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() -> None: + """ + Test the ``migrate_chart`` command when importing a chart. + """ + chart_config = { + "slice_name": "Birth names by state", + "description": None, + "certified_by": None, + "certification_details": None, + "viz_type": "area", + "params": json.dumps( + { + "datasource": "21__table", + "viz_type": "area", + "granularity_sqla": "ds", + "time_grain_sqla": "P1D", + "time_range": "No filter", + "metrics": ["count"], + "adhoc_filters": [], + "groupby": ["state"], + "order_desc": True, + "row_limit": 10000, + "show_brush": "auto", + "show_legend": True, + "line_interpolation": "linear", + "stacked_style": "stack", + "color_scheme": "supersetColors", + "rich_tooltip": True, + "bottom_margin": "auto", + "x_ticks_layout": "auto", + "x_axis_format": "smart_date", + "y_axis_format": "SMART_NUMBER", + "y_axis_bounds": [None, None], + "rolling_type": "None", + "comparison_type": "values", + "annotation_layers": [], + "extra_form_data": {}, + "dashboards": [], + } + ), + "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": "Birth names by state", + "description": None, + "certified_by": None, + "certification_details": None, + "viz_type": "echarts_area", + "params": json.dumps( + { + "datasource": "21__table", + "viz_type": "echarts_area", + "time_grain_sqla": "P1D", + "metrics": ["count"], + "adhoc_filters": [ + { + "clause": "WHERE", + "subject": "ds", + "operator": "TEMPORAL_RANGE", + "comparator": "No filter", + "expressionType": "SIMPLE", + } + ], + "groupby": ["state"], + "order_desc": True, + "row_limit": 10000, + "show_brush": "auto", + "show_legend": True, + "line_interpolation": "linear", + "color_scheme": "supersetColors", + "rich_tooltip": True, + "bottom_margin": "auto", + "x_ticks_layout": "auto", + "x_axis_format": "smart_date", + "y_axis_format": "SMART_NUMBER", + "y_axis_bounds": [None, None], + "rolling_type": "None", + "comparison_type": "values", + "annotation_layers": [], + "extra_form_data": {}, + "dashboards": [], + "show_extra_controls": True, + "stack": "Stack", + "x_axis": "ds", + } + ), + "cache_timeout": None, + "uuid": "ffd15af2-2188-425c-b6b4-df28aac45872", + "version": "1.0.0", + "dataset_uuid": "a18b9cb0-b8d3-42ed-bd33-0f0fadbf0f6d", + } From 30c4d966f832c32f33545b3e944618f52a46db9c Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Mon, 17 Jul 2023 13:51:43 -0700 Subject: [PATCH 2/3] Skip the area migration --- .../charts/commands/importers/v1/utils.py | 6 +- .../migrations/shared/migrate_viz/base.py | 10 ++ .../shared/migrate_viz/processors.py | 8 +- .../commands/importers/v1/utils_test.py | 137 ++++++++++++------ 4 files changed, 113 insertions(+), 48 deletions(-) diff --git a/superset/charts/commands/importers/v1/utils.py b/superset/charts/commands/importers/v1/utils.py index 079376fab67f5..e258abe269ff6 100644 --- a/superset/charts/commands/importers/v1/utils.py +++ b/superset/charts/commands/importers/v1/utils.py @@ -26,7 +26,7 @@ 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.migrations.shared.migrate_viz.base import ChartMigratorStatus, MigrateViz from superset.models.slice import Slice @@ -70,7 +70,9 @@ def migrate_chart(config: dict[str, Any]) -> dict[str, Any]: migrators = { class_.source_viz_type: class_ for class_ in processors.__dict__.values() - if isclass(class_) and issubclass(class_, MigrateViz) and class_ != MigrateViz + if isclass(class_) + and issubclass(class_, MigrateViz) + and class_.status == ChartMigratorStatus.COMPLETE } output = copy.deepcopy(config) diff --git a/superset/migrations/shared/migrate_viz/base.py b/superset/migrations/shared/migrate_viz/base.py index e277fcd208e81..e10f56cac540f 100644 --- a/superset/migrations/shared/migrate_viz/base.py +++ b/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) diff --git a/superset/migrations/shared/migrate_viz/processors.py b/superset/migrations/shared/migrate_viz/processors.py index 4db308e04c093..cb97e6850415c 100644 --- a/superset/migrations/shared/migrate_viz/processors.py +++ b/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" @@ -83,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 @@ -104,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 diff --git a/tests/unit_tests/charts/commands/importers/v1/utils_test.py b/tests/unit_tests/charts/commands/importers/v1/utils_test.py index ab8fd0d1f7ded..c99ecaf6d6da7 100644 --- a/tests/unit_tests/charts/commands/importers/v1/utils_test.py +++ b/tests/unit_tests/charts/commands/importers/v1/utils_test.py @@ -20,9 +20,11 @@ from superset.charts.commands.importers.v1.utils import migrate_chart -def test_migrate_chart() -> None: +def test_migrate_chart_area() -> None: """ - Test the ``migrate_chart`` command when importing a chart. + 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", @@ -32,32 +34,32 @@ def test_migrate_chart() -> None: "viz_type": "area", "params": json.dumps( { + "adhoc_filters": [], + "annotation_layers": [], + "bottom_margin": "auto", + "color_scheme": "supersetColors", + "comparison_type": "values", + "dashboards": [], "datasource": "21__table", - "viz_type": "area", + "extra_form_data": {}, "granularity_sqla": "ds", - "time_grain_sqla": "P1D", - "time_range": "No filter", - "metrics": ["count"], - "adhoc_filters": [], "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, - "line_interpolation": "linear", "stacked_style": "stack", - "color_scheme": "supersetColors", - "rich_tooltip": True, - "bottom_margin": "auto", - "x_ticks_layout": "auto", + "time_grain_sqla": "P1D", + "time_range": "No filter", + "viz_type": "area", "x_axis_format": "smart_date", - "y_axis_format": "SMART_NUMBER", + "x_ticks_layout": "auto", "y_axis_bounds": [None, None], - "rolling_type": "None", - "comparison_type": "values", - "annotation_layers": [], - "extra_form_data": {}, - "dashboards": [], + "y_axis_format": "SMART_NUMBER", } ), "cache_timeout": None, @@ -66,19 +68,85 @@ def test_migrate_chart() -> None: "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": "Birth names by state", + "slice_name": "Pivot Table", "description": None, "certified_by": None, "certification_details": None, - "viz_type": "echarts_area", + "viz_type": "pivot_table_v2", "params": json.dumps( { - "datasource": "21__table", - "viz_type": "echarts_area", - "time_grain_sqla": "P1D", - "metrics": ["count"], + "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", @@ -88,27 +156,6 @@ def test_migrate_chart() -> None: "expressionType": "SIMPLE", } ], - "groupby": ["state"], - "order_desc": True, - "row_limit": 10000, - "show_brush": "auto", - "show_legend": True, - "line_interpolation": "linear", - "color_scheme": "supersetColors", - "rich_tooltip": True, - "bottom_margin": "auto", - "x_ticks_layout": "auto", - "x_axis_format": "smart_date", - "y_axis_format": "SMART_NUMBER", - "y_axis_bounds": [None, None], - "rolling_type": "None", - "comparison_type": "values", - "annotation_layers": [], - "extra_form_data": {}, - "dashboards": [], - "show_extra_controls": True, - "stack": "Stack", - "x_axis": "ds", } ), "cache_timeout": None, From 83bb63fc0b7c6d08b2ca431277eab88ca9922ada Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Thu, 20 Jul 2023 18:55:42 -0700 Subject: [PATCH 3/3] Add comment about area chart --- superset/charts/commands/importers/v1/utils.py | 5 +++-- superset/migrations/shared/migrate_viz/base.py | 10 ---------- .../migrations/shared/migrate_viz/processors.py | 17 ++++++++++------- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/superset/charts/commands/importers/v1/utils.py b/superset/charts/commands/importers/v1/utils.py index e258abe269ff6..589ae76a310a0 100644 --- a/superset/charts/commands/importers/v1/utils.py +++ b/superset/charts/commands/importers/v1/utils.py @@ -26,7 +26,7 @@ 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.migrations.shared.migrate_viz.base import MigrateViz from superset.models.slice import Slice @@ -72,7 +72,8 @@ def migrate_chart(config: dict[str, Any]) -> dict[str, Any]: for class_ in processors.__dict__.values() if isclass(class_) and issubclass(class_, MigrateViz) - and class_.status == ChartMigratorStatus.COMPLETE + and hasattr(class_, "source_viz_type") + and class_ != processors.MigrateAreaChart # incomplete } output = copy.deepcopy(config) diff --git a/superset/migrations/shared/migrate_viz/base.py b/superset/migrations/shared/migrate_viz/base.py index e10f56cac540f..e277fcd208e81 100644 --- a/superset/migrations/shared/migrate_viz/base.py +++ b/superset/migrations/shared/migrate_viz/base.py @@ -18,7 +18,6 @@ import copy import json -from enum import Enum from typing import Any from alembic import op @@ -32,12 +31,6 @@ Base = declarative_base() -class ChartMigratorStatus(Enum): - NOT_IMPLEMENTED = "NOT_IMPLEMENTED" - INCOMPLETE = "INCOMPLETE" - COMPLETE = "COMPLETE" - - class Slice(Base): # type: ignore __tablename__ = "slices" @@ -58,9 +51,6 @@ 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) diff --git a/superset/migrations/shared/migrate_viz/processors.py b/superset/migrations/shared/migrate_viz/processors.py index cb97e6850415c..70c3c2705542c 100644 --- a/superset/migrations/shared/migrate_viz/processors.py +++ b/superset/migrations/shared/migrate_viz/processors.py @@ -16,7 +16,7 @@ # under the License. from typing import Any -from .base import ChartMigratorStatus, MigrateViz +from .base import MigrateViz class MigrateTreeMap(MigrateViz): @@ -35,12 +35,19 @@ 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"} - status = ChartMigratorStatus.INCOMPLETE - def _pre_action(self) -> None: if self.data.get("contribution"): self.data["contributionMode"] = "row" @@ -85,8 +92,6 @@ 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 @@ -108,8 +113,6 @@ class MigrateDualLine(MigrateViz): } remove_keys = {"metric", "metric_2"} - status = ChartMigratorStatus.COMPLETE - def _pre_action(self) -> None: self.data["yAxisIndex"] = 0 self.data["yAxisIndexB"] = 1