query_context.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
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)],
    }