fix: save dataset and repopulate state (#20965)

* save dataset and repopulate state

* disable dashboard button if dataset is missing

* fix error message

* fix tests

superset-frontend/packages/superset-ui-core/src/api/types/core.ts

@@ -0,0 +1,31 @@
+/**
+ * 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.
+ */
+
+// /superset/sqllab_viz
+interface SqlLabPostRequest {
+  data: {
+    schema: string;
+    sql: string;
+    dbId: number;
+    templateParams?: string | undefined;
+    datasourceName: string;
+    metrics?: string[];
+    columns?: string[];
+  };
+}

superset-frontend/src/explore/actions/datasourcesActions.test.ts

@@ -17,10 +17,14 @@
  * under the License.
  */
 import { DatasourceType } from '@superset-ui/core';
+import fetchMock from 'fetch-mock';
 import {
   setDatasource,
   changeDatasource,
+  saveDataset,
 } from 'src/explore/actions/datasourcesActions';
+import sinon from 'sinon';
+import * as ClientError from 'src/utils/getClientErrorObject';
 import datasourcesReducer from '../reducers/datasourcesReducer';
 import { updateFormDataByDatasource } from './exploreActions';
 
@@ -51,10 +55,21 @@ const NEW_DATASOURCE = {
   description: null,
 };
 
+const SAVE_DATASET_POST_ARGS = {
+  schema: 'foo',
+  sql: 'select * from bar',
+  database: { id: 1 },
+  templateParams: undefined,
+  datasourceName: 'new dataset',
+  columns: [],
+};
+
 const defaultDatasourcesReducerState = {
   [CURRENT_DATASOURCE.uid]: CURRENT_DATASOURCE,
 };
 
+const saveDatasetEndpoint = `glob:*/superset/sqllab_viz/`;
+
 test('sets new datasource', () => {
   const newState = datasourcesReducer(
     defaultDatasourcesReducerState,
@@ -83,3 +98,42 @@ test('change datasource action', () => {
     updateFormDataByDatasource(CURRENT_DATASOURCE, NEW_DATASOURCE),
   );
 });
+
+test('saveDataset handles success', async () => {
+  const datasource = { id: 1 };
+  const saveDatasetResponse = {
+    data: datasource,
+  };
+  fetchMock.reset();
+  fetchMock.post(saveDatasetEndpoint, saveDatasetResponse);
+  const dispatch = sinon.spy();
+  const getState = sinon.spy(() => ({ explore: { datasource } }));
+  const dataset = await saveDataset(SAVE_DATASET_POST_ARGS)(dispatch);
+
+  expect(fetchMock.calls(saveDatasetEndpoint)).toHaveLength(1);
+  expect(dispatch.callCount).toBe(1);
+  const thunk = dispatch.getCall(0).args[0];
+  thunk(dispatch, getState);
+  expect(dispatch.getCall(1).args[0].type).toEqual('SET_DATASOURCE');
+
+  expect(dataset).toEqual(datasource);
+});
+
+test('updateSlice with add to existing dashboard handles failure', async () => {
+  fetchMock.reset();
+  const sampleError = new Error('sampleError');
+  fetchMock.post(saveDatasetEndpoint, { throws: sampleError });
+  const dispatch = sinon.spy();
+  const errorSpy = jest.spyOn(ClientError, 'getClientErrorObject');
+
+  let caughtError;
+  try {
+    await saveDataset(SAVE_DATASET_POST_ARGS)(dispatch);
+  } catch (error) {
+    caughtError = error;
+  }
+
+  expect(caughtError).toEqual(sampleError);
+  expect(fetchMock.calls(saveDatasetEndpoint)).toHaveLength(4);
+  expect(errorSpy).toHaveBeenCalledWith(sampleError);
+});

superset-frontend/src/explore/actions/datasourcesActions.ts

@@ -17,8 +17,12 @@
  * under the License.
  */
 
-import { Dispatch } from 'redux';
+import { Dispatch, AnyAction } from 'redux';
+import { ThunkDispatch } from 'redux-thunk';
 import { Dataset } from '@superset-ui/chart-controls';
+import { SupersetClient } from '@superset-ui/core';
+import { addDangerToast } from 'src/components/MessageToasts/actions';
+import { getClientErrorObject } from 'src/utils/getClientErrorObject';
 import { updateFormDataByDatasource } from './exploreActions';
 import { ExplorePageState } from '../types';
 
@@ -31,6 +35,45 @@ export function setDatasource(datasource: Dataset) {
   return { type: SET_DATASOURCE, datasource };
 }
 
+export function saveDataset({
+  schema,
+  sql,
+  database,
+  templateParams,
+  datasourceName,
+  columns,
+}: Omit<SqlLabPostRequest['data'], 'dbId'> & { database: { id: number } }) {
+  return async function (dispatch: ThunkDispatch<any, undefined, AnyAction>) {
+    // Create a dataset object
+    try {
+      const {
+        json: { data },
+      } = await SupersetClient.post({
+        endpoint: '/superset/sqllab_viz/',
+        postPayload: {
+          data: {
+            schema,
+            sql,
+            dbId: database?.id,
+            templateParams,
+            datasourceName,
+            metrics: [],
+            columns,
+          },
+        },
+      });
+      // Update form_data to point to new dataset
+      dispatch(changeDatasource(data));
+      return data;
+    } catch (error) {
+      getClientErrorObject(error).then(e => {
+        dispatch(addDangerToast(e.error));
+      });
+      throw error;
+    }
+  };
+}
+
 export function changeDatasource(newDatasource: Dataset) {
   return function (dispatch: Dispatch, getState: () => ExplorePageState) {
     const {
@@ -44,6 +87,7 @@ export function changeDatasource(newDatasource: Dataset) {
 export const datasourcesActions = {
   setDatasource,
   changeDatasource,
+  saveDataset,
 };
 
 export type AnyDatasourcesAction = SetDatasource;

superset-frontend/src/explore/actions/exploreActions.ts

@@ -104,6 +104,11 @@ export function setExploreControls(formData: QueryFormData) {
   return { type: SET_EXPLORE_CONTROLS, formData };
 }
 
+export const SET_FORM_DATA = 'UPDATE_FORM_DATA';
+export function setFormData(formData: QueryFormData) {
+  return { type: SET_FORM_DATA, formData };
+}
+
 export const UPDATE_CHART_TITLE = 'UPDATE_CHART_TITLE';
 export function updateChartTitle(sliceName: string) {
   return { type: UPDATE_CHART_TITLE, sliceName };

superset-frontend/src/explore/actions/hydrateExplore.test.ts

@@ -82,7 +82,77 @@ test('creates hydrate action from initial data', () => {
           controls: expect.any(Object),
           form_data: exploreInitialData.form_data,
           slice: exploreInitialData.slice,
-          controlsTransferred: [],
+          standalone: null,
+          force: null,
+        },
+      },
+    }),
+  );
+});
+
+test('creates hydrate action with existing state', () => {
+  const dispatch = jest.fn();
+  const getState = jest.fn(() => ({
+    user: {},
+    charts: {},
+    datasources: {},
+    common: {},
+    explore: { controlsTransferred: ['all_columns'] },
+  }));
+  // ignore type check - we dont need exact explore state for this test
+  // @ts-ignore
+  hydrateExplore(exploreInitialData)(dispatch, getState);
+  expect(dispatch).toHaveBeenCalledWith(
+    expect.objectContaining({
+      type: HYDRATE_EXPLORE,
+      data: {
+        charts: {
+          371: {
+            id: 371,
+            chartAlert: null,
+            chartStatus: null,
+            chartStackTrace: null,
+            chartUpdateEndTime: null,
+            chartUpdateStartTime: 0,
+            latestQueryFormData: {
+              cache_timeout: undefined,
+              datasource: '8__table',
+              slice_id: 371,
+              url_params: undefined,
+              viz_type: 'table',
+            },
+            sliceFormData: {
+              cache_timeout: undefined,
+              datasource: '8__table',
+              slice_id: 371,
+              url_params: undefined,
+              viz_type: 'table',
+            },
+            queryController: null,
+            queriesResponse: null,
+            triggerQuery: false,
+            lastRendered: 0,
+          },
+        },
+        datasources: {
+          '8__table': exploreInitialData.dataset,
+        },
+        saveModal: {
+          dashboards: [],
+          saveModalAlert: null,
+        },
+        explore: {
+          can_add: false,
+          can_download: false,
+          can_overwrite: false,
+          isDatasourceMetaLoading: false,
+          isStarred: false,
+          triggerRender: false,
+          datasource: exploreInitialData.dataset,
+          controls: expect.any(Object),
+          controlsTransferred: ['all_columns'],
+          form_data: exploreInitialData.form_data,
+          slice: exploreInitialData.slice,
           standalone: null,
           force: null,
         },

superset-frontend/src/explore/actions/hydrateExplore.ts

@@ -49,7 +49,8 @@ export const HYDRATE_EXPLORE = 'HYDRATE_EXPLORE';
 export const hydrateExplore =
   ({ form_data, slice, dataset }: ExplorePageInitialData) =>
   (dispatch: Dispatch, getState: () => ExplorePageState) => {
-    const { user, datasources, charts, sliceEntities, common } = getState();
+    const { user, datasources, charts, sliceEntities, common, explore } =
+      getState();
 
     const sliceId = getUrlParam(URL_PARAMS.sliceId);
     const dashboardId = getUrlParam(URL_PARAMS.dashboardId);
@@ -119,7 +120,7 @@ export const hydrateExplore =
       controls: initialControls,
       form_data: initialFormData,
       slice: initialSlice,
-      controlsTransferred: [],
+      controlsTransferred: explore.controlsTransferred,
       standalone: getUrlParam(URL_PARAMS.standalone),
       force: getUrlParam(URL_PARAMS.force),
       sliceDashboards: initialFormData.dashboards,

superset-frontend/src/explore/actions/saveModalActions.js

@@ -135,8 +135,13 @@ const addToasts = (isNewSlice, sliceName, addedToDashboard) => {
 
 //  Update existing slice
 export const updateSlice =
-  ({ slice_id: sliceId, owners }, sliceName, formData, addedToDashboard) =>
-  async dispatch => {
+  ({ slice_id: sliceId, owners }, sliceName, addedToDashboard) =>
+  async (dispatch, getState) => {
+    const {
+      explore: {
+        form_data: { url_params: _, ...formData },
+      },
+    } = getState();
     try {
       const response = await SupersetClient.put({
         endpoint: `/api/v1/chart/${sliceId}`,
@@ -154,7 +159,12 @@ export const updateSlice =
 
 //  Create new slice
 export const createSlice =
-  (sliceName, formData, addedToDashboard) => async dispatch => {
+  (sliceName, addedToDashboard) => async (dispatch, getState) => {
+    const {
+      explore: {
+        form_data: { url_params: _, ...formData },
+      },
+    } = getState();
     try {
       const response = await SupersetClient.post({
         endpoint: `/api/v1/chart/`,

superset-frontend/src/explore/actions/saveModalActions.test.js

@@ -75,6 +75,7 @@ const formData = {
   datasource: `${datasourceId}__${datasourceType}`,
   dashboards,
 };
+const mockExploreState = { explore: { form_data: formData } };
 
 const sliceResponsePayload = {
   id: 10,
@@ -95,11 +96,11 @@ test('updateSlice handles success', async () => {
   fetchMock.reset();
   fetchMock.put(updateSliceEndpoint, sliceResponsePayload);
   const dispatch = sinon.spy();
-  const slice = await updateSlice(
-    { slice_id: sliceId },
-    sliceName,
-    formData,
-  )(dispatch);
+  const getState = sinon.spy(() => mockExploreState);
+  const slice = await updateSlice({ slice_id: sliceId }, sliceName)(
+    dispatch,
+    getState,
+  );
 
   expect(fetchMock.calls(updateSliceEndpoint)).toHaveLength(1);
   expect(dispatch.callCount).toBe(2);
@@ -117,9 +118,10 @@ test('updateSlice handles failure', async () => {
   fetchMock.reset();
   fetchMock.put(updateSliceEndpoint, { throws: sampleError });
   const dispatch = sinon.spy();
+  const getState = sinon.spy(() => mockExploreState);
   let caughtError;
   try {
-    await updateSlice({ slice_id: sliceId }, sliceName, formData)(dispatch);
+    await updateSlice({ slice_id: sliceId }, sliceName)(dispatch, getState);
   } catch (error) {
     caughtError = error;
   }
@@ -139,7 +141,8 @@ test('createSlice handles success', async () => {
   fetchMock.reset();
   fetchMock.post(createSliceEndpoint, sliceResponsePayload);
   const dispatch = sinon.spy();
-  const slice = await createSlice(sliceName, formData)(dispatch);
+  const getState = sinon.spy(() => mockExploreState);
+  const slice = await createSlice(sliceName)(dispatch, getState);
   expect(fetchMock.calls(createSliceEndpoint)).toHaveLength(1);
   expect(dispatch.callCount).toBe(2);
   expect(dispatch.getCall(0).args[0].type).toBe(SAVE_SLICE_SUCCESS);
@@ -156,9 +159,10 @@ test('createSlice handles failure', async () => {
   fetchMock.reset();
   fetchMock.post(createSliceEndpoint, { throws: sampleError });
   const dispatch = sinon.spy();
+  const getState = sinon.spy(() => mockExploreState);
   let caughtError;
   try {
-    await createSlice(sliceName, formData)(dispatch);
+    await createSlice(sliceName)(dispatch, getState);
   } catch (error) {
     caughtError = error;
   }
@@ -243,10 +247,11 @@ test('updateSlice with add to new dashboard handles success', async () => {
   fetchMock.reset();
   fetchMock.put(updateSliceEndpoint, sliceResponsePayload);
   const dispatch = sinon.spy();
-  const slice = await updateSlice({ slice_id: sliceId }, sliceName, formData, {
+  const getState = sinon.spy(() => mockExploreState);
+  const slice = await updateSlice({ slice_id: sliceId }, sliceName, {
     new: true,
     title: dashboardName,
-  })(dispatch);
+  })(dispatch, getState);
 
   expect(fetchMock.calls(updateSliceEndpoint)).toHaveLength(1);
   expect(dispatch.callCount).toBe(3);
@@ -269,10 +274,11 @@ test('updateSlice with add to existing dashboard handles success', async () => {
   fetchMock.reset();
   fetchMock.put(updateSliceEndpoint, sliceResponsePayload);
   const dispatch = sinon.spy();
-  const slice = await updateSlice({ slice_id: sliceId }, sliceName, formData, {
+  const getState = sinon.spy(() => mockExploreState);
+  const slice = await updateSlice({ slice_id: sliceId }, sliceName, {
     new: false,
     title: dashboardName,
-  })(dispatch);
+  })(dispatch, getState);
 
   expect(fetchMock.calls(updateSliceEndpoint)).toHaveLength(1);
   expect(dispatch.callCount).toBe(3);

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

@@ -21,7 +21,7 @@ import React from 'react';
 import { Input } from 'src/components/Input';
 import { Form, FormItem } from 'src/components/Form';
 import Alert from 'src/components/Alert';
-import { t, styled, SupersetClient, DatasourceType } from '@superset-ui/core';
+import { t, styled, DatasourceType } from '@superset-ui/core';
 import ReactMarkdown from 'react-markdown';
 import Modal from 'src/components/Modal';
 import { Radio } from 'src/components/Radio';
@@ -31,7 +31,6 @@ import { SelectValue } from 'antd/lib/select';
 import { connect } from 'react-redux';
 import { withRouter, RouteComponentProps } from 'react-router-dom';
 import { InfoTooltipWithTrigger } from '@superset-ui/chart-controls';
-import { getClientErrorObject } from 'src/utils/getClientErrorObject';
 import Loading from 'src/components/Loading';
 
 // Session storage key for recent dashboard
@@ -87,7 +86,6 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
       alert: null,
       action: this.canOverwriteSlice() ? 'overwrite' : 'saveas',
       isLoading: false,
-      saveStatus: null,
     };
     this.onDashboardSelectChange = this.onDashboardSelectChange.bind(this);
     this.onSliceNameChange = this.onSliceNameChange.bind(this);
@@ -154,7 +152,6 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
     this.setState({ alert: null, isLoading: true });
     this.props.actions.removeSaveModalAlert();
     // eslint-disable-next-line @typescript-eslint/no-unused-vars
-    const { url_params, ...formData } = this.props.form_data || {};
 
     let promise = Promise.resolve();
 
@@ -170,38 +167,22 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
       const { templateParams } = this.props.datasource;
       const columns = this.props.datasource?.columns || [];
 
-      // Create a dataset object
-      await SupersetClient.post({
-        endpoint: '/superset/sqllab_viz/',
-        postPayload: {
-          data: {
-            schema,
-            sql,
-            dbId: database?.id,
-            templateParams,
-            datasourceName: this.state.datasetName,
-            metrics: [],
-            columns,
-          },
-        },
-      })
-        .then(({ json }) => json)
-        .then(async (data: { table_id: number }) => {
-          // Update form_data to point to new dataset
-          formData.datasource = `${data.table_id}__table`;
-          this.setState({ saveStatus: 'succeed' });
-        })
-        .catch(response =>
-          getClientErrorObject(response).then(e => {
-            this.setState({ isLoading: false, saveStatus: 'failed' });
-            this.props.addDangerToast(e.error);
-          }),
-        );
+      try {
+        await this.props.actions.saveDataset({
+          schema,
+          sql,
+          database,
+          templateParams,
+          datasourceName: this.state.datasetName,
+          columns,
+        });
+      } catch {
+        // Don't continue since server was unable to create dataset
+        this.setState({ isLoading: false });
+        return;
+      }
     }
 
-    // Don't continue since server was unable to create dataset
-    if (this.state.saveStatus === 'failed') return;
-
     let dashboard: DashboardGetResponse | null = null;
     if (this.state.newDashboardName || this.state.saveToDashboardId) {
       let saveToDashboardId = this.state.saveToDashboardId || null;
@@ -221,7 +202,11 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
           dashboard = response.result;
           const dashboards = new Set<number>(this.props.sliceDashboards);
           dashboards.add(dashboard.id);
-          formData.dashboards = Array.from(dashboards);
+          const { url_params, ...formData } = this.props.form_data || {};
+          this.props.actions.setFormData({
+            ...formData,
+            dashboards: Array.from(dashboards),
+          });
         });
     }
 
@@ -231,7 +216,6 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
         this.props.actions.updateSlice(
           this.props.slice,
           this.state.newSliceName,
-          formData,
           dashboard
             ? {
                 title: dashboard.dashboard_title,
@@ -244,7 +228,6 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
       promise = promise.then(() =>
         this.props.actions.createSlice(
           this.state.newSliceName,
-          formData,
           dashboard
             ? {
                 title: dashboard.dashboard_title,
@@ -386,7 +369,9 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
         buttonSize="small"
         disabled={
           !this.state.newSliceName ||
-          (!this.state.saveToDashboardId && !this.state.newDashboardName)
+          (!this.state.saveToDashboardId && !this.state.newDashboardName) ||
+          (this.props.datasource?.type !== DatasourceType.Table &&
+            !this.state.datasetName)
         }
         onClick={() => this.saveOrOverwrite(true)}
       >
@@ -399,7 +384,12 @@ class SaveModal extends React.Component<SaveModalProps, SaveModalState> {
         buttonSize="small"
         buttonStyle="primary"
         onClick={() => this.saveOrOverwrite(false)}
-        disabled={this.state.isLoading || !this.state.newSliceName}
+        disabled={
+          this.state.isLoading ||
+          !this.state.newSliceName ||
+          (this.props.datasource?.type !== DatasourceType.Table &&
+            !this.state.datasetName)
+        }
         data-test="btn-modal-save"
       >
         {!this.canOverwriteSlice() && this.props.slice

superset-frontend/src/explore/reducers/exploreReducer.js

@@ -62,6 +62,8 @@ export default function exploreReducer(state = {}, action) {
         // reset time range filter to default
         newFormData.time_range = DEFAULT_TIME_RANGE;
 
+        newFormData.datasource = newDatasource.uid;
+
         // reset control values for column/metric related controls
         Object.entries(controls).forEach(([controlName, controlState]) => {
           if (
@@ -215,6 +217,12 @@ export default function exploreReducer(state = {}, action) {
         controls: getControlsState(state, action.formData),
       };
     },
+    [actions.SET_FORM_DATA]() {
+      return {
+        ...state,
+        form_data: action.formData,
+      };
+    },
     [actions.UPDATE_CHART_TITLE]() {
       return {
         ...state,

superset/views/core.py

@@ -2135,7 +2135,12 @@ class Superset(BaseSupersetView):  # pylint: disable=too-many-public-methods
         table.columns = cols
         table.metrics = [SqlMetric(metric_name="count", expression="count(*)")]
         db.session.commit()
-        return json_success(json.dumps({"table_id": table.id}))
+
+        return json_success(
+            json.dumps(
+                {"table_id": table.id, "data": sanitize_datasource_data(table.data)}
+            )
+        )
 
     @has_access
     @expose("/extra_table_metadata/<int:database_id>/<table_name>/<schema>/")