chore: consolidate sqllab store into SPA store (#25088)

superset-frontend/src/SqlLab/App.jsx

@@ -17,7 +17,6 @@
  * under the License.
  */
 import React from 'react';
-import persistState from 'redux-localstorage';
 import { Provider } from 'react-redux';
 import { hot } from 'react-hot-loader/root';
 import {
@@ -30,16 +29,11 @@ import { GlobalStyles } from 'src/GlobalStyles';
 import { setupStore, userReducer } from 'src/views/store';
 import setupExtensions from 'src/setup/setupExtensions';
 import getBootstrapData from 'src/utils/getBootstrapData';
-import { tableApiUtil } from 'src/hooks/apiResources/tables';
+import { persistSqlLabStateEnhancer } from 'src/SqlLab/middlewares/persistSqlLabStateEnhancer';
 import getInitialState from './reducers/getInitialState';
 import { reducers } from './reducers/index';
 import App from './components/App';
-import {
-  emptyTablePersistData,
-  emptyQueryResults,
-  clearQueryEditors,
-} from './utils/reduxStateToLocalStorageHelper';
-import { BYTES_PER_CHAR, KB_STORAGE } from './constants';
+import { rehydratePersistedState } from './utils/reduxStateToLocalStorageHelper';
 import setupApp from '../setup/setupApp';
 
 import '../assets/stylesheets/reactable-pagination.less';
@@ -54,90 +48,16 @@ const bootstrapData = getBootstrapData();
 initFeatureFlags(bootstrapData.common.feature_flags);
 
 const initialState = getInitialState(bootstrapData);
-const sqlLabPersistStateConfig = {
-  paths: ['sqlLab'],
-  config: {
-    slicer: paths => state => {
-      const subset = {};
-      paths.forEach(path => {
-        // this line is used to remove old data from browser localStorage.
-        // we used to persist all redux state into localStorage, but
-        // 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],
-            tables: emptyTablePersistData(state[path].tables),
-            queries: emptyQueryResults(state[path].queries),
-            queryEditors: clearQueryEditors(state[path].queryEditors),
-            unsavedQueryEditor: clearQueryEditors([
-              state[path].unsavedQueryEditor,
-            ])[0],
-          };
-        }
-      });
-
-      const data = JSON.stringify(subset);
-      // 2 digit precision
-      const currentSize =
-        Math.round(((data.length * BYTES_PER_CHAR) / KB_STORAGE) * 100) / 100;
-      if (state.localStorageUsageInKilobytes !== currentSize) {
-        state.localStorageUsageInKilobytes = currentSize; // eslint-disable-line no-param-reassign
-      }
-
-      return subset;
-    },
-    merge: (initialState, persistedState = {}) => {
-      const result = {
-        ...initialState,
-        ...persistedState,
-        sqlLab: {
-          ...(persistedState?.sqlLab || {}),
-          // Overwrite initialState over persistedState for sqlLab
-          // since a logic in getInitialState overrides the value from persistedState
-          ...initialState.sqlLab,
-        },
-      };
-      return result;
-    },
-  },
-};
 
 export const store = setupStore({
   initialState,
   rootReducers: { ...reducers, user: userReducer },
   ...(!isFeatureEnabled(FeatureFlag.SQLLAB_BACKEND_PERSISTENCE) && {
-    enhancers: [
-      persistState(
-        sqlLabPersistStateConfig.paths,
-        sqlLabPersistStateConfig.config,
-      ),
-    ],
+    enhancers: [persistSqlLabStateEnhancer],
   }),
 });
 
-// Rehydrate server side persisted table metadata
-initialState.sqlLab.tables.forEach(
-  ({ name: table, schema, dbId, persistData }) => {
-    if (dbId && schema && table && persistData?.columns) {
-      store.dispatch(
-        tableApiUtil.upsertQueryData(
-          'tableMetadata',
-          { dbId, schema, table },
-          persistData,
-        ),
-      );
-      store.dispatch(
-        tableApiUtil.upsertQueryData(
-          'tableExtendedMetadata',
-          { dbId, schema, table },
-          {},
-        ),
-      );
-    }
-  },
-);
+rehydratePersistedState(store.dispatch, initialState);
 
 // Highlight the navbar menu
 const menus = document.querySelectorAll('.nav.navbar-nav li.dropdown');

superset-frontend/src/SqlLab/actions/sqlLab.js

