fix(cache): Add cache warmup for non-legacy charts (#24671)

superset/charts/commands/export.py

@@ -34,7 +34,7 @@ logger = logging.getLogger(__name__)
 
 
 # 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):

superset/charts/commands/warm_up_cache.py

@@ -21,12 +21,17 @@ from typing import Any, Optional, Union
 import simplejson as json
 from flask import g
 
-from superset.charts.commands.exceptions import WarmUpCacheChartNotFoundError
+from superset.charts.commands.exceptions import (
+    ChartInvalidError,
+    WarmUpCacheChartNotFoundError,
+)
+from superset.charts.data.commands.get_data_command import ChartDataCommand
 from superset.commands.base import BaseCommand
 from superset.extensions import db
 from superset.models.slice import Slice
 from superset.utils.core import error_msg_from_exception
 from superset.views.utils import get_dashboard_extra_filters, get_form_data, get_viz
+from superset.viz import viz_types
 
 
 class ChartWarmUpCacheCommand(BaseCommand):
@@ -43,31 +48,51 @@ class ChartWarmUpCacheCommand(BaseCommand):
     def run(self) -> dict[str, Any]:
         self.validate()
         chart: Slice = self._chart_or_id  # type: ignore
+
         try:
             form_data = get_form_data(chart.id, use_slice_data=True)[0]
-            if self._dashboard_id:
-                form_data["extra_filters"] = (
-                    json.loads(self._extra_filters)
-                    if self._extra_filters
-                    else get_dashboard_extra_filters(chart.id, self._dashboard_id)
-                )
 
-            if not chart.datasource:
-                raise Exception("Chart's datasource does not exist")
+            if form_data.get("viz_type") in viz_types:
+                # Legacy visualizations.
+                if not chart.datasource:
+                    raise ChartInvalidError("Chart's datasource does not exist")
 
-            obj = get_viz(
-                datasource_type=chart.datasource.type,
-                datasource_id=chart.datasource.id,
-                form_data=form_data,
-                force=True,
-            )
+                if self._dashboard_id:
+                    form_data["extra_filters"] = (
+                        json.loads(self._extra_filters)
+                        if self._extra_filters
+                        else get_dashboard_extra_filters(chart.id, self._dashboard_id)
+                    )
 
-            # 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"]
+                g.form_data = form_data  # pylint: disable=assigning-non-slot
+                payload = get_viz(
+                    datasource_type=chart.datasource.type,
+                    datasource_id=chart.datasource.id,
+                    form_data=form_data,
+                    force=True,
+                ).get_payload()
+                delattr(g, "form_data")
+                error = payload["errors"] or None
+                status = payload["status"]
+            else:
+                # Non-legacy visualizations.
+                query_context = chart.get_query_context()
+
+                if not query_context:
+                    raise ChartInvalidError("Chart's query context does not exist")
+
+                query_context.force = True
+                command = ChartDataCommand(query_context)
+                command.validate()
+                payload = command.run()
+
+                # Report the first error.
+                for query in payload["queries"]:
+                    error = query["error"]
+                    status = query["status"]
+
+                    if error is not None:
+                        break
         except Exception as ex:  # pylint: disable=broad-except
             error = error_msg_from_exception(ex)
             status = None

superset/views/core.py

@@ -43,6 +43,7 @@ from superset import (
     security_manager,
 )
 from superset.charts.commands.exceptions import ChartNotFoundError
+from superset.charts.commands.warm_up_cache import ChartWarmUpCacheCommand
 from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType
 from superset.connectors.base.models import BaseDatasource
 from superset.connectors.sqla.models import SqlaTable
@@ -72,6 +73,7 @@ from superset.utils import core as utils
 from superset.utils.async_query_manager import AsyncQueryTokenException
 from superset.utils.cache import etag_cache
 from superset.utils.core import (
+    base_json_conv,
     DatasourceType,
     get_user_id,
     get_username,
@@ -95,7 +97,6 @@ from superset.views.utils import (
     check_datasource_perms,
     check_explore_cache_perms,
     check_resource_permissions,
-    get_dashboard_extra_filters,
     get_datasource_info,
     get_form_data,
     get_viz,
@@ -769,7 +770,8 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
     @api
     @has_access_api
     @expose("/warm_up_cache/", methods=("GET",))
-    def warm_up_cache(  # pylint: disable=too-many-locals,no-self-use
+    @deprecated(new_target="api/v1/chart/warm_up_cache/")
+    def warm_up_cache(  # pylint: disable=no-self-use
         self,
     ) -> FlaskResponse:
         """Warms up the cache for the slice or table.
@@ -825,43 +827,22 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
                 .all()
             )
 
-        result = []
-
-        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)
-                    )
-
-                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,
-                )
-
-                # 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
-                error = utils.error_msg_from_exception(ex)
-                status = None
-
-            result.append(
-                {"slice_id": slc.id, "viz_error": error, "viz_status": status}
-            )
-
-        return json_success(json.dumps(result))
+        return json_success(
+            json.dumps(
+                [
+                    {
+                        "slice_id" if key == "chart_id" else key: value
+                        for key, value in ChartWarmUpCacheCommand(
+                            slc, dashboard_id, extra_filters
+                        )
+                        .run()
+                        .items()
+                    }
+                    for slc in slices
+                ],
+                default=base_json_conv,
+            ),
+        )
 
     @has_access
     @expose("/dashboard/<dashboard_id_or_slug>/")

tests/integration_tests/charts/api_tests.py

@@ -14,37 +14,41 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# isort:skip_file
 """Unit tests for Superset"""
 import json
 from io import BytesIO
+from unittest import mock
 from zipfile import is_zipfile, ZipFile
 
 import prison
 import pytest
 import yaml
+from flask_babel import lazy_gettext as _
+from parameterized import parameterized
 from sqlalchemy import and_
 from sqlalchemy.sql import func
 
+from superset.charts.commands.exceptions import ChartDataQueryFailedError
+from superset.charts.data.commands.get_data_command import ChartDataCommand
 from superset.connectors.sqla.models import SqlaTable
 from superset.extensions import cache_manager, db, security_manager
 from superset.models.core import Database, FavStar, FavStarClassName
 from superset.models.dashboard import Dashboard
-from superset.reports.models import ReportSchedule, ReportScheduleType
 from superset.models.slice import Slice
+from superset.reports.models import ReportSchedule, ReportScheduleType
 from superset.utils.core import get_example_default_schema
 from superset.utils.database import get_example_database
-
-from tests.integration_tests.conftest import with_feature_flags
+from superset.viz import viz_types
 from tests.integration_tests.base_api_tests import ApiOwnersTestCaseMixin
 from tests.integration_tests.base_tests import SupersetTestCase
+from tests.integration_tests.conftest import with_feature_flags
 from tests.integration_tests.fixtures.birth_names_dashboard import (
     load_birth_names_dashboard_with_slices,
     load_birth_names_data,
 )
 from tests.integration_tests.fixtures.energy_dashboard import (
-    load_energy_table_with_slice,
     load_energy_table_data,
+    load_energy_table_with_slice,
 )
 from tests.integration_tests.fixtures.importexport import (
     chart_config,
@@ -1710,12 +1714,16 @@ class TestChartApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixin):
         assert data["result"][0]["slice_name"] == "name0"
         assert data["result"][0]["datasource_id"] == 1
 
-    @pytest.mark.usefixtures(
-        "load_energy_table_with_slice", "load_birth_names_dashboard_with_slices"
+    @parameterized.expand(
+        [
+            "Top 10 Girl Name Share",  # Legacy chart
+            "Pivot Table v2",  # Non-legacy chart
+        ],
     )
-    def test_warm_up_cache(self):
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    def test_warm_up_cache(self, slice_name):
         self.login()
-        slc = self.get_slice("Top 10 Girl Name Share", db.session)
+        slc = self.get_slice(slice_name, db.session)
         rv = self.client.put("/api/v1/chart/warm_up_cache", json={"chart_id": slc.id})
         self.assertEqual(rv.status_code, 200)
         data = json.loads(rv.data.decode("utf-8"))
@@ -1780,7 +1788,6 @@ class TestChartApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixin):
         )
         self.assertEqual(rv.status_code, 400)
         data = json.loads(rv.data.decode("utf-8"))
-        print(data)
         self.assertEqual(
             data,
             {
@@ -1791,3 +1798,81 @@ class TestChartApi(SupersetTestCase, ApiOwnersTestCaseMixin, InsertChartMixin):
                 }
             },
         )
+
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    def test_warm_up_cache_error(self) -> None:
+        self.login()
+        slc = self.get_slice("Pivot Table v2", db.session)
+
+        with mock.patch.object(ChartDataCommand, "run") as mock_run:
+            mock_run.side_effect = ChartDataQueryFailedError(
+                _(
+                    "Error: %(error)s",
+                    error=_("Empty query?"),
+                )
+            )
+
+            assert json.loads(
+                self.client.put(
+                    "/api/v1/chart/warm_up_cache",
+                    json={"chart_id": slc.id},
+                ).data
+            ) == {
+                "result": [
+                    {
+                        "chart_id": slc.id,
+                        "viz_error": "Error: Empty query?",
+                        "viz_status": None,
+                    },
+                ],
+            }
+
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    def test_warm_up_cache_no_query_context(self) -> None:
+        self.login()
+        slc = self.get_slice("Pivot Table v2", db.session)
+
+        with mock.patch.object(Slice, "get_query_context") as mock_get_query_context:
+            mock_get_query_context.return_value = None
+
+            assert json.loads(
+                self.client.put(
+                    f"/api/v1/chart/warm_up_cache",
+                    json={"chart_id": slc.id},
+                ).data
+            ) == {
+                "result": [
+                    {
+                        "chart_id": slc.id,
+                        "viz_error": "Chart's query context does not exist",
+                        "viz_status": None,
+                    },
+                ],
+            }
+
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    def test_warm_up_cache_no_datasource(self) -> None:
+        self.login()
+        slc = self.get_slice("Top 10 Girl Name Share", db.session)
+
+        with mock.patch.object(
+            Slice,
+            "datasource",
+            new_callable=mock.PropertyMock,
+        ) as mock_datasource:
+            mock_datasource.return_value = None
+
+            assert json.loads(
+                self.client.put(
+                    f"/api/v1/chart/warm_up_cache",
+                    json={"chart_id": slc.id},
+                ).data
+            ) == {
+                "result": [
+                    {
+                        "chart_id": slc.id,
+                        "viz_error": "Chart's datasource does not exist",
+                        "viz_status": None,
+                    },
+                ],
+            }

tests/integration_tests/charts/commands_tests.py

@@ -96,52 +96,7 @@ class TestExportChartsCommand(SupersetTestCase):
             "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")
@@ -187,6 +142,7 @@ class TestExportChartsCommand(SupersetTestCase):
             "certification_details",
             "viz_type",
             "params",
+            "query_context",
             "cache_timeout",
             "uuid",
             "version",

tests/integration_tests/core_tests.py

@@ -14,72 +14,67 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# isort:skip_file
 """Unit tests for Superset"""
 import datetime
 import doctest
 import html
 import json
 import logging
-from urllib.parse import quote
-
-import prison
-import superset.utils.database
-from superset.utils.core import backend
-from tests.integration_tests.fixtures.public_role import public_role_like_gamma
-from tests.integration_tests.fixtures.birth_names_dashboard import (
-    load_birth_names_dashboard_with_slices,
-    load_birth_names_data,
-)
-from sqlalchemy import Table
-
-import pytest
-import pytz
 import random
 import unittest
 from unittest import mock
+from urllib.parse import quote
 
 import pandas as pd
+import prison
+import pytest
+import pytz
 import sqlalchemy as sqla
+from flask_babel import lazy_gettext as _
+from sqlalchemy import Table
 from sqlalchemy.exc import SQLAlchemyError
-from superset.models.cache import CacheKey
-from superset.utils.database import get_example_database
-from tests.integration_tests.conftest import with_feature_flags
-from tests.integration_tests.fixtures.energy_dashboard import (
-    load_energy_table_with_slice,
-    load_energy_table_data,
-)
-from tests.integration_tests.insert_chart_mixin import InsertChartMixin
-from tests.integration_tests.test_app import app
+
+import superset.utils.database
 import superset.views.utils
-from superset import (
-    dataframe,
-    db,
-    security_manager,
-    sql_lab,
-)
+from superset import dataframe, db, security_manager, sql_lab
+from superset.charts.commands.exceptions import ChartDataQueryFailedError
+from superset.charts.data.commands.get_data_command import ChartDataCommand
 from superset.common.db_query_status import QueryStatus
 from superset.connectors.sqla.models import SqlaTable
 from superset.db_engine_specs.base import BaseEngineSpec
 from superset.db_engine_specs.mssql import MssqlEngineSpec
-from superset.exceptions import SupersetException
+from superset.exceptions import QueryObjectValidationError, SupersetException
 from superset.extensions import async_query_manager, cache_manager
 from superset.models import core as models
 from superset.models.annotations import Annotation, AnnotationLayer
+from superset.models.cache import CacheKey
 from superset.models.dashboard import Dashboard
 from superset.models.slice import Slice
 from superset.models.sql_lab import Query
 from superset.result_set import SupersetResultSet
 from superset.utils import core as utils
+from superset.utils.core import backend
+from superset.utils.database import get_example_database
 from superset.views import core as views
 from superset.views.database.views import DatabaseView
-
-from .base_tests import SupersetTestCase
+from tests.integration_tests.conftest import CTAS_SCHEMA_NAME, with_feature_flags
+from tests.integration_tests.fixtures.birth_names_dashboard import (
+    load_birth_names_dashboard_with_slices,
+    load_birth_names_data,
+)
+from tests.integration_tests.fixtures.energy_dashboard import (
+    load_energy_table_data,
+    load_energy_table_with_slice,
+)
+from tests.integration_tests.fixtures.public_role import public_role_like_gamma
 from tests.integration_tests.fixtures.world_bank_dashboard import (
     load_world_bank_dashboard_with_slices,
     load_world_bank_data,
 )
-from tests.integration_tests.conftest import CTAS_SCHEMA_NAME
+from tests.integration_tests.insert_chart_mixin import InsertChartMixin
+from tests.integration_tests.test_app import app
+
+from .base_tests import SupersetTestCase
 
 logger = logging.getLogger(__name__)
 
@@ -389,7 +384,8 @@ class TestCore(SupersetTestCase, InsertChartMixin):
         db.session.commit()
 
     @pytest.mark.usefixtures(
-        "load_energy_table_with_slice", "load_birth_names_dashboard_with_slices"
+        "load_birth_names_dashboard_with_slices",
+        "load_energy_table_with_slice",
     )
     def test_warm_up_cache(self):
         self.login()
@@ -415,6 +411,29 @@ class TestCore(SupersetTestCase, InsertChartMixin):
             + quote(json.dumps([{"col": "name", "op": "in", "val": ["Jennifer"]}]))
         ) == [{"slice_id": slc.id, "viz_error": None, "viz_status": "success"}]
 
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    def test_warm_up_cache_error(self) -> None:
+        self.login()
+        slc = self.get_slice("Pivot Table v2", db.session)
+
+        with mock.patch.object(
+            ChartDataCommand,
+            "run",
+            side_effect=ChartDataQueryFailedError(
+                _(
+                    "Error: %(error)s",
+                    error=_("Empty query?"),
+                )
+            ),
+        ):
+            assert self.get_json_resp(f"/superset/warm_up_cache?slice_id={slc.id}") == [
+                {
+                    "slice_id": slc.id,
+                    "viz_error": "Error: Empty query?",
+                    "viz_status": None,
+                }
+            ]
+
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     def test_cache_logging(self):
         self.login("admin")