ktmud commented on code in PR #20761:
URL: https://github.com/apache/superset/pull/20761#discussion_r924045280
##########
superset/migrations/shared/migrate_viz/base.py:
##########
@@ -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")
Review Comment:
I'm not sure if we should raise an error here. Maybe just silently override?
##########
superset/migrations/shared/migrate_viz/base.py:
##########
@@ -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"
Review Comment:
Not need for such enum since we will import different stuff for different
migrations anyway.
##########
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)
Review Comment:
Data may be corrupted, let's always try/catch just to be safe.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
