feat: expand new chart data endpoint coverage (#9903)

* feat: implement new chart API for additional components

* Fix python tests

* Fix tests

* Fix lint

* fix camel case error in requestParams

* lint

* fix samples row limit

* Add samples row limit to config

* remove unnecessary code

* lint

* Address review comments
This commit is contained in:
Ville Brofeldt
2020-06-02 10:47:28 +03:00
committed by GitHub
parent c7618ee54b
commit 38a6bd79da
12 changed files with 196 additions and 126 deletions

View File

@@ -16,6 +16,7 @@
# under the License.
import copy
import logging
import math
import pickle as pkl
from datetime import datetime, timedelta
from typing import Any, ClassVar, Dict, List, Optional, Union
@@ -50,8 +51,8 @@ class QueryContext:
queries: List[QueryObject]
force: bool
custom_cache_timeout: Optional[int]
response_type: utils.ChartDataResponseType
response_format: utils.ChartDataResponseFormat
result_type: utils.ChartDataResultType
result_format: utils.ChartDataResultFormat
# TODO: Type datasource and query_object dictionary with TypedDict when it becomes
# a vanilla python type https://github.com/python/mypy/issues/5288
@@ -61,8 +62,8 @@ class QueryContext:
queries: List[Dict[str, Any]],
force: bool = False,
custom_cache_timeout: Optional[int] = None,
response_format: Optional[utils.ChartDataResponseFormat] = None,
response_type: Optional[utils.ChartDataResponseType] = None,
result_type: Optional[utils.ChartDataResultType] = None,
result_format: Optional[utils.ChartDataResultFormat] = None,
) -> None:
self.datasource = ConnectorRegistry.get_datasource(
str(datasource["type"]), int(datasource["id"]), db.session
@@ -70,8 +71,8 @@ class QueryContext:
self.queries = [QueryObject(**query_obj) for query_obj in queries]
self.force = force
self.custom_cache_timeout = custom_cache_timeout
self.response_format = response_format or utils.ChartDataResponseFormat.JSON
self.response_type = response_type or utils.ChartDataResponseType.RESULTS
self.result_type = result_type or utils.ChartDataResultType.FULL
self.result_format = result_format or utils.ChartDataResultFormat.JSON
def get_query_result(self, query_object: QueryObject) -> Dict[str, Any]:
"""Returns a pandas dataframe based on the query object"""
@@ -134,7 +135,7 @@ class QueryContext:
def get_data(
self, df: pd.DataFrame,
) -> Union[str, List[Dict[str, Any]]]: # pylint: disable=no-self-use
if self.response_format == utils.ChartDataResponseFormat.CSV:
if self.result_format == utils.ChartDataResultFormat.CSV:
include_index = not isinstance(df.index, pd.RangeIndex)
result = df.to_csv(index=include_index, **config["CSV_EXPORT"])
return result or ""
@@ -143,18 +144,18 @@ class QueryContext:
def get_single_payload(self, query_obj: QueryObject) -> Dict[str, Any]:
"""Returns a payload of metadata and data"""
if self.response_type == utils.ChartDataResponseType.QUERY:
if self.result_type == utils.ChartDataResultType.QUERY:
return {
"query": self.datasource.get_query_str(query_obj.to_dict()),
"language": self.datasource.query_language,
}
if self.response_type == utils.ChartDataResponseType.SAMPLES:
row_limit = query_obj.row_limit or 1000
if self.result_type == utils.ChartDataResultType.SAMPLES:
row_limit = query_obj.row_limit or math.inf
query_obj = copy.copy(query_obj)
query_obj.groupby = []
query_obj.metrics = []
query_obj.post_processing = []
query_obj.row_limit = row_limit
query_obj.row_limit = min(row_limit, config["SAMPLES_ROW_LIMIT"])
query_obj.columns = [o.column_name for o in self.datasource.columns]
payload = self.get_df_payload(query_obj)