[superset] branch master updated: fix: send CSV pivoted in reports (#16347)

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new ec8d3b0  fix: send CSV pivoted in reports (#16347)
ec8d3b0 is described below

commit ec8d3b03e4e14451ee5047a0a46a2a941ee13b4b
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Wed Aug 18 19:36:48 2021 -0700

    fix: send CSV pivoted in reports (#16347)
---
 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 3fa48c0..7a21fe2 100644
--- a/superset/charts/api.py
+++ b/superset/charts/api.py
@@ -499,6 +499,12 @@ class ChartRestApi(BaseSupersetModelRestApi):
         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 @@ class ChartRestApi(BaseSupersetModelRestApi):
             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 96d1e7c..9ef976b 100644
--- a/superset/charts/post_processing.py
+++ b/superset/charts/post_processing.py
@@ -26,11 +26,17 @@ In order to do that, we reproduce the post-processing in Python
 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 ea86104..349d62d 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