fix(db): use paginated_update for viz migration

superset/migrations/shared/migrate_viz/__init__.py

@@ -0,0 +1,17 @@
+# 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 .processors import *

superset/utils/migrate_viz.py → superset/migrations/shared/migrate_viz/base.py

@@ -17,21 +17,39 @@
 from __future__ import annotations
 
 import json
-from enum import Enum
-from typing import Dict, Set, Type, TYPE_CHECKING
+from typing import Dict, Set
 
-if TYPE_CHECKING:
-    from superset.models.slice import Slice
+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.migrations.shared.utils import paginated_update, try_load_json
+
+Base = declarative_base()
+
+
+class Slice(Base):  # type: ignore
+    __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)
+
+
+FORM_DATA_BAK_FIELD_NAME = "form_data_bak"
 
 
 class MigrateViz:
     remove_keys: Set[str] = set()
-    mapping_keys: Dict[str, str] = {}
+    rename_keys: Dict[str, str] = {}
     source_viz_type: str
     target_viz_type: str
 
     def __init__(self, form_data: str) -> None:
-        self.data = json.loads(form_data)
+        self.data = try_load_json(form_data)
 
     def _pre_action(self) -> None:
         """some actions before migrate"""
@@ -45,11 +63,11 @@ def _migrate(self) -> None:
 
         rv_data = {}
         for key, value in self.data.items():
-            if key in self.mapping_keys and self.mapping_keys[key] in rv_data:
+            if key in self.rename_keys and self.rename_keys[key] in rv_data:
                 raise ValueError("Duplicate key in target viz")
 
-            if key in self.mapping_keys:
-                rv_data[self.mapping_keys[key]] = value
+            if key in self.rename_keys:
+                rv_data[self.rename_keys[key]] = value
 
             if key in self.remove_keys:
                 continue
@@ -62,7 +80,7 @@ def _post_action(self) -> None:
         """some actions after migrate"""
 
     @classmethod
-    def upgrade(cls, slc: Slice) -> Slice:
+    def upgrade_slice(cls, slc: Slice) -> Slice:
         clz = cls(slc.params)
         slc.viz_type = cls.target_viz_type
         form_data_bak = clz.data.copy()
@@ -72,77 +90,56 @@ def upgrade(cls, slc: Slice) -> Slice:
         clz._post_action()
 
         # only backup params
-        slc.params = json.dumps({**clz.data, "form_data_bak": form_data_bak})
+        slc.params = json.dumps({**clz.data, FORM_DATA_BAK_FIELD_NAME: form_data_bak})
 
-        query_context = json.loads(slc.query_context or "{}")
+        query_context = try_load_json(slc.query_context)
         if "form_data" in query_context:
             query_context["form_data"] = clz.data
             slc.query_context = json.dumps(query_context)
         return slc
 
     @classmethod
-    def downgrade(cls, slc: Slice) -> Slice:
-        form_data = json.loads(slc.params)
-        if "form_data_bak" in form_data and "viz_type" in form_data.get(
-            "form_data_bak"
-        ):
-            form_data_bak = form_data["form_data_bak"]
+    def downgrade_slice(cls, slc: Slice) -> Slice:
+        form_data = try_load_json(slc.params)
+        form_data_bak = form_data.get(FORM_DATA_BAK_FIELD_NAME, {})
+        if "viz_type" in form_data_bak:
             slc.params = json.dumps(form_data_bak)
             slc.viz_type = form_data_bak.get("viz_type")
-
-            query_context = json.loads(slc.query_context or "{}")
+            query_context = try_load_json(slc.query_context)
             if "form_data" in query_context:
                 query_context["form_data"] = form_data_bak
                 slc.query_context = json.dumps(query_context)
         return slc
 
-
-class MigrateTreeMap(MigrateViz):
-    source_viz_type = "treemap"
-    target_viz_type = "treemap_v2"
-    remove_keys = {"metrics"}
-
-    def _pre_action(self) -> None:
-        if (
-            "metrics" in self.data
-            and isinstance(self.data["metrics"], list)
-            and len(self.data["metrics"]) > 0
+    @classmethod
+    def upgrade(cls) -> None:
+        bind = op.get_bind()
+        session = db.Session(bind=bind)
+        slices = session.query(Slice).filter(Slice.viz_type == cls.source_viz_type)
+        for slc in paginated_update(
+            slices,
+            lambda current, total: print(
+                f"  Updating {current}/{total} charts", end="\r"
+            ),
         ):
-            self.data["metric"] = self.data["metrics"][0]
-
+            new_viz = cls.upgrade_slice(slc)
+            session.merge(new_viz)
 
-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,
-}
+    @classmethod
+    def downgrade(cls) -> None:
+        bind = op.get_bind()
+        session = db.Session(bind=bind)
+        slices = session.query(Slice).filter(
+            and_(
+                Slice.viz_type == cls.target_viz_type,
+                Slice.params.like(f"%{FORM_DATA_BAK_FIELD_NAME}%"),
+            )
+        )
+        for slc in paginated_update(
+            slices,
+            lambda current, total: print(
+                f"  Downgrading {current}/{total} charts", end="\r"
+            ),
+        ):
+            new_viz = cls.downgrade_slice(slc)
+            session.merge(new_viz)

superset/migrations/shared/migrate_viz/processors.py

@@ -0,0 +1,55 @@
+# 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 .base import MigrateViz
+
+
+class MigrateTreeMap(MigrateViz):
+    source_viz_type = "treemap"
+    target_viz_type = "treemap_v2"
+    remove_keys = {"metrics"}
+
+    def _pre_action(self) -> None:
+        if (
+            "metrics" in self.data
+            and isinstance(self.data["metrics"], list)
+            and len(self.data["metrics"]) > 0
+        ):
+            self.data["metric"] = self.data["metrics"][0]
+
+
+class MigrateAreaChart(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

superset/migrations/shared/utils.py

@@ -14,10 +14,11 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import json
 import logging
 import os
 import time
-from typing import Any, Callable, Iterator, Optional, Union
+from typing import Any, Callable, Dict, Iterator, Optional, Union
 from uuid import uuid4
 
 from alembic import op
@@ -115,3 +116,11 @@ def paginated_update(
         if print_page_progress:
             print_page_progress(end, count)
         start += batch_size
+
+
+def try_load_json(data: Optional[str]) -> Dict[str, Any]:
+    try:
+        return data and json.loads(data) or {}
+    except json.decoder.JSONDecodeError:
+        print(f"Failed to parse: {data}")
+        return {}

superset/migrations/versions/2022-07-07_13-00_c747c78868b6_migrating_legacy_treemap.py

@@ -21,82 +21,16 @@
 Create Date: 2022-06-30 22:04:17.686635
 
 """
+from superset.migrations.shared.migrate_viz import MigrateTreeMap
 
 # revision identifiers, used by Alembic.
-
 revision = "c747c78868b6"
 down_revision = "cdcf3d64daf4"
 
-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
-
-treemap_processor = get_migrate_class[MigrateVizEnum.treemap]
-
-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 == treemap_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 = treemap_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()
+    MigrateTreeMap.upgrade()
 
 
 def downgrade():
-    bind = op.get_bind()
-    session = db.Session(bind=bind)
-
-    slices = session.query(Slice).filter(
-        and_(
-            Slice.viz_type == treemap_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 = treemap_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()
+    MigrateTreeMap.downgrade()