From 3e8d1a1513d7a038131ada4774970cef7d8b3bfd Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Wed, 18 Aug 2021 18:03:40 -0700
Subject: [PATCH] fix: send CSV pivoted in reports

---
 superset/charts/api.py             | 12 ++++++------
 superset/charts/post_processing.py | 24 +++++++++++++++++++++---
 superset/common/query_actions.py   |  1 +
 3 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/superset/charts/api.py b/superset/charts/api.py
index 3fa48c076b177..7a21fe2676d05 100644
--- 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
index 96d1e7ca0ed44..9ef976bdec62b 100644
--- 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
index ea8610433993d..349d62d38227e 100644
--- 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