feat(SIP-39): Async query support for charts

setup.cfg

@@ -30,7 +30,7 @@ combine_as_imports = true
 include_trailing_comma = true
 line_length = 88
 known_first_party = superset
-known_third_party =alembic,apispec,backoff,bleach,cachelib,celery,click,colorama,contextlib2,croniter,cryptography,dateutil,flask,flask_appbuilder,flask_babel,flask_caching,flask_compress,flask_login,flask_migrate,flask_sqlalchemy,flask_talisman,flask_testing,flask_wtf,freezegun,geohash,geopy,humanize,isodate,jinja2,markdown,markupsafe,marshmallow,msgpack,numpy,pandas,parameterized,parsedatetime,pathlib2,polyline,prison,pyarrow,pyhive,pytest,pytz,retry,selenium,setuptools,simplejson,slack,sqlalchemy,sqlalchemy_utils,sqlparse,werkzeug,wtforms,wtforms_json,yaml
+known_third_party =alembic,apispec,backoff,bleach,cachelib,celery,click,colorama,contextlib2,croniter,cryptography,dateutil,flask,flask_appbuilder,flask_babel,flask_caching,flask_compress,flask_login,flask_migrate,flask_sqlalchemy,flask_talisman,flask_testing,flask_wtf,freezegun,geohash,geopy,humanize,isodate,jinja2,jwt,markdown,markupsafe,marshmallow,msgpack,numpy,pandas,parameterized,parsedatetime,pathlib2,polyline,prison,pyarrow,pyhive,pytest,pytz,redis,retry,selenium,setuptools,simplejson,slack,sqlalchemy,sqlalchemy_utils,sqlparse,werkzeug,wtforms,wtforms_json,yaml
 multi_line_output = 3
 order_by_type = false
 

superset/app.py

@@ -30,6 +30,7 @@
     _event_logger,
     APP_DIR,
     appbuilder,
+    async_query_manager,
     cache_manager,
     celery_app,
     csrf,
@@ -485,6 +486,7 @@ def init_app_in_ctx(self) -> None:
         self.configure_url_map_converters()
         self.configure_data_sources()
         self.configure_auth_provider()
+        self.configure_async_queries()
 
         # Hook that provides administrators a handle on the Flask APP
         # after initialization
@@ -639,6 +641,11 @@ def configure_wtf(self) -> None:
             for ex in csrf_exempt_list:
                 csrf.exempt(ex)
 
+    def configure_async_queries(self) -> None:
+        if feature_flag_manager.is_feature_enabled("GLOBAL_ASYNC_QUERIES"):
+            logger.info("*************** Init async queries")
+            async_query_manager.init_app(self.flask_app)
+
     def register_blueprints(self) -> None:
         for bp in self.config["BLUEPRINTS"]:
             try:

superset/charts/api.py

@@ -33,10 +33,13 @@
 from superset import is_feature_enabled, thumbnail_cache
 from superset.charts.commands.bulk_delete import BulkDeleteChartCommand
 from superset.charts.commands.create import CreateChartCommand
+from superset.charts.commands.data import ChartDataCommand
 from superset.charts.commands.delete import DeleteChartCommand
 from superset.charts.commands.exceptions import (
     ChartBulkDeleteFailedError,
     ChartCreateFailedError,
+    ChartDataCacheLoadError,
+    ChartDataQueryFailedError,
     ChartDeleteFailedError,
     ChartForbiddenError,
     ChartInvalidError,
@@ -49,7 +52,6 @@
 from superset.charts.filters import ChartAllTextFilter, ChartFavoriteFilter, ChartFilter
 from superset.charts.schemas import (
     CHART_SCHEMAS,
-    ChartDataQueryContextSchema,
     ChartPostSchema,
     ChartPutSchema,
     get_delete_ids_schema,
@@ -61,9 +63,10 @@
 )
 from superset.constants import RouteMethod
 from superset.exceptions import SupersetSecurityException
-from superset.extensions import event_logger
+from superset.extensions import async_query_manager, event_logger
 from superset.models.slice import Slice
 from superset.tasks.thumbnails import cache_chart_thumbnail
+from superset.utils.async_query_manager import AsyncQueryTokenException
 from superset.utils.core import ChartDataResultFormat, json_int_dttm_ser
 from superset.utils.screenshots import ChartScreenshot
 from superset.utils.urls import get_url_path
@@ -89,6 +92,8 @@ class ChartRestApi(BaseSupersetModelRestApi):
         RouteMethod.RELATED,
         "bulk_delete",  # not using RouteMethod since locally defined
         "data",
+        "data_from_cache",
+        "viz_types",
         "favorite_status",
     }
     class_permission_name = "SliceModelView"
