fix: auto-complete of tables and names are not working in SQL lab

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

@@ -171,17 +171,20 @@ class AceEditorWrapper extends React.PureComponent<Props, State> {
       meta: 'schema',
     }));
     const columns = {};
-    const tables = props.extendedTables || props.tables || [];
+
+    const tables = props.tables || [];
+    const extendedTables = props.extendedTables || [];
 
     const tableWords = tables.map(t => {
-      const tableName = t.name;
-      const cols = t.columns || [];
+      const tableName = t.value;
+      const extendedTable = extendedTables.find(et => et.name === tableName);
+      const cols = (extendedTable && extendedTable.columns) || [];
       cols.forEach(col => {
         columns[col.name] = null; // using an object as a unique set
       });
 
       return {
-        name: tableName,
+        name: t.label,
         value: tableName,
         score: TABLE_AUTOCOMPLETE_SCORE,
         meta: 'table',

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

@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React from 'react';
+import React, { useEffect, useRef, useCallback } from 'react';
 import Button from 'src/components/Button';
 import { t, styled, css, SupersetTheme } from '@superset-ui/core';
 import Collapse from 'src/components/Collapse';
@@ -41,7 +41,10 @@ interface actionsTypes {
   addDangerToast: (msg: string) => void;
   queryEditorSetSchema: (queryEditor: QueryEditor, schema?: string) => void;
   queryEditorSetSchemaOptions: () => void;
-  queryEditorSetTableOptions: (options: Array<any>) => void;
+  queryEditorSetTableOptions: (
+    queryEditor: QueryEditor,
+    options: Array<any>,
+  ) => void;
   resetState: () => void;
 }
 
@@ -86,6 +89,13 @@ export default function SqlEditorLeftBar({
   tables = [],
   height = 500,
 }: SqlEditorLeftBarProps) {
+  // Ref needed to avoid infinite rerenders on handlers
+  // that require and modify the queryEditor
+  const queryEditorRef = useRef<QueryEditor>(queryEditor);
+  useEffect(() => {
+    queryEditorRef.current = queryEditor;
+  }, [queryEditor]);
+
   const onDbChange = ({ id: dbId }: { id: number }) => {
     actions.queryEditorSetDb(queryEditor, dbId);
     actions.queryEditorSetFunctionNames(queryEditor, dbId);
@@ -132,9 +142,23 @@ export default function SqlEditorLeftBar({
   const shouldShowReset = window.location.search === '?reset=1';
   const tableMetaDataHeight = height - 130; // 130 is the height of the selects above
 
-  const onSchemaChange = (schema: string) => {
-    actions.queryEditorSetSchema(queryEditor, schema);
-  };
+  const handleSchemaChange = useCallback(
+    (schema: string) => {
+      if (queryEditorRef.current) {
+        actions.queryEditorSetSchema(queryEditorRef.current, schema);
+      }
+    },
+    [actions],
+  );
+
+  const handleTablesLoad = React.useCallback(
+    (options: Array<any>) => {
+      if (queryEditorRef.current) {
+        actions.queryEditorSetTableOptions(queryEditorRef.current, options);
+      }
+    },
+    [actions],
+  );
 
   return (
     <div className="SqlEditorLeftBar">
@@ -143,10 +167,10 @@ export default function SqlEditorLeftBar({
         getDbList={actions.setDatabases}
         handleError={actions.addDangerToast}
         onDbChange={onDbChange}
-        onSchemaChange={onSchemaChange}
+        onSchemaChange={handleSchemaChange}
         onSchemasLoad={actions.queryEditorSetSchemaOptions}
         onTableChange={onTableChange}
-        onTablesLoad={actions.queryEditorSetTableOptions}
+        onTablesLoad={handleTablesLoad}
         schema={queryEditor.schema}
         sqlLabMode
       />

superset-frontend/src/components/TableSelector/TableSelector.test.tsx

@@ -20,12 +20,13 @@
 import React from 'react';
 import { render, screen, waitFor } from 'spec/helpers/testing-library';
 import { SupersetClient } from '@superset-ui/core';
+import { act } from 'react-dom/test-utils';
 import userEvent from '@testing-library/user-event';
 import TableSelector from '.';
 
 const SupersetClientGet = jest.spyOn(SupersetClient, 'get');
 
-const createProps = () => ({
+const createProps = (props = {}) => ({
   database: {
     id: 1,
     database_name: 'main',
@@ -34,23 +35,33 @@ const createProps = () => ({
   },
   schema: 'test_schema',
   handleError: jest.fn(),
+  ...props,
 });
 
-beforeAll(() => {
-  SupersetClientGet.mockImplementation(
-    async () =>
-      ({
-        json: {
-          options: [
-            { label: 'table_a', value: 'table_a' },
-            { label: 'table_b', value: 'table_b' },
-          ],
-        },
-      } as any),
-  );
+afterEach(() => {
+  jest.clearAllMocks();
 });
 
+const getSchemaMockFunction = async () =>
+  ({
+    json: {
+      result: ['schema_a', 'schema_b'],
+    },
+  } as any);
+
+const getTableMockFunction = async () =>
+  ({
+    json: {
+      options: [
+        { label: 'table_a', value: 'table_a' },
+        { label: 'table_b', value: 'table_b' },
+      ],
+    },
+  } as any);
+
 test('renders with default props', async () => {
+  SupersetClientGet.mockImplementation(getTableMockFunction);
+
   const props = createProps();
   render(<TableSelector {...props} />, { useRedux: true });
   const databaseSelect = screen.getByRole('combobox', {
@@ -70,6 +81,8 @@ test('renders with default props', async () => {
 });
 
 test('renders table options', async () => {
+  SupersetClientGet.mockImplementation(getTableMockFunction);
+
   const props = createProps();
   render(<TableSelector {...props} />, { useRedux: true });
   const tableSelect = screen.getByRole('combobox', {
@@ -85,6 +98,8 @@ test('renders table options', async () => {
 });
 
 test('renders disabled without schema', async () => {
+  SupersetClientGet.mockImplementation(getTableMockFunction);
+
   const props = createProps();
   render(<TableSelector {...props} schema={undefined} />, { useRedux: true });
   const tableSelect = screen.getByRole('combobox', {
@@ -94,3 +109,42 @@ test('renders disabled without schema', async () => {
     expect(tableSelect).toBeDisabled();
   });
 });
+
+test('table options are notified after schema selection', async () => {
+  SupersetClientGet.mockImplementation(getSchemaMockFunction);
+
+  const callback = jest.fn();
+  const props = createProps({
+    onTablesLoad: callback,
+    schema: undefined,
+  });
+  render(<TableSelector {...props} />, { useRedux: true });
+
+  const schemaSelect = screen.getByRole('combobox', {
+    name: 'Select schema or type schema name',
+  });
+  expect(schemaSelect).toBeInTheDocument();
+  expect(callback).not.toHaveBeenCalled();
+
+  userEvent.click(schemaSelect);
+
+  expect(
+    await screen.findByRole('option', { name: 'schema_a' }),
+  ).toBeInTheDocument();
+  expect(
+    await screen.findByRole('option', { name: 'schema_b' }),
+  ).toBeInTheDocument();
+
+  SupersetClientGet.mockImplementation(getTableMockFunction);
+
+  act(() => {
+    userEvent.click(screen.getAllByText('schema_a')[1]);
+  });
+
+  await waitFor(() => {
+    expect(callback).toHaveBeenCalledWith([
+      { label: 'table_a', value: 'table_a' },
+      { label: 'table_b', value: 'table_b' },
+    ]);
+  });
+});

superset-frontend/src/components/TableSelector/index.tsx

@@ -208,15 +208,14 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
               currentTable = option;
             }
           });
-          if (onTablesLoad) {
-            onTablesLoad(json.options);
-          }
+
+          onTablesLoad?.(json.options);
           setTableOptions(options);
           setCurrentTable(currentTable);
           setLoadingTables(false);
           if (forceRefresh) addSuccessToast('List updated');
         })
-        .catch(e => {
+        .catch(() => {
           setLoadingTables(false);
           handleError(t('There was an error loading the tables'));
         });