feat(SqlLab): Change Save Dataset Button to Split Save Query Button IV (#20852)

* Moving entire split save btn PR

* Addressed review comments

* Remove arbitrary div from ErrorBoundary in Chart

* Added accidentally removed comment

* Fix act errors in SaveQuery tests

* Fix SaveDatasetActionButton test

* SaveDatasetModal test almost working

* SaveDatasetModal tests all passing

* Clean SaveDatasetModal test

* Fix create chart button and SaveDatasetModal text in SQL Lab

* Fix untitled dataset name on SaveDatasetModal in explore

* Fix styling on split save button

superset-frontend/src/SqlLab/components/SaveQuery/SaveQuery.test.jsx

@@ -17,60 +17,87 @@
  * under the License.
  */
 import React from 'react';
-import { shallow } from 'enzyme';
-import * as sinon from 'sinon';
+import { render, screen, waitFor } from 'spec/helpers/testing-library';
+import userEvent from '@testing-library/user-event';
 import SaveQuery from 'src/SqlLab/components/SaveQuery';
-import Modal from 'src/components/Modal';
-import Button from 'src/components/Button';
-import { FormItem } from 'src/components/Form';
+import { databases } from 'src/SqlLab/fixtures';
+
+const mockedProps = {
+  query: {
+    dbId: 1,
+    schema: 'main',
+    sql: 'SELECT * FROM t',
+  },
+  defaultLabel: 'untitled',
+  animation: false,
+  database: databases.result[0],
+  onUpdate: () => {},
+  onSave: () => {},
+};
+
+const splitSaveBtnProps = {
+  ...mockedProps,
+  database: {
+    ...mockedProps.database,
+    allows_virtual_table_explore: true,
+  },
+};
 
 describe('SavedQuery', () => {
-  const mockedProps = {
-    query: {
-      dbId: 1,
-      schema: 'main',
-      sql: 'SELECT * FROM t',
-    },
-    defaultLabel: 'untitled',
-    animation: false,
-  };
-  it('is valid', () => {
-    expect(React.isValidElement(<SaveQuery />)).toBe(true);
-  });
-  it('is valid with props', () => {
-    expect(React.isValidElement(<SaveQuery {...mockedProps} />)).toBe(true);
-  });
-  it('has a Modal', () => {
-    const wrapper = shallow(<SaveQuery {...mockedProps} />);
-    expect(wrapper.find(Modal)).toExist();
-  });
-  // TODO: eschutho convert test to RTL
-  // eslint-disable-next-line jest/no-disabled-tests
-  it.skip('has a cancel button', () => {
-    const wrapper = shallow(<SaveQuery {...mockedProps} />);
-    const modal = wrapper.find(Modal);
+  it('renders a non-split save button when allows_virtual_table_explore is not enabled', () => {
+    render(<SaveQuery {...mockedProps} />, { useRedux: true });
 
-    expect(modal.find('[data-test="cancel-query"]')).toHaveLength(1);
-  });
-  it('has 2 FormItem', () => {
-    const wrapper = shallow(<SaveQuery {...mockedProps} />);
-    const modal = wrapper.find(Modal);
+    const saveBtn = screen.getByRole('button', { name: /save/i });
 
-    expect(modal.find(FormItem)).toHaveLength(2);
+    expect(saveBtn).toBeVisible();
   });
-  // eslint-disable-next-line jest/no-disabled-tests
-  it.skip('has a save button if this is a new query', () => {
-    const saveSpy = sinon.spy();
-    const wrapper = shallow(<SaveQuery {...mockedProps} onSave={saveSpy} />);
-    const modal = wrapper.find(Modal);
 
-    expect(modal.find(Button)).toHaveLength(2);
-    modal.find(Button).at(0).simulate('click');
-    expect(saveSpy.calledOnce).toBe(true);
+  it('renders a save query modal when user clicks save button', () => {
+    render(<SaveQuery {...mockedProps} />, { useRedux: true });
+
+    const saveBtn = screen.getByRole('button', { name: /save/i });
+    userEvent.click(saveBtn);
+
+    const saveQueryModalHeader = screen.getByRole('heading', {
+      name: /save query/i,
+    });
+
+    expect(saveQueryModalHeader).toBeVisible();
   });
-  // eslint-disable-next-line jest/no-disabled-tests
-  it.skip('has an update button if this is an existing query', () => {
-    const updateSpy = sinon.spy();
+
+  it('renders the save query modal UI', () => {
+    render(<SaveQuery {...mockedProps} />, { useRedux: true });
+
+    const saveBtn = screen.getByRole('button', { name: /save/i });
+    userEvent.click(saveBtn);
+
+    const closeBtn = screen.getByRole('button', { name: /close/i });
+    const saveQueryModalHeader = screen.getByRole('heading', {
+      name: /save query/i,
+    });
+    const nameLabel = screen.getByText(/name/i);
+    const descriptionLabel = screen.getByText(/description/i);
+    const textBoxes = screen.getAllByRole('textbox');
+    const nameTextbox = textBoxes[0];
+    const descriptionTextbox = textBoxes[1];
+    // There are now two save buttons, the initial save button and the modal save button
+    const saveBtns = screen.getAllByRole('button', { name: /save/i });
+    const cancelBtn = screen.getByRole('button', { name: /cancel/i });
+
+    expect(closeBtn).toBeVisible();
+    expect(saveQueryModalHeader).toBeVisible();
+    expect(nameLabel).toBeVisible();
+    expect(descriptionLabel).toBeVisible();
+    expect(textBoxes.length).toBe(2);
+    expect(nameTextbox).toBeVisible();
+    expect(descriptionTextbox).toBeVisible();
+    expect(saveBtns.length).toBe(2);
+    expect(saveBtns[0]).toBeVisible();
+    expect(saveBtns[1]).toBeVisible();
+    expect(cancelBtn).toBeVisible();
+  });
+
+  it('renders a "save as new" and "update" button if query already exists', () => {
     const props = {
       ...mockedProps,
       query: {
@@ -78,11 +105,81 @@ describe('SavedQuery', () => {
         remoteId: '42',
       },
     };
-    const wrapper = shallow(<SaveQuery {...props} onUpdate={updateSpy} />);
-    const modal = wrapper.find(Modal);
+    render(<SaveQuery {...props} />, { useRedux: true });
 
-    expect(modal.find(Button)).toHaveLength(3);
-    modal.find(Button).at(0).simulate('click');
-    expect(updateSpy.calledOnce).toBe(true);
+    const saveBtn = screen.getByRole('button', { name: /save/i });
+    userEvent.click(saveBtn);
+
+    const saveAsNewBtn = screen.getByRole('button', { name: /save as new/i });
+    const updateBtn = screen.getByRole('button', { name: /update/i });
+
+    expect(saveAsNewBtn).toBeVisible();
+    expect(updateBtn).toBeVisible();
+  });
+
+  it('renders a split save button when allows_virtual_table_explore is enabled', async () => {
+    render(<SaveQuery {...splitSaveBtnProps} />, { useRedux: true });
+
+    await waitFor(() => {
+      const saveBtn = screen.getByRole('button', { name: /save/i });
+      const caretBtn = screen.getByRole('button', { name: /caret-down/i });
+
+      expect(saveBtn).toBeVisible();
+      expect(caretBtn).toBeVisible();
+    });
+  });
+
+  it('renders a save dataset modal when user clicks "save dataset" menu item', async () => {
+    render(<SaveQuery {...splitSaveBtnProps} />, { useRedux: true });
+
+    await waitFor(() => {
+      const caretBtn = screen.getByRole('button', { name: /caret-down/i });
+      userEvent.click(caretBtn);
+
+      const saveDatasetMenuItem = screen.getByText(/save dataset/i);
+      userEvent.click(saveDatasetMenuItem);
+    });
+
+    const saveDatasetHeader = screen.getByText(/save or overwrite dataset/i);
+
+    expect(saveDatasetHeader).toBeVisible();
+  });
+
+  it('renders the save dataset modal UI', async () => {
+    render(<SaveQuery {...splitSaveBtnProps} />, { useRedux: true });
+
+    await waitFor(() => {
+      const caretBtn = screen.getByRole('button', { name: /caret-down/i });
+      userEvent.click(caretBtn);
+
+      const saveDatasetMenuItem = screen.getByText(/save dataset/i);
+      userEvent.click(saveDatasetMenuItem);
+    });
+
+    const closeBtn = screen.getByRole('button', { name: /close/i });
+    const saveDatasetHeader = screen.getByText(/save or overwrite dataset/i);
+    const saveRadio = screen.getByRole('radio', {
+      name: /save as new untitled dataset/i,
+    });
+    const saveLabel = screen.getByText(/save as new/i);
+    const saveTextbox = screen.getByRole('textbox');
+    const overwriteRadio = screen.getByRole('radio', {
+      name: /overwrite existing/i,
+    });
+    const overwriteLabel = screen.getByText(/overwrite existing/i);
+    const overwriteCombobox = screen.getByRole('combobox');
+    const overwritePlaceholderText = screen.getByText(
+      /select or type dataset name/i,
+    );
+
+    expect(saveDatasetHeader).toBeVisible();
+    expect(closeBtn).toBeVisible();
+    expect(saveRadio).toBeVisible();
+    expect(saveLabel).toBeVisible();
+    expect(saveTextbox).toBeVisible();
+    expect(overwriteRadio).toBeVisible();
+    expect(overwriteLabel).toBeVisible();
+    expect(overwriteCombobox).toBeVisible();
+    expect(overwritePlaceholderText).toBeVisible();
   });
 });

superset-frontend/src/SqlLab/components/SaveQuery/index.tsx

@@ -21,21 +21,11 @@ import { Row, Col } from 'src/components';
 import { Input, TextArea } from 'src/components/Input';
 import { t, styled } from '@superset-ui/core';
 import Button from 'src/components/Button';
+import { Menu } from 'src/components/Menu';
 import { Form, FormItem } from 'src/components/Form';
 import Modal from 'src/components/Modal';
-import Icons from 'src/components/Icons';
-
-const Styles = styled.span`
-  span[role='img'] {
-    display: flex;
-    margin: 0;
-    color: ${({ theme }) => theme.colors.grayscale.base};
-    svg {
-      vertical-align: -${({ theme }) => theme.gridUnit * 1.25}px;
-      margin: 0;
-    }
-  }
-`;
+import SaveDatasetActionButton from 'src/SqlLab/components/SaveDatasetActionButton';
+import { SaveDatasetModal } from 'src/SqlLab/components/SaveDatasetModal';
 
 interface SaveQueryProps {
   query: any;
@@ -43,6 +33,7 @@ interface SaveQueryProps {
   onSave: (arg0: QueryPayload) => void;
   onUpdate: (arg0: QueryPayload) => void;
   saveQueryWarning: string | null;
+  database: Record<string, any>;
 }
 
 type QueryPayload = {
@@ -68,37 +59,57 @@ type QueryPayload = {
     type: string;
     value: string;
   }>;
-  title: string;
+  name: string;
 };
 
+const Styles = styled.span`
+  span[role='img'] {
+    display: flex;
+    margin: 0;
+    color: ${({ theme }) => theme.colors.grayscale.base};
+    svg {
+      vertical-align: -${({ theme }) => theme.gridUnit * 1.25}px;
+      margin: 0;
+    }
+  }
+`;
+
 export default function SaveQuery({
   query,
   defaultLabel = t('Undefined'),
   onSave = () => {},
   onUpdate,
   saveQueryWarning = null,
+  database,
 }: SaveQueryProps) {
   const [description, setDescription] = useState<string>(
     query.description || '',
   );
   const [label, setLabel] = useState<string>(defaultLabel);
   const [showSave, setShowSave] = useState<boolean>(false);
+  const [showSaveDatasetModal, setShowSaveDatasetModal] = useState(false);
   const isSaved = !!query.remoteId;
+  const canExploreDatabase = !!database?.allows_virtual_table_explore;
+
+  const overlayMenu = (
+    <Menu>
+      <Menu.Item onClick={() => setShowSaveDatasetModal(true)}>
+        {t('Save dataset')}
+      </Menu.Item>
+    </Menu>
+  );
 
   const queryPayload = () => ({
     ...query,
-    title: label,
+    name: label,
     description,
   });
 
   useEffect(() => {
-    if (!isSaved) {
-      setLabel(defaultLabel);
-    }
+    if (!isSaved) setLabel(defaultLabel);
   }, [defaultLabel]);
-  const close = () => {
-    setShowSave(false);
-  };
+
+  const close = () => setShowSave(false);
 
   const onSaveWrapper = () => {
     onSave(queryPayload());
@@ -118,10 +129,6 @@ export default function SaveQuery({
     setDescription(e.target.value);
   };
 
-  const toggleSave = () => {
-    setShowSave(!showSave);
-  };
-
   const renderModalBody = () => (
     <Form layout="vertical">
       <Row>
@@ -161,10 +168,17 @@ export default function SaveQuery({
 
   return (
     <Styles className="SaveQuery">
-      <Button buttonSize="small" onClick={toggleSave}>
-        <Icons.Save iconSize="xl" />
-        {isSaved ? t('Save') : t('Save as')}
-      </Button>
+      <SaveDatasetActionButton
+        setShowSave={setShowSave}
+        overlayMenu={canExploreDatabase ? overlayMenu : null}
+      />
+      <SaveDatasetModal
+        visible={showSaveDatasetModal}
+        onHide={() => setShowSaveDatasetModal(false)}
+        buttonTextOnSave={t('Save & Explore')}
+        buttonTextOnOverwrite={t('Overwrite & Explore')}
+        datasource={query}
+      />
       <Modal
         className="save-query-modal"
         onHandledPrimaryAction={onSaveWrapper}