@@ -468,8 +473,12 @@ def data(self) -> Response:
                 application/json:
                   schema:
                     $ref: "#/components/schemas/ChartDataResponseSchema"
+            202:
+              $ref: '#/components/responses/202'
             400:
               $ref: '#/components/responses/400'
+            401:
+              $ref: '#/components/responses/401'
             500:
               $ref: '#/components/responses/500'
         """
@@ -480,41 +489,138 @@ def data(self) -> Response:
             json_body = json.loads(request.form["form_data"])
         else:
             return self.response_400(message="Request is not JSON")
+
         try:
-            query_context = ChartDataQueryContextSchema().load(json_body)
-        except KeyError:
-            return self.response_400(message="Request is incorrect")
+            command = ChartDataCommand()
+            command.set_query_context(json_body)
+            command.validate()
         except ValidationError as error:
             return self.response_400(
                 message=_("Request is incorrect: %(error)s", error=error.messages)
             )
-        try:
-            query_context.raise_for_access()
         except SupersetSecurityException:
             return self.response_401()
-        payload = query_context.get_payload()
-        for query in payload:
-            if query.get("error"):
-                return self.response_400(message=f"Error: {query['error']}")
-        result_format = query_context.result_format
 
+        if is_feature_enabled("GLOBAL_ASYNC_QUERIES"):
+            try:
+                command.validate_request(request)
+            except AsyncQueryTokenException:
+                return self.response_401()
+
+            result = command.run_async()
+            return self.response(202, **result)
+
+        # TODO: DRY
+        try:
+            result = command.run()
+        except ChartDataQueryFailedError as exc:
+            return self.response_400(message=exc.message)
+
+        result_format = result["query_context"].result_format
+        response = self.response_400(
+            message=f"Unsupported result_format: {result_format}"
+        )
+
+        if result_format == ChartDataResultFormat.CSV:
+            # return the first result
+            data = result["queries"][0]["data"]
+            response = CsvResponse(
+                data,
+                status=200,
+                headers=generate_download_headers("csv"),
+                mimetype="application/csv",
+            )
+
+        if result_format == ChartDataResultFormat.JSON:
+            response_data = simplejson.dumps(
+                {"result": result["queries"]},
+                default=json_int_dttm_ser,
+                ignore_nan=True,
+            )
+            resp = make_response(response_data, 200)
+            resp.headers["Content-Type"] = "application/json; charset=utf-8"
+            response = resp
+
+        return response
+
+    @expose("/data/<cache_key>", methods=["GET"])
+    @event_logger.log_this
+    @protect()
+    @safe
+    @statsd_metrics
+    def data_from_cache(self, cache_key: str) -> Response:
+        """
+        Takes a query context cache key returns payload
+        data response for the given query.
+        ---
+        get:
+          description: >-
+            Takes a query context constructed in the client and returns payload data
+            response for the given query.
+          parameters:
+          - in: path
+            schema:
+              type: string
+            name: cache_key
+          responses:
+            200:
+              description: Query result
+              content:
+                application/json:
+                  schema:
+                    $ref: "#/components/schemas/ChartDataResponseSchema"
+            400:
+              $ref: '#/components/responses/400'
+            401:
+              $ref: '#/components/responses/401'
+            404:
+              $ref: '#/components/responses/404'
+            500:
+              $ref: '#/components/responses/500'
+        """
+        command = ChartDataCommand()
+        try:
+            cached_data = command.load_query_context_from_cache(cache_key)
+            command.set_query_context(cached_data)
+            command.validate()
+        except ChartDataCacheLoadError:
+            return self.response_404()
+        except ValidationError as error:
+            return self.response_400(
+                message=_("Request is incorrect: %(error)s", error=error.messages)
+            )
+        except SupersetSecurityException as exc:
+            logger.info(exc)
+            return self.response_401()
+
+        # TODO: DRY
+        try:
+            result = command.run()
+        except ChartDataCacheLoadError as exc:
+            return self.response_400(message=exc.message)
+        except ChartDataQueryFailedError as exc:
+            return self.response_400(message=exc.message)
+
+        result_format = result["query_context"].result_format
         response = self.response_400(
             message=f"Unsupported result_format: {result_format}"
         )
 
         if result_format == ChartDataResultFormat.CSV:
             # return the first result
-            result = payload[0]["data"]
+            data = result["queries"][0]["data"]
             response = CsvResponse(
-                result,
+                data,
                 status=200,
                 headers=generate_download_headers("csv"),
                 mimetype="application/csv",
             )
 
         if result_format == ChartDataResultFormat.JSON:
             response_data = simplejson.dumps(
-                {"result": payload}, default=json_int_dttm_ser, ignore_nan=True
+                {"result": result["queries"]},
+                default=json_int_dttm_ser,
+                ignore_nan=True,
             )
             resp = make_response(response_data, 200)
             resp.headers["Content-Type"] = "application/json; charset=utf-8"

superset/charts/commands/data.py

@@ -0,0 +1,96 @@
+# 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.
+import logging
+from typing import Any, Dict, Optional
+
+from flask import Request
+from marshmallow import ValidationError
+
+from superset import cache
+from superset.charts.commands.exceptions import (
+    ChartDataCacheLoadError,
+    ChartDataQueryFailedError,
+)
+from superset.charts.schemas import ChartDataQueryContextSchema
+from superset.commands.base import BaseCommand
+from superset.common.query_context import QueryContext
+from superset.exceptions import CacheLoadError
+from superset.extensions import async_query_manager
+from superset.tasks.async_queries import load_chart_data_into_cache
+
+logger = logging.getLogger(__name__)
+
+
+class ChartDataCommand(BaseCommand):
+    def __init__(self) -> None:
+        self._form_data: Dict[str, Any]
+        self._query_context: QueryContext
+        self._async_channel_id: str
+
+    def run(self, **kwargs: Any) -> Dict[str, Any]:
+        # caching is handled in query_context.get_df_payload (also evals `force` property)
+        cache_query_context = kwargs["cache"] if "cache" in kwargs else False
+        force_cached = kwargs["force_cached"] if "force_cached" in kwargs else False
+        try:
+            payload = self._query_context.get_payload(
+                cache_query_context=cache_query_context, force_cached=force_cached
+            )
+        except CacheLoadError:
+            raise ChartDataCacheLoadError()
+
+        for query in payload["queries"]:
+            if query.get("error"):
+                raise ChartDataQueryFailedError(f"Error: {query['error']}")
+
+        return_value = {
+            "query_context": self._query_context,
+            "queries": payload["queries"],
+        }
+        if cache_query_context:
+            return_value.update(cache_key=payload["cache_key"])
+
+        return return_value
+
+    def run_async(self) -> Dict[str, Any]:
+        # TODO: confirm cache backend is configured
+        job_metadata = async_query_manager.init_job(self._async_channel_id)
+        load_chart_data_into_cache.delay(job_metadata, self._form_data)
+
+        return job_metadata
+
+    def set_query_context(self, form_data: Dict[str, Any]) -> None:
+        self._form_data = form_data
+        try:
+            self._query_context = ChartDataQueryContextSchema().load(self._form_data)
+        except KeyError:
+            raise ValidationError("Request is incorrect")
+        except ValidationError as error:
+            raise error
+
+    def validate(self) -> None:
+        self._query_context.raise_for_access()
+
+    def validate_request(self, request: Request) -> None:
+        jwt_data = async_query_manager.parse_jwt_from_request(request)
+        self._async_channel_id = jwt_data["channel"]
+
+    def load_query_context_from_cache(self, cache_key: str) -> Dict[str, Any]:
+        cache_value = cache.get(cache_key)
+        if not cache_value:
+            raise ChartDataCacheLoadError("Cached data not found")
+
+        return cache_value["data"]

superset/charts/commands/exceptions.py

@@ -83,3 +83,11 @@ class ChartForbiddenError(ForbiddenError):
 
 class ChartBulkDeleteFailedError(CreateFailedError):
     message = _("Charts could not be deleted.")
+
+
+class ChartDataQueryFailedError(CommandException):
+    pass
+
+
+class ChartDataCacheLoadError(CommandException):
+    pass