fix: Adjust viz migrations to also migrate the queries object (#33285)

Co-authored-by: Michael S. Molina <michael.s.molina@gmail.com>
Co-authored-by: Michael S. Molina <70410625+michael-s-molina@users.noreply.github.com>
(cherry picked from commit 57183da315)

tests/integration_tests/migrations/c747c78868b6_migrating_legacy_treemap__tests.py

@@ -1,91 +0,0 @@
-# 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 superset.app import SupersetApp
-from superset.migrations.shared.migrate_viz import MigrateTreeMap
-from superset.utils import json
-
-treemap_form_data = """{
-  "adhoc_filters": [
-    {
-      "clause": "WHERE",
-      "comparator": [
-        "Edward"
-      ],
-      "expressionType": "SIMPLE",
-      "filterOptionName": "filter_xhbus6irfa_r10k9nwmwy",
-      "isExtra": false,
-      "isNew": false,
-      "operator": "IN",
-      "operatorId": "IN",
-      "sqlExpression": null,
-      "subject": "name"
-    }
-  ],
-  "color_scheme": "bnbColors",
-  "datasource": "2__table",
-  "extra_form_data": {},
-  "granularity_sqla": "ds",
-  "groupby": [
-    "state",
-    "gender"
-  ],
-  "metrics": [
-    "sum__num"
-  ],
-  "number_format": ",d",
-  "order_desc": true,
-  "row_limit": 10,
-  "time_range": "No filter",
-  "timeseries_limit_metric": "sum__num",
-  "treemap_ratio": 1.618033988749895,
-  "viz_type": "treemap"
-}
-"""
-
-
-def test_treemap_migrate(app_context: SupersetApp) -> None:
-    from superset.models.slice import Slice
-
-    slc = Slice(
-        viz_type=MigrateTreeMap.source_viz_type,
-        datasource_type="table",
-        params=treemap_form_data,
-        query_context=f'{{"form_data": {treemap_form_data}}}',
-    )
-
-    MigrateTreeMap.upgrade_slice(slc)
-    assert slc.viz_type == MigrateTreeMap.target_viz_type
-    # verify form_data
-    new_form_data = json.loads(slc.params)
-    assert new_form_data["metric"] == "sum__num"
-    assert new_form_data["viz_type"] == "treemap_v2"
-    assert "metrics" not in new_form_data
-    assert json.dumps(new_form_data["form_data_bak"], sort_keys=True) == json.dumps(
-        json.loads(treemap_form_data), sort_keys=True
-    )
-
-    # verify query_context
-    new_query_context = json.loads(slc.query_context)
-    assert new_query_context["form_data"]["viz_type"] == "treemap_v2"
-
-    # downgrade
-    MigrateTreeMap.downgrade_slice(slc)
-    assert slc.viz_type == MigrateTreeMap.source_viz_type
-    assert json.dumps(json.loads(slc.params), sort_keys=True) == json.dumps(
-        json.loads(treemap_form_data), sort_keys=True
-    )

tests/unit_tests/migrations/viz/dual_line_to_mixed_chart_test.py

@@ -30,6 +30,7 @@ ADHOC_FILTERS = [
 ]
 
 SOURCE_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "metric": "num_boys",
     "y_axis_format": ",d",
     "y_axis_bounds": [50, 100],
@@ -44,6 +45,7 @@ SOURCE_FORM_DATA: dict[str, Any] = {
 }
 
 TARGET_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "metrics": ["num_boys"],
     "y_axis_format": ",d",
     "y_axis_bounds": [50, 100],

tests/unit_tests/migrations/viz/heatmap_v1_v2_test.py

@@ -20,10 +20,15 @@ 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] = {
+    "datasource": "1__table",
     "any_other_key": "untouched",
-    "all_columns_x": ["category"],
-    "all_columns_y": ["product"],
-    "metric": ["sales"],
+    "all_columns_x": "category",
+    "all_columns_y": "product",
+    "metric": {
+        "label": "sales",
+        "expressionType": "SQL",
+        "sqlExpression": "max(sales)",
+    },
     "adhoc_filters": [],
     "row_limit": 100,
     "sort_by_metric": True,
@@ -47,10 +52,15 @@ SOURCE_FORM_DATA: dict[str, Any] = {
 }
 
 TARGET_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "any_other_key": "untouched",
-    "x_axis": ["category"],
-    "groupby": ["product"],
-    "metric": ["sales"],
+    "x_axis": "category",
+    "groupby": "product",
+    "metric": {
+        "label": "sales",
+        "expressionType": "SQL",
+        "sqlExpression": "max(sales)",
+    },
     "adhoc_filters": [],
     "row_limit": 100,
     "legend_type": "continuous",

tests/unit_tests/migrations/viz/histogram_v1_v2_test.py

@@ -20,6 +20,7 @@ from superset.migrations.shared.migrate_viz import MigrateHistogramChart
 from tests.unit_tests.migrations.viz.utils import migrate_and_assert
 
 SOURCE_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "all_columns_x": ["category"],
     "adhoc_filters": [],
     "cumulative": True,
@@ -33,6 +34,7 @@ SOURCE_FORM_DATA: dict[str, Any] = {
 }
 
 TARGET_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "adhoc_filters": [],
     "bins": 5,
     "column": "category",

tests/unit_tests/migrations/viz/nvd3_bubble_chart_to_echarts_test.py

@@ -20,6 +20,7 @@ from superset.migrations.shared.migrate_viz import MigrateBubbleChart
 from tests.unit_tests.migrations.viz.utils import migrate_and_assert
 
 SOURCE_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "adhoc_filters": [],
     "bottom_margin": 20,
     "color_scheme": "default",
@@ -29,7 +30,7 @@ SOURCE_FORM_DATA: dict[str, Any] = {
     "max_bubble_size": 50,
     "series": ["region"],
     "show_legend": True,
-    "size": 75,
+    "size": {"label": "sales", "expressionType": "SQL", "sqlExpression": "max(sales)"},
     "viz_type": "bubble",
     "x": "year",
     "x_axis_format": "SMART_DATE",
@@ -46,6 +47,7 @@ SOURCE_FORM_DATA: dict[str, Any] = {
 }
 
 TARGET_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "adhoc_filters": [],
     "color_scheme": "default",
     "entity": "count",
@@ -56,7 +58,7 @@ TARGET_FORM_DATA: dict[str, Any] = {
     "row_limit": 100,
     "series": ["region"],
     "show_legend": True,
-    "size": 75,
+    "size": {"label": "sales", "expressionType": "SQL", "sqlExpression": "max(sales)"},
     "truncateYAxis": True,
     "viz_type": "bubble_v2",
     "x": "year",

tests/unit_tests/migrations/viz/pivot_table_v1_v2_test.py

@@ -20,6 +20,7 @@ from superset.migrations.shared.migrate_viz import MigratePivotTable
 from tests.unit_tests.migrations.viz.utils import migrate_and_assert
 
 SOURCE_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "any_other_key": "untouched",
     "columns": ["state"],
     "combine_metric": True,
@@ -33,6 +34,7 @@ SOURCE_FORM_DATA: dict[str, Any] = {
 }
 
 TARGET_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "any_other_key": "untouched",
     "aggregateFunction": "Sum",
     "colTotals": True,

tests/unit_tests/migrations/viz/time_related_fields_test.py

@@ -20,12 +20,14 @@ from superset.migrations.shared.migrate_viz import MigratePivotTable
 from tests.unit_tests.migrations.viz.utils import migrate_and_assert
 
 SOURCE_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "granularity_sqla": "ds",
     "time_range": "2024-04-24T00:00:00 : 2025-10-18T00:00:00",
     "viz_type": "pivot_table",
 }
 
 TARGET_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "form_data_bak": SOURCE_FORM_DATA,
     "granularity_sqla": "ds",
     "rowOrder": "value_z_to_a",

tests/unit_tests/migrations/viz/utils.py

@@ -20,6 +20,7 @@ from superset.migrations.shared.migrate_viz import MigrateViz
 from superset.utils import json
 
 TIMESERIES_SOURCE_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "bottom_margin": 20,
     "comparison_type": "absolute",
     "contribution": True,
@@ -42,6 +43,7 @@ TIMESERIES_SOURCE_FORM_DATA: dict[str, Any] = {
 }
 
 TIMESERIES_TARGET_FORM_DATA: dict[str, Any] = {
+    "datasource": "1__table",
     "comparison_type": "difference",
     "contributionMode": "row",
     "logAxis": True,
@@ -75,7 +77,7 @@ def migrate_and_assert(
         viz_type=cls.source_viz_type,
         datasource_type="table",
         params=dumped_form_data,
-        query_context=f'{{"form_data": {dumped_form_data}}}',
+        query_context=f'{{"form_data": {dumped_form_data}, "queries": []}}',
     )
 
     # upgrade
@@ -83,6 +85,7 @@ def migrate_and_assert(
 
     # verify form_data
     new_form_data = json.loads(slc.params)
+    new_form_data.pop("queries_bak", None)
     assert new_form_data == target
     assert new_form_data["form_data_bak"] == source