# 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": ["", '"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)