feat: Area viz migration (apache#20359)

charan1314 · Jul 19, 2022 · d36c0e0 · d36c0e0
1 parent f443452
commit d36c0e0
Show file tree

Hide file tree

Showing 5 changed files with 236 additions and 10 deletions.
diff --git a/...et/migrations/versions/2022-07-07_00-00_cdcf3d64daf4_add_user_id_dttm_idx_to_log_model.py b/...et/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/..._c747c78868b6_migrating_legacy_treemap.py → ..._c747c78868b6_migrating_legacy_treemap.py b/..._c747c78868b6_migrating_legacy_treemap.py → ..._c747c78868b6_migrating_legacy_treemap.py
@@ -17,15 +17,15 @@
 """Migrating legacy TreeMap
 
 Revision ID: c747c78868b6
-Revises: e786798587de
+Revises: cdcf3d64daf4
 Create Date: 2022-06-30 22:04:17.686635
 
 """
 
 # 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
@@ -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
@@ -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
@@ -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
+    )