From 4e6e87f62ee749e8326fa734225dde955c4d0025 Mon Sep 17 00:00:00 2001 From: Yongjie Zhao Date: Fri, 8 Jul 2022 19:56:03 +0800 Subject: [PATCH] feat: Area viz migration (#20359) --- ...4daf4_add_user_id_dttm_idx_to_log_model.py | 4 +- ..._c747c78868b6_migrating_legacy_treemap.py} | 4 +- ...4-00_06e1e70058c7_migrating_legacy_area.py | 101 ++++++++++++++++++ superset/utils/migrate_viz.py | 38 +++++-- .../viz_migration/area_migration_test.py | 99 +++++++++++++++++ 5 files changed, 236 insertions(+), 10 deletions(-) rename superset/migrations/versions/{2022-06-30_22-04_c747c78868b6_migrating_legacy_treemap.py => 2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py} (98%) create mode 100644 superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py create mode 100644 tests/unit_tests/utils/viz_migration/area_migration_test.py diff --git a/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py b/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py index 207cffb4179c..1122571e17a8 100644 --- a/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py +++ b/superset/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py @@ -17,14 +17,14 @@ """Add user_id and dttm composite index to Log model Revision ID: cdcf3d64daf4 -Revises: b0d0249074e4 +Revises: 7fb8bca906d2 Create Date: 2022-04-05 13:27:06.028908 """ # revision identifiers, used by Alembic. revision = "cdcf3d64daf4" -down_revision = "c747c78868b6" +down_revision = "7fb8bca906d2" from alembic import op diff --git a/superset/migrations/versions/2022-06-30_22-04_c747c78868b6_migrating_legacy_treemap.py b/superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py similarity index 98% rename from superset/migrations/versions/2022-06-30_22-04_c747c78868b6_migrating_legacy_treemap.py rename to superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py index c420af5fca12..5f93e7cd752c 100644 --- a/superset/migrations/versions/2022-06-30_22-04_c747c78868b6_migrating_legacy_treemap.py +++ b/superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py @@ -17,7 +17,7 @@ """Migrating legacy TreeMap Revision ID: c747c78868b6 -Revises: e786798587de +Revises: cdcf3d64daf4 Create Date: 2022-06-30 22:04:17.686635 """ @@ -25,7 +25,7 @@ # revision identifiers, used by Alembic. revision = "c747c78868b6" -down_revision = "7fb8bca906d2" +down_revision = "cdcf3d64daf4" from alembic import op from sqlalchemy import and_, Column, Integer, String, Text diff --git a/superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py b/superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py new file mode 100644 index 000000000000..3def02268d9b --- /dev/null +++ b/superset/migrations/versions/2022-07-07_14-00_06e1e70058c7_migrating_legacy_area.py @@ -0,0 +1,101 @@ +# 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. +"""Migrating legacy Area + +Revision ID: 06e1e70058c7 +Revises: c747c78868b6 +Create Date: 2022-06-13 14:17:51.872706 + +""" + +# revision identifiers, used by Alembic. +revision = "06e1e70058c7" +down_revision = "c747c78868b6" + +from alembic import op +from sqlalchemy import and_, Column, Integer, String, Text +from sqlalchemy.ext.declarative import declarative_base + +from superset import db +from superset.utils.migrate_viz import get_migrate_class, MigrateVizEnum + +area_processor = get_migrate_class[MigrateVizEnum.area] + +Base = declarative_base() + + +class Slice(Base): + __tablename__ = "slices" + + id = Column(Integer, primary_key=True) + slice_name = Column(String(250)) + viz_type = Column(String(250)) + params = Column(Text) + query_context = Column(Text) + + +def upgrade(): + bind = op.get_bind() + session = db.Session(bind=bind) + + slices = session.query(Slice).filter( + Slice.viz_type == area_processor.source_viz_type + ) + total = slices.count() + idx = 0 + for slc in slices.yield_per(1000): + try: + idx += 1 + print(f"Upgrading ({idx}/{total}): {slc.slice_name}#{slc.id}") + new_viz = area_processor.upgrade(slc) + session.merge(new_viz) + except Exception as exc: + print( + "Error while processing migration: '{}'\nError: {}\n".format( + slc.slice_name, str(exc) + ) + ) + session.commit() + session.close() + + +def downgrade(): + bind = op.get_bind() + session = db.Session(bind=bind) + + slices = session.query(Slice).filter( + and_( + Slice.viz_type == area_processor.target_viz_type, + Slice.params.like("%form_data_bak%"), + ) + ) + total = slices.count() + idx = 0 + for slc in slices.yield_per(1000): + try: + idx += 1 + print(f"Downgrading ({idx}/{total}): {slc.slice_name}#{slc.id}") + new_viz = area_processor.downgrade(slc) + session.merge(new_viz) + except Exception as exc: + print( + "Error while processing migration: '{}'\nError: {}\n".format( + slc.slice_name, str(exc) + ) + ) + session.commit() + session.close() diff --git a/superset/utils/migrate_viz.py b/superset/utils/migrate_viz.py index 65ae467cb803..6e59f1257fcb 100644 --- a/superset/utils/migrate_viz.py +++ b/superset/utils/migrate_viz.py @@ -24,12 +24,6 @@ from superset.models.slice import Slice -# pylint: disable=invalid-name -class MigrateVizEnum(str, Enum): - # the Enum member name is viz_type in database - treemap = "treemap" - - class MigrateViz: remove_keys: Set[str] = set() mapping_keys: Dict[str, str] = {} @@ -117,6 +111,38 @@ def _pre_action(self) -> None: self.data["metric"] = self.data["metrics"][0] +class MigrateArea(MigrateViz): + source_viz_type = "area" + target_viz_type = "echarts_area" + remove_keys = {"contribution", "stacked_style", "x_axis_label"} + + def _pre_action(self) -> None: + if self.data.get("contribution"): + self.data["contributionMode"] = "row" + + stacked = self.data.get("stacked_style") + if stacked: + stacked_map = { + "expand": "Expand", + "stack": "Stack", + } + self.data["show_extra_controls"] = True + self.data["stack"] = stacked_map.get(stacked) + + x_axis_label = self.data.get("x_axis_label") + if x_axis_label: + self.data["x_axis_title"] = x_axis_label + self.data["x_axis_title_margin"] = 30 + + +# pylint: disable=invalid-name +class MigrateVizEnum(str, Enum): + # the Enum member name is viz_type in database + treemap = "treemap" + area = "area" + + get_migrate_class: Dict[MigrateVizEnum, Type[MigrateViz]] = { MigrateVizEnum.treemap: MigrateTreeMap, + MigrateVizEnum.area: MigrateArea, } diff --git a/tests/unit_tests/utils/viz_migration/area_migration_test.py b/tests/unit_tests/utils/viz_migration/area_migration_test.py new file mode 100644 index 000000000000..8857a96c940a --- /dev/null +++ b/tests/unit_tests/utils/viz_migration/area_migration_test.py @@ -0,0 +1,99 @@ +# 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.app import SupersetApp +from superset.utils.migrate_viz import get_migrate_class, MigrateVizEnum + +area_form_data = """{ + "adhoc_filters": [], + "annotation_layers": [], + "bottom_margin": "auto", + "color_scheme": "lyftColors", + "comparison_type": "values", + "contribution": true, + "datasource": "2__table", + "extra_form_data": {}, + "granularity_sqla": "ds", + "groupby": [ + "gender" + ], + "line_interpolation": "linear", + "metrics": [ + "sum__num" + ], + "order_desc": true, + "rich_tooltip": true, + "rolling_type": "None", + "row_limit": 10000, + "show_brush": "auto", + "show_controls": true, + "show_legend": true, + "slice_id": 165, + "stacked_style": "stack", + "time_grain_sqla": "P1D", + "time_range": "No filter", + "viz_type": "area", + "x_axis_format": "smart_date", + "x_axis_label": "x asix label", + "x_axis_showminmax": false, + "x_ticks_layout": "auto", + "y_axis_bounds": [ + null, + null + ], + "y_axis_format": "SMART_NUMBER" +} +""" + +area_processor = get_migrate_class[MigrateVizEnum.area] + + +def test_area_migrate(app_context: SupersetApp) -> None: + from superset.models.slice import Slice + + slc = Slice( + viz_type="area", + datasource_type="table", + params=area_form_data, + query_context=f'{{"form_data": {area_form_data}}}', + ) + + slc = area_processor.upgrade(slc) + assert slc.viz_type == area_processor.target_viz_type + # verify form_data + new_form_data = json.loads(slc.params) + assert new_form_data["contributionMode"] == "row" + assert "contribution" not in new_form_data + assert new_form_data["show_extra_controls"] is True + assert new_form_data["stack"] == "Stack" + assert new_form_data["x_axis_title"] == "x asix label" + assert new_form_data["x_axis_title_margin"] == 30 + assert json.dumps(new_form_data["form_data_bak"], sort_keys=True) == json.dumps( + json.loads(area_form_data), sort_keys=True + ) + + # verify query_context + new_query_context = json.loads(slc.query_context) + assert new_query_context["form_data"]["viz_type"] == area_processor.target_viz_type + + # downgrade + slc = area_processor.downgrade(slc) + assert slc.viz_type == area_processor.source_viz_type + assert json.dumps(json.loads(slc.params), sort_keys=True) == json.dumps( + json.loads(area_form_data), sort_keys=True + )