feat(chart-data): add rowcount, timegrain and column result types (#1…

…3271)

* feat(chart-data): add rowcount, timegrain and column result types

* break out actions from query_context

* rename module

superset/charts/schemas.py

@@ -887,10 +887,22 @@ class AnnotationLayerSchema(Schema):
     )
 
 
+class ChartDataDatasourceSchema(Schema):
+    description = "Chart datasource"
+    id = fields.Integer(description="Datasource id", required=True,)
+    type = fields.String(
+        description="Datasource type",
+        validate=validate.OneOf(choices=("druid", "table")),
+    )
+
+
 class ChartDataQueryObjectSchema(Schema):
     class Meta:  # pylint: disable=too-few-public-methods
         unknown = EXCLUDE
 
+    datasource = fields.Nested(ChartDataDatasourceSchema, allow_none=True)
+    result_type = EnumField(ChartDataResultType, by_value=True, allow_none=True)
+
     annotation_layers = fields.List(
         fields.Nested(AnnotationLayerSchema),
         description="Annotation layers to apply to chart",
@@ -971,10 +983,10 @@ class Meta:  # pylint: disable=too-few-public-methods
         description="Metric used to limit timeseries queries by.", allow_none=True,
     )
     row_limit = fields.Integer(
-        description='Maximum row count. Default: `config["ROW_LIMIT"]`',
+        description='Maximum row count (0=disabled). Default: `config["ROW_LIMIT"]`',
         allow_none=True,
         validate=[
-            Range(min=1, error=_("`row_limit` must be greater than or equal to 1"))
+            Range(min=0, error=_("`row_limit` must be greater than or equal to 0"))
         ],
     )
     row_offset = fields.Integer(
@@ -1038,14 +1050,9 @@ class Meta:  # pylint: disable=too-few-public-methods
         values=fields.String(description="The value of the query parameter"),
         allow_none=True,
     )
-
-
-class ChartDataDatasourceSchema(Schema):
-    description = "Chart datasource"
-    id = fields.Integer(description="Datasource id", required=True,)
-    type = fields.String(
-        description="Datasource type",
-        validate=validate.OneOf(choices=("druid", "table")),
+    is_rowcount = fields.Boolean(
+        description="Should the rowcount of the actual query be returned",
+        allow_none=True,
     )
 
 

superset/common/query_actions.py

@@ -0,0 +1,182 @@
+# 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 copy
+import math
+from typing import Any, Callable, cast, Dict, List, Optional, TYPE_CHECKING
+
+from flask_babel import _
+
+from superset import app
+from superset.connectors.base.models import BaseDatasource
+from superset.exceptions import QueryObjectValidationError
+from superset.utils.core import (
+    ChartDataResultType,
+    extract_column_dtype,
+    extract_dataframe_dtypes,
+    get_time_filter_status,
+    QueryStatus,
+)
+
+if TYPE_CHECKING:
+    from superset.common.query_context import QueryContext
+    from superset.common.query_object import QueryObject
+
+config = app.config
+
+
+def _get_datasource(
+    query_context: "QueryContext", query_obj: "QueryObject"
+) -> BaseDatasource:
+    return query_obj.datasource or query_context.datasource
+
+
+def _get_columns(
+    query_context: "QueryContext", query_obj: "QueryObject", _: bool
+) -> Dict[str, Any]:
+    datasource = _get_datasource(query_context, query_obj)
+    return {
+        "data": [
+            {
+                "column_name": col.column_name,
+                "verbose_name": col.verbose_name,
+                "dtype": extract_column_dtype(col),
+            }
+            for col in datasource.columns
+        ]
+    }
+
+
+def _get_timegrains(
+    query_context: "QueryContext", query_obj: "QueryObject", _: bool
+) -> Dict[str, Any]:
+    datasource = _get_datasource(query_context, query_obj)
+    return {
+        "data": [
+            {
+                "name": grain.name,
+                "function": grain.function,
+                "duration": grain.duration,
+            }
+            for grain in datasource.database.grains()
+        ]
+    }
+
+
+def _get_query(
+    query_context: "QueryContext", query_obj: "QueryObject", _: bool,
+) -> Dict[str, Any]:
+    datasource = _get_datasource(query_context, query_obj)
+    return {
+        "query": datasource.get_query_str(query_obj.to_dict()),
+        "language": datasource.query_language,
+    }
+
+
+def _get_full(
+    query_context: "QueryContext",
+    query_obj: "QueryObject",
+    force_cached: Optional[bool] = False,
+) -> Dict[str, Any]:
+    datasource = _get_datasource(query_context, query_obj)
+    result_type = query_obj.result_type or query_context.result_type
+    payload = query_context.get_df_payload(query_obj, force_cached=force_cached)
+    df = payload["df"]
+    status = payload["status"]
+    if status != QueryStatus.FAILED:
+        payload["colnames"] = list(df.columns)
+        payload["coltypes"] = extract_dataframe_dtypes(df)
+        payload["data"] = query_context.get_data(df)
+    del payload["df"]
+
+    filters = query_obj.filter
+    filter_columns = cast(List[str], [flt.get("col") for flt in filters])
+    columns = set(datasource.column_names)
+    applied_time_columns, rejected_time_columns = get_time_filter_status(
+        datasource, query_obj.applied_time_extras
+    )
+    payload["applied_filters"] = [
+        {"column": col} for col in filter_columns if col in columns
+    ] + applied_time_columns
+    payload["rejected_filters"] = [
+        {"reason": "not_in_datasource", "column": col}
+        for col in filter_columns
+        if col not in columns
+    ] + rejected_time_columns
+
+    if result_type == ChartDataResultType.RESULTS and status != QueryStatus.FAILED:
+        return {"data": payload["data"]}
+    return payload
+
+
+def _get_samples(
+    query_context: "QueryContext", query_obj: "QueryObject", force_cached: bool = False
+) -> Dict[str, Any]:
+    datasource = _get_datasource(query_context, query_obj)
+    row_limit = query_obj.row_limit or math.inf
+    query_obj = copy.copy(query_obj)
+    query_obj.is_timeseries = False
+    query_obj.orderby = []
+    query_obj.groupby = []
+    query_obj.metrics = []
+    query_obj.post_processing = []
+    query_obj.row_limit = min(row_limit, config["SAMPLES_ROW_LIMIT"])
+    query_obj.row_offset = 0
+    query_obj.columns = [o.column_name for o in datasource.columns]
+    return _get_full(query_context, query_obj, force_cached)
+
+
+def _get_results(
+    query_context: "QueryContext", query_obj: "QueryObject", force_cached: bool = False
+) -> Dict[str, Any]:
+    payload = _get_full(query_context, query_obj, force_cached)
+    return {"data": payload["data"]}
+
+
+_result_type_functions: Dict[
+    ChartDataResultType, Callable[["QueryContext", "QueryObject", bool], Dict[str, Any]]
+] = {
+    ChartDataResultType.COLUMNS: _get_columns,
+    ChartDataResultType.TIMEGRAINS: _get_timegrains,
+    ChartDataResultType.QUERY: _get_query,
+    ChartDataResultType.SAMPLES: _get_samples,
+    ChartDataResultType.FULL: _get_full,
+    ChartDataResultType.RESULTS: _get_results,
+}
+
+
+def get_query_results(
+    result_type: ChartDataResultType,
+    query_context: "QueryContext",
+    query_obj: "QueryObject",
+    force_cached: bool,
+) -> Dict[str, Any]:
+    """
+    Return result payload for a chart data request.
+
+    :param result_type: the type of result to return
+    :param query_context: query context to which the query object belongs
+    :param query_obj: query object for which to retrieve the results
+    :param force_cached: should results be forcefully retrieved from cache
+    :raises QueryObjectValidationError: if an unsupported result type is requested
+    :return: JSON serializable result payload
+    """
+    result_func = _result_type_functions.get(result_type)
+    if result_func:
+        return result_func(query_context, query_obj, force_cached)
+    raise QueryObjectValidationError(
+        _("Invalid result type: %(result_type)", result_type=result_type)
+    )