feat(explore): export csv data pivoted for Pivot Table [ID-9] (#17512)

* feat(explore): export csv data pivoted for Pivot Table

* Implement dropdown with download csv options

* Change label to "Original"

* Add tests

* Add form data to query context

* Add form data to query context generator

* Explicitly make form_data optional

superset-frontend/packages/superset-ui-core/src/query/buildQueryContext.ts

@@ -73,6 +73,7 @@ export default function buildQueryContext(
         ...hooks,
       },
     }),
+    form_data: formData,
     result_format: formData.result_format || 'json',
     result_type: formData.result_type || 'full',
   };

superset-frontend/packages/superset-ui-core/src/query/types/Query.ts

@@ -21,7 +21,12 @@ import { DatasourceType } from './Datasource';
 import { BinaryOperator, SetOperator, UnaryOperator } from './Operator';
 import { AppliedTimeExtras, TimeRange, TimeRangeEndpoints } from './Time';
 import { AnnotationLayer } from './AnnotationLayer';
-import { QueryFields, QueryFormColumn, QueryFormMetric } from './QueryFormData';
+import {
+  QueryFields,
+  QueryFormColumn,
+  QueryFormData,
+  QueryFormMetric,
+} from './QueryFormData';
 import { Maybe } from '../../types';
 import { PostProcessingRule } from './PostProcessing';
 import { JsonObject } from '../../connection';
@@ -158,6 +163,7 @@ export interface QueryContext {
   /** Response format */
   result_format: string;
   queries: QueryObject[];
+  form_data?: QueryFormData;
 }
 
 export default {};

superset-frontend/spec/javascripts/explore/components/ExploreActionButtons_spec.jsx

@@ -82,4 +82,36 @@ describe('ExploreActionButtons', () => {
       spyExportChart.restore();
     });
   });
+
+  describe('Dropdown csv button when viz type is pivot table', () => {
+    let wrapper;
+    const defaultProps = {
+      actions: {},
+      canDownloadCSV: false,
+      latestQueryFormData: { viz_type: 'pivot_table_v2' },
+      queryEndpoint: 'localhost',
+      chartHeight: '30px',
+    };
+
+    beforeEach(() => {
+      wrapper = mount(
+        <ThemeProvider theme={supersetTheme}>
+          <ExploreActionButtons {...defaultProps} />
+        </ThemeProvider>,
+        {
+          wrappingComponent: Provider,
+          wrappingComponentProps: {
+            store: mockStore,
+          },
+        },
+      );
+    });
+
+    it('should render a dropdown button when viz type is pivot table', () => {
+      const csvTrigger = wrapper.find(
+        'div[role="button"] span[aria-label="caret-down"]',
+      );
+      expect(csvTrigger).toExist();
+    });
+  });
 });

superset-frontend/src/explore/components/ExploreActionButtons.tsx

@@ -16,9 +16,9 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { useState } from 'react';
+import React, { ReactElement, useState } from 'react';
 import cx from 'classnames';
-import { t } from '@superset-ui/core';
+import { QueryFormData, t } from '@superset-ui/core';
 import Icons from 'src/components/Icons';
 import { Tooltip } from 'src/components/Tooltip';
 import copyTextToClipboard from 'src/utils/copy';
@@ -27,13 +27,15 @@ import { useUrlShortener } from 'src/common/hooks/useUrlShortener';
 import EmbedCodeButton from './EmbedCodeButton';
 import { exportChart, getExploreLongUrl } from '../exploreUtils';
 import ExploreAdditionalActionsMenu from './ExploreAdditionalActionsMenu';
+import { ExportToCSVDropdown } from './ExportToCSVDropdown';
 
 type ActionButtonProps = {
-  icon: React.ReactElement;
+  prefixIcon: React.ReactElement;
-  text?: string;
+  suffixIcon?: React.ReactElement;
+  text?: string | ReactElement;
   tooltip: string;
   className?: string;
-  onClick: React.MouseEventHandler<HTMLElement>;
+  onClick?: React.MouseEventHandler<HTMLElement>;
   onTooltipVisibilityChange?: (visible: boolean) => void;
   'data-test'?: string;
 };
@@ -42,18 +44,27 @@ type ExploreActionButtonsProps = {
   actions: { redirectSQLLab: () => void; openPropertiesModal: () => void };
   canDownloadCSV: boolean;
   chartStatus: string;
-  latestQueryFormData: {};
+  latestQueryFormData: QueryFormData;
   queriesResponse: {};
   slice: { slice_name: string };
   addDangerToast: Function;
 };
 
+const VIZ_TYPES_PIVOTABLE = ['pivot_table', 'pivot_table_v2'];
+
 const ActionButton = (props: ActionButtonProps) => {
-  const { icon, text, tooltip, className, onTooltipVisibilityChange, ...rest } =
+  const {
-    props;
+    prefixIcon,
+    suffixIcon,
+    text,
+    tooltip,
+    className,
+    onTooltipVisibilityChange,
+    ...rest
+  } = props;
   return (
     <Tooltip
-      id={`${icon}-tooltip`}
+      id={`${prefixIcon}-tooltip`}
       placement="top"
       title={tooltip}
       trigger={['hover']}
@@ -71,8 +82,9 @@ const ActionButton = (props: ActionButtonProps) => {
         style={{ height: 30 }}
         {...rest}
       >
-        {icon}
+        {prefixIcon}
         {text && <span style={{ marginLeft: 5 }}>{text}</span>}
+        {suffixIcon}
       </div>
     </Tooltip>
   );
@@ -123,6 +135,14 @@ const ExploreActionButtons = (props: ExploreActionButtonsProps) => {
       })
     : null;
 
+  const doExportCSVPivoted = canDownloadCSV
+    ? exportChart.bind(this, {
+        formData: latestQueryFormData,
+        resultType: 'post_processed',
+        resultFormat: 'csv',
+      })
+    : null;
+
   const doExportJson = exportChart.bind(this, {
     formData: latestQueryFormData,
     resultType: 'results',
@@ -142,7 +162,7 @@ const ExploreActionButtons = (props: ExploreActionButtonsProps) => {
       {latestQueryFormData && (
         <>
           <ActionButton
-            icon={<Icons.Link iconSize="l" />}
+            prefixIcon={<Icons.Link iconSize="l" />}
             tooltip={copyTooltip}
             onClick={doCopyLink}
             data-test="short-link-button"
@@ -151,24 +171,47 @@ const ExploreActionButtons = (props: ExploreActionButtonsProps) => {
             }
           />
           <ActionButton
-            icon={<Icons.Email iconSize="l" />}
+            prefixIcon={<Icons.Email iconSize="l" />}
             tooltip={t('Share chart by email')}
             onClick={doShareEmail}
           />
           <EmbedCodeButton latestQueryFormData={latestQueryFormData} />
           <ActionButton
-            icon={<Icons.FileTextOutlined iconSize="m" />}
+            prefixIcon={<Icons.FileTextOutlined iconSize="m" />}
             text=".JSON"
             tooltip={t('Export to .JSON format')}
             onClick={doExportJson}
           />
-          <ActionButton
+          {VIZ_TYPES_PIVOTABLE.includes(latestQueryFormData.viz_type) ? (
-            icon={<Icons.FileExcelOutlined iconSize="m" />}
+            <ExportToCSVDropdown
-            text=".CSV"
+              exportCSVOriginal={doExportCSV}
-            tooltip={t('Export to .CSV format')}
+              exportCSVPivoted={doExportCSVPivoted}
-            onClick={doExportCSV}
+            >
-            className={exportToCSVClasses}
+              <ActionButton
-          />
+                prefixIcon={<Icons.FileExcelOutlined iconSize="m" />}
+                suffixIcon={
+                  <Icons.CaretDown
+                    iconSize="l"
+                    css={theme => `
+                    margin-left: ${theme.gridUnit}px;
+                    margin-right: ${-theme.gridUnit}px;
+                  `}
+                  />
+                }
+                text=".CSV"
+                tooltip={t('Export to .CSV format')}
+                className={exportToCSVClasses}
+              />
+            </ExportToCSVDropdown>
+          ) : (
+            <ActionButton
+              prefixIcon={<Icons.FileExcelOutlined iconSize="m" />}
+              text=".CSV"
+              tooltip={t('Export to .CSV format')}
+              onClick={doExportCSV}
+              className={exportToCSVClasses}
+            />
+          )}
         </>
       )}
       <ExploreAdditionalActionsMenu

superset-frontend/src/explore/components/ExportToCSVDropdown/ExportToCSVDropdown.test.tsx

@@ -0,0 +1,75 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import React from 'react';
+import userEvent from '@testing-library/user-event';
+import { render, screen } from 'spec/helpers/testing-library';
+import { ExportToCSVDropdown } from './index';
+
+const exportCSVOriginal = jest.fn();
+const exportCSVPivoted = jest.fn();
+
+test('Dropdown button with menu renders', () => {
+  render(
+    <ExportToCSVDropdown
+      exportCSVOriginal={exportCSVOriginal}
+      exportCSVPivoted={exportCSVPivoted}
+    >
+      <div>.CSV</div>
+    </ExportToCSVDropdown>,
+  );
+
+  expect(screen.getByText('.CSV')).toBeVisible();
+
+  userEvent.click(screen.getByText('.CSV'));
+  expect(screen.getByRole('menu')).toBeInTheDocument();
+  expect(screen.getByText('Original')).toBeInTheDocument();
+  expect(screen.getByText('Pivoted')).toBeInTheDocument();
+});
+
+test('Call export csv original on click', () => {
+  render(
+    <ExportToCSVDropdown
+      exportCSVOriginal={exportCSVOriginal}
+      exportCSVPivoted={exportCSVPivoted}
+    >
+      <div>.CSV</div>
+    </ExportToCSVDropdown>,
+  );
+
+  userEvent.click(screen.getByText('.CSV'));
+  userEvent.click(screen.getByText('Original'));
+
+  expect(exportCSVOriginal).toHaveBeenCalled();
+});
+
+test('Call export csv pivoted on click', () => {
+  render(
+    <ExportToCSVDropdown
+      exportCSVOriginal={exportCSVOriginal}
+      exportCSVPivoted={exportCSVPivoted}
+    >
+      <div>.CSV</div>
+    </ExportToCSVDropdown>,
+  );
+
+  userEvent.click(screen.getByText('.CSV'));
+  userEvent.click(screen.getByText('Pivoted'));
+
+  expect(exportCSVPivoted).toHaveBeenCalled();
+});

superset-frontend/src/explore/components/ExportToCSVDropdown/index.tsx

@@ -0,0 +1,90 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+import React, { ReactChild, useCallback } from 'react';
+import { t, styled } from '@superset-ui/core';
+import Icons from 'src/components/Icons';
+import { Dropdown, Menu } from 'src/common/components';
+
+enum MENU_KEYS {
+  EXPORT_ORIGINAL = 'export_original',
+  EXPORT_PIVOTED = 'export_pivoted',
+}
+
+interface ExportToCSVButtonProps {
+  exportCSVOriginal: () => void;
+  exportCSVPivoted: () => void;
+  children: ReactChild;
+}
+
+const MenuItemContent = styled.div`
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+
+  span[role='img'] {
+    font-size: ${({ theme }) => theme.typography.sizes.l}px;
+    margin-left: ${({ theme }) => theme.gridUnit * 4}px;
+  }
+`;
+
+export const ExportToCSVDropdown = ({
+  exportCSVOriginal,
+  exportCSVPivoted,
+  children,
+}: ExportToCSVButtonProps) => {
+  const handleMenuClick = useCallback(
+    ({ key }: { key: React.Key }) => {
+      switch (key) {
+        case MENU_KEYS.EXPORT_ORIGINAL:
+          exportCSVOriginal();
+          break;
+        case MENU_KEYS.EXPORT_PIVOTED:
+          exportCSVPivoted();
+          break;
+        default:
+          break;
+      }
+    },
+    [exportCSVPivoted, exportCSVOriginal],
+  );
+
+  return (
+    <Dropdown
+      trigger={['click']}
+      overlay={
+        <Menu onClick={handleMenuClick} selectable={false}>
+          <Menu.Item key={MENU_KEYS.EXPORT_ORIGINAL}>
+            <MenuItemContent>
+              {t('Original')}
+              <Icons.Download />
+            </MenuItemContent>
+          </Menu.Item>
+          <Menu.Item key={MENU_KEYS.EXPORT_PIVOTED}>
+            <MenuItemContent>
+              {t('Pivoted')}
+              <Icons.Download />
+            </MenuItemContent>
+          </Menu.Item>
+        </Menu>
+      }
+    >
+      {children}
+    </Dropdown>
+  );
+};

superset/charts/data/api.py

@@ -242,7 +242,8 @@ class ChartDataRestApi(ChartRestApi):
         ):
             return self._run_async(json_body, command)
 
-        return self._get_data_response(command)
+        form_data = json_body.get("form_data")
+        return self._get_data_response(command, form_data=form_data)
 
     @expose("/data/<cache_key>", methods=["GET"])
     @protect()

superset/charts/schemas.py

@@ -1160,6 +1160,8 @@ class ChartDataQueryContextSchema(Schema):
     result_type = EnumField(ChartDataResultType, by_value=True)
     result_format = EnumField(ChartDataResultFormat, by_value=True)
 
+    form_data = fields.Raw(allow_none=True, required=False)
+
     # pylint: disable=unused-argument
     @post_load
     def make_query_context(self, data: Dict[str, Any], **kwargs: Any) -> QueryContext:

superset/common/query_context.py

@@ -47,6 +47,7 @@ class QueryContext:
 
     datasource: BaseDatasource
     queries: List[QueryObject]
+    form_data: Optional[Dict[str, Any]]
     result_type: ChartDataResultType
     result_format: ChartDataResultFormat
     force: bool
@@ -63,6 +64,7 @@ class QueryContext:
         *,
         datasource: BaseDatasource,
         queries: List[QueryObject],
+        form_data: Optional[Dict[str, Any]],
         result_type: ChartDataResultType,
         result_format: ChartDataResultFormat,
         force: bool = False,
@@ -73,6 +75,7 @@ class QueryContext:
         self.result_type = result_type
         self.result_format = result_format
         self.queries = queries
+        self.form_data = form_data
         self.force = force
         self.custom_cache_timeout = custom_cache_timeout
         self.cache_values = cache_values

superset/common/query_context_factory.py

@@ -46,6 +46,7 @@ class QueryContextFactory:  # pylint: disable=too-few-public-methods
         *,
         datasource: DatasourceDict,
         queries: List[Dict[str, Any]],
+        form_data: Optional[Dict[str, Any]] = None,
         result_type: Optional[ChartDataResultType] = None,
         result_format: Optional[ChartDataResultFormat] = None,
         force: bool = False,
@@ -69,6 +70,7 @@ class QueryContextFactory:  # pylint: disable=too-few-public-methods
         return QueryContext(
             datasource=datasource_model_instance,
             queries=queries_,
+            form_data=form_data,
             result_type=result_type,
             result_format=result_format,
             force=force,

tests/common/query_context_generator.py

@@ -16,7 +16,7 @@
 # under the License.
 import copy
 import dataclasses
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Optional
 
 from superset.common.chart_data import ChartDataResultType
 from superset.utils.core import AnnotationType, DTTM_ALIAS, TimeRangeEndpoint
@@ -242,7 +242,9 @@ class QueryContextGenerator:
         add_time_offsets: bool = False,
         table_id=1,
         table_type="table",
+        form_data: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
+        form_data = form_data or {}
         table_name = query_name.split(":")[0]
         table = self.get_table(table_name, table_id, table_type)
         return {
@@ -253,7 +255,8 @@ class QueryContextGenerator:
                 )
             ],
             "result_type": ChartDataResultType.FULL,
+            "form_data": form_data,
         }
 
-    def get_table(self, name, id, type):
+    def get_table(self, name, id_, type_):
-        return Table(id, type, name)
+        return Table(id_, type_, name)

tests/integration_tests/fixtures/query_context.py

@@ -14,14 +14,14 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from tests.common.query_context_generator import QueryContextGenerator
 from tests.integration_tests.base_tests import SupersetTestCase
 
 
 class QueryContextGeneratorInteg(QueryContextGenerator):
-    def get_table(self, name, id, type):
+    def get_table(self, name, id_, type_):
         return SupersetTestCase.get_table(name=name)
 
 
@@ -29,6 +29,7 @@ def get_query_context(
     query_name: str,
     add_postprocessing_operations: bool = False,
     add_time_offsets: bool = False,
+    form_data: Optional[Dict[str, Any]] = None,
 ) -> Dict[str, Any]:
     """
     Create a request payload for retrieving a QueryContext object via the
@@ -40,8 +41,12 @@ def get_query_context(
     :param datasource_type: type of datasource to query.
     :param add_postprocessing_operations: Add post-processing operations to QueryObject
     :param add_time_offsets: Add time offsets to QueryObject(advanced analytics)
+    :param form_data: chart metadata
     :return: Request payload
     """
     return QueryContextGeneratorInteg().generate(
-        query_name, add_postprocessing_operations, add_time_offsets
+        query_name=query_name,
+        add_postprocessing_operations=add_postprocessing_operations,
+        add_time_offsets=add_time_offsets,
+        form_data=form_data,
     )