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

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;
-  text?: string;
+  prefixIcon: React.ReactElement;
+  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 } =
-    props;
+  const {
+    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
-            icon={<Icons.FileExcelOutlined iconSize="m" />}
-            text=".CSV"
-            tooltip={t('Export to .CSV format')}
-            onClick={doExportCSV}
-            className={exportToCSVClasses}
-          />
+          {VIZ_TYPES_PIVOTABLE.includes(latestQueryFormData.viz_type) ? (
+            <ExportToCSVDropdown
+              exportCSVOriginal={doExportCSV}
+              exportCSVPivoted={doExportCSVPivoted}
+            >
+              <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 @@ def data(self) -> Response:
         ):
             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()