mirror of
https://github.com/apache/superset.git
synced 2026-04-19 16:14:52 +00:00
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:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user