mirror of
https://github.com/apache/superset.git
synced 2026-04-22 17:45:21 +00:00
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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user