def data(self) -> Response: # pylint: disable=too-many-return-statements
"""
Takes a query context constructed in the client and returns payload
data response for the given query.
---
post:
description: >-
Takes a query context constructed in the client and returns payload data
response for the given query.
requestBody:
description: >-
A query context consists of a datasource from which to fetch data
and one or many query objects.
required: true
content:
application/json:
schema:
$ref: "#/components/schemas/ChartDataQueryContextSchema"
responses:
200:
description: Query result
content:
application/json:
schema:
$ref: "#/components/schemas/ChartDataResponseSchema"
400:
$ref: '#/components/responses/400'
500:
$ref: '#/components/responses/500'
"""
if request.is_json:
json_body = request.json
elif request.form.get("form_data"):
# CSV export submits regular form data
json_body = json.loads(request.form["form_data"])
else:
return self.response_400(message="Request is not JSON")
try:
query_context = ChartDataQueryContextSchema().load(json_body)
except KeyError:
return self.response_400(message="Request is incorrect")
except ValidationError as error:
return self.response_400(message=_(
"Request is incorrect: %(error)s", error=error.messages))
try:
query_context.raise_for_access()
except SupersetSecurityException:
return self.response_401()
payload = query_context.get_payload()
for query in payload:
if query.get("error"):
return self.response_400(message=f"Error: {query['error']}")
result_format = query_context.result_format
if result_format == ChartDataResultFormat.CSV:
# return the first result
result = payload[0]["data"]
return CsvResponse(
result,
status=200,
headers=generate_download_headers("csv"),
mimetype="application/csv",
)
if result_format == ChartDataResultFormat.JSON:
response_data = simplejson.dumps({"result": payload},
default=json_int_dttm_ser,
ignore_nan=True)
resp = make_response(response_data, 200)
resp.headers["Content-Type"] = "application/json; charset=utf-8"
return resp
return self.response_400(
message=f"Unsupported result_format: {result_format}")
class ChartDataCommand(BaseCommand):
def __init__(self) -> None:
self._form_data: Dict[str, Any]
self._query_context: QueryContext
self._async_channel_id: str
def run(self, **kwargs: Any) -> Dict[str, Any]:
# caching is handled in query_context.get_df_payload
# (also evals `force` property)
cache_query_context = kwargs.get("cache", False)
force_cached = kwargs.get("force_cached", False)
try:
payload = self._query_context.get_payload(
cache_query_context=cache_query_context, force_cached=force_cached
)
except CacheLoadError as ex:
raise ChartDataCacheLoadError(ex.message) from ex
# TODO: QueryContext should support SIP-40 style errors
for query in payload["queries"]:
if query.get("error"):
raise ChartDataQueryFailedError(f"Error: {query['error']}")
return_value = {
"query_context": self._query_context,
"queries": payload["queries"],
}
if cache_query_context:
return_value.update(cache_key=payload["cache_key"])
return return_value
def run_async(self, user_id: Optional[str]) -> Dict[str, Any]:
job_metadata = async_query_manager.init_job(self._async_channel_id, user_id)
load_chart_data_into_cache.delay(job_metadata, self._form_data)
return job_metadata
def set_query_context(self, form_data: Dict[str, Any]) -> QueryContext:
self._form_data = form_data
try:
self._query_context = ChartDataQueryContextSchema().load(self._form_data)
except KeyError as ex:
raise ValidationError("Request is incorrect") from ex
except ValidationError as error:
raise error
return self._query_context
def validate(self) -> None:
self._query_context.raise_for_access()
def validate_async_request(self, request: Request) -> None:
jwt_data = async_query_manager.parse_jwt_from_request(request)
self._async_channel_id = jwt_data["channel"]
def load_query_context_from_cache( # pylint: disable=no-self-use
self, cache_key: str
) -> Dict[str, Any]:
cache_value = cache.get(cache_key)
if not cache_value:
raise ChartDataCacheLoadError("Cached data not found")
return cache_value["data"]
