refactor(sql_lab): SQL Lab Persistent Saved State (#17771)

* a lot of console logs * testing * test * added saved_query to remoteId * created useEffect so that title properly changes in modal * Update superset-frontend/src/SqlLab/actions/sqlLab.js Co-authored-by: Lyndsi Kay Williams <55605634+lyndsiWilliams@users.noreply.github.com> Co-authored-by: Lyndsi Kay Williams <55605634+lyndsiWilliams@users.noreply.github.com>
apache · Jan 18, 2022 · 88db2cc · 88db2cc
1 parent 5bfe2d4
commit 88db2cc
Show file tree

Hide file tree

Showing 6 changed files with 52 additions and 8 deletions.
diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.js b/superset-frontend/src/SqlLab/actions/sqlLab.js
@@ -636,12 +636,13 @@ export function switchQueryEditor(queryEditor, displayLimit) {
             title: json.label,
             sql: json.sql,
             selectedText: null,
-            latestQueryId: json.latest_query ? json.latest_query.id : null,
+            latestQueryId: json.latest_query?.id,
             autorun: json.autorun,
             dbId: json.database_id,
             templateParams: json.template_params,
             schema: json.schema,
             queryLimit: json.query_limit,
+            remoteId: json.saved_query?.id,
             validationResult: {
               id: null,
               errors: [],
@@ -864,19 +865,38 @@ export function saveQuery(query) {
       stringify: false,
     })
       .then(result => {
+        const savedQuery = convertQueryToClient(result.json.item);
         dispatch({
           type: QUERY_EDITOR_SAVED,
           query,
-          result: convertQueryToClient(result.json.item),
+          result: savedQuery,
         });
-        dispatch(addSuccessToast(t('Your query was saved')));
         dispatch(queryEditorSetTitle(query, query.title));
+        return savedQuery;
       })
       .catch(() =>
         dispatch(addDangerToast(t('Your query could not be saved'))),
       );
 }
 
+export const addSavedQueryToTabState =
+  (queryEditor, savedQuery) => dispatch => {
+    const sync = isFeatureEnabled(FeatureFlag.SQLLAB_BACKEND_PERSISTENCE)
+      ? SupersetClient.put({
+          endpoint: `/tabstateview/${queryEditor.id}`,
+          postPayload: { saved_query_id: savedQuery.remoteId },
+        })
+      : Promise.resolve();
+
+    return sync
+      .catch(() => {
+        dispatch(addDangerToast(t('Your query was not properly saved')));
+      })
+      .then(() => {
+        dispatch(addSuccessToast(t('Your query was saved')));
+      });
+  };
+
 export function updateSavedQuery(query) {
   return dispatch =>
     SupersetClient.put({

diff --git a/superset-frontend/src/SqlLab/actions/sqlLab.test.js b/superset-frontend/src/SqlLab/actions/sqlLab.test.js
@@ -23,7 +23,6 @@ import configureMockStore from 'redux-mock-store';
 import thunk from 'redux-thunk';
 import shortid from 'shortid';
 import * as featureFlags from 'src/featureFlags';
-import { ADD_TOAST } from 'src/components/MessageToasts/actions';
 import * as actions from 'src/SqlLab/actions/sqlLab';
 import { defaultQueryEditor, query } from '../fixtures';
 
@@ -93,7 +92,7 @@ describe('async actions', () => {
       expect.assertions(1);
 
       return makeRequest().then(() => {
-        expect(dispatch.callCount).toBe(3);
+        expect(dispatch.callCount).toBe(2);
       });
     });
 
@@ -111,7 +110,6 @@ describe('async actions', () => {
       const store = mockStore({});
       const expectedActionTypes = [
         actions.QUERY_EDITOR_SAVED,
-        ADD_TOAST,
         actions.QUERY_EDITOR_SET_TITLE,
       ];
       return store.dispatch(actions.saveQuery(query)).then(() => {

diff --git a/superset-frontend/src/SqlLab/components/SaveQuery/index.tsx b/superset-frontend/src/SqlLab/components/SaveQuery/index.tsx
@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { useState } from 'react';
+import React, { useState, useEffect } from 'react';
 import { Row, Col, Input, TextArea } from 'src/common/components';
 import { t, styled } from '@superset-ui/core';
 import Button from 'src/components/Button';
@@ -90,6 +90,11 @@ export default function SaveQuery({
     description,
   });
 
+  useEffect(() => {
+    if (!isSaved) {
+      setLabel(defaultLabel);
+    }
+  }, [defaultLabel]);
   const close = () => {
     setShowSave(false);
   };

diff --git a/superset-frontend/src/SqlLab/components/SqlEditor/index.jsx b/superset-frontend/src/SqlLab/components/SqlEditor/index.jsx
@@ -52,6 +52,7 @@ import {
   queryEditorSetTemplateParams,
   runQuery,
   saveQuery,
+  addSavedQueryToTabState,
   scheduleQuery,
   setActiveSouthPaneTab,
   updateSavedQuery,
@@ -192,6 +193,7 @@ class SqlEditor extends React.PureComponent {
     this.canValidateQuery = this.canValidateQuery.bind(this);
     this.runQuery = this.runQuery.bind(this);
     this.stopQuery = this.stopQuery.bind(this);
+    this.saveQuery = this.saveQuery.bind(this);
     this.onSqlChanged = this.onSqlChanged.bind(this);
     this.setQueryEditorSql = this.setQueryEditorSql.bind(this);
     this.setQueryEditorSqlWithDebounce = debounce(
@@ -592,6 +594,12 @@ class SqlEditor extends React.PureComponent {
     );
   }
 
+  async saveQuery(query) {
+    const { queryEditor: qe, actions } = this.props;
+    const savedQuery = await actions.saveQuery(query);
+    actions.addSavedQueryToTabState(qe, savedQuery);
+  }
+
   renderEditorBottomBar() {
     const { queryEditor: qe } = this.props;
 
@@ -630,6 +638,7 @@ class SqlEditor extends React.PureComponent {
         )}
       </Menu>
     );
+
     return (
       <StyledToolbar className="sql-toolbar" id="js-sql-toolbar">
         <div className="leftItems">
@@ -693,7 +702,7 @@ class SqlEditor extends React.PureComponent {
             <SaveQuery
               query={qe}
               defaultLabel={qe.title || qe.description}
-              onSave={this.props.actions.saveQuery}
+              onSave={this.saveQuery}
               onUpdate={this.props.actions.updateSavedQuery}
               saveQueryWarning={this.props.saveQueryWarning}
             />
@@ -809,6 +818,7 @@ function mapDispatchToProps(dispatch) {
       queryEditorSetTemplateParams,
       runQuery,
       saveQuery,
+      addSavedQueryToTabState,
       scheduleQuery,
       setActiveSouthPaneTab,
       updateSavedQuery,

diff --git a/superset-frontend/src/SqlLab/reducers/getInitialState.js b/superset-frontend/src/SqlLab/reducers/getInitialState.js
@@ -79,6 +79,7 @@ export default function getInitialState({
         latestQueryId: activeTab.latest_query
           ? activeTab.latest_query.id
           : null,
+        remoteId: activeTab.saved_query?.id,
         autorun: activeTab.autorun,
         templateParams: activeTab.template_params || undefined,
         dbId: activeTab.database_id,

diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py
@@ -211,6 +211,11 @@ class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin, ImportExportMixin):
     def __repr__(self) -> str:
         return str(self.label)
 
+    def to_dict(self) -> Dict[str, Any]:
+        return {
+            "id": self.id,
+        }
+
     @property
     def pop_tab_link(self) -> Markup:
         return Markup(
@@ -285,6 +290,10 @@ class TabState(Model, AuditMixinNullable, ExtraJSONMixin):
     template_params = Column(Text)
     hide_left_bar = Column(Boolean, default=False)
 
+    # any saved queries that are associated with the Tab State
+    saved_query_id = Column(Integer, ForeignKey("saved_query.id"), nullable=True)
+    saved_query = relationship("SavedQuery", foreign_keys=[saved_query_id])
+
     def to_dict(self) -> Dict[str, Any]:
         return {
             "id": self.id,
@@ -300,6 +309,7 @@ def to_dict(self) -> Dict[str, Any]:
             "autorun": self.autorun,
             "template_params": self.template_params,
             "hide_left_bar": self.hide_left_bar,
+            "saved_query": self.saved_query.to_dict() if self.saved_query else None,
         }