@@ -37,8 +37,11 @@ import {
 import { getClientErrorObject } from 'src/utils/getClientErrorObject';
 import COMMON_ERR_MESSAGES from 'src/utils/errorMessages';
 import { LOG_ACTIONS_SQLLAB_FETCH_FAILED_QUERY } from 'src/logger/LogUtils';
+import getBootstrapData from 'src/utils/getBootstrapData';
 import { logEvent } from 'src/logger/actions';
 import { newQueryTabName } from '../utils/newQueryTabName';
+import getInitialState from '../reducers/getInitialState';
+import { rehydratePersistedState } from '../utils/reduxStateToLocalStorageHelper';
 
 export const RESET_STATE = 'RESET_STATE';
 export const ADD_QUERY_EDITOR = 'ADD_QUERY_EDITOR';
@@ -136,8 +139,21 @@ export function getUpToDateQuery(rootState, queryEditor, key) {
   };
 }
 
-export function resetState() {
-  return { type: RESET_STATE };
+export function resetState(data) {
+  return (dispatch, getState) => {
+    const { common } = getState();
+    const initialState = getInitialState({
+      ...getBootstrapData(),
+      common,
+      ...data,
+    });
+
+    dispatch({
+      type: RESET_STATE,
+      sqlLabInitialState: initialState.sqlLab,
+    });
+    rehydratePersistedState(dispatch, initialState);
+  };
 }
 
 export function updateQueryEditor(alterations) {

superset-frontend/src/SqlLab/middlewares/persistSqlLabStateEnhancer.js

@@ -0,0 +1,85 @@
+/**
+ * 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.
+ */
+// TODO: requires redux-localstorage > 1.0 for typescript support
+import persistState from 'redux-localstorage';
+import {
+  emptyTablePersistData,
+  emptyQueryResults,
+  clearQueryEditors,
+} from '../utils/reduxStateToLocalStorageHelper';
+import { BYTES_PER_CHAR, KB_STORAGE } from '../constants';
+
+const CLEAR_ENTITY_HELPERS_MAP = {
+  tables: emptyTablePersistData,
+  queries: emptyQueryResults,
+  queryEditors: clearQueryEditors,
+  unsavedQueryEditor: qe => clearQueryEditors([qe])[0],
+};
+
+const sqlLabPersistStateConfig = {
+  paths: ['sqlLab'],
+  config: {
+    slicer: paths => state => {
+      const subset = {};
+      paths.forEach(path => {
+        // this line is used to remove old data from browser localStorage.
+        // we used to persist all redux state into localStorage, but
+        // 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] = Object.fromEntries(
+            Object.entries(state[path]).map(([key, value]) => [
+              key,
+              CLEAR_ENTITY_HELPERS_MAP[key]?.(value) ?? value,
+            ]),
+          );
+        }
+      });
+
+      const data = JSON.stringify(subset);
+      // 2 digit precision
+      const currentSize =
+        Math.round(((data.length * BYTES_PER_CHAR) / KB_STORAGE) * 100) / 100;
+      if (state.localStorageUsageInKilobytes !== currentSize) {
+        state.localStorageUsageInKilobytes = currentSize; // eslint-disable-line no-param-reassign
+      }
+
+      return subset;
+    },
+    merge: (initialState, persistedState = {}) => {
+      const result = {
+        ...initialState,
+        ...persistedState,
+        sqlLab: {
+          ...(persistedState?.sqlLab || {}),
+          // Overwrite initialState over persistedState for sqlLab
+          // since a logic in getInitialState overrides the value from persistedState
+          ...initialState.sqlLab,
+        },
+      };
+      return result;
+    },
+  },
+};
+
+export const persistSqlLabStateEnhancer = persistState(
+  sqlLabPersistStateConfig.paths,
+  sqlLabPersistStateConfig.config,
+);

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

@@ -41,7 +41,7 @@ export default function getInitialState({
   tab_state_ids: tabStateIds = [],
   databases,
   queries: queries_,
-  user,
+  ...otherBootstrapData
 }: BootstrapData & Partial<InitialState>) {
   /**
    * Before YYYY-MM-DD, the state for SQL Lab was stored exclusively in the
@@ -205,10 +205,7 @@ export default function getInitialState({
       (common || {})?.flash_messages || [],
     ),
     localStorageUsageInKilobytes: 0,
-    common: {
-      flash_messages: common.flash_messages,
-      conf: common.conf,
-    },
-    user,
+    common,
+    ...otherBootstrapData,
   };
 }

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

@@ -17,7 +17,6 @@
  * under the License.
  */
 import { normalizeTimestamp, QueryState, t } from '@superset-ui/core';
-import getInitialState from './getInitialState';
 import * as actions from '../actions/sqlLab';
 import { now } from '../../utils/dates';
 import {
@@ -165,7 +164,7 @@ export default function sqlLabReducer(state = {}, action) {
       return { ...state, queries: newQueries };
     },
     [actions.RESET_STATE]() {
-      return { ...getInitialState() };
+      return { ...action.sqlLabInitialState };
     },
     [actions.MERGE_TABLE]() {
       const at = { ...action.table };

superset-frontend/src/SqlLab/utils/reduxStateToLocalStorageHelper.js

@@ -17,6 +17,7 @@
  * under the License.
  */
 import pick from 'lodash/pick';
+import { tableApiUtil } from 'src/hooks/apiResources/tables';
 import {
   BYTES_PER_CHAR,
   KB_STORAGE,
@@ -96,3 +97,25 @@ export function clearQueryEditors(queryEditors) {
       ),
   );
 }
+
+export function rehydratePersistedState(dispatch, state) {
+  // Rehydrate server side persisted table metadata
+  state.sqlLab.tables.forEach(({ name: table, schema, dbId, persistData }) => {
+    if (dbId && schema && table && persistData?.columns) {
+      dispatch(
+        tableApiUtil.upsertQueryData(
+          'tableMetadata',
+          { dbId, schema, table },
+          persistData,
+        ),
+      );
+      dispatch(
+        tableApiUtil.upsertQueryData(
+          'tableExtendedMetadata',
+          { dbId, schema, table },
+          {},
+        ),
+      );
+    }
+  });
+}

superset-frontend/src/views/store.ts

@@ -16,7 +16,11 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import { configureStore, ConfigureStoreOptions, Store } from '@reduxjs/toolkit';
+import {
+  configureStore,
+  ConfigureStoreOptions,
+  StoreEnhancer,
+} from '@reduxjs/toolkit';
 import thunk from 'redux-thunk';
 import { api } from 'src/hooks/apiResources/queryApi';
 import messageToastReducer from 'src/components/MessageToasts/reducers';
@@ -34,6 +38,11 @@ import logger from 'src/middleware/loggerMiddleware';
 import saveModal from 'src/explore/reducers/saveModalReducer';
 import explore from 'src/explore/reducers/exploreReducer';
 import exploreDatasources from 'src/explore/reducers/datasourcesReducer';
+import { FeatureFlag, isFeatureEnabled } from '@superset-ui/core';
+
+import { persistSqlLabStateEnhancer } from 'src/SqlLab/middlewares/persistSqlLabStateEnhancer';
+import sqlLabReducer from 'src/SqlLab/reducers/sqlLab';
+import getInitialState from 'src/SqlLab/reducers/getInitialState';
 import { DatasourcesState } from 'src/dashboard/types';
 import {
   DatasourcesActionPayload,
@@ -113,6 +122,8 @@ const CombinedDatasourceReducers = (
 };
 
 const reducers = {
+  sqlLab: sqlLabReducer,
+  localStorageUsage: noopReducer(0),
   messageToasts: messageToastReducer,
   common: noopReducer(bootstrapData.common),
   user: userReducer,
@@ -140,14 +151,14 @@ const reducers = {
  */
 export function setupStore({
   disableDebugger = false,
-  initialState = {},
+  initialState = getInitialState(bootstrapData),
   rootReducers = reducers,
   ...overrides
 }: {
   disableDebugger?: boolean;
   initialState?: ConfigureStoreOptions['preloadedState'];
   rootReducers?: ConfigureStoreOptions['reducer'];
-} & Partial<ConfigureStoreOptions> = {}): Store {
+} & Partial<ConfigureStoreOptions> = {}) {
   return configureStore({
     preloadedState: initialState,
     reducer: {
@@ -156,9 +167,12 @@ export function setupStore({
     },
     middleware: getMiddleware,
     devTools: process.env.WEBPACK_MODE === 'development' && !disableDebugger,
+    ...(!isFeatureEnabled(FeatureFlag.SQLLAB_BACKEND_PERSISTENCE) && {
+      enhancers: [persistSqlLabStateEnhancer as StoreEnhancer],
+    }),
     ...overrides,
   });
 }
 
-export const store: Store = setupStore();
+export const store = setupStore();
 export type RootState = ReturnType<typeof store.getState>;