| # 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 |
| from typing import Any, Dict, List |
| |
| QUERY_OBJECTS = { |
| "birth_names": { |
| "extras": {"where": "", "time_range_endpoints": ["inclusive", "exclusive"]}, |
| "granularity": "ds", |
| "groupby": ["name"], |
| "is_timeseries": False, |
| "metrics": [{"label": "sum__num"}], |
| "order_desc": True, |
| "orderby": [], |
| "row_limit": 100, |
| "time_range": "100 years ago : now", |
| "timeseries_limit": 0, |
| "timeseries_limit_metric": None, |
| "filters": [{"col": "gender", "op": "==", "val": "boy"}], |
| "having": "", |
| "having_filters": [], |
| "where": "", |
| } |
| } |
| |
| POSTPROCESSING_OPERATIONS = { |
| "birth_names": [ |
| { |
| "operation": "aggregate", |
| "options": { |
| "groupby": ["gender"], |
| "aggregates": { |
| "q1": { |
| "operator": "percentile", |
| "column": "sum__num", |
| "options": {"q": 25}, |
| }, |
| "median": {"operator": "median", "column": "sum__num",}, |
| }, |
| }, |
| }, |
| {"operation": "sort", "options": {"columns": {"q1": False, "gender": True},},}, |
| ] |
| } |
| |
| |
| def _get_query_object( |
| datasource_name: str, add_postprocessing_operations: bool |
| ) -> Dict[str, Any]: |
| if datasource_name not in QUERY_OBJECTS: |
| raise Exception( |
| f"QueryObject fixture not defined for datasource: {datasource_name}" |
| ) |
| query_object = copy.deepcopy(QUERY_OBJECTS[datasource_name]) |
| if add_postprocessing_operations: |
| query_object["post_processing"] = _get_postprocessing_operation(datasource_name) |
| return query_object |
| |
| |
| def _get_postprocessing_operation(datasource_name: str) -> List[Dict[str, Any]]: |
| if datasource_name not in QUERY_OBJECTS: |
| raise Exception( |
| f"Post-processing fixture not defined for datasource: {datasource_name}" |
| ) |
| return copy.deepcopy(POSTPROCESSING_OPERATIONS[datasource_name]) |
| |
| |
| def get_query_context( |
| datasource_name: str = "birth_names", |
| datasource_id: int = 0, |
| datasource_type: str = "table", |
| add_postprocessing_operations: bool = False, |
| ) -> Dict[str, Any]: |
| """ |
| Create a request payload for retrieving a QueryContext object via the |
| `api/v1/chart/data` endpoint. By default returns a payload corresponding to one |
| generated by the "Boy Name Cloud" chart in the examples. |
| |
| :param datasource_name: name of datasource to query. Different datasources require |
| different parameters in the QueryContext. |
| :param datasource_id: id of datasource to query. |
| :param datasource_type: type of datasource to query. |
| :param add_postprocessing_operations: Add post-processing operations to QueryObject |
| :return: Request payload |
| """ |
| return { |
| "datasource": {"id": datasource_id, "type": datasource_type}, |
| "queries": [_get_query_object(datasource_name, add_postprocessing_operations)], |
| } |