feat(chart-data-api): download multiple csvs as zip (#18618)

* feat(chart-data-api): download multiple csvs as zip

* break out util

* check for empty request
This commit is contained in:
Ville Brofeldt
2022-02-08 21:25:06 +02:00
committed by GitHub
parent 9c08bc0ffc
commit 125be78ee6
3 changed files with 63 additions and 7 deletions

View File

@@ -21,7 +21,7 @@ import logging
from typing import Any, Dict, Optional, TYPE_CHECKING
import simplejson
from flask import g, make_response, request
from flask import current_app, g, make_response, request, Response
from flask_appbuilder.api import expose, protect
from flask_babel import gettext as _
from marshmallow import ValidationError
@@ -44,13 +44,11 @@ from superset.connectors.base.models import BaseDatasource
from superset.exceptions import QueryObjectValidationError
from superset.extensions import event_logger
from superset.utils.async_query_manager import AsyncQueryTokenException
from superset.utils.core import json_int_dttm_ser
from superset.utils.core import create_zip, json_int_dttm_ser
from superset.views.base import CsvResponse, generate_download_headers
from superset.views.base_api import statsd_metrics
if TYPE_CHECKING:
from flask import Response
from superset.common.query_context import QueryContext
logger = logging.getLogger(__name__)
@@ -350,9 +348,25 @@ class ChartDataRestApi(ChartRestApi):
if not security_manager.can_access("can_csv", "Superset"):
return self.response_403()
# return the first result
data = result["queries"][0]["data"]
return CsvResponse(data, headers=generate_download_headers("csv"))
if not result["queries"]:
return self.response_400(_("Empty query result"))
if len(result["queries"]) == 1:
# return single query results csv format
data = result["queries"][0]["data"]
return CsvResponse(data, headers=generate_download_headers("csv"))
# return multi-query csv results bundled as a zip file
encoding = current_app.config["CSV_EXPORT"].get("encoding", "utf-8")
files = {
f"query_{idx + 1}.csv": result["data"].encode(encoding)
for idx, result in enumerate(result["queries"])
}
return Response(
create_zip(files),
headers=generate_download_headers("zip"),
mimetype="application/zip",
)
if result_format == ChartDataResultFormat.JSON:
response_data = simplejson.dumps(