fix(sqllab): invalid table metadata request (#21304)

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

@@ -1137,8 +1137,9 @@ function getTableExtendedMetadata(table, query, dispatch) {
     );
 }
 
-export function addTable(query, database, tableName, schemaName) {
-  return function (dispatch) {
+export function addTable(queryEditor, database, tableName, schemaName) {
+  return function (dispatch, getState) {
+    const query = getUpToDateQuery(getState(), queryEditor, queryEditor.id);
     const table = {
       dbId: query.dbId,
       queryEditorId: query.id,

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

@@ -736,29 +736,71 @@ describe('async actions', () => {
         const database = { disable_data_preview: true };
         const tableName = 'table';
         const schemaName = 'schema';
-        const store = mockStore({});
+        const store = mockStore(initialState);
         const expectedActionTypes = [
           actions.MERGE_TABLE, // addTable
           actions.MERGE_TABLE, // getTableMetadata
           actions.MERGE_TABLE, // getTableExtendedMetadata
           actions.MERGE_TABLE, // addTable
         ];
-        return store
-          .dispatch(actions.addTable(query, database, tableName, schemaName))
-          .then(() => {
-            expect(store.getActions().map(a => a.type)).toEqual(
-              expectedActionTypes,
-            );
-            expect(store.getActions()[0].prepend).toBeTruthy();
-            expect(fetchMock.calls(updateTableSchemaEndpoint)).toHaveLength(1);
-            expect(fetchMock.calls(getTableMetadataEndpoint)).toHaveLength(1);
-            expect(fetchMock.calls(getExtraTableMetadataEndpoint)).toHaveLength(
-              1,
-            );
+        const request = actions.addTable(
+          query,
+          database,
+          tableName,
+          schemaName,
+        );
+        return request(store.dispatch, store.getState).then(() => {
+          expect(store.getActions().map(a => a.type)).toEqual(
+            expectedActionTypes,
+          );
+          expect(store.getActions()[0].prepend).toBeTruthy();
+          expect(fetchMock.calls(updateTableSchemaEndpoint)).toHaveLength(1);
+          expect(fetchMock.calls(getTableMetadataEndpoint)).toHaveLength(1);
+          expect(fetchMock.calls(getExtraTableMetadataEndpoint)).toHaveLength(
+            1,
+          );
 
-            // tab state is not updated, since no query was run
-            expect(fetchMock.calls(updateTabStateEndpoint)).toHaveLength(0);
-          });
+          // tab state is not updated, since no query was run
+          expect(fetchMock.calls(updateTabStateEndpoint)).toHaveLength(0);
+        });
+      });
+
+      it('fetches table schema state from unsaved change', () => {
+        const database = { disable_data_preview: true };
+        const tableName = 'table';
+        const schemaName = 'schema';
+        const expectedDbId = 473892;
+        const store = mockStore({
+          ...initialState,
+          sqlLab: {
+            ...initialState.sqlLab,
+            unsavedQueryEditor: {
+              id: query.id,
+              dbId: expectedDbId,
+            },
+          },
+        });
+        const request = actions.addTable(
+          query,
+          database,
+          tableName,
+          schemaName,
+        );
+        return request(store.dispatch, store.getState).then(() => {
+          expect(
+            fetchMock.calls(
+              `glob:**/api/v1/database/${expectedDbId}/table/*/*/`,
+            ),
+          ).toHaveLength(1);
+          expect(
+            fetchMock.calls(
+              `glob:**/api/v1/database/${expectedDbId}/table_extra/*/*/`,
+            ),
+          ).toHaveLength(1);
+
+          // tab state is not updated, since no query was run
+          expect(fetchMock.calls(updateTabStateEndpoint)).toHaveLength(0);
+        });
       });
 
       it('updates and runs data preview query when configured', () => {

superset-frontend/src/SqlLab/components/SqlEditorLeftBar/SqlEditorLeftBar.test.jsx

@@ -19,14 +19,12 @@
 import React from 'react';
 import configureStore from 'redux-mock-store';
 import fetchMock from 'fetch-mock';
-import { shallow } from 'enzyme';
 import { render, screen } from 'spec/helpers/testing-library';
 import userEvent from '@testing-library/user-event';
 import { Provider } from 'react-redux';
 import '@testing-library/jest-dom/extend-expect';
 import thunk from 'redux-thunk';
 import SqlEditorLeftBar from 'src/SqlLab/components/SqlEditorLeftBar';
-import TableElement from 'src/SqlLab/components/TableElement';
 import { supersetTheme, ThemeProvider } from '@superset-ui/core';
 import {
   table,
@@ -46,74 +44,86 @@ const mockedProps = {
 const middlewares = [thunk];
 const mockStore = configureStore(middlewares);
 const store = mockStore(initialState);
-fetchMock.get('glob:*/api/v1/database/*/schemas/?*', { result: [] });
-describe('SqlEditorLeftBar', () => {
-  let wrapper;
-
-  beforeEach(() => {
-    wrapper = shallow(<SqlEditorLeftBar {...mockedProps} />, {
-      context: { store },
-    });
-  });
 
-  afterEach(() => {
-    wrapper.unmount();
-  });
+fetchMock.get('glob:*/api/v1/database/*/schemas/?*', { result: [] });
+fetchMock.get('glob:*/superset/tables/**', {
+  json: {
+    options: [
+      {
+        label: 'ab_user',
+        value: 'ab_user',
+      },
+    ],
+    tableLength: 1,
+  },
+});
 
+describe('Left Panel Expansion', () => {
   it('is valid', () => {
-    expect(React.isValidElement(<SqlEditorLeftBar {...mockedProps} />)).toBe(
-      true,
-    );
+    expect(
+      React.isValidElement(
+        <Provider store={store}>
+          <SqlEditorLeftBar {...mockedProps} />
+        </Provider>,
+      ),
+    ).toBe(true);
   });
 
   it('renders a TableElement', () => {
-    expect(wrapper.find(TableElement)).toExist();
-  });
-});
-
-describe('Left Panel Expansion', () => {
-  it('table should be visible when expanded is true', () => {
-    const { container } = render(
+    const { queryAllByTestId } = render(
       <ThemeProvider theme={supersetTheme}>
         <Provider store={store}>
           <SqlEditorLeftBar {...mockedProps} />
         </Provider>
       </ThemeProvider>,
     );
-    const dbSelect = screen.getByRole('combobox', {
-      name: 'Select database or type database name',
-    });
-    const schemaSelect = screen.getByRole('combobox', {
-      name: 'Select schema or type schema name',
-    });
-    const dropdown = screen.getByText(/Select table or type table name/i);
-    const abUser = screen.getByText(/ab_user/i);
-    expect(dbSelect).toBeInTheDocument();
-    expect(schemaSelect).toBeInTheDocument();
-    expect(dropdown).toBeInTheDocument();
-    expect(abUser).toBeInTheDocument();
-    expect(
-      container.querySelector('.ant-collapse-content-active'),
-    ).toBeInTheDocument();
+    expect(queryAllByTestId('table-element').length).toBeGreaterThanOrEqual(1);
   });
 
-  it('should toggle the table when the header is clicked', async () => {
-    const collapseMock = jest.fn();
-    render(
-      <ThemeProvider theme={supersetTheme}>
-        <Provider store={store}>
-          <SqlEditorLeftBar
-            actions={{ ...mockedActions, collapseTable: collapseMock }}
-            tables={[table]}
-            queryEditor={defaultQueryEditor}
-            database={databases}
-            height={0}
-          />
-        </Provider>
-      </ThemeProvider>,
-    );
-    const header = screen.getByText(/ab_user/);
-    userEvent.click(header);
-    expect(collapseMock).toHaveBeenCalled();
+  describe('Left Panel Expansion', () => {
+    it('table should be visible when expanded is true', () => {
+      const { container } = render(
+        <ThemeProvider theme={supersetTheme}>
+          <Provider store={store}>
+            <SqlEditorLeftBar {...mockedProps} />
+          </Provider>
+        </ThemeProvider>,
+      );
+      const dbSelect = screen.getByRole('combobox', {
+        name: 'Select database or type database name',
+      });
+      const schemaSelect = screen.getByRole('combobox', {
+        name: 'Select schema or type schema name',
+      });
+      const dropdown = screen.getByText(/Select table or type table name/i);
+      const abUser = screen.getByText(/ab_user/i);
+      expect(dbSelect).toBeInTheDocument();
+      expect(schemaSelect).toBeInTheDocument();
+      expect(dropdown).toBeInTheDocument();
+      expect(abUser).toBeInTheDocument();
+      expect(
+        container.querySelector('.ant-collapse-content-active'),
+      ).toBeInTheDocument();
+    });
+
+    it('should toggle the table when the header is clicked', async () => {
+      const collapseMock = jest.fn();
+      render(
+        <ThemeProvider theme={supersetTheme}>
+          <Provider store={store}>
+            <SqlEditorLeftBar
+              actions={{ ...mockedActions, collapseTable: collapseMock }}
+              tables={[table]}
+              queryEditor={defaultQueryEditor}
+              database={databases}
+              height={0}
+            />
+          </Provider>
+        </ThemeProvider>,
+      );
+      const header = screen.getByText(/ab_user/);
+      userEvent.click(header);
+      expect(collapseMock).toHaveBeenCalled();
+    });
   });
 });

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

@@ -25,14 +25,15 @@ 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';
 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 } from 'src/SqlLab/types';
+import { QueryEditor, SchemaOption, SqlLabRootState } from 'src/SqlLab/types';
 import { DatabaseObject } from 'src/components/DatabaseSelector';
 import { EmptyStateSmall } from 'src/components/EmptyState';
 import {
@@ -116,6 +117,15 @@ export default function SqlEditorLeftBar({
   const [userSelectedDb, setUserSelected] = useState<DatabaseObject | null>(
     null,
   );
+  const schema = useSelector<SqlLabRootState, string>(
+    ({ sqlLab: { unsavedQueryEditor } }) => {
+      const updatedQueryEditor = {
+        ...queryEditor,
+        ...(unsavedQueryEditor.id === queryEditor.id && unsavedQueryEditor),
+      };
+      return updatedQueryEditor.schema;
+    },
+  );
 
   useEffect(() => {
     const bool = querystring.parse(window.location.search).db;
@@ -263,7 +273,7 @@ export default function SqlEditorLeftBar({
         onSchemasLoad={handleSchemasLoad}
         onTableSelectChange={onTablesChange}
         onTablesLoad={handleTablesLoad}
-        schema={queryEditor.schema}
+        schema={schema}
         tableValue={selectedTableNames}
         sqlLabMode
       />

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

@@ -262,6 +262,9 @@ class TabbedSqlEditors extends React.PureComponent {
     const qeid = this.props.tabHistory[this.props.tabHistory.length - 1];
     if (key !== qeid) {
       const queryEditor = this.props.queryEditors.find(qe => qe.id === key);
+      if (!queryEditor) {
+        return;
+      }
       this.props.actions.switchQueryEditor(
         queryEditor,
         this.props.displayLimit,

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

@@ -270,6 +270,7 @@ const TableElement = ({ table, actions, ...props }: TableElementProps) => {
 
     const metadata = (
       <div
+        data-test="table-element"
         onMouseEnter={() => setHover(true)}
         onMouseLeave={() => setHover(false)}
         css={{ paddingTop: 6 }}

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

@@ -110,7 +110,12 @@ export default function sqlLabReducer(state = {}, action) {
       );
     },
     [actions.REMOVE_QUERY_EDITOR]() {
-      let newState = removeFromArr(state, 'queryEditors', action.queryEditor);
+      const queryEditor = {
+        ...action.queryEditor,
+        ...(action.queryEditor.id === state.unsavedQueryEditor.id &&
+          state.unsavedQueryEditor),
+      };
+      let newState = removeFromArr(state, 'queryEditors', queryEditor);
       // List of remaining queryEditor ids
       const qeIds = newState.queryEditors.map(qe => qe.id);
 
@@ -127,10 +132,19 @@ export default function sqlLabReducer(state = {}, action) {
 
       // Remove associated table schemas
       const tables = state.tables.filter(
-        table => table.queryEditorId !== action.queryEditor.id,
+        table => table.queryEditorId !== queryEditor.id,
       );
 
-      newState = { ...newState, tabHistory, tables, queries };
+      newState = {
+        ...newState,
+        tabHistory,
+        tables,
+        queries,
+        unsavedQueryEditor: {
+          ...(action.queryEditor.id !== state.unsavedQueryEditor.id &&
+            state.unsavedQueryEditor),
+        },
+      };
       return newState;
     },
     [actions.REMOVE_QUERY]() {

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

@@ -75,6 +75,28 @@ describe('sqlLabReducer', () => {
         initialState.queryEditors.length,
       );
     });
+    it('should remove a query editor including unsaved changes', () => {
+      expect(newState.queryEditors).toHaveLength(
+        initialState.queryEditors.length + 1,
+      );
+      let action = {
+        type: actions.QUERY_EDITOR_SETDB,
+        queryEditor: qe,
+        dbId: 123,
+      };
+      newState = sqlLabReducer(newState, action);
+      expect(newState.unsavedQueryEditor.dbId).toEqual(action.dbId);
+      action = {
+        type: actions.REMOVE_QUERY_EDITOR,
+        queryEditor: qe,
+      };
+      newState = sqlLabReducer(newState, action);
+      expect(newState.queryEditors).toHaveLength(
+        initialState.queryEditors.length,
+      );
+      expect(newState.unsavedQueryEditor.dbId).toBeUndefined();
+      expect(newState.unsavedQueryEditor.id).toBeUndefined();
+    });
     it('should set q query editor active', () => {
       const expectedTitle = 'new updated title';
       const addQueryEditorAction = {