chore: Migrate warm up cache endpoint to api v1

superset/cachekeys/api.py

@@ -24,7 +24,9 @@
 from marshmallow.exceptions import ValidationError
 from sqlalchemy.exc import SQLAlchemyError
 
+from superset.cachekeys.commands.warm_up_cache import WarmUpCacheCommand
 from superset.cachekeys.schemas import CacheInvalidationRequestSchema
+from superset.commands.exceptions import CommandException
 from superset.connectors.sqla.models import SqlaTable
 from superset.extensions import cache_manager, db, event_logger, stats_logger_manager
 from superset.models.cache import CacheKey
@@ -40,10 +42,98 @@ class CacheRestApi(BaseSupersetModelRestApi):
     class_permission_name = "CacheRestApi"
     include_route_methods = {
         "invalidate",
+        "warm_up_cache",
     }
 
     openapi_spec_component_schemas = (CacheInvalidationRequestSchema,)
 
+    @expose("/warm_up_cache", methods=["GET"])
+    @protect()
+    @safe
+    @statsd_metrics
+    @event_logger.log_this_with_context(
+        action=lambda self, *args, **kwargs: f"{self.__class__.__name__}"
+        f".warm_up_cache",
+        log_to_statsd=False,
+    )
+    def warm_up_cache(self) -> Response:
+        """Warms up the cache for the slice or table.
+
+        Note for slices a force refresh occurs.
+
+        In terms of the `extra_filters` these can be obtained from records in the JSON
+        encoded `logs.json` column associated with the `explore_json` action.
+
+        ---
+        get:
+          description: >-
+            Warms up the cache for the slice or table
+          parameters:
+          - in: query
+            name: slice_id
+            schema:
+              type: integer
+            description: The ID of the chart to warm up cache for
+          - in: query
+            name: dashboard_id
+            schema:
+              type: integer
+            description: The ID of the dashboard to get filters for when warming cache
+          - in: query
+            name: table_name
+            schema:
+              type: string
+            description: The name of the table to warm up cache for
+          - in: query
+            name: db_name
+            schema:
+              type: string
+            description: The name of the database where the table is located
+          - in: query
+            name: extra_filters
+            schema:
+              type: string
+            description: Extra filters to apply when warming up cache
+          responses:
+            200:
+              description: Each chart's warmup status
+              content:
+                application/json:
+                  schema:
+                    type: object
+                    properties:
+                      result:
+                        type: array
+                        items:
+                          type: object
+                          properties:
+                            slice_id:
+                              type: integer
+                            viz_error:
+                              type: string
+                            viz_status:
+                              type: string
+            400:
+              $ref: '#/components/responses/400'
+            404:
+              $ref: '#/components/responses/404'
+            500:
+              $ref: '#/components/responses/500'
+        """
+        slice_id = request.args.get("slice_id")
+        dashboard_id = request.args.get("dashboard_id")
+        table_name = request.args.get("table_name")
+        db_name = request.args.get("db_name")
+        extra_filters = request.args.get("extra_filters")
+
+        try:
+            payload = WarmUpCacheCommand(
+                slice_id, dashboard_id, table_name, db_name, extra_filters
+            ).run()
+            return self.response(200, result=payload)
+        except CommandException as ex:
+            return self.response(ex.status, message=ex.message)
+
     @expose("/invalidate", methods=["POST"])
     @protect()
     @safe

superset/cachekeys/commands/__init__.py

@@ -0,0 +1,16 @@
+# 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.

superset/cachekeys/commands/exceptions.py

@@ -0,0 +1,44 @@
+# 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 flask_babel import lazy_gettext as _
+
+from superset.commands.exceptions import CommandException
+
+
+class WarmUpCacheParametersExpectedError(CommandException):
+    status = 400
+    message = _(
+        "Malformed request. slice_id or table_name "
+        "and db_name arguments are expected"
+    )
+
+
+class WarmUpCacheChartNotFoundError(CommandException):
+    def __init__(self, chart_id: int):
+        message = f"Chart {chart_id} not found"
+        super().__init__(message)
+
+    status = 404
+
+
+class WarmUpCacheTableNotFoundError(CommandException):
+    def __init__(self, table_name: str, db_name: str):
+        message = f"Table {table_name} wasn't found in the database {db_name}"
+        super().__init__(message)
+
+    status = 404

superset/cachekeys/commands/warm_up_cache.py

@@ -0,0 +1,116 @@
+# 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 typing import Any, Dict, List, Optional
+
+import simplejson as json
+from flask import g
+
+from superset.cachekeys.commands.exceptions import (
+    WarmUpCacheChartNotFoundError,
+    WarmUpCacheParametersExpectedError,
+    WarmUpCacheTableNotFoundError,
+)
+from superset.commands.base import BaseCommand
+from superset.connectors.sqla.models import SqlaTable
+from superset.extensions import db
+from superset.models.core import Database
+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
+
+
+class WarmUpCacheCommand(BaseCommand):
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        slice_id: Optional[int],
+        dashboard_id: Optional[int],
+        table_name: Optional[str],
+        db_name: Optional[str],
+        extra_filters: Optional[str],
+    ):
+        self._slice_id = slice_id
+        self._dashboard_id = dashboard_id
+        self._table_name = table_name
+        self._db_name = db_name
+        self._extra_filters = extra_filters
+        self._slices: List[Slice] = []
+
+    def run(self) -> List[Dict[str, Any]]:
+        self.validate()
+        result: List[Dict[str, Any]] = []
+        for slc in self._slices:
+            try:
+                form_data = get_form_data(slc.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(slc.id, self._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 = error_msg_from_exception(ex)
+                status = None
+
+            result.append(
+                {"slice_id": slc.id, "viz_error": error, "viz_status": status}
+            )
+
+        return result
+
+    def validate(self) -> None:
+        if not self._slice_id and not (self._table_name and self._db_name):
+            raise WarmUpCacheParametersExpectedError()
+        if self._slice_id:
+            self._slices = db.session.query(Slice).filter_by(id=self._slice_id).all()
+            if not self._slices:
+                raise WarmUpCacheChartNotFoundError(self._slice_id)
+        elif self._table_name and self._db_name:
+            table = (
+                db.session.query(SqlaTable)
+                .join(Database)
+                .filter(
+                    Database.database_name == self._db_name,
+                    SqlaTable.table_name == self._table_name,
+                )
+            ).one_or_none()
+            if not table:
+                raise WarmUpCacheTableNotFoundError(self._table_name, self._db_name)
+            self._slices = (
+                db.session.query(Slice)
+                .filter_by(datasource_id=table.id, datasource_type=table.type)
+                .all()
+            )

superset/tasks/cache.py

@@ -40,7 +40,7 @@ def get_url(chart: Slice, dashboard: Optional[Dashboard] = None) -> str:
     """Return external URL for warming up a given chart/table cache."""
     with app.test_request_context():
         baseurl = "{WEBDRIVER_BASEURL}".format(**app.config)
-        url = f"{baseurl}superset/warm_up_cache/?slice_id={chart.id}"
+        url = f"{baseurl}api/v1/cachekey/warm_up_cache/?slice_id={chart.id}"
         if dashboard:
             url += f"&dashboard_id={dashboard.id}"
         return url
@@ -51,7 +51,7 @@ class Strategy:  # pylint: disable=too-few-public-methods
     A cache warm up strategy.
 
     Each strategy defines a `get_urls` method that returns a list of URLs to
-    be fetched from the `/superset/warm_up_cache/` endpoint.
+    be fetched from the `/api/v1/cachekey/warm_up_cache` endpoint.
 
     Strategies can be configured in `superset/config.py`:
 

superset/views/core.py

@@ -1697,6 +1697,7 @@ def fave_slices(self, user_id: Optional[int] = None) -> FlaskResponse:
     @api
     @has_access_api
     @expose("/warm_up_cache/", methods=["GET"])
+    @deprecated(new_target="api/v1/cachekey/warm_up_cache")
     def warm_up_cache(  # pylint: disable=too-many-locals,no-self-use
         self,
     ) -> FlaskResponse:

tests/integration_tests/cachekeys/api_tests.py

@@ -16,17 +16,28 @@
 # under the License.
 # isort:skip_file
 """Unit tests for Superset"""
+import json
 from typing import Dict, Any
+from urllib.parse import quote
 
 import pytest
 
 from superset.extensions import cache_manager, db
 from superset.models.cache import CacheKey
 from superset.utils.core import get_example_default_schema
+from superset.utils.database import get_example_database
 from tests.integration_tests.base_tests import (
     SupersetTestCase,
     post_assert_metric,
 )
+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,
+)
 
 
 @pytest.fixture
@@ -165,3 +176,40 @@ def test_invalidate_existing_caches(invalidate):
         .datasource_uid
         == "X__table"
     )
+
+
+class TestWarmUpCache(SupersetTestCase):
+    @pytest.mark.usefixtures(
+        "load_energy_table_with_slice", "load_birth_names_dashboard_with_slices"
+    )
+    def test_warm_up_cache(self):
+        self.login()
+        slc = self.get_slice("Girls", db.session)
+        data = self.get_json_resp(
+            "/api/v1/cachekey/warm_up_cache?slice_id={}".format(slc.id)
+        )
+        self.assertEqual(
+            data["result"],
+            [{"slice_id": slc.id, "viz_error": None, "viz_status": "success"}],
+        )
+
+        data = self.get_json_resp(
+            "/api/v1/cachekey/warm_up_cache?table_name=energy_usage"
+            f"&db_name={get_example_database().database_name}"
+        )
+        assert len(data) > 0
+
+        dashboard = self.get_dash_by_slug("births")
+
+        assert self.get_json_resp(
+            f"/api/v1/cachekey/warm_up_cache?dashboard_id={dashboard.id}&slice_id={slc.id}"
+        )["result"] == [
+            {"slice_id": slc.id, "viz_error": None, "viz_status": "success"}
+        ]
+
+        assert self.get_json_resp(
+            f"/api/v1/cachekey/warm_up_cache?dashboard_id={dashboard.id}&slice_id={slc.id}&extra_filters="
+            + quote(json.dumps([{"col": "name", "op": "in", "val": ["Jennifer"]}]))
+        )["result"] == [
+            {"slice_id": slc.id, "viz_error": None, "viz_status": "success"}
+        ]