query_context_generator.py

# 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 dataclasses
from typing import Any, Optional

from superset.common.chart_data import ChartDataResultType
from superset.utils.core import AnnotationType, DTTM_ALIAS

query_birth_names = {
    "extras": {"where": "", "time_grain_sqla": "P1D"},
    "columns": ["name"],
    "metrics": [{"label": "sum__num"}],
    "orderby": [("sum__num", False)],
    "row_limit": 100,
    "granularity": "ds",
    "time_range": "100 years ago : now",
    "timeseries_limit": 0,
    "timeseries_limit_metric": None,
    "order_desc": True,
    "filters": [
        {"col": "gender", "op": "==", "val": "boy"},
        {"col": "num", "op": "IS NOT NULL"},
        {"col": "name", "op": "NOT IN", "val": ["<NULL>", '"abc"']},
    ],
    "having": "",
    "where": "",
}

QUERY_OBJECTS: dict[str, dict[str, object]] = {
    "birth_names": query_birth_names,
    # `:suffix` are overrides only
    "birth_names:include_time": {
        "groupby": [DTTM_ALIAS, "name"],
    },
    "birth_names:orderby_dup_alias": {
        "metrics": [
            {
                "expressionType": "SIMPLE",
                "column": {"column_name": "num_girls", "type": "BIGINT(20)"},
                "aggregate": "SUM",
                "label": "num_girls",
            },
            {
                "expressionType": "SIMPLE",
                "column": {"column_name": "num_boys", "type": "BIGINT(20)"},
                "aggregate": "SUM",
                "label": "num_boys",
            },
        ],
        "orderby": [
            [
                {
                    "expressionType": "SIMPLE",
                    "column": {"column_name": "num_girls", "type": "BIGINT(20)"},
                    "aggregate": "SUM",
                    # the same underlying expression, but different label
                    "label": "SUM(num_girls)",
                },
                False,
            ],
            # reference the ambiguous alias in SIMPLE metric
            [
                {
                    "expressionType": "SIMPLE",
                    "column": {"column_name": "num_boys", "type": "BIGINT(20)"},
                    "aggregate": "AVG",
                    "label": "AVG(num_boys)",
                },
                False,
            ],
            # reference the ambiguous alias in CUSTOM SQL metric
            [
                {
                    "expressionType": "SQL",
                    "sqlExpression": "MAX(CASE WHEN num_boys > 0 THEN 1 ELSE 0 END)",
                    "label": "MAX(CASE WHEN...",
                },
                True,
            ],
        ],
    },
    "birth_names:only_orderby_has_metric": {
        "metrics": [],
    },
}

ANNOTATION_LAYERS = {
    AnnotationType.FORMULA: {
        "annotationType": "FORMULA",
        "color": "#ff7f44",
        "hideLine": False,
        "name": "my formula",
        "opacity": "",
        "overrides": {"time_range": None},
        "show": True,
        "showMarkers": False,
        "sourceType": "",
        "style": "solid",
        "value": "3+x",
        "width": 5,
    },
    AnnotationType.EVENT: {
        "name": "my event",
        "annotationType": "EVENT",
        "sourceType": "NATIVE",
        "color": "#e04355",
        "opacity": "",
        "style": "solid",
        "width": 5,
        "showMarkers": False,
        "hideLine": False,
        "value": 1,
        "overrides": {"time_range": None},
        "show": True,
        "titleColumn": "",
        "descriptionColumns": [],
        "timeColumn": "",
        "intervalEndColumn": "",
    },
    AnnotationType.INTERVAL: {
        "name": "my interval",
        "annotationType": "INTERVAL",
        "sourceType": "NATIVE",
        "color": "#e04355",
        "opacity": "",
        "style": "solid",
        "width": 1,
        "showMarkers": False,
        "hideLine": False,
        "value": 1,
        "overrides": {"time_range": None},
        "show": True,
        "titleColumn": "",
        "descriptionColumns": [],
        "timeColumn": "",
        "intervalEndColumn": "",
    },
    AnnotationType.TIME_SERIES: {
        "annotationType": "TIME_SERIES",
        "color": None,
        "descriptionColumns": [],
        "hideLine": False,
        "intervalEndColumn": "",
        "name": "my line",
        "opacity": "",
        "overrides": {"time_range": None},
        "show": True,
        "showMarkers": False,
        "sourceType": "line",
        "style": "dashed",
        "timeColumn": "",
        "titleColumn": "",
        "value": 837,
        "width": 5,
    },
}

POSTPROCESSING_OPERATIONS = {
    "birth_names": [
        {
            "operation": "aggregate",
            "options": {
                "groupby": ["name"],
                "aggregates": {
                    "q1": {
                        "operator": "percentile",
                        "column": "sum__num",
                        # todo: rename "interpolation" to "method" when we updated
                        #  numpy.
                        #  https://numpy.org/doc/stable/reference/generated/numpy.percentile.html
                        "options": {"q": 25, "interpolation": "lower"},
                    },
                    "median": {
                        "operator": "median",
                        "column": "sum__num",
                    },
                },
            },
        },
        {
            "operation": "sort",
            "options": {"by": ["q1", "name"], "ascending": [False, True]},
        },
    ]
}


def get_query_object(
    query_name: str,
    add_postprocessing_operations: bool,
    add_time_offsets: bool,
) -> dict[str, Any]:
    if query_name not in QUERY_OBJECTS:
        raise Exception(f"QueryObject fixture not defined for datasource: {query_name}")
    obj = QUERY_OBJECTS[query_name]

    # apply overrides
    if ":" in query_name:
        parent_query_name = query_name.split(":")[0]
        obj = {
            **QUERY_OBJECTS[parent_query_name],
            **obj,
        }

    query_object = copy.deepcopy(obj)
    if add_postprocessing_operations:
        query_object["post_processing"] = _get_postprocessing_operation(query_name)
    if add_time_offsets:
        query_object["time_offsets"] = ["1 year ago"]

    return query_object


def _get_postprocessing_operation(query_name: str) -> list[dict[str, Any]]:
    if query_name not in QUERY_OBJECTS:
        raise Exception(
            f"Post-processing fixture not defined for datasource: {query_name}"
        )
    return copy.deepcopy(POSTPROCESSING_OPERATIONS[query_name])


@dataclasses.dataclass
class Table:
    id: int
    type: str
    name: str


class QueryContextGenerator:
    def generate(
        self,
        query_name: str,
        add_postprocessing_operations: bool = False,
        add_time_offsets: bool = False,
        table_id=1,
        table_type="table",
        form_data: Optional[dict[str, Any]] = None,
    ) -> dict[str, Any]:
        form_data = form_data or {}
        table_name = query_name.split(":")[0]
        table = self.get_table(table_name, table_id, table_type)
        return {
            "datasource": {"id": table.id, "type": table.type},
            "queries": [
                get_query_object(
                    query_name,
                    add_postprocessing_operations,
                    add_time_offsets,
                )
            ],
            "result_type": ChartDataResultType.FULL,
            "form_data": form_data,
        }

    def get_table(self, name, id_, type_):
        return Table(id_, type_, name)