fix(sqllab): Invalid schema fetch for deprecated value (#22695)

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

@@ -128,10 +128,22 @@ export function getUpToDateQuery(rootState, queryEditor, key) {
     sqlLab: { unsavedQueryEditor },
   } = rootState;
   const id = key ?? queryEditor.id;
-  return {
+  const updatedQueryEditor = {
     ...queryEditor,
     ...(id === unsavedQueryEditor.id && unsavedQueryEditor),
   };
+
+  if (
+    'schema' in updatedQueryEditor &&
+    !updatedQueryEditor.schemaOptions?.find(
+      ({ value }) => value === updatedQueryEditor.schema,
+    )
+  ) {
+    // remove the deprecated schema option
+    delete updatedQueryEditor.schema;
+  }
+
+  return updatedQueryEditor;
 }
 
 export function resetState() {

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

@@ -316,6 +316,54 @@ describe('async actions', () => {
     });
   });
 
+  describe('getUpToDateQuery', () => {
+    const id = 'randomidvalue2';
+    const unsavedQueryEditor = {
+      id,
+      sql: 'some query here',
+      schemaOptions: [{ value: 'testSchema' }, { value: 'schema2' }],
+    };
+
+    it('returns payload with the updated queryEditor', () => {
+      const queryEditor = {
+        id,
+        name: 'Dummy query editor',
+        schema: 'testSchema',
+      };
+      const state = {
+        sqlLab: {
+          tabHistory: [id],
+          queryEditors: [queryEditor],
+          unsavedQueryEditor,
+        },
+      };
+      expect(actions.getUpToDateQuery(state, queryEditor)).toEqual({
+        ...queryEditor,
+        ...unsavedQueryEditor,
+      });
+    });
+
+    it('ignores the deprecated schema option', () => {
+      const queryEditor = {
+        id,
+        name: 'Dummy query editor',
+        schema: 'oldSchema',
+      };
+      const state = {
+        sqlLab: {
+          tabHistory: [id],
+          queryEditors: [queryEditor],
+          unsavedQueryEditor,
+        },
+      };
+      expect(actions.getUpToDateQuery(state, queryEditor)).toEqual({
+        ...queryEditor,
+        ...unsavedQueryEditor,
+        schema: undefined,
+      });
+    });
+  });
+
   describe('postStopQuery', () => {
     const stopQueryEndpoint = 'glob:*/api/v1/query/stop';
     fetchMock.post(stopQueryEndpoint, {});

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

@@ -43,6 +43,7 @@ const mockQueryEditor = {
   autorun: false,
   sql: 'SELECT * FROM ...',
   remoteId: 999,
+  schemaOptions: [{ value: 'query_schema' }, { value: 'query_schema_updated' }],
 };
 const disabled = {
   id: 'disabledEditorId',
@@ -82,6 +83,7 @@ const standardProviderWithUnsaved: React.FC = ({ children }) => (
         ...initialState,
         sqlLab: {
           ...initialState.sqlLab,
+          queryEditors: [mockQueryEditor],
           unsavedQueryEditor,
         },
       })}
@@ -123,7 +125,7 @@ describe('ShareSqlLabQuery', () => {
         });
       });
       const button = screen.getByRole('button');
-      const { id, remoteId, ...expected } = mockQueryEditor;
+      const { id, remoteId, schemaOptions, ...expected } = mockQueryEditor;
       const storeQuerySpy = jest.spyOn(utils, 'storeQuery');
       userEvent.click(button);
       expect(storeQuerySpy.mock.calls).toHaveLength(1);

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

@@ -41,15 +41,22 @@ const middlewares = [thunk];
 const mockStore = configureStore(middlewares);
 const store = mockStore(initialState);
 
-fetchMock.get('glob:*/api/v1/database/*/schemas/?*', { result: [] });
+beforeEach(() => {
-fetchMock.get('glob:*/superset/tables/**', {
+  fetchMock.get('glob:*/api/v1/database/**', { result: [] });
-  options: [
+  fetchMock.get('glob:*/api/v1/database/*/schemas/?*', { result: [] });
-    {
+  fetchMock.get('glob:*/superset/tables/**', {
-      label: 'ab_user',
+    options: [
-      value: 'ab_user',
+      {
-    },
+        label: 'ab_user',
-  ],
+        value: 'ab_user',
-  tableLength: 1,
+      },
+    ],
+    tableLength: 1,
+  });
+});
+
+afterEach(() => {
+  fetchMock.restore();
 });
 
 const renderAndWait = (props, store) =>
@@ -110,8 +117,9 @@ test('should toggle the table when the header is clicked', async () => {
   userEvent.click(header);
 
   await waitFor(() => {
-    expect(store.getActions()).toHaveLength(4);
+    expect(store.getActions()[store.getActions().length - 1].type).toEqual(
-    expect(store.getActions()[3].type).toEqual('COLLAPSE_TABLE');
+      'COLLAPSE_TABLE',
+    );
   });
 });
 
@@ -129,14 +137,77 @@ test('When changing database the table list must be updated', async () => {
         database_name: 'new_db',
         backend: 'postgresql',
       }}
-      queryEditor={{ ...mockedProps.queryEditor, schema: 'new_schema' }}
+      queryEditorId={defaultQueryEditor.id}
       tables={[{ ...mockedProps.tables[0], dbId: 2, name: 'new_table' }]}
     />,
     {
       useRedux: true,
-      initialState,
+      store: mockStore({
+        ...initialState,
+        sqlLab: {
+          ...initialState.sqlLab,
+          unsavedQueryEditor: {
+            id: defaultQueryEditor.id,
+            schema: 'new_schema',
+          },
+        },
+      }),
     },
   );
   expect(await screen.findByText(/new_db/i)).toBeInTheDocument();
   expect(await screen.findByText(/new_table/i)).toBeInTheDocument();
 });
+
+test('ignore schema api when current schema is deprecated', async () => {
+  const invalidSchemaName = 'None';
+  await renderAndWait(
+    mockedProps,
+    mockStore({
+      ...initialState,
+      sqlLab: {
+        ...initialState.sqlLab,
+        unsavedQueryEditor: {
+          id: defaultQueryEditor.id,
+          schema: invalidSchemaName,
+        },
+      },
+    }),
+  );
+
+  expect(await screen.findByText(/Database/i)).toBeInTheDocument();
+  expect(screen.queryByText(/None/i)).not.toBeInTheDocument();
+  expect(fetchMock.calls()).not.toContainEqual(
+    expect.arrayContaining([
+      expect.stringContaining(
+        `/tables/${mockedProps.database.id}/${invalidSchemaName}/`,
+      ),
+    ]),
+  );
+});
+
+test('fetches schema api when current schema is among the list', async () => {
+  const validSchema = 'schema1';
+  await renderAndWait(
+    mockedProps,
+    mockStore({
+      ...initialState,
+      sqlLab: {
+        ...initialState.sqlLab,
+        unsavedQueryEditor: {
+          id: defaultQueryEditor.id,
+          schema: validSchema,
+          schemaOptions: [{ name: validSchema, value: validSchema }],
+        },
+      },
+    }),
+  );
+
+  expect(await screen.findByText(/schema1/i)).toBeInTheDocument();
+  expect(fetchMock.calls()).toContainEqual(
+    expect.arrayContaining([
+      expect.stringContaining(
+        `/tables/${mockedProps.database.id}/${validSchema}/`,
+      ),
+    ]),
+  );
+});

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

@@ -117,7 +117,6 @@ const SqlEditorLeftBar = ({
 }: SqlEditorLeftBarProps) => {
   const dispatch = useDispatch();
   const queryEditor = useQueryEditor(queryEditorId, ['dbId', 'schema']);
-
   const [emptyResultsWithSearch, setEmptyResultsWithSearch] = useState(false);
   const [userSelectedDb, setUserSelected] = useState<DatabaseObject | null>(
     null,

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

@@ -16,6 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+import { useMemo } from 'react';
 import pick from 'lodash/pick';
 import { shallowEqual, useSelector } from 'react-redux';
 import { SqlLabRootState, QueryEditor } from 'src/SqlLab/types';
@@ -24,7 +25,10 @@ export default function useQueryEditor<T extends keyof QueryEditor>(
   sqlEditorId: string,
   attributes: ReadonlyArray<T>,
 ) {
-  return useSelector<SqlLabRootState, Pick<QueryEditor, T | 'id'>>(
+  const queryEditor = useSelector<
+    SqlLabRootState,
+    Pick<QueryEditor, T | 'id' | 'schema'>
+  >(
     ({ sqlLab: { unsavedQueryEditor, queryEditors } }) =>
       pick(
         {
@@ -32,7 +36,32 @@ export default function useQueryEditor<T extends keyof QueryEditor>(
           ...(sqlEditorId === unsavedQueryEditor.id && unsavedQueryEditor),
         },
         ['id'].concat(attributes),
-      ) as Pick<QueryEditor, T | 'id'>,
+      ) as Pick<QueryEditor, T | 'id' | 'schema'>,
     shallowEqual,
   );
+  const { schema, schemaOptions } = useSelector<
+    SqlLabRootState,
+    Pick<QueryEditor, 'schema' | 'schemaOptions'>
+  >(
+    ({ sqlLab: { unsavedQueryEditor, queryEditors } }) =>
+      pick(
+        {
+          ...queryEditors.find(({ id }) => id === sqlEditorId),
+          ...(sqlEditorId === unsavedQueryEditor.id && unsavedQueryEditor),
+        },
+        ['schema', 'schemaOptions'],
+      ) as Pick<QueryEditor, T | 'schema' | 'schemaOptions'>,
+    shallowEqual,
+  );
+
+  const schemaOptionsMap = useMemo(
+    () => new Set(schemaOptions?.map(({ value }) => value)),
+    [schemaOptions],
+  );
+
+  if ('schema' in queryEditor && schema && !schemaOptionsMap.has(schema)) {
+    delete queryEditor.schema;
+  }
+
+  return queryEditor as Pick<QueryEditor, T | 'id'>;
 }

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

@@ -70,7 +70,7 @@ test('includes id implicitly', () => {
 test('returns updated values from unsaved change', () => {
   const expectedSql = 'SELECT updated_column\nFROM updated_table\nWHERE';
   const { result } = renderHook(
-    () => useQueryEditor(defaultQueryEditor.id, ['id', 'sql']),
+    () => useQueryEditor(defaultQueryEditor.id, ['id', 'sql', 'schema']),
     {
       wrapper: createWrapper({
         useRedux: true,
@@ -88,5 +88,31 @@ test('returns updated values from unsaved change', () => {
     },
   );
   expect(result.current.id).toEqual(defaultQueryEditor.id);
+  expect(result.current.schema).toEqual(defaultQueryEditor.schema);
   expect(result.current.sql).toEqual(expectedSql);
 });
+
+test('skips the deprecated schema option', () => {
+  const expectedSql = 'SELECT updated_column\nFROM updated_table\nWHERE';
+  const { result } = renderHook(
+    () => useQueryEditor(defaultQueryEditor.id, ['schema']),
+    {
+      wrapper: createWrapper({
+        useRedux: true,
+        store: mockStore({
+          ...initialState,
+          sqlLab: {
+            ...initialState.sqlLab,
+            unsavedQueryEditor: {
+              id: defaultQueryEditor.id,
+              sql: expectedSql,
+              schema: 'deprecated schema',
+            },
+          },
+        }),
+      }),
+    },
+  );
+  expect(result.current.schema).not.toEqual(defaultQueryEditor.schema);
+  expect(result.current.schema).toBeUndefined();
+});

superset-frontend/src/SqlLab/types.ts

@@ -35,7 +35,7 @@ export interface QueryEditor {
   id: string;
   dbId?: number;
   name: string;
-  schema: string;
+  schema?: string;
   autorun: boolean;
   sql: string;
   remoteId: number | null;