fix(sqllab): perf regression on #21532 refactor (#21632)

superset-frontend/spec/helpers/testing-library.tsx

@@ -47,7 +47,7 @@ type Options = Omit<RenderOptions, 'queries'> & {
   store?: Store;
 };
 
-function createWrapper(options?: Options) {
+export function createWrapper(options?: Options) {
   const {
     useDnd,
     useRedux,

superset-frontend/src/SqlLab/components/AceEditorWrapper/AceEditorWrapper.test.tsx

@@ -48,7 +48,7 @@ jest.mock('src/components/AsyncAceEditor', () => ({
 const setup = (queryEditor: QueryEditor, store?: Store) =>
   render(
     <AceEditorWrapper
-      queryEditor={queryEditor}
+      queryEditorId={queryEditor.id}
       height="100px"
       hotkeys={[]}
       database={{}}

superset-frontend/src/SqlLab/components/AceEditorWrapper/index.tsx

@@ -37,7 +37,7 @@ import {
   AceCompleterKeyword,
   FullSQLEditor as AceEditor,
 } from 'src/components/AsyncAceEditor';
-import { QueryEditor } from 'src/SqlLab/types';
+import useQueryEditor from 'src/SqlLab/hooks/useQueryEditor';
 
 type HotKey = {
   key: string;
@@ -50,7 +50,7 @@ type AceEditorWrapperProps = {
   autocomplete: boolean;
   onBlur: (sql: string) => void;
   onChange: (sql: string) => void;
-  queryEditor: QueryEditor;
+  queryEditorId: string;
   database: any;
   extendedTables?: Array<{ name: string; columns: any[] }>;
   height: string;
@@ -61,7 +61,7 @@ const AceEditorWrapper = ({
   autocomplete,
   onBlur = () => {},
   onChange = () => {},
-  queryEditor,
+  queryEditorId,
   database,
   extendedTables = [],
   height,
@@ -69,7 +69,17 @@ const AceEditorWrapper = ({
 }: AceEditorWrapperProps) => {
   const dispatch = useDispatch();
 
-  const { sql: currentSql } = queryEditor;
+  const queryEditor = useQueryEditor(queryEditorId, [
+    'id',
+    'dbId',
+    'sql',
+    'functionNames',
+    'schemaOptions',
+    'tableOptions',
+    'validationResult',
+    'schema',
+  ]);
+  const currentSql = queryEditor.sql ?? '';
   const functionNames = queryEditor.functionNames ?? [];
   const schemas = queryEditor.schemaOptions ?? [];
   const tables = queryEditor.tableOptions ?? [];

superset-frontend/src/SqlLab/components/QueryLimitSelect/QueryLimitSelect.test.tsx

@@ -51,7 +51,7 @@ const defaultQueryLimit = 100;
 const setup = (props?: Partial<QueryLimitSelectProps>, store?: Store) =>
   render(
     <QueryLimitSelect
-      queryEditor={defaultQueryEditor}
+      queryEditorId={defaultQueryEditor.id}
       maxRow={100000}
       defaultQueryLimit={defaultQueryLimit}
       {...props}
@@ -67,12 +67,20 @@ describe('QueryLimitSelect', () => {
     const queryLimit = 10;
     const { getByText } = setup(
       {
-        queryEditor: {
+        queryEditorId: defaultQueryEditor.id,
-          ...defaultQueryEditor,
-          queryLimit,
-        },
       },
-      mockStore(initialState),
+      mockStore({
+        ...initialState,
+        sqlLab: {
+          ...initialState.sqlLab,
+          queryEditors: [
+            {
+              ...defaultQueryEditor,
+              queryLimit,
+            },
+          ],
+        },
+      }),
     );
     expect(getByText(queryLimit)).toBeInTheDocument();
   });
@@ -129,7 +137,9 @@ describe('QueryLimitSelect', () => {
         {
           type: 'QUERY_EDITOR_SET_QUERY_LIMIT',
           queryLimit: LIMIT_DROPDOWN[expectedIndex],
-          queryEditor: defaultQueryEditor,
+          queryEditor: {
+            id: defaultQueryEditor.id,
+          },
         },
       ]),
     );

superset-frontend/src/SqlLab/components/QueryLimitSelect/index.tsx

@@ -17,16 +17,16 @@
  * under the License.
  */
 import React from 'react';
-import { useSelector, useDispatch } from 'react-redux';
+import { useDispatch } from 'react-redux';
 import { styled, useTheme } from '@superset-ui/core';
 import { AntdDropdown } from 'src/components';
 import { Menu } from 'src/components/Menu';
 import Icons from 'src/components/Icons';
-import { SqlLabRootState, QueryEditor } from 'src/SqlLab/types';
 import { queryEditorSetQueryLimit } from 'src/SqlLab/actions/sqlLab';
+import useQueryEditor from 'src/SqlLab/hooks/useQueryEditor';
 
 export interface QueryLimitSelectProps {
-  queryEditor: QueryEditor;
+  queryEditorId: string;
   maxRow: number;
   defaultQueryLimit: number;
 }
@@ -79,19 +79,12 @@ function renderQueryLimit(
 }
 
 const QueryLimitSelect = ({
-  queryEditor,
+  queryEditorId,
   maxRow,
   defaultQueryLimit,
 }: QueryLimitSelectProps) => {
-  const queryLimit = useSelector<SqlLabRootState, number>(
+  const queryEditor = useQueryEditor(queryEditorId, ['id', 'queryLimit']);
-    ({ sqlLab: { unsavedQueryEditor } }) => {
+  const queryLimit = queryEditor.queryLimit || defaultQueryLimit;
-      const updatedQueryEditor = {
-        ...queryEditor,
-        ...(unsavedQueryEditor.id === queryEditor.id && unsavedQueryEditor),
-      };
-      return updatedQueryEditor.queryLimit || defaultQueryLimit;
-    },
-  );
   const dispatch = useDispatch();
   const setQueryLimit = (updatedQueryLimit: number) =>
     dispatch(queryEditorSetQueryLimit(queryEditor, updatedQueryLimit));

superset-frontend/src/SqlLab/components/RunQueryActionButton/RunQueryActionButton.test.tsx

@@ -41,13 +41,13 @@ jest.mock('src/components/Select/AsyncSelect', () => () => (
 ));
 
 const defaultProps = {
-  queryEditor: defaultQueryEditor,
+  queryEditorId: defaultQueryEditor.id,
   allowAsync: false,
   dbId: 1,
   queryState: 'ready',
-  runQuery: jest.fn(),
+  runQuery: () => {},
   selectedText: null,
-  stopQuery: jest.fn(),
+  stopQuery: () => {},
   overlayCreateAsMenu: null,
 };
 
@@ -57,95 +57,104 @@ const setup = (props?: Partial<Props>, store?: Store) =>
     ...(store && { store }),
   });
 
-describe('RunQueryActionButton', () => {
+it('renders a single Button', () => {
-  beforeEach(() => {
+  const { getByRole } = setup({}, mockStore(initialState));
-    defaultProps.runQuery.mockReset();
+  expect(getByRole('button')).toBeInTheDocument();
-    defaultProps.stopQuery.mockReset();
+});
-  });
+
-
+it('renders a label for Run Query', () => {
-  it('renders a single Button', () => {
+  const { getByText } = setup({}, mockStore(initialState));
-    const { getByRole } = setup({}, mockStore(initialState));
+  expect(getByText('Run')).toBeInTheDocument();
-    expect(getByRole('button')).toBeInTheDocument();
+});
-  });
+
-
+it('renders a label for Selected Query', () => {
-  it('renders a label for Run Query', () => {
+  const { getByText } = setup(
-    const { getByText } = setup({}, mockStore(initialState));
+    {},
-    expect(getByText('Run')).toBeInTheDocument();
+    mockStore({
-  });
+      ...initialState,
-
+      sqlLab: {
-  it('renders a label for Selected Query', () => {
+        ...initialState.sqlLab,
-    const { getByText } = setup(
+        unsavedQueryEditor: {
-      {},
+          id: defaultQueryEditor.id,
-      mockStore({
+          selectedText: 'select * from\n-- this is comment\nwhere',
-        ...initialState,
+        },
-        sqlLab: {
+      },
-          ...initialState.sqlLab,
+    }),
-          unsavedQueryEditor: {
+  );
-            id: defaultQueryEditor.id,
+  expect(getByText('Run selection')).toBeInTheDocument();
-            selectedText: 'FROM',
+});
-          },
+
-        },
+it('disable button when sql from unsaved changes is empty', () => {
-      }),
+  const { getByRole } = setup(
-    );
+    {},
-    expect(getByText('Run selection')).toBeInTheDocument();
+    mockStore({
-  });
+      ...initialState,
-
+      sqlLab: {
-  it('disable button when sql from unsaved changes is empty', () => {
+        ...initialState.sqlLab,
-    const { getByRole } = setup(
+        unsavedQueryEditor: {
-      {},
+          id: defaultQueryEditor.id,
-      mockStore({
+          sql: '',
-        ...initialState,
+        },
-        sqlLab: {
+      },
-          ...initialState.sqlLab,
+    }),
-          unsavedQueryEditor: {
+  );
-            id: defaultQueryEditor.id,
+  const button = getByRole('button');
-            sql: '',
+  expect(button).toBeDisabled();
-          },
+});
-        },
+
-      }),
+it('disable button when selectedText only contains blank contents', () => {
-    );
+  const { getByRole } = setup(
-    const button = getByRole('button');
+    {},
-    expect(button).toBeDisabled();
+    mockStore({
-  });
+      ...initialState,
-
+      sqlLab: {
-  it('enable default button for unrelated unsaved changes', () => {
+        ...initialState.sqlLab,
-    const { getByRole } = setup(
+        unsavedQueryEditor: {
-      {},
+          id: defaultQueryEditor.id,
-      mockStore({
+          selectedText: '-- this is comment\n\n     \t',
-        ...initialState,
+        },
-        sqlLab: {
+      },
-          ...initialState.sqlLab,
+    }),
-          unsavedQueryEditor: {
+  );
-            id: `${defaultQueryEditor.id}-other`,
+  const button = getByRole('button');
-            sql: '',
+  expect(button).toBeDisabled();
-          },
+});
-        },
+
-      }),
+it('enable default button for unrelated unsaved changes', () => {
-    );
+  const { getByRole } = setup(
-    const button = getByRole('button');
+    {},
-    expect(button).toBeEnabled();
+    mockStore({
-  });
+      ...initialState,
-
+      sqlLab: {
-  it('dispatch runQuery on click', async () => {
+        ...initialState.sqlLab,
-    const { getByRole } = setup({}, mockStore(initialState));
+        unsavedQueryEditor: {
-    const button = getByRole('button');
+          id: `${defaultQueryEditor.id}-other`,
-    expect(defaultProps.runQuery).toHaveBeenCalledTimes(0);
+          sql: '',
-    fireEvent.click(button);
+        },
-    await waitFor(() => expect(defaultProps.runQuery).toHaveBeenCalledTimes(1));
+      },
-  });
+    }),
-
+  );
-  describe('on running state', () => {
+  const button = getByRole('button');
-    it('dispatch stopQuery on click', async () => {
+  expect(button).toBeEnabled();
-      const { getByRole } = setup(
+});
-        { queryState: 'running' },
+
-        mockStore(initialState),
+it('dispatch runQuery on click', async () => {
-      );
+  const runQuery = jest.fn();
-      const button = getByRole('button');
+  const { getByRole } = setup({ runQuery }, mockStore(initialState));
-      expect(defaultProps.stopQuery).toHaveBeenCalledTimes(0);
+  const button = getByRole('button');
-      fireEvent.click(button);
+  expect(runQuery).toHaveBeenCalledTimes(0);
-      await waitFor(() =>
+  fireEvent.click(button);
-        expect(defaultProps.stopQuery).toHaveBeenCalledTimes(1),
+  await waitFor(() => expect(runQuery).toHaveBeenCalledTimes(1));
-      );
+});
-    });
+
-  });
+it('dispatch stopQuery on click while running state', async () => {
+  const stopQuery = jest.fn();
+  const { getByRole } = setup(
+    { queryState: 'running', stopQuery },
+    mockStore(initialState),
+  );
+  const button = getByRole('button');
+  expect(stopQuery).toHaveBeenCalledTimes(0);
+  fireEvent.click(button);
+  await waitFor(() => expect(stopQuery).toHaveBeenCalledTimes(1));
 });

superset-frontend/src/SqlLab/components/RunQueryActionButton/index.tsx

@@ -24,16 +24,11 @@ import Button from 'src/components/Button';
 import Icons from 'src/components/Icons';
 import { DropdownButton } from 'src/components/DropdownButton';
 import { detectOS } from 'src/utils/common';
-import { shallowEqual, useSelector } from 'react-redux';
+import { QueryButtonProps } from 'src/SqlLab/types';
-import {
+import useQueryEditor from 'src/SqlLab/hooks/useQueryEditor';
-  QueryEditor,
-  SqlLabRootState,
-  QueryButtonProps,
-} from 'src/SqlLab/types';
-import { getUpToDateQuery } from 'src/SqlLab/actions/sqlLab';
 
 export interface Props {
-  queryEditor: QueryEditor;
+  queryEditorId: string;
   allowAsync: boolean;
   queryState?: string;
   runQuery: (c?: boolean) => void;
@@ -86,29 +81,21 @@ const StyledButton = styled.span`
   }
 `;
 
-const RunQueryActionButton = ({
+const RunQueryActionButton: React.FC<Props> = ({
   allowAsync = false,
-  queryEditor,
+  queryEditorId,
   queryState,
   overlayCreateAsMenu,
   runQuery,
   stopQuery,
-}: Props) => {
+}) => {
   const theme = useTheme();
   const userOS = detectOS();
-  const { selectedText, sql } = useSelector<
+
-    SqlLabRootState,
+  const { selectedText, sql } = useQueryEditor(queryEditorId, [
-    Pick<QueryEditor, 'selectedText' | 'sql'>
+    'selectedText',
-  >(rootState => {
+    'sql',
-    const currentQueryEditor = getUpToDateQuery(
+  ]);
-      rootState,
-      queryEditor,
-    ) as unknown as QueryEditor;
-    return {
-      selectedText: currentQueryEditor.selectedText,
-      sql: currentQueryEditor.sql,
-    };
-  }, shallowEqual);
 
   const shouldShowStopBtn =
     !!queryState && ['running', 'pending'].indexOf(queryState) > -1;
@@ -117,7 +104,10 @@ const RunQueryActionButton = ({
     ? (DropdownButton as React.FC)
     : Button;
 
-  const isDisabled = !sql || !sql.trim();
+  const sqlContent = selectedText || sql || '';
+  const isDisabled =
+    !sqlContent ||
+    !sqlContent.replace(/(\/\*[^*]*\*\/)|(\/\/[^*]*)|(--[^.].*)/gm, '').trim();
 
   const stopButtonTooltipText = useMemo(
     () =>

superset-frontend/src/SqlLab/components/SaveQuery/SaveQuery.test.tsx

@@ -25,15 +25,28 @@ import SaveQuery from 'src/SqlLab/components/SaveQuery';
 import { initialState, databases } from 'src/SqlLab/fixtures';
 
 const mockedProps = {
-  queryEditor: {
+  queryEditorId: '123',
-    dbId: 1,
-    schema: 'main',
-    sql: 'SELECT * FROM t',
-  },
   animation: false,
   database: databases.result[0],
   onUpdate: () => {},
   onSave: () => {},
+  saveQueryWarning: null,
+  columns: [],
+};
+
+const mockState = {
+  ...initialState,
+  sqlLab: {
+    ...initialState.sqlLab,
+    queryEditors: [
+      {
+        id: mockedProps.queryEditorId,
+        dbId: 1,
+        schema: 'main',
+        sql: 'SELECT * FROM t',
+      },
+    ],
+  },
 };
 
 const splitSaveBtnProps = {
@@ -51,7 +64,7 @@ describe('SavedQuery', () => {
   it('renders a non-split save button when allows_virtual_table_explore is not enabled', () => {
     render(<SaveQuery {...mockedProps} />, {
       useRedux: true,
-      store: mockStore(initialState),
+      store: mockStore(mockState),
     });
 
     const saveBtn = screen.getByRole('button', { name: /save/i });
@@ -62,7 +75,7 @@ describe('SavedQuery', () => {
   it('renders a save query modal when user clicks save button', () => {
     render(<SaveQuery {...mockedProps} />, {
       useRedux: true,
-      store: mockStore(initialState),
+      store: mockStore(mockState),
     });
 
     const saveBtn = screen.getByRole('button', { name: /save/i });
@@ -78,7 +91,7 @@ describe('SavedQuery', () => {
   it('renders the save query modal UI', () => {
     render(<SaveQuery {...mockedProps} />, {
       useRedux: true,
-      store: mockStore(initialState),
+      store: mockStore(mockState),
     });
 
     const saveBtn = screen.getByRole('button', { name: /save/i });
@@ -111,16 +124,18 @@ describe('SavedQuery', () => {
   });
 
   it('renders a "save as new" and "update" button if query already exists', () => {
-    const props = {
+    render(<SaveQuery {...mockedProps} />, {
-      ...mockedProps,
-      queryEditor: {
-        ...mockedProps.query,
-        remoteId: '42',
-      },
-    };
-    render(<SaveQuery {...props} />, {
       useRedux: true,
-      store: mockStore(initialState),
+      store: mockStore({
+        ...mockState,
+        sqlLab: {
+          ...mockState.sqlLab,
+          unsavedQueryEditor: {
+            id: mockedProps.queryEditorId,
+            remoteId: '42',
+          },
+        },
+      }),
     });
 
     const saveBtn = screen.getByRole('button', { name: /save/i });
@@ -136,7 +151,7 @@ describe('SavedQuery', () => {
   it('renders a split save button when allows_virtual_table_explore is enabled', async () => {
     render(<SaveQuery {...splitSaveBtnProps} />, {
       useRedux: true,
-      store: mockStore(initialState),
+      store: mockStore(mockState),
     });
 
     await waitFor(() => {
@@ -151,7 +166,7 @@ describe('SavedQuery', () => {
   it('renders a save dataset modal when user clicks "save dataset" menu item', async () => {
     render(<SaveQuery {...splitSaveBtnProps} />, {
       useRedux: true,
-      store: mockStore(initialState),
+      store: mockStore(mockState),
     });
 
     await waitFor(() => {
@@ -170,7 +185,7 @@ describe('SavedQuery', () => {
   it('renders the save dataset modal UI', async () => {
     render(<SaveQuery {...splitSaveBtnProps} />, {
       useRedux: true,
-      store: mockStore(initialState),
+      store: mockStore(mockState),
     });
 
     await waitFor(() => {

superset-frontend/src/SqlLab/components/SaveQuery/index.tsx

@@ -16,8 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { useState, useEffect } from 'react';
+import React, { useState, useEffect, useMemo } from 'react';
-import { useSelector, shallowEqual } from 'react-redux';
 import { Row, Col } from 'src/components';
 import { Input, TextArea } from 'src/components/Input';
 import { t, styled } from '@superset-ui/core';
@@ -31,10 +30,11 @@ import {
   ISaveableDatasource,
 } from 'src/SqlLab/components/SaveDatasetModal';
 import { getDatasourceAsSaveableDataset } from 'src/utils/datasourceUtils';
-import { QueryEditor, SqlLabRootState } from 'src/SqlLab/types';
+import useQueryEditor from 'src/SqlLab/hooks/useQueryEditor';
+import { QueryEditor } from 'src/SqlLab/types';
 
 interface SaveQueryProps {
-  queryEditor: QueryEditor;
+  queryEditorId: string;
   columns: ISaveableDatasource['columns'];
   onSave: (arg0: QueryPayload) => void;
   onUpdate: (arg0: QueryPayload) => void;
@@ -43,30 +43,22 @@ interface SaveQueryProps {
 }
 
 type QueryPayload = {
-  autorun: boolean;
+  name: string;
-  dbId: number;
   description?: string;
   id?: string;
-  latestQueryId: string;
+} & Pick<
-  queryLimit: number;
+  QueryEditor,
-  remoteId: number;
+  | 'autorun'
-  schema: string;
+  | 'dbId'
-  schemaOptions: Array<{
+  | 'schema'
-    label: string;
+  | 'sql'
-    title: string;
+  | 'selectedText'
-    value: string;
+  | 'remoteId'
-  }>;
+  | 'latestQueryId'
-  selectedText: string | null;
+  | 'queryLimit'
-  sql: string;
+  | 'tableOptions'
-  tableOptions: Array<{
+  | 'schemaOptions'
-    label: string;
+>;
-    schema: string;
-    title: string;
-    type: string;
-    value: string;
-  }>;
-  name: string;
-};
 
 const Styles = styled.span`
   span[role='img'] {
@@ -81,20 +73,33 @@ const Styles = styled.span`
 `;
 
 export default function SaveQuery({
-  queryEditor,
+  queryEditorId,
   onSave = () => {},
   onUpdate,
   saveQueryWarning = null,
   database,
   columns,
 }: SaveQueryProps) {
-  const query = useSelector<SqlLabRootState, QueryEditor>(
+  const queryEditor = useQueryEditor(queryEditorId, [
-    ({ sqlLab: { unsavedQueryEditor } }) => ({
+    'autorun',
+    'name',
+    'description',
+    'remoteId',
+    'dbId',
+    'latestQueryId',
+    'queryLimit',
+    'schema',
+    'schemaOptions',
+    'selectedText',
+    'sql',
+    'tableOptions',
+  ]);
+  const query = useMemo(
+    () => ({
       ...queryEditor,
-      ...(queryEditor.id === unsavedQueryEditor.id && unsavedQueryEditor),
       columns,
     }),
-    shallowEqual,
+    [queryEditor, columns],
   );
   const defaultLabel = query.name || query.description || t('Undefined');
   const [description, setDescription] = useState<string>(
@@ -114,12 +119,12 @@ export default function SaveQuery({
     </Menu>
   );
 
-  const queryPayload = () =>
+  const queryPayload = () => ({
-    ({
+    ...query,
-      ...query,
+    name: label,
-      name: label,
+    description,
-      description,
+    dbId: query.dbId ?? 0,
-    } as any as QueryPayload);
+  });
 
   useEffect(() => {
     if (!isSaved) setLabel(defaultLabel);

superset-frontend/src/SqlLab/components/ShareSqlLabQuery/ShareSqlLabQuery.test.tsx

@@ -31,30 +31,42 @@ import ShareSqlLabQuery from 'src/SqlLab/components/ShareSqlLabQuery';
 import { initialState } from 'src/SqlLab/fixtures';
 
 const mockStore = configureStore([thunk]);
-const store = mockStore(initialState);
+const defaultProps = {
-let isFeatureEnabledMock;
+  queryEditorId: 'qe1',
+  addDangerToast: jest.fn(),
+};
+const mockQueryEditor = {
+  id: defaultProps.queryEditorId,
+  dbId: 0,
+  name: 'query title',
+  schema: 'query_schema',
+  autorun: false,
+  sql: 'SELECT * FROM ...',
+  remoteId: 999,
+};
+const disabled = {
+  id: 'disabledEditorId',
+  remoteId: undefined,
+};
 
-const standardProvider = ({ children }) => (
+const mockState = {
+  ...initialState,
+  sqlLab: {
+    ...initialState.sqlLab,
+    queryEditors: [mockQueryEditor, disabled],
+  },
+};
+const store = mockStore(mockState);
+let isFeatureEnabledMock: jest.SpyInstance;
+
+const standardProvider: React.FC = ({ children }) => (
   <ThemeProvider theme={supersetTheme}>
     <Provider store={store}>{children}</Provider>
   </ThemeProvider>
 );
 
-const defaultProps = {
-  queryEditor: {
-    id: 'qe1',
-    dbId: 0,
-    name: 'query title',
-    schema: 'query_schema',
-    autorun: false,
-    sql: 'SELECT * FROM ...',
-    remoteId: 999,
-  },
-  addDangerToast: jest.fn(),
-};
-
 const unsavedQueryEditor = {
-  id: defaultProps.queryEditor.id,
+  id: defaultProps.queryEditorId,
   dbId: 9888,
   name: 'query title changed',
   schema: 'query_schema_updated',
@@ -62,7 +74,7 @@ const unsavedQueryEditor = {
   autorun: true,
 };
 
-const standardProviderWithUnsaved = ({ children }) => (
+const standardProviderWithUnsaved: React.FC = ({ children }) => (
   <ThemeProvider theme={supersetTheme}>
     <Provider
       store={mockStore({
@@ -100,7 +112,7 @@ describe('ShareSqlLabQuery', () => {
     });
 
     afterAll(() => {
-      isFeatureEnabledMock.restore();
+      isFeatureEnabledMock.mockReset();
     });
 
     it('calls storeQuery() with the query when getCopyUrl() is called', async () => {
@@ -110,7 +122,7 @@ describe('ShareSqlLabQuery', () => {
         });
       });
       const button = screen.getByRole('button');
-      const { id, remoteId, ...expected } = defaultProps.queryEditor;
+      const { id, remoteId, ...expected } = mockQueryEditor;
       const storeQuerySpy = jest.spyOn(utils, 'storeQuery');
       userEvent.click(button);
       expect(storeQuerySpy.mock.calls).toHaveLength(1);
@@ -142,7 +154,7 @@ describe('ShareSqlLabQuery', () => {
     });
 
     afterAll(() => {
-      isFeatureEnabledMock.restore();
+      isFeatureEnabledMock.mockReset();
     });
 
     it('does not call storeQuery() with the query when getCopyUrl() is called and feature is not enabled', async () => {
@@ -160,10 +172,7 @@ describe('ShareSqlLabQuery', () => {
 
     it('button is disabled and there is a request to save the query', async () => {
       const updatedProps = {
-        queryEditor: {
+        queryEditorId: disabled.id,
-          ...defaultProps.queryEditor,
-          remoteId: undefined,
-        },
       };
 
       render(<ShareSqlLabQuery {...updatedProps} />, {

superset-frontend/src/SqlLab/components/ShareSqlLabQuery/index.tsx

@@ -17,7 +17,6 @@
  * under the License.
  */
 import React from 'react';
-import { shallowEqual, useSelector } from 'react-redux';
 import { t, useTheme, styled } from '@superset-ui/core';
 import Button from 'src/components/Button';
 import Icons from 'src/components/Icons';
@@ -26,10 +25,10 @@ import CopyToClipboard from 'src/components/CopyToClipboard';
 import { storeQuery } from 'src/utils/common';
 import { getClientErrorObject } from 'src/utils/getClientErrorObject';
 import { FeatureFlag, isFeatureEnabled } from 'src/featureFlags';
-import { QueryEditor, SqlLabRootState } from 'src/SqlLab/types';
+import useQueryEditor from 'src/SqlLab/hooks/useQueryEditor';
 
 interface ShareSqlLabQueryPropTypes {
-  queryEditor: QueryEditor;
+  queryEditorId: string;
   addDangerToast: (msg: string) => void;
 }
 
@@ -44,21 +43,15 @@ const StyledIcon = styled(Icons.Link)`
 `;
 
 function ShareSqlLabQuery({
-  queryEditor,
+  queryEditorId,
   addDangerToast,
 }: ShareSqlLabQueryPropTypes) {
   const theme = useTheme();
 
-  const { dbId, name, schema, autorun, sql, remoteId } = useSelector<
+  const { dbId, name, schema, autorun, sql, remoteId } = useQueryEditor(
-    SqlLabRootState,
+    queryEditorId,
-    Partial<QueryEditor>
+    ['dbId', 'name', 'schema', 'autorun', 'sql', 'remoteId'],
-  >(({ sqlLab: { unsavedQueryEditor } }) => {
+  );
-    const { dbId, name, schema, autorun, sql, remoteId } = {
-      ...queryEditor,
-      ...(unsavedQueryEditor.id === queryEditor.id && unsavedQueryEditor),
-    };
-    return { dbId, name, schema, autorun, sql, remoteId };
-  }, shallowEqual);
 
   const getCopyUrlForKvStore = (callback: Function) => {
     const sharedQuery = { dbId, name, schema, autorun, sql };

superset-frontend/src/SqlLab/components/SqlEditor/index.jsx

@@ -163,13 +163,8 @@ const SqlEditor = ({
   const theme = useTheme();
   const dispatch = useDispatch();
 
-  const { currentQueryEditor, database, latestQuery, hideLeftBar } =
+  const { database, latestQuery, hideLeftBar } = useSelector(
-    useSelector(({ sqlLab: { unsavedQueryEditor, databases, queries } }) => {
+    ({ sqlLab: { unsavedQueryEditor, databases, queries } }) => {
-      const currentQueryEditor = {
-        ...queryEditor,
-        ...(queryEditor.id === unsavedQueryEditor.id && unsavedQueryEditor),
-      };
-
       let { dbId, latestQueryId, hideLeftBar } = queryEditor;
       if (unsavedQueryEditor.id === queryEditor.id) {
         dbId = unsavedQueryEditor.dbId || dbId;
@@ -177,12 +172,12 @@ const SqlEditor = ({
         hideLeftBar = unsavedQueryEditor.hideLeftBar || hideLeftBar;
       }
       return {
-        currentQueryEditor,
         database: databases[dbId],
         latestQuery: queries[latestQueryId],
         hideLeftBar,
       };
-    });
+    },
+  );
 
   const queryEditors = useSelector(({ sqlLab }) => sqlLab.queryEditors);
 
@@ -540,7 +535,7 @@ const SqlEditor = ({
           <span>
             <RunQueryActionButton
               allowAsync={database ? database.allow_run_async : false}
-              queryEditor={queryEditor}
+              queryEditorId={queryEditor.id}
               queryState={latestQuery?.state}
               runQuery={startQuery}
               stopQuery={stopQuery}
@@ -559,7 +554,7 @@ const SqlEditor = ({
             )}
           <span>
             <QueryLimitSelect
-              queryEditor={queryEditor}
+              queryEditorId={queryEditor.id}
               maxRow={maxRow}
               defaultQueryLimit={defaultQueryLimit}
             />
@@ -576,7 +571,7 @@ const SqlEditor = ({
         <div className="rightItems">
           <span>
             <SaveQuery
-              queryEditor={queryEditor}
+              queryEditorId={queryEditor.id}
               columns={latestQuery?.results?.columns || []}
               onSave={onSaveQuery}
               onUpdate={query => dispatch(updateSavedQuery(query))}
@@ -585,7 +580,7 @@ const SqlEditor = ({
             />
           </span>
           <span>
-            <ShareSqlLabQuery queryEditor={queryEditor} />
+            <ShareSqlLabQuery queryEditorId={queryEditor.id} />
           </span>
           <AntdDropdown overlay={renderDropdown()} trigger="click">
             <Icons.MoreHoriz iconColor={theme.colors.grayscale.base} />
@@ -616,7 +611,7 @@ const SqlEditor = ({
             autocomplete={autocompleteEnabled}
             onBlur={setQueryEditorAndSaveSql}
             onChange={onSqlChanged}
-            queryEditor={currentQueryEditor}
+            queryEditorId={queryEditor.id}
             database={database}
             extendedTables={tables}
             height={`${aceEditorHeight}px`}

superset-frontend/src/SqlLab/components/SqlEditorLeftBar/index.tsx

@@ -25,7 +25,6 @@ import React, {
   Dispatch,
   SetStateAction,
 } from 'react';
-import { useSelector } from 'react-redux';
 import querystring from 'query-string';
 import Button from 'src/components/Button';
 import { t, styled, css, SupersetTheme } from '@superset-ui/core';
@@ -33,7 +32,8 @@ import Collapse from 'src/components/Collapse';
 import Icons from 'src/components/Icons';
 import { TableSelectorMultiple } from 'src/components/TableSelector';
 import { IconTooltip } from 'src/components/IconTooltip';
-import { QueryEditor, SchemaOption, SqlLabRootState } from 'src/SqlLab/types';
+import { QueryEditor, SchemaOption } from 'src/SqlLab/types';
+import useQueryEditor from 'src/SqlLab/hooks/useQueryEditor';
 import { DatabaseObject } from 'src/components/DatabaseSelector';
 import { EmptyStateSmall } from 'src/components/EmptyState';
 import {
@@ -117,15 +117,7 @@ export default function SqlEditorLeftBar({
   const [userSelectedDb, setUserSelected] = useState<DatabaseObject | null>(
     null,
   );
-  const schema = useSelector<SqlLabRootState, string>(
+  const { schema } = useQueryEditor(queryEditor.id, ['schema']);
-    ({ sqlLab: { unsavedQueryEditor } }) => {
-      const updatedQueryEditor = {
-        ...queryEditor,
-        ...(unsavedQueryEditor.id === queryEditor.id && unsavedQueryEditor),
-      };
-      return updatedQueryEditor.schema;
-    },
-  );
 
   useEffect(() => {
     const bool = querystring.parse(window.location.search).db;

superset-frontend/src/SqlLab/hooks/useQueryEditor/index.ts

@@ -0,0 +1,38 @@
+/**
+ * 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.
+ */
+import pick from 'lodash/pick';
+import { shallowEqual, useSelector } from 'react-redux';
+import { SqlLabRootState, QueryEditor } from 'src/SqlLab/types';
+
+export default function useQueryEditor<T extends keyof QueryEditor>(
+  sqlEditorId: string,
+  attributes: ReadonlyArray<T>,
+) {
+  return useSelector<SqlLabRootState, Pick<QueryEditor, T | 'id'>>(
+    ({ sqlLab: { unsavedQueryEditor, queryEditors } }) =>
+      pick(
+        {
+          ...queryEditors.find(({ id }) => id === sqlEditorId),
+          ...(sqlEditorId === unsavedQueryEditor.id && unsavedQueryEditor),
+        },
+        ['id'].concat(attributes),
+      ) as Pick<QueryEditor, T | 'id'>,
+    shallowEqual,
+  );
+}

superset-frontend/src/SqlLab/hooks/useQueryEditor/useQueryEditor.test.ts

@@ -0,0 +1,92 @@
+/**
+ * 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.
+ */
+import configureStore from 'redux-mock-store';
+import thunk from 'redux-thunk';
+import { initialState, defaultQueryEditor } from 'src/SqlLab/fixtures';
+import { renderHook } from '@testing-library/react-hooks';
+import { createWrapper } from 'spec/helpers/testing-library';
+
+import useQueryEditor from '.';
+
+const middlewares = [thunk];
+const mockStore = configureStore(middlewares);
+
+test('returns selected queryEditor values', () => {
+  const { result } = renderHook(
+    () =>
+      useQueryEditor(defaultQueryEditor.id, [
+        'id',
+        'name',
+        'dbId',
+        'schemaOptions',
+      ]),
+    {
+      wrapper: createWrapper({
+        useRedux: true,
+        store: mockStore(initialState),
+      }),
+    },
+  );
+  expect(result.current).toEqual({
+    id: defaultQueryEditor.id,
+    name: defaultQueryEditor.name,
+    dbId: defaultQueryEditor.dbId,
+    schemaOptions: defaultQueryEditor.schemaOptions,
+  });
+});
+
+test('includes id implicitly', () => {
+  const { result } = renderHook(
+    () => useQueryEditor(defaultQueryEditor.id, ['name']),
+    {
+      wrapper: createWrapper({
+        useRedux: true,
+        store: mockStore(initialState),
+      }),
+    },
+  );
+  expect(result.current).toEqual({
+    id: defaultQueryEditor.id,
+    name: defaultQueryEditor.name,
+  });
+});
+
+test('returns updated values from unsaved change', () => {
+  const expectedSql = 'SELECT updated_column\nFROM updated_table\nWHERE';
+  const { result } = renderHook(
+    () => useQueryEditor(defaultQueryEditor.id, ['id', 'sql']),
+    {
+      wrapper: createWrapper({
+        useRedux: true,
+        store: mockStore({
+          ...initialState,
+          sqlLab: {
+            ...initialState.sqlLab,
+            unsavedQueryEditor: {
+              id: defaultQueryEditor.id,
+              sql: expectedSql,
+            },
+          },
+        }),
+      }),
+    },
+  );
+  expect(result.current.id).toEqual(defaultQueryEditor.id);
+  expect(result.current.sql).toEqual(expectedSql);
+});

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

@@ -449,7 +449,7 @@ export default function sqlLabReducer(state = {}, action) {
         );
         return {
           ...(action.queryEditor.id === state.unsavedQueryEditor.id
-            ? alterInObject(
+            ? alterInArr(
                 mergeUnsavedState,
                 'queryEditors',
                 action.queryEditor,