Allow overwriting a SQLLab query that has previously been saved (#8298)

* ignore direnv

* allow overwriting saved queries

* simplify state management a little bit

* fix tests and linting

.gitignore

@@ -24,6 +24,7 @@
 .coverage
 .DS_Store
 .eggs
+.envrc
 .idea
 .mypy_cache
 .python-version

superset/assets/package.json

@@ -79,9 +79,9 @@
     "@superset-ui/legacy-preset-chart-deckgl": "^0.1.0",
     "@superset-ui/legacy-preset-chart-nvd3": "^0.11.0",
     "@superset-ui/number-format": "^0.12.1",
-    "@superset-ui/query": "^0.12.2",
     "@superset-ui/plugin-chart-table": "^0.11.0",
     "@superset-ui/preset-chart-xy": "^0.11.0",
+    "@superset-ui/query": "^0.12.2",
     "@superset-ui/time-format": "^0.12.1",
     "@superset-ui/translation": "^0.12.0",
     "@types/react-json-tree": "^0.6.11",

superset/assets/spec/javascripts/sqllab/SaveQuery_spec.jsx

@@ -19,14 +19,18 @@
 import React from 'react';
 import { FormControl } from 'react-bootstrap';
 import { shallow } from 'enzyme';
+import * as sinon from 'sinon';
 import SaveQuery from '../../../src/SqlLab/components/SaveQuery';
 import ModalTrigger from '../../../src/components/ModalTrigger';
+import Button from '../../../src/components/Button';
 
 describe('SavedQuery', () => {
   const mockedProps = {
-    dbId: 1,
-    schema: 'main',
-    sql: 'SELECT * FROM t',
+    query: {
+      dbId: 1,
+      schema: 'main',
+      sql: 'SELECT * FROM t',
+    },
     defaultLabel: 'untitled',
     animation: false,
   };
@@ -54,4 +58,27 @@ describe('SavedQuery', () => {
     const modal = shallow(wrapper.instance().renderModalBody());
     expect(modal.find(FormControl)).toHaveLength(2);
   });
+  it('has a save button if this is a new query', () => {
+    const saveSpy = sinon.spy();
+    const wrapper = shallow(<SaveQuery {...mockedProps} onSave={saveSpy} />);
+    const modal = shallow(wrapper.instance().renderModalBody());
+    expect(modal.find(Button)).toHaveLength(2);
+    modal.find(Button).at(0).simulate('click');
+    expect(saveSpy.calledOnce).toBe(true);
+  });
+  it('has an update button if this is an existing query', () => {
+    const updateSpy = sinon.spy();
+    const props = {
+      ...mockedProps,
+      query: {
+        ...mockedProps.query,
+        remoteId: '42',
+      },
+    };
+    const wrapper = shallow(<SaveQuery {...props} onUpdate={updateSpy} />);
+    const modal = shallow(wrapper.instance().renderModalBody());
+    expect(modal.find(Button)).toHaveLength(3);
+    modal.find(Button).at(0).simulate('click');
+    expect(updateSpy.calledOnce).toBe(true);
+  });
 });

superset/assets/src/SqlLab/actions/sqlLab.js

@@ -20,6 +20,8 @@ import shortid from 'shortid';
 import JSONbig from 'json-bigint';
 import { t } from '@superset-ui/translation';
 import { SupersetClient } from '@superset-ui/connection';
+import invert from 'lodash/invert';
+import mapKeys from 'lodash/mapKeys';
 
 import { now } from '../../modules/dates';
 import {
@@ -32,6 +34,7 @@ import COMMON_ERR_MESSAGES from '../../utils/errorMessages';
 
 export const RESET_STATE = 'RESET_STATE';
 export const ADD_QUERY_EDITOR = 'ADD_QUERY_EDITOR';
+export const UPDATE_QUERY_EDITOR = 'UPDATE_QUERY_EDITOR';
 export const CLONE_QUERY_TO_NEW_TAB = 'CLONE_QUERY_TO_NEW_TAB';
 export const REMOVE_QUERY_EDITOR = 'REMOVE_QUERY_EDITOR';
 export const MERGE_TABLE = 'MERGE_TABLE';
@@ -82,6 +85,24 @@ export const addInfoToast = addInfoToastAction;
 export const addSuccessToast = addSuccessToastAction;
 export const addDangerToast = addDangerToastAction;
 
+// a map of SavedQuery field names to the different names used client-side,
+// because for now making the names consistent is too complicated
+// so it might as well only happen in one place
+const queryClientMapping = {
+  id: 'remoteId',
+  db_id: 'dbId',
+  client_id: 'id',
+  label: 'title',
+};
+const queryServerMapping = invert(queryClientMapping);
+
+// uses a mapping like those above to convert object key names to another style
+const fieldConverter = mapping => obj =>
+  mapKeys(obj, (value, key) => key in mapping ? mapping[key] : key);
+
+const convertQueryToServer = fieldConverter(queryServerMapping);
+const convertQueryToClient = fieldConverter(queryClientMapping);
+
 export function resetState() {
   return { type: RESET_STATE };
 }
@@ -105,13 +126,29 @@ export function saveQuery(query) {
   return dispatch =>
     SupersetClient.post({
       endpoint: '/savedqueryviewapi/api/create',
-      postPayload: query,
+      postPayload: convertQueryToServer(query),
       stringify: false,
     })
       .then(() => dispatch(addSuccessToast(t('Your query was saved'))))
       .catch(() => dispatch(addDangerToast(t('Your query could not be saved'))));
 }
 
+export function updateQueryEditor(alterations) {
+  return { type: UPDATE_QUERY_EDITOR, alterations };
+}
+
+export function updateSavedQuery(query) {
+  return dispatch =>
+    SupersetClient.put({
+      endpoint: `/savedqueryviewapi/api/update/${query.remoteId}`,
+      postPayload: convertQueryToServer(query),
+      stringify: false,
+    })
+      .then(() => dispatch(addSuccessToast(t('Your query was updated'))))
+      .catch(() => dispatch(addDangerToast(t('Your query could not be updated'))))
+      .then(() => dispatch(updateQueryEditor(query)));
+}
+
 export function scheduleQuery(query) {
   return dispatch =>
     SupersetClient.post({
@@ -504,13 +541,9 @@ export function popSavedQuery(saveQueryId) {
   return function (dispatch) {
     return SupersetClient.get({ endpoint: `/savedqueryviewapi/api/get/${saveQueryId}` })
       .then(({ json }) => {
-        const { result } = json;
         const queryEditorProps = {
-          title: result.label,
-          dbId: result.db_id ? parseInt(result.db_id, 10) : null,
-          schema: result.schema,
+          ...convertQueryToClient(json.result),
           autorun: false,
-          sql: result.sql,
         };
         return dispatch(addQueryEditor(queryEditorProps));
       })

superset/assets/src/SqlLab/components/SaveQuery.jsx

@@ -25,12 +25,11 @@ import Button from '../../components/Button';
 import ModalTrigger from '../../components/ModalTrigger';
 
 const propTypes = {
+  query: PropTypes.object,
   defaultLabel: PropTypes.string,
-  sql: PropTypes.string,
-  schema: PropTypes.string,
-  dbId: PropTypes.number,
   animation: PropTypes.bool,
   onSave: PropTypes.func,
+  onUpdate: PropTypes.func,
   saveQueryWarning: PropTypes.string,
 };
 const defaultProps = {
@@ -50,23 +49,21 @@ class SaveQuery extends React.PureComponent {
     };
     this.toggleSave = this.toggleSave.bind(this);
     this.onSave = this.onSave.bind(this);
+    this.onUpdate = this.onUpdate.bind(this);
     this.onCancel = this.onCancel.bind(this);
     this.onLabelChange = this.onLabelChange.bind(this);
     this.onDescriptionChange = this.onDescriptionChange.bind(this);
   }
   onSave() {
-    const query = {
-      label: this.state.label,
-      description: this.state.description,
-      db_id: this.props.dbId,
-      schema: this.props.schema,
-      sql: this.props.sql,
-    };
-    this.props.onSave(query);
-    this.saveModal.close();
+    this.props.onSave(this.queryPayload());
+    this.close();
+  }
+  onUpdate() {
+    this.props.onUpdate(this.queryPayload());
+    this.close();
   }
   onCancel() {
-    this.saveModal.close();
+    this.close();
   }
   onLabelChange(e) {
     this.setState({ label: e.target.value });
@@ -74,10 +71,21 @@ class SaveQuery extends React.PureComponent {
   onDescriptionChange(e) {
     this.setState({ description: e.target.value });
   }
+  queryPayload() {
+    return {
+      ...this.props.query,
+      title: this.state.label,
+      description: this.state.description,
+    };
+  }
+  close() {
+    if (this.saveModal) this.saveModal.close();
+  }
   toggleSave(e) {
     this.setState({ target: e.target, showSave: !this.state.showSave });
   }
   renderModalBody() {
+    const isSaved = !!this.props.query.remoteId;
     return (
       <FormGroup bsSize="small">
         <Row>
@@ -124,12 +132,21 @@ class SaveQuery extends React.PureComponent {
         )}
         <Row>
           <Col md={12}>
+            {isSaved && (
+              <Button
+                bsStyle="primary"
+                onClick={this.onUpdate}
+                className="m-r-3"
+              >
+                {t('Update')}
+              </Button>
+            )}
             <Button
-              bsStyle="primary"
+              bsStyle={isSaved ? undefined : 'primary'}
               onClick={this.onSave}
               className="m-r-3"
             >
-              {t('Save')}
+              {isSaved ? t('Save New') : t('Save')}
             </Button>
             <Button onClick={this.onCancel} className="cancelQuery">
               {t('Cancel')}

superset/assets/src/SqlLab/components/SqlEditor.jsx

@@ -432,12 +432,11 @@ class SqlEditor extends React.PureComponent {
             }
             <span className="m-r-5">
               <SaveQuery
-                defaultLabel={qe.title}
-                sql={qe.sql}
+                query={qe}
+                defaultLabel={qe.description == null ? qe.title : qe.description}
                 className="m-r-5"
                 onSave={this.props.actions.saveQuery}
-                schema={qe.schema}
-                dbId={qe.dbId}
+                onUpdate={this.props.actions.updateSavedQuery}
                 saveQueryWarning={this.props.saveQueryWarning}
               />
             </span>

superset/assets/src/SqlLab/reducers/sqlLab.js

@@ -39,11 +39,18 @@ export default function sqlLabReducer(state = {}, action) {
       const newState = Object.assign({}, state, { tabHistory });
       return addToArr(newState, 'queryEditors', action.queryEditor);
     },
+    [actions.UPDATE_QUERY_EDITOR]() {
+      const id = action.alterations.remoteId;
+      const existing = state.queryEditors.find(qe => qe.remoteId === id);
+      if (existing == null) return state;
+      return alterInArr(state, 'queryEditors', existing, action.alterations, 'remoteId');
+    },
     [actions.CLONE_QUERY_TO_NEW_TAB]() {
       const progenitor = state.queryEditors.find(
         qe => qe.id === state.tabHistory[state.tabHistory.length - 1],
       );
       const qe = {
+        remoteId: progenitor.remoteId,
         id: shortid.generate(),
         title: t('Copy of %s', progenitor.title),
         dbId: action.query.dbId ? action.query.dbId : null,

superset/assets/src/SqlLab/utils/reduxStateToLocalStorageHelper.js

@@ -19,6 +19,7 @@
 import { LOCALSTORAGE_MAX_QUERY_AGE_MS } from '../constants';
 
 const PERSISTENT_QUERY_EDITOR_KEYS = new Set([
+  'remoteId',
   'autorun',
   'dbId',
   'height',

superset/views/core.py

@@ -2678,7 +2678,7 @@ class Superset(BaseSupersetView):
     @expose("/sql_json/", methods=["POST"])
     @event_logger.log_this
     def sql_json(self):
-        """Runs arbitrary sql and returns and json"""
+        """Runs arbitrary sql and returns data as json"""
         # Collect Values
         database_id: int = request.json.get("database_id")
         schema: str = request.json.get("schema")

superset/views/sql_lab.py

@@ -156,7 +156,15 @@ class SavedQueryViewApi(SavedQueryView):
         "sql",
         "extra_json",
     ]
-    show_columns = ["label", "db_id", "schema", "description", "sql", "extra_json"]
+    show_columns = [
+        "id",
+        "label",
+        "db_id",
+        "schema",
+        "description",
+        "sql",
+        "extra_json",
+    ]
     add_columns = show_columns
     edit_columns = add_columns