From 356b0d8ee56f42fe2de651b1a8ee08f1c4857172 Mon Sep 17 00:00:00 2001 From: "Michael S. Molina" <70410625+michael-s-molina@users.noreply.github.com> Date: Fri, 29 Mar 2024 12:08:43 -0300 Subject: [PATCH] feat: Adds Heatmap chart migration logic (#27771) --- requirements/development.txt | 15 +++- superset/cli/viz_migrations.py | 3 + .../shared/migrate_viz/processors.py | 15 ++++ .../migrations/viz/heatmap_v1_v2_test.py | 78 +++++++++++++++++++ 4 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 tests/unit_tests/migrations/viz/heatmap_v1_v2_test.py diff --git a/requirements/development.txt b/requirements/development.txt index 23c356c06aa97..48302d4996f2f 100644 --- a/requirements/development.txt +++ b/requirements/development.txt @@ -10,6 +10,8 @@ # via # -r requirements/base.in # -r requirements/development.in +appnope==0.1.4 + # via ipython astroid==2.15.8 # via pylint asttokens==2.2.1 @@ -193,6 +195,10 @@ ptyprocess==0.7.0 # via pexpect pure-eval==0.2.2 # via stack-data +pure-sasl==0.6.2 + # via + # pyhive + # thrift-sasl pydata-google-auth==1.7.0 # via pandas-gbq pydruid==0.6.6 @@ -201,7 +207,7 @@ pyee==11.0.1 # via playwright pyfakefs==5.2.2 # via -r requirements/development.in -pyhive[presto]==0.7.0 +pyhive[hive_pure_sasl]==0.7.0 # via apache-superset pyinstrument==4.4.0 # via -r requirements/development.in @@ -243,7 +249,12 @@ statsd==4.0.1 tableschema==1.20.10 # via apache-superset thrift==0.20.0 - # via apache-superset + # via + # apache-superset + # pyhive + # thrift-sasl +thrift-sasl==0.4.3 + # via pyhive tomli==2.0.1 # via # build diff --git a/superset/cli/viz_migrations.py b/superset/cli/viz_migrations.py index f24dd8f444cbf..c80c3a5136945 100644 --- a/superset/cli/viz_migrations.py +++ b/superset/cli/viz_migrations.py @@ -27,6 +27,7 @@ class VizType(str, Enum): AREA = "area" BUBBLE = "bubble" DUAL_LINE = "dual_line" + HEATMAP = "heatmap" LINE = "line" PIVOT_TABLE = "pivot_table" SUNBURST = "sunburst" @@ -79,6 +80,7 @@ def migrate(viz_type: VizType, is_downgrade: bool = False) -> None: MigrateAreaChart, MigrateBubbleChart, MigrateDualLine, + MigrateHeatmapChart, MigrateLineChart, MigratePivotTable, MigrateSunburst, @@ -89,6 +91,7 @@ def migrate(viz_type: VizType, is_downgrade: bool = False) -> None: VizType.AREA: MigrateAreaChart, VizType.BUBBLE: MigrateBubbleChart, VizType.DUAL_LINE: MigrateDualLine, + VizType.HEATMAP: MigrateHeatmapChart, VizType.LINE: MigrateLineChart, VizType.PIVOT_TABLE: MigratePivotTable, VizType.SUNBURST: MigrateSunburst, diff --git a/superset/migrations/shared/migrate_viz/processors.py b/superset/migrations/shared/migrate_viz/processors.py index 5fbd624aa8bf9..e1a77bd5bc87d 100644 --- a/superset/migrations/shared/migrate_viz/processors.py +++ b/superset/migrations/shared/migrate_viz/processors.py @@ -213,3 +213,18 @@ def _pre_action(self) -> None: # Truncate y-axis by default to preserve layout self.data["y_axis_showminmax"] = True + + +class MigrateHeatmapChart(MigrateViz): + source_viz_type = "heatmap" + target_viz_type = "heatmap_v2" + rename_keys = { + "all_columns_x": "x_axis", + "all_columns_y": "groupby", + "y_axis_bounds": "value_bounds", + "show_perc": "show_percentage", + } + remove_keys = {"sort_by_metric", "canvas_image_rendering"} + + def _pre_action(self) -> None: + self.data["legend_type"] = "continuous" diff --git a/tests/unit_tests/migrations/viz/heatmap_v1_v2_test.py b/tests/unit_tests/migrations/viz/heatmap_v1_v2_test.py new file mode 100644 index 0000000000000..e343df904f7eb --- /dev/null +++ b/tests/unit_tests/migrations/viz/heatmap_v1_v2_test.py @@ -0,0 +1,78 @@ +# 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 typing import Any + +from superset.migrations.shared.migrate_viz import MigrateHeatmapChart +from tests.unit_tests.migrations.viz.utils import migrate_and_assert + +SOURCE_FORM_DATA: dict[str, Any] = { + "any_other_key": "untouched", + "all_columns_x": ["category"], + "all_columns_y": ["product"], + "metric": ["sales"], + "adhoc_filters": [], + "row_limit": 100, + "sort_by_metric": True, + "linear_color_scheme": "blue", + "xscale_interval": 2, + "yscale_interval": 2, + "canvas_image_rendering": "auto", + "normalize_across": "x", + "left_margin": 50, + "bottom_margin": 50, + "y_axis_bounds": [0, 100], + "y_axis_format": "SMART_NUMBER", + "currency_format": "USD", + "sort_x_axis": "alpha_asc", + "sort_y_axis": "alpha_asc", + "show_legend": True, + "show_perc": True, + "show_values": True, + "normalized": True, + "viz_type": "heatmap", +} + +TARGET_FORM_DATA: dict[str, Any] = { + "any_other_key": "untouched", + "x_axis": ["category"], + "groupby": ["product"], + "metric": ["sales"], + "adhoc_filters": [], + "row_limit": 100, + "legend_type": "continuous", + "linear_color_scheme": "blue", + "xscale_interval": 2, + "yscale_interval": 2, + "normalize_across": "x", + "left_margin": 50, + "bottom_margin": 50, + "value_bounds": [0, 100], + "y_axis_format": "SMART_NUMBER", + "currency_format": "USD", + "sort_x_axis": "alpha_asc", + "sort_y_axis": "alpha_asc", + "show_legend": True, + "show_percentage": True, + "show_values": True, + "normalized": True, + "viz_type": "heatmap_v2", + "form_data_bak": SOURCE_FORM_DATA, +} + + +def test_migration() -> None: + migrate_and_assert(MigrateHeatmapChart, SOURCE_FORM_DATA, TARGET_FORM_DATA)