[explore view] fix long query issue from Run in SQL LAB Button (#9345)

* [explore view] fix long query issue from Run in SQL LAB Button

* SQL Lab page needs to take the post form data, too

* fix variable names

* updated payload dict, rename hidden form

Co-authored-by: Jesse Yang <jesse.yang@airbnb.com>

superset-frontend/src/SqlLab/App.jsx

@@ -46,7 +46,9 @@ setupApp();
 
 const appContainer = document.getElementById('app');
 const bootstrapData = JSON.parse(appContainer.getAttribute('data-bootstrap'));
+
 initFeatureFlags(bootstrapData.common.feature_flags);
+
 const initialState = getInitialState(bootstrapData);
 const sqlLabPersistStateConfig = {
   paths: ['sqlLab'],
@@ -59,7 +61,6 @@ const sqlLabPersistStateConfig = {
         // it caused configurations passed from server-side got override.
         // see PR 6257 for details
         delete state[path].common; // eslint-disable-line no-param-reassign
-
         if (path === 'sqlLab') {
           subset[path] = {
             ...state[path],

superset-frontend/src/SqlLab/components/TabbedSqlEditors.jsx

@@ -39,6 +39,7 @@ const propTypes = {
   databases: PropTypes.object.isRequired,
   queries: PropTypes.object.isRequired,
   queryEditors: PropTypes.array,
+  requestedQuery: PropTypes.object,
   tabHistory: PropTypes.array.isRequired,
   tables: PropTypes.array.isRequired,
   offline: PropTypes.bool,
@@ -48,6 +49,7 @@ const propTypes = {
 const defaultProps = {
   queryEditors: [],
   offline: false,
+  requestedQuery: null,
   saveQueryWarning: null,
   scheduleQueryWarning: null,
 };
@@ -99,7 +101,12 @@ class TabbedSqlEditors extends React.PureComponent {
         });
     }
 
-    const query = URI(window.location).search(true);
+    // merge post form data with GET search params
+    const query = {
+      ...this.props.requestedQuery,
+      ...URI(window.location).search(true),
+    };
+
     // Popping a new tab based on the querystring
     if (query.id || query.sql || query.savedQueryId || query.datasourceKey) {
       if (query.id) {
@@ -374,7 +381,7 @@ class TabbedSqlEditors extends React.PureComponent {
 TabbedSqlEditors.propTypes = propTypes;
 TabbedSqlEditors.defaultProps = defaultProps;
 
-function mapStateToProps({ sqlLab, common }) {
+function mapStateToProps({ sqlLab, common, requestedQuery }) {
   return {
     databases: sqlLab.databases,
     queryEditors: sqlLab.queryEditors,
@@ -388,6 +395,7 @@ function mapStateToProps({ sqlLab, common }) {
     maxRow: common.conf.SQL_MAX_ROW,
     saveQueryWarning: common.conf.SQLLAB_SAVE_WARNING_MESSAGE,
     scheduleQueryWarning: common.conf.SQLLAB_SCHEDULE_WARNING_MESSAGE,
+    requestedQuery,
   };
 }
 function mapDispatchToProps(dispatch) {

superset-frontend/src/SqlLab/reducers/getInitialState.js

@@ -19,8 +19,16 @@
 import { t } from '@superset-ui/translation';
 import getToastsFromPyFlashMessages from '../../messageToasts/utils/getToastsFromPyFlashMessages';
 
-export default function getInitialState({ defaultDbId, ...restBootstrapData }) {
-  /*
+export default function getInitialState({
+  defaultDbId,
+  common,
+  active_tab: activeTab,
+  tab_state_ids: tabStateIds = [],
+  databases,
+  queries: queries_,
+  requested_query: requestedQuery,
+}) {
+  /**
    * Before YYYY-MM-DD, the state for SQL Lab was stored exclusively in the
    * browser's localStorage. The feature flag `SQLLAB_BACKEND_PERSISTENCE`
    * moves the state to the backend instead, migrating it from local storage.
@@ -39,7 +47,7 @@ export default function getInitialState({ defaultDbId, ...restBootstrapData }) {
     autorun: false,
     templateParams: null,
     dbId: defaultDbId,
-    queryLimit: restBootstrapData.common.conf.DEFAULT_SQLLAB_LIMIT,
+    queryLimit: common.conf.DEFAULT_SQLLAB_LIMIT,
     validationResult: {
       id: null,
       errors: [],
@@ -52,11 +60,11 @@ export default function getInitialState({ defaultDbId, ...restBootstrapData }) {
     },
   };
 
-  /* Load state from the backend. This will be empty if the feature flag
+  /**
+   * Load state from the backend. This will be empty if the feature flag
    * `SQLLAB_BACKEND_PERSISTENCE` is off.
    */
-  const activeTab = restBootstrapData.active_tab;
-  restBootstrapData.tab_state_ids.forEach(({ id, label }) => {
+  tabStateIds.forEach(({ id, label }) => {
     let queryEditor;
     if (activeTab && activeTab.id === id) {
       queryEditor = {
@@ -92,7 +100,6 @@ export default function getInitialState({ defaultDbId, ...restBootstrapData }) {
   });
 
   const tabHistory = activeTab ? [activeTab.id.toString()] : [];
-
   const tables = [];
   if (activeTab) {
     activeTab.table_schemas
@@ -126,9 +133,10 @@ export default function getInitialState({ defaultDbId, ...restBootstrapData }) {
       });
   }
 
-  const { databases, queries } = restBootstrapData;
+  const queries = { ...queries_ };
 
-  /* If the `SQLLAB_BACKEND_PERSISTENCE` feature flag is off, or if the user
+  /**
+   * If the `SQLLAB_BACKEND_PERSISTENCE` feature flag is off, or if the user
    * hasn't used SQL Lab after it has been turned on, the state will be stored
    * in the browser's local storage.
    */
@@ -173,13 +181,14 @@ export default function getInitialState({ defaultDbId, ...restBootstrapData }) {
       tables,
       queriesLastUpdate: Date.now(),
     },
+    requestedQuery,
     messageToasts: getToastsFromPyFlashMessages(
-      (restBootstrapData.common || {}).flash_messages || [],
+      (common || {}).flash_messages || [],
     ),
     localStorageUsageInKilobytes: 0,
     common: {
-      flash_messages: restBootstrapData.common.flash_messages,
-      conf: restBootstrapData.common.conf,
+      flash_messages: common.flash_messages,
+      conf: common.conf,
     },
   };
 }

superset-frontend/src/chart/chartAction.js

@@ -25,6 +25,7 @@ import { isFeatureEnabled, FeatureFlag } from 'src/featureFlags';
 import {
   getExploreUrlAndPayload,
   getAnnotationJsonUrl,
+  postForm,
 } from '../explore/exploreUtils';
 import {
   requiresQuery,
@@ -358,14 +359,12 @@ export function redirectSQLLab(formData) {
       postPayload: { form_data: formData },
     })
       .then(({ json }) => {
-        const redirectUrl = new URL(window.location);
-        redirectUrl.pathname = '/superset/sqllab';
-        for (const key of redirectUrl.searchParams.keys()) {
-          redirectUrl.searchParams.delete(key);
-        }
-        redirectUrl.searchParams.set('datasourceKey', formData.datasource);
-        redirectUrl.searchParams.set('sql', json.query);
-        window.open(redirectUrl.href, '_blank');
+        const redirectUrl = '/superset/sqllab';
+        const payload = {
+          datasourceKey: formData.datasource,
+          sql: json.query,
+        };
+        postForm(redirectUrl, payload);
       })
       .catch(() =>
         dispatch(addDangerToast(t('An error occurred while loading the SQL'))),

superset-frontend/src/explore/exploreUtils.js

@@ -190,29 +190,36 @@ export function getExploreUrlAndPayload({
   };
 }
 
+export function postForm(url, payload, target = '_blank') {
+  if (!url) {
+    return;
+  }
+
+  const hiddenForm = document.createElement('form');
+  hiddenForm.action = url;
+  hiddenForm.method = 'POST';
+  hiddenForm.target = target;
+  const token = document.createElement('input');
+  token.type = 'hidden';
+  token.name = 'csrf_token';
+  token.value = (document.getElementById('csrf_token') || {}).value;
+  hiddenForm.appendChild(token);
+  const data = document.createElement('input');
+  data.type = 'hidden';
+  data.name = 'form_data';
+  data.value = safeStringify(payload);
+  hiddenForm.appendChild(data);
+
+  document.body.appendChild(hiddenForm);
+  hiddenForm.submit();
+  document.body.removeChild(hiddenForm);
+}
+
 export function exportChart(formData, endpointType) {
   const { url, payload } = getExploreUrlAndPayload({
     formData,
     endpointType,
     allowDomainSharding: false,
   });
-
-  const exploreForm = document.createElement('form');
-  exploreForm.action = url;
-  exploreForm.method = 'POST';
-  exploreForm.target = '_blank';
-  const token = document.createElement('input');
-  token.type = 'hidden';
-  token.name = 'csrf_token';
-  token.value = (document.getElementById('csrf_token') || {}).value;
-  exploreForm.appendChild(token);
-  const data = document.createElement('input');
-  data.type = 'hidden';
-  data.name = 'form_data';
-  data.value = safeStringify(payload);
-  exploreForm.appendChild(data);
-
-  document.body.appendChild(exploreForm);
-  exploreForm.submit();
-  document.body.removeChild(exploreForm);
+  postForm(url, payload);
 }

superset/views/core.py

@@ -2713,7 +2713,7 @@ class Superset(BaseSupersetView):
         )
 
     @staticmethod
-    def _get_sqllab_payload(user_id: int) -> Dict[str, Any]:
+    def _get_sqllab_tabs(user_id: int) -> Dict[str, Any]:
         # send list of tab state ids
         tabs_state = (
             db.session.query(TabState.id, TabState.label)
@@ -2753,8 +2753,6 @@ class Superset(BaseSupersetView):
             }
 
         return {
-            "defaultDbId": config["SQLLAB_DEFAULT_DBID"],
-            "common": common_bootstrap_payload(),
             "tab_state_ids": tabs_state,
             "active_tab": active_tab.to_dict() if active_tab else None,
             "databases": databases,
@@ -2762,10 +2760,21 @@ class Superset(BaseSupersetView):
         }
 
     @has_access
-    @expose("/sqllab")
+    @expose("/sqllab", methods=["GET", "POST"])
     def sqllab(self):
         """SQL Editor"""
-        payload = self._get_sqllab_payload(g.user.get_id())
+        payload = {
+            "defaultDbId": config["SQLLAB_DEFAULT_DBID"],
+            "common": common_bootstrap_payload(),
+            **self._get_sqllab_tabs(g.user.get_id()),
+        }
+
+        form_data = request.form.get("form_data")
+        if form_data:
+            try:
+                payload["requested_query"] = json.loads(form_data)
+            except json.JSONDecodeError:
+                pass
         bootstrap_data = json.dumps(
             payload, default=utils.pessimistic_json_iso_dttm_ser
         )

tests/core_tests.py

@@ -1177,7 +1177,7 @@ class CoreTests(SupersetTestCase):
 
         # we should have only 1 query returned, since the second one is not
         # associated with any tabs
-        payload = views.Superset._get_sqllab_payload(user_id=user_id)
+        payload = views.Superset._get_sqllab_tabs(user_id=user_id)
         self.assertEqual(len(payload["queries"]), 1)