fix(sql lab): table selector should display all the selected tables

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

@@ -16,12 +16,12 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { useEffect, useRef, useCallback } from 'react';
+import React, { useEffect, useRef, useCallback, useMemo } from 'react';
 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 TableSelector from 'src/components/TableSelector';
+import { TableSelectorMultiple } from 'src/components/TableSelector';
 import { IconTooltip } from 'src/components/IconTooltip';
 import { QueryEditor } from 'src/SqlLab/types';
 import { DatabaseObject } from 'src/components/DatabaseSelector';
@@ -101,10 +101,32 @@ export default function SqlEditorLeftBar({
     actions.queryEditorSetFunctionNames(queryEditor, dbId);
   };
 
-  const onTableChange = (tableName: string, schemaName: string) => {
-    if (tableName && schemaName) {
-      actions.addTable(queryEditor, database, tableName, schemaName);
+  const selectedTableNames = useMemo(
+    () => tables?.map(table => table.name) || [],
+    [tables],
+  );
+
+  const onTablesChange = (tableNames: string[], schemaName: string) => {
+    if (!schemaName) {
+      return;
     }
+
+    const currentTables = [...tables];
+    const tablesToAdd = tableNames.filter(name => {
+      const index = currentTables.findIndex(table => table.name === name);
+      if (index >= 0) {
+        currentTables.splice(index, 1);
+        return false;
+      }
+
+      return true;
+    });
+
+    tablesToAdd.forEach(tableName =>
+      actions.addTable(queryEditor, database, tableName, schemaName),
+    );
+
+    currentTables.forEach(table => actions.removeTable(table));
   };
 
   const onToggleTable = (updatedTables: string[]) => {
@@ -162,16 +184,17 @@ export default function SqlEditorLeftBar({
 
   return (
     <div className="SqlEditorLeftBar">
-      <TableSelector
+      <TableSelectorMultiple
         database={database}
         getDbList={actions.setDatabases}
         handleError={actions.addDangerToast}
         onDbChange={onDbChange}
         onSchemaChange={handleSchemaChange}
         onSchemasLoad={actions.queryEditorSetSchemaOptions}
-        onTableChange={onTableChange}
+        onTableSelectChange={onTablesChange}
         onTablesLoad={handleTablesLoad}
         schema={queryEditor.schema}
+        tableValue={selectedTableNames}
         sqlLabMode
       />
       <div className="divider" />

superset-frontend/src/components/Datasource/DatasourceEditor.jsx

@@ -1008,7 +1008,7 @@ class DatasourceEditor extends React.PureComponent {
                         handleError={this.props.addDangerToast}
                         schema={datasource.schema}
                         sqlLabMode={false}
-                        tableName={datasource.table_name}
+                        tableValue={datasource.table_name}
                         onSchemaChange={
                           this.state.isEditMode
                             ? schema =>
@@ -1024,7 +1024,7 @@ class DatasourceEditor extends React.PureComponent {
                                 )
                             : undefined
                         }
-                        onTableChange={
+                        onTableSelectChange={
                           this.state.isEditMode
                             ? table =>
                                 this.onDatasourcePropChange('table_name', table)

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

@@ -18,11 +18,11 @@
  */
 
 import React from 'react';
-import { render, screen, waitFor } from 'spec/helpers/testing-library';
+import { render, screen, waitFor, within } 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 '.';
+import TableSelector, { TableSelectorMultiple } from '.';
 
 const SupersetClientGet = jest.spyOn(SupersetClient, 'get');
 
@@ -55,10 +55,17 @@ const getTableMockFunction = async () =>
       options: [
         { label: 'table_a', value: 'table_a' },
         { label: 'table_b', value: 'table_b' },
+        { label: 'table_c', value: 'table_c' },
+        { label: 'table_d', value: 'table_d' },
       ],
     },
   } as any);
 
+const getSelectItemContainer = (select: HTMLElement) =>
+  select.parentElement?.parentElement?.getElementsByClassName(
+    'ant-select-selection-item',
+  );
+
 test('renders with default props', async () => {
   SupersetClientGet.mockImplementation(getTableMockFunction);
 
@@ -145,6 +152,96 @@ test('table options are notified after schema selection', async () => {
     expect(callback).toHaveBeenCalledWith([
       { label: 'table_a', value: 'table_a' },
       { label: 'table_b', value: 'table_b' },
+      { label: 'table_c', value: 'table_c' },
+      { label: 'table_d', value: 'table_d' },
     ]);
   });
 });
+
+test('table select retain value if not in SQL Lab mode', async () => {
+  SupersetClientGet.mockImplementation(getTableMockFunction);
+
+  const callback = jest.fn();
+  const props = createProps({
+    onTableSelectChange: callback,
+    sqlLabMode: false,
+  });
+
+  render(<TableSelector {...props} />, { useRedux: true });
+
+  const tableSelect = screen.getByRole('combobox', {
+    name: 'Select table or type table name',
+  });
+
+  expect(screen.queryByText('table_a')).not.toBeInTheDocument();
+  expect(getSelectItemContainer(tableSelect)).toHaveLength(0);
+
+  userEvent.click(tableSelect);
+
+  expect(
+    await screen.findByRole('option', { name: 'table_a' }),
+  ).toBeInTheDocument();
+
+  act(() => {
+    userEvent.click(screen.getAllByText('table_a')[1]);
+  });
+
+  expect(callback).toHaveBeenCalled();
+
+  const selectedValueContainer = getSelectItemContainer(tableSelect);
+
+  expect(selectedValueContainer).toHaveLength(1);
+  expect(
+    await within(selectedValueContainer?.[0] as HTMLElement).findByText(
+      'table_a',
+    ),
+  ).toBeInTheDocument();
+});
+
+test('table multi select retain all the values selected', async () => {
+  SupersetClientGet.mockImplementation(getTableMockFunction);
+
+  const callback = jest.fn();
+  const props = createProps({
+    onTableSelectChange: callback,
+  });
+
+  render(<TableSelectorMultiple {...props} />, { useRedux: true });
+
+  const tableSelect = screen.getByRole('combobox', {
+    name: 'Select table or type table name',
+  });
+
+  expect(screen.queryByText('table_a')).not.toBeInTheDocument();
+  expect(getSelectItemContainer(tableSelect)).toHaveLength(0);
+
+  userEvent.click(tableSelect);
+
+  expect(
+    await screen.findByRole('option', { name: 'table_a' }),
+  ).toBeInTheDocument();
+
+  act(() => {
+    const item = screen.getAllByText('table_a');
+    userEvent.click(item[item.length - 1]);
+  });
+
+  act(() => {
+    const item = screen.getAllByText('table_c');
+    userEvent.click(item[item.length - 1]);
+  });
+
+  const selectedValueContainer = getSelectItemContainer(tableSelect);
+
+  expect(selectedValueContainer).toHaveLength(2);
+  expect(
+    await within(selectedValueContainer?.[0] as HTMLElement).findByText(
+      'table_a',
+    ),
+  ).toBeInTheDocument();
+  expect(
+    await within(selectedValueContainer?.[1] as HTMLElement).findByText(
+      'table_c',
+    ),
+  ).toBeInTheDocument();
+});

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

@@ -23,6 +23,8 @@ import React, {
   useMemo,
   useEffect,
 } from 'react';
+import { SelectValue } from 'antd/lib/select';
+
 import { styled, SupersetClient, t } from '@superset-ui/core';
 import { Select } from 'src/components';
 import { FormLabel } from 'src/components/Form';
@@ -87,12 +89,13 @@ interface TableSelectorProps {
   onDbChange?: (db: DatabaseObject) => void;
   onSchemaChange?: (schema?: string) => void;
   onSchemasLoad?: () => void;
-  onTableChange?: (tableName?: string, schema?: string) => void;
   onTablesLoad?: (options: Array<any>) => void;
   readOnly?: boolean;
   schema?: string;
   sqlLabMode?: boolean;
-  tableName?: string;
+  tableValue?: string | string[];
+  onTableSelectChange?: (value?: string | string[], schema?: string) => void;
+  tableSelectMode?: 'single' | 'multiple';
 }
 
 interface Table {
@@ -150,34 +153,52 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
   onDbChange,
   onSchemaChange,
   onSchemasLoad,
-  onTableChange,
   onTablesLoad,
   readOnly = false,
   schema,
   sqlLabMode = true,
-  tableName,
+  tableSelectMode = 'single',
+  tableValue = undefined,
+  onTableSelectChange,
 }) => {
   const [currentDatabase, setCurrentDatabase] = useState<
     DatabaseObject | undefined
   >(database);
   const [currentSchema, setCurrentSchema] = useState<string | undefined>(
     schema,
   );
-  const [currentTable, setCurrentTable] = useState<TableOption | undefined>();
+
+  const [tableOptions, setTableOptions] = useState<TableOption[]>([]);
+  const [tableSelectValue, setTableSelectValue] = useState<
+    SelectValue | undefined
+  >(undefined);
   const [refresh, setRefresh] = useState(0);
   const [previousRefresh, setPreviousRefresh] = useState(0);
   const [loadingTables, setLoadingTables] = useState(false);
-  const [tableOptions, setTableOptions] = useState<TableOption[]>([]);
   const { addSuccessToast } = useToasts();
 
   useEffect(() => {
     // reset selections
     if (database === undefined) {
       setCurrentDatabase(undefined);
       setCurrentSchema(undefined);
-      setCurrentTable(undefined);
+      setTableSelectValue(undefined);
     }
-  }, [database]);
+  }, [database, tableSelectMode]);
+
+  useEffect(() => {
+    if (tableSelectMode === 'single') {
+      setTableSelectValue(
+        tableOptions.find(option => option.value === tableValue),
+      );
+    } else {
+      setTableSelectValue(
+        tableOptions?.filter(
+          option => option && tableValue?.includes(option.value),
+        ) || [],
+      );
+    }
+  }, [tableOptions, tableValue, tableSelectMode]);
 
   useEffect(() => {
     if (currentDatabase && currentSchema) {
@@ -195,23 +216,18 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
 
       SupersetClient.get({ endpoint })
         .then(({ json }) => {
-          const options: TableOption[] = [];
-          let currentTable;
-          json.options.forEach((table: Table) => {
-            const option = {
+          const options: TableOption[] = json.options.map((table: Table) => {
+            const option: TableOption = {
               value: table.value,
               label: <TableOption table={table} />,
               text: table.label,
             };
-            options.push(option);
-            if (table.label === tableName) {
-              currentTable = option;
-            }
+
+            return option;
           });
 
           onTablesLoad?.(json.options);
           setTableOptions(options);
-          setCurrentTable(currentTable);
           setLoadingTables(false);
           if (forceRefresh) addSuccessToast('List updated');
         })
@@ -223,7 +239,7 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
     // We are using the refresh state to re-trigger the query
     // previousRefresh should be out of dependencies array
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [currentDatabase, currentSchema, onTablesLoad, refresh]);
+  }, [currentDatabase, currentSchema, onTablesLoad, setTableOptions, refresh]);
 
   function renderSelectRow(select: ReactNode, refreshBtn: ReactNode) {
     return (
@@ -234,10 +250,18 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
     );
   }
 
-  const internalTableChange = (table?: TableOption) => {
-    setCurrentTable(table);
-    if (onTableChange && currentSchema) {
-      onTableChange(table?.value, currentSchema);
+  const internalTableChange = (
+    selectedOptions: TableOption | TableOption[] | undefined,
+  ) => {
+    if (currentSchema) {
+      onTableSelectChange?.(
+        Array.isArray(selectedOptions)
+          ? selectedOptions.map(option => option?.value)
+          : selectedOptions?.value,
+        currentSchema,
+      );
+    } else {
+      setTableSelectValue(selectedOptions);
     }
   };
 
@@ -253,6 +277,7 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
     if (onSchemaChange) {
       onSchemaChange(schema);
     }
+
     internalTableChange(undefined);
   };
 
@@ -305,11 +330,15 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
         lazyLoading={false}
         loading={loadingTables}
         name="select-table"
-        onChange={(table: TableOption) => internalTableChange(table)}
+        onChange={(options: TableOption | TableOption[]) =>
+          internalTableChange(options)
+        }
         options={tableOptions}
         placeholder={t('Select table or type table name')}
         showSearch
-        value={currentTable}
+        mode={tableSelectMode}
+        value={tableSelectValue}
+        allowClear={tableSelectMode === 'multiple'}
       />
     );
 
@@ -332,4 +361,7 @@ const TableSelector: FunctionComponent<TableSelectorProps> = ({
   );
 };
 
+export const TableSelectorMultiple: FunctionComponent<TableSelectorProps> =
+  props => <TableSelector tableSelectMode="multiple" {...props} />;
+
 export default TableSelector;