fix: cache warmup solution non legacy charts. (#23012)

apache · Feb 14, 2023 · e755b4f · e755b4f
1 parent 0ec1e6e
commit e755b4f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 66 deletions.
diff --git a/superset/charts/commands/export.py b/superset/charts/commands/export.py
@@ -34,7 +34,7 @@
 
 
 # keys present in the standard export that are not needed
-REMOVE_KEYS = ["datasource_type", "datasource_name", "query_context", "url_params"]
+REMOVE_KEYS = ["datasource_type", "datasource_name", "url_params"]
 
 
 class ExportChartsCommand(ExportModelsCommand):

diff --git a/superset/views/core.py b/superset/views/core.py
@@ -56,6 +56,7 @@
 )
 from superset.charts.commands.exceptions import ChartNotFoundError
 from superset.charts.dao import ChartDAO
+from superset.charts.data.commands.get_data_command import ChartDataCommand
 from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.common.db_query_status import QueryStatus
 from superset.connectors.base.models import BaseDatasource
@@ -1740,28 +1741,35 @@ def warm_up_cache(  # pylint: disable=too-many-locals,no-self-use
 
         for slc in slices:
             try:
-                form_data = get_form_data(slc.id, use_slice_data=True)[0]
-                if dashboard_id:
-                    form_data["extra_filters"] = (
-                        json.loads(extra_filters)
-                        if extra_filters
-                        else get_dashboard_extra_filters(slc.id, dashboard_id)
-                    )
+                query_context = slc.get_query_context()
+                if query_context:
+                    query_context.force = True
+                    command = ChartDataCommand(query_context)
+                    command.validate()
+                    payload = command.run()
+                else:
+                    form_data = get_form_data(slc.id, use_slice_data=True)[0]
+                    if dashboard_id:
+                        form_data["extra_filters"] = (
+                            json.loads(extra_filters)
+                            if extra_filters
+                            else get_dashboard_extra_filters(slc.id, dashboard_id)
+                        )
 
-                if not slc.datasource:
-                    raise Exception("Slice's datasource does not exist")
+                    if not slc.datasource:
+                        raise Exception("Slice's datasource does not exist")
 
-                obj = get_viz(
-                    datasource_type=slc.datasource.type,
-                    datasource_id=slc.datasource.id,
-                    form_data=form_data,
-                    force=True,
-                )
+                    obj = get_viz(
+                        datasource_type=slc.datasource.type,
+                        datasource_id=slc.datasource.id,
+                        form_data=form_data,
+                        force=True,
+                    )
+                    # pylint: disable=assigning-non-slot
+                    g.form_data = form_data
+                    payload = obj.get_payload()
+                    delattr(g, "form_data")
 
-                # pylint: disable=assigning-non-slot
-                g.form_data = form_data
-                payload = obj.get_payload()
-                delattr(g, "form_data")
                 error = payload["errors"] or None
                 status = payload["status"]
             except Exception as ex:  # pylint: disable=broad-except

diff --git a/tests/integration_tests/charts/commands_tests.py b/tests/integration_tests/charts/commands_tests.py
@@ -88,52 +88,7 @@ def test_export_chart_command(self, mock_g):
             "dataset_uuid": str(example_chart.table.uuid),
             "uuid": str(example_chart.uuid),
             "version": "1.0.0",
-        }
-
-    @patch("superset.security.manager.g")
-    @pytest.mark.usefixtures("load_energy_table_with_slice")
-    def test_export_chart_with_query_context(self, mock_g):
-        """Test that charts that have a query_context are exported correctly"""
-
-        mock_g.user = security_manager.find_user("alpha")
-        example_chart = db.session.query(Slice).filter_by(slice_name="Heatmap").one()
-        command = ExportChartsCommand([example_chart.id])
-
-        contents = dict(command.run())
-
-        expected = [
-            "metadata.yaml",
-            f"charts/Heatmap_{example_chart.id}.yaml",
-            "datasets/examples/energy_usage.yaml",
-            "databases/examples.yaml",
-        ]
-        assert expected == list(contents.keys())
-
-        metadata = yaml.safe_load(contents[f"charts/Heatmap_{example_chart.id}.yaml"])
-
-        assert metadata == {
-            "slice_name": "Heatmap",
-            "description": None,
-            "certified_by": None,
-            "certification_details": None,
-            "viz_type": "heatmap",
-            "params": {
-                "all_columns_x": "source",
-                "all_columns_y": "target",
-                "canvas_image_rendering": "pixelated",
-                "collapsed_fieldsets": "",
-                "linear_color_scheme": "blue_white_yellow",
-                "metric": "sum__value",
-                "normalize_across": "heatmap",
-                "slice_name": "Heatmap",
-                "viz_type": "heatmap",
-                "xscale_interval": "1",
-                "yscale_interval": "1",
-            },
-            "cache_timeout": None,
-            "dataset_uuid": str(example_chart.table.uuid),
-            "uuid": str(example_chart.uuid),
-            "version": "1.0.0",
+            "query_context": None,
         }
 
     @patch("superset.security.manager.g")
@@ -179,6 +134,7 @@ def test_export_chart_command_key_order(self, mock_g):
             "certification_details",
             "viz_type",
             "params",
+            "query_context",
             "cache_timeout",
             "uuid",
             "version",