feat(explore): Dataset Panel Options when Source = Query II (#20299)

* Created/tested query dataset dropdown

* Add isValidDatasourceType to @superset-ui/core and hide query dropdown

* Removed isValidDatasourceType check

* Remove the rest of isValidDatasourceType check

superset-frontend/src/explore/components/controls/DatasourceControl/DatasourceControl.test.tsx

@@ -20,7 +20,7 @@
 import React from 'react';
 import { render, screen, act } from 'spec/helpers/testing-library';
 import userEvent from '@testing-library/user-event';
-import { SupersetClient } from '@superset-ui/core';
+import { SupersetClient, DatasourceType } from '@superset-ui/core';
 import * as Utils from 'src/explore/exploreUtils';
 import DatasourceControl from '.';
 
@@ -158,3 +158,52 @@ test('Click on View in SQL Lab', async () => {
 
   expect(postFormSpy).toBeCalledTimes(1);
 });
+
+test('Should open a different menu when datasource=query', () => {
+  const props = createProps();
+  const queryProps = {
+    ...props,
+    datasource: {
+      ...props.datasource,
+      type: DatasourceType.Query,
+    },
+  };
+  render(<DatasourceControl {...queryProps} />);
+
+  expect(screen.queryByText('Query preview')).not.toBeInTheDocument();
+  expect(screen.queryByText('View in SQL Lab')).not.toBeInTheDocument();
+  expect(screen.queryByText('Save as dataset')).not.toBeInTheDocument();
+
+  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+
+  expect(screen.getByText('Query preview')).toBeInTheDocument();
+  expect(screen.getByText('View in SQL Lab')).toBeInTheDocument();
+  expect(screen.getByText('Save as dataset')).toBeInTheDocument();
+});
+
+test('Click on Save as dataset', () => {
+  const props = createProps();
+  const queryProps = {
+    ...props,
+    datasource: {
+      ...props.datasource,
+      type: DatasourceType.Query,
+    },
+  };
+
+  render(<DatasourceControl {...queryProps} />, { useRedux: true });
+  userEvent.click(screen.getByTestId('datasource-menu-trigger'));
+  userEvent.click(screen.getByText('Save as dataset'));
+
+  // Renders a save dataset modal
+  const saveRadioBtn = screen.getByRole('radio', {
+    name: /save as new undefined/i,
+  });
+  const overwriteRadioBtn = screen.getByRole('radio', {
+    name: /overwrite existing select or type dataset name/i,
+  });
+  expect(saveRadioBtn).toBeVisible();
+  expect(overwriteRadioBtn).toBeVisible();
+  expect(screen.getByRole('button', { name: /save/i })).toBeVisible();
+  expect(screen.getByRole('button', { name: /close/i })).toBeVisible();
+});

superset-frontend/src/explore/components/controls/DatasourceControl/index.jsx

@@ -19,7 +19,7 @@
  */
 import React from 'react';
 import PropTypes from 'prop-types';
-import { t, styled, withTheme } from '@superset-ui/core';
+import { t, styled, withTheme, DatasourceType } from '@superset-ui/core';
 import { getUrlParam } from 'src/utils/urlUtils';
 
 import { AntdDropdown } from 'src/components';
@@ -30,6 +30,7 @@ import {
   ChangeDatasourceModal,
   DatasourceModal,
 } from 'src/components/Datasource';
+import { SaveDatasetModal } from 'src/SqlLab/components/SaveDatasetModal';
 import { postForm } from 'src/explore/exploreUtils';
 import Button from 'src/components/Button';
 import ErrorAlert from 'src/components/ErrorMessage/ErrorAlert';
@@ -112,6 +113,8 @@ const Styles = styled.div`
 const CHANGE_DATASET = 'change_dataset';
 const VIEW_IN_SQL_LAB = 'view_in_sql_lab';
 const EDIT_DATASET = 'edit_dataset';
+const QUERY_PREVIEW = 'query_preview';
+const SAVE_AS_DATASET = 'save_as_dataset';
 
 class DatasourceControl extends React.PureComponent {
   constructor(props) {
@@ -126,6 +129,7 @@ class DatasourceControl extends React.PureComponent {
     this.toggleEditDatasourceModal = this.toggleEditDatasourceModal.bind(this);
     this.toggleShowDatasource = this.toggleShowDatasource.bind(this);
     this.handleMenuItemClick = this.handleMenuItemClick.bind(this);
+    this.toggleSaveDatasetModal = this.toggleSaveDatasetModal.bind(this);
   }
 
   onDatasourceSave(datasource) {
@@ -166,25 +170,51 @@ class DatasourceControl extends React.PureComponent {
     }));
   }
 
+  toggleSaveDatasetModal() {
+    this.setState(({ showSaveDatasetModal }) => ({
+      showSaveDatasetModal: !showSaveDatasetModal,
+    }));
+  }
+
   handleMenuItemClick({ key }) {
-    if (key === CHANGE_DATASET) {
-      this.toggleChangeDatasourceModal();
-    }
-    if (key === EDIT_DATASET) {
-      this.toggleEditDatasourceModal();
-    }
-    if (key === VIEW_IN_SQL_LAB) {
-      const { datasource } = this.props;
-      const payload = {
-        datasourceKey: `${datasource.id}__${datasource.type}`,
-        sql: datasource.sql,
-      };
-      postForm('/superset/sqllab/', payload);
+    switch (key) {
+      case CHANGE_DATASET:
+        this.toggleChangeDatasourceModal();
+        break;
+
+      case EDIT_DATASET:
+        this.toggleEditDatasourceModal();
+        break;
+
+      case VIEW_IN_SQL_LAB:
+        {
+          const { datasource } = this.props;
+          const payload = {
+            datasourceKey: `${datasource.id}__${datasource.type}`,
+            sql: datasource.sql,
+          };
+          postForm('/superset/sqllab/', payload);
+        }
+        break;
+
+      case QUERY_PREVIEW:
+        break;
+
+      case SAVE_AS_DATASET:
+        this.toggleSaveDatasetModal();
+        break;
+
+      default:
+        break;
     }
   }
 
   render() {
-    const { showChangeDatasourceModal, showEditDatasourceModal } = this.state;
+    const {
+      showChangeDatasourceModal,
+      showEditDatasourceModal,
+      showSaveDatasetModal,
+    } = this.state;
     const { datasource, onChange, theme } = this.props;
     const isMissingDatasource = datasource.id == null;
     let isMissingParams = false;
@@ -205,7 +235,7 @@ class DatasourceControl extends React.PureComponent {
 
     const editText = t('Edit dataset');
 
-    const datasourceMenu = (
+    const defaultDatasourceMenu = (
       <Menu onClick={this.handleMenuItemClick}>
         {this.props.isEditable && (
           <Menu.Item
@@ -233,6 +263,14 @@ class DatasourceControl extends React.PureComponent {
       </Menu>
     );
 
+    const queryDatasourceMenu = (
+      <Menu onClick={this.handleMenuItemClick}>
+        <Menu.Item key={QUERY_PREVIEW}>{t('Query preview')}</Menu.Item>
+        <Menu.Item key={VIEW_IN_SQL_LAB}>{t('View in SQL Lab')}</Menu.Item>
+        <Menu.Item key={SAVE_AS_DATASET}>{t('Save as dataset')}</Menu.Item>
+      </Menu>
+    );
+
     const { health_check_message: healthCheckMessage } = datasource;
 
     let extra = {};
@@ -265,7 +303,11 @@ class DatasourceControl extends React.PureComponent {
             <WarningIconWithTooltip warningMarkdown={extra.warning_markdown} />
           )}
           <AntdDropdown
-            overlay={datasourceMenu}
+            overlay={
+              datasource.type === DatasourceType.Query
+                ? queryDatasourceMenu
+                : defaultDatasourceMenu
+            }
             trigger={['click']}
             data-test="datasource-menu"
           >
@@ -340,6 +382,18 @@ class DatasourceControl extends React.PureComponent {
             onChange={onChange}
           />
         )}
+        {showSaveDatasetModal && (
+          <SaveDatasetModal
+            visible={showSaveDatasetModal}
+            onHide={this.toggleSaveDatasetModal}
+            buttonTextOnSave={t('Save & Explore')}
+            buttonTextOnOverwrite={t('Overwrite & Explore')}
+            modalDescription={t(
+              'Save this query as a virtual dataset to continue exploring',
+            )}
+            datasource={datasource}
+          />
+        )}
       </Styles>
     );
   }