fix: send CSV pivoted in reports (apache#16347)

opus-42 · Nov 14, 2021 · 4bf9038 · 4bf9038
1 parent c16a040
commit 4bf9038
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 9 deletions.
diff --git a/superset/charts/api.py b/superset/charts/api.py
@@ -499,6 +499,12 @@ def send_chart_response(
         result_type = result["query_context"].result_type
         result_format = result["query_context"].result_format
 
+        # Post-process the data so it matches the data presented in the chart.
+        # This is needed for sending reports based on text charts that do the
+        # post-processing of data, eg, the pivot table.
+        if result_type == ChartDataResultType.POST_PROCESSED:
+            result = apply_post_process(result, form_data)
+
         if result_format == ChartDataResultFormat.CSV:
             # Verify user has permission to export CSV file
             if not security_manager.can_access("can_csv", "Superset"):
@@ -509,12 +515,6 @@ def send_chart_response(
             return CsvResponse(data, headers=generate_download_headers("csv"))
 
         if result_format == ChartDataResultFormat.JSON:
-            # Post-process the data so it matches the data presented in the chart.
-            # This is needed for sending reports based on text charts that do the
-            # post-processing of data, eg, the pivot table.
-            if result_type == ChartDataResultType.POST_PROCESSED:
-                result = apply_post_process(result, form_data)
-
             response_data = simplejson.dumps(
                 {"result": result["queries"]},
                 default=json_int_dttm_ser,

diff --git a/superset/charts/post_processing.py b/superset/charts/post_processing.py
@@ -26,11 +26,17 @@
 for these chart types.
 """
 
+from io import StringIO
 from typing import Any, Dict, List, Optional, Tuple
 
 import pandas as pd
 
-from superset.utils.core import DTTM_ALIAS, extract_dataframe_dtypes, get_metric_name
+from superset.utils.core import (
+    ChartDataResultFormat,
+    DTTM_ALIAS,
+    extract_dataframe_dtypes,
+    get_metric_name,
+)
 
 
 def get_column_key(label: Tuple[str, ...], metrics: List[str]) -> Tuple[Any, ...]:
@@ -276,7 +282,13 @@ def apply_post_process(
     post_processor = post_processors[viz_type]
 
     for query in result["queries"]:
-        df = pd.DataFrame.from_dict(query["data"])
+        if query["result_format"] == ChartDataResultFormat.JSON:
+            df = pd.DataFrame.from_dict(query["data"])
+        elif query["result_format"] == ChartDataResultFormat.CSV:
+            df = pd.read_csv(StringIO(query["data"]))
+        else:
+            raise Exception(f"Result format {query['result_format']} not supported")
+
         processed_df = post_processor(df, form_data)
 
         query["colnames"] = list(processed_df.columns)
@@ -298,6 +310,12 @@ def apply_post_process(
             for index in processed_df.index
         ]
 
-        query["data"] = processed_df.to_dict()
+        if query["result_format"] == ChartDataResultFormat.JSON:
+            query["data"] = processed_df.to_dict()
+        elif query["result_format"] == ChartDataResultFormat.CSV:
+            buf = StringIO()
+            processed_df.to_csv(buf)
+            buf.seek(0)
+            query["data"] = buf.getvalue()
 
     return result
diff --git a/superset/common/query_actions.py b/superset/common/query_actions.py
@@ -104,6 +104,7 @@ def _get_full(
         payload["indexnames"] = list(df.index)
         payload["coltypes"] = extract_dataframe_dtypes(df)
         payload["data"] = query_context.get_data(df)
+        payload["result_format"] = query_context.result_format
     del payload["df"]
 
     filters = query_obj.filter