From e660de69364c5f40b646fa54cf16f493f8ccf1e8 Mon Sep 17 00:00:00 2001
From: "Michael S. Molina" <70410625+michael-s-molina@users.noreply.github.com>
Date: Thu, 22 Jul 2021 15:17:31 -0300
Subject: [PATCH] chore: Adds lazy loading to the Select component (#15799)

---
 .../src/components/Select/Select.stories.tsx  |  19 ++-
 .../src/components/Select/Select.tsx          | 150 ++++++++++++++----
 .../FiltersConfigForm/DatasetSelect.tsx       |  24 +--
 .../FiltersConfigForm/DefaultValue.tsx        |  12 +-
 .../FiltersConfigForm/FiltersConfigForm.tsx   |  68 ++++----
 .../getControlItemsMap.test.tsx               |   1 +
 .../FiltersConfigForm/getControlItemsMap.tsx  |   3 +-
 .../FiltersConfigModal.test.tsx               |  16 +-
 8 files changed, 195 insertions(+), 98 deletions(-)

diff --git a/superset-frontend/src/components/Select/Select.stories.tsx b/superset-frontend/src/components/Select/Select.stories.tsx
index 5955176f536c..e6dd867d972c 100644
--- a/superset-frontend/src/components/Select/Select.stories.tsx
+++ b/superset-frontend/src/components/Select/Select.stories.tsx
@@ -144,6 +144,11 @@ InteractiveSelect.argTypes = {
       disable: true,
     },
   },
+  fetchOnlyOnSearch: {
+    table: {
+      disable: true,
+    },
+  },
 };
 
 InteractiveSelect.story = {
@@ -296,10 +301,12 @@ const USERS = [
 
 export const AsyncSelect = ({
   withError,
+  withInitialValue,
   responseTime,
   ...rest
 }: SelectProps & {
   withError: boolean;
+  withInitialValue: boolean;
   responseTime: number;
 }) => {
   const [requests, setRequests] = useState<ReactNode[]>([]);
@@ -375,6 +382,11 @@ export const AsyncSelect = ({
         <Select
           {...rest}
           options={withError ? fetchUserListError : fetchUserListPage}
+          value={
+            withInitialValue
+              ? { label: 'Valentina', value: 'Valentina' }
+              : undefined
+          }
         />
       </div>
       <div
@@ -398,9 +410,11 @@ export const AsyncSelect = ({
 };
 
 AsyncSelect.args = {
-  withError: false,
-  pageSize: 10,
   allowNewOptions: false,
+  fetchOnlyOnSearch: false,
+  pageSize: 10,
+  withError: false,
+  withInitialValue: false,
 };
 
 AsyncSelect.argTypes = {
@@ -431,6 +445,7 @@ AsyncSelect.argTypes = {
       type: 'range',
       min: 0.5,
       max: 5,
+      step: 0.5,
     },
   },
 };
diff --git a/superset-frontend/src/components/Select/Select.tsx b/superset-frontend/src/components/Select/Select.tsx
index e24547b0b9ab..be38aef8fd30 100644
--- a/superset-frontend/src/components/Select/Select.tsx
+++ b/superset-frontend/src/components/Select/Select.tsx
@@ -28,16 +28,17 @@ import React, {
   useCallback,
 } from 'react';
 import { styled, t } from '@superset-ui/core';
-import { Select as AntdSelect } from 'antd';
-import Icons from 'src/components/Icons';
-import {
+import AntdSelect, {
   SelectProps as AntdSelectProps,
   SelectValue as AntdSelectValue,
   LabeledValue as AntdLabeledValue,
 } from 'antd/lib/select';
+import { DownOutlined, SearchOutlined } from '@ant-design/icons';
 import debounce from 'lodash/debounce';
-import { getClientErrorObject } from 'src/utils/getClientErrorObject';
 import { isEqual } from 'lodash';
+import { Spin } from 'antd';
+import Icons from 'src/components/Icons';
+import { getClientErrorObject } from 'src/utils/getClientErrorObject';
 import { hasOption } from './utils';
 
 type AntdSelectAllProps = AntdSelectProps<AntdSelectValue>;
@@ -47,7 +48,6 @@ type PickedSelectProps = Pick<
   | 'allowClear'
   | 'autoFocus'
   | 'value'
-  | 'defaultValue'
   | 'disabled'
   | 'filterOption'
   | 'notFoundContent'
@@ -79,6 +79,7 @@ export interface SelectProps extends PickedSelectProps {
   options: OptionsType | OptionsPagePromise;
   pageSize?: number;
   invertSelection?: boolean;
+  fetchOnlyOnSearch?: boolean;
 }
 
 const StyledContainer = styled.div`
@@ -122,6 +123,10 @@ const StyledError = styled.div`
   `}
 `;
 
+const StyledSpin = styled(Spin)`
+  margin-top: ${({ theme }) => -theme.gridUnit}px;
+`;
+
 const MAX_TAG_COUNT = 4;
 const TOKEN_SEPARATORS = [',', '\n', '\t', ';'];
 const DEBOUNCE_TIMEOUT = 500;
@@ -137,15 +142,16 @@ const Error = ({ error }: { error: string }) => (
 const Select = ({
   allowNewOptions = false,
   ariaLabel,
+  fetchOnlyOnSearch,
   filterOption = true,
   header = null,
+  invertSelection = false,
   mode = 'single',
   name,
+  options,
   pageSize = DEFAULT_PAGE_SIZE,
   placeholder = t('Select ...'),
-  options,
   showSearch,
-  invertSelection = false,
   value,
   ...props
 }: SelectProps) => {
@@ -164,7 +170,8 @@ const Select = ({
   const [isDropdownVisible, setIsDropdownVisible] = useState(false);
   const [page, setPage] = useState(0);
   const [totalCount, setTotalCount] = useState(0);
-  const fetchedQueries = useRef(new Set<string>());
+  const [loadingEnabled, setLoadingEnabled] = useState(false);
+  const fetchedQueries = useRef(new Map<string, number>());
   const mappedMode = isSingleMode
     ? undefined
     : allowNewOptions
@@ -177,6 +184,26 @@ const Select = ({
     );
   }, [options]);
 
+  useEffect(() => {
+    if (isAsync && value) {
+      const array: AntdLabeledValue[] = Array.isArray(value)
+        ? (value as AntdLabeledValue[])
+        : [value as AntdLabeledValue];
+      const options: AntdLabeledValue[] = [];
+      array.forEach(element => {
+        const found = selectOptions.find(
+          option => option.value === element.value,
+        );
+        if (!found) {
+          options.push(element);
+        }
+      });
+      if (options.length > 0) {
+        setSelectOptions([...selectOptions, ...options]);
+      }
+    }
+  }, [isAsync, selectOptions, value]);
+
   useEffect(() => {
     setSelectValue(value);
   }, [value]);
@@ -185,24 +212,56 @@ const Select = ({
     (selectedValue: AntdSelectValue | undefined) => {
       // bringing selected options to the top of the list
       if (selectedValue) {
-        const currentValue = selectedValue as string[] | string;
-        const topOptions = selectOptions.filter(opt =>
-          Array.isArray(currentValue)
-            ? currentValue.includes(opt.value)
-            : currentValue === opt.value,
-        );
-        const otherOptions = selectOptions.filter(
-          opt => !topOptions.find(tOpt => tOpt.value === opt.value),
-        );
+        const topOptions: OptionsType = [];
+        const otherOptions: OptionsType = [];
+
+        selectOptions.forEach(opt => {
+          let found = false;
+          if (Array.isArray(selectedValue)) {
+            if (isAsync) {
+              found =
+                (selectedValue as AntdLabeledValue[]).find(
+                  element => element.value === opt.value,
+                ) !== undefined;
+            } else {
+              found = selectedValue.includes(opt.value);
+            }
+          } else {
+            found = isAsync
+              ? (selectedValue as AntdLabeledValue).value === opt.value
+              : selectedValue === opt.value;
+          }
+
+          if (found) {
+            topOptions.push(opt);
+          } else {
+            otherOptions.push(opt);
+          }
+        });
+
         // fallback for custom options in tags mode as they
         // do not appear in the selectOptions state
-        if (!isSingleMode && Array.isArray(currentValue)) {
-          // eslint-disable-next-line no-restricted-syntax
-          for (const val of currentValue) {
-            if (!topOptions.find(tOpt => tOpt.value === val)) {
-              topOptions.push({ label: val, value: val });
+        if (!isSingleMode && Array.isArray(selectedValue)) {
+          selectedValue.forEach((val: string | number | AntdLabeledValue) => {
+            if (
+              !topOptions.find(
+                tOpt =>
+                  tOpt.value ===
+                  (isAsync ? (val as AntdLabeledValue)?.value : val),
+              )
+            ) {
+              if (isAsync) {
+                const labelValue = val as AntdLabeledValue;
+                topOptions.push({
+                  label: labelValue.label,
+                  value: labelValue.value,
+                });
+              } else {
+                const value = val as string | number;
+                topOptions.push({ label: String(value), value });
+              }
             }
-          }
+          });
         }
 
         const sortedOptions = [...topOptions, ...otherOptions];
@@ -211,7 +270,7 @@ const Select = ({
         }
       }
     },
-    [isSingleMode, selectOptions],
+    [isAsync, isSingleMode, selectOptions],
   );
 
   const handleOnSelect = (
@@ -220,7 +279,11 @@ const Select = ({
     if (isSingleMode) {
       setSelectValue(selectedValue);
     } else {
-      const currentSelected = Array.isArray(selectValue) ? selectValue : [];
+      const currentSelected = selectValue
+        ? Array.isArray(selectValue)
+          ? selectValue
+          : [selectValue]
+        : [];
       if (
         typeof selectedValue === 'number' ||
         typeof selectedValue === 'string'
@@ -271,7 +334,9 @@ const Select = ({
   const handlePaginatedFetch = useMemo(
     () => (value: string, page: number, pageSize: number) => {
       const key = `${value};${page};${pageSize}`;
-      if (fetchedQueries.current.has(key)) {
+      const cachedCount = fetchedQueries.current.get(key);
+      if (cachedCount) {
+        setTotalCount(cachedCount);
         return;
       }
       setLoading(true);
@@ -279,7 +344,7 @@ const Select = ({
       fetchOptions(value, page, pageSize)
         .then(({ data, totalCount }: OptionsTypePage) => {
           handleData(data);
-          fetchedQueries.current.add(key);
+          fetchedQueries.current.set(key, totalCount);
           setTotalCount(totalCount);
         })
         .catch(onError)
@@ -351,6 +416,11 @@ const Select = ({
 
   const handleOnDropdownVisibleChange = (isDropdownVisible: boolean) => {
     setIsDropdownVisible(isDropdownVisible);
+
+    if (isAsync && !loadingEnabled) {
+      setLoadingEnabled(true);
+    }
+
     // multiple or tags mode keep the dropdown visible while selecting options
     // this waits for the dropdown to be closed before sorting the top options
     if (!isSingleMode && !isDropdownVisible) {
@@ -359,13 +429,20 @@ const Select = ({
   };
 
   useEffect(() => {
-    const foundOption = hasOption(searchedValue, selectOptions);
-    if (isAsync && !foundOption) {
+    const allowFetch = !fetchOnlyOnSearch || searchedValue;
+    if (isAsync && loadingEnabled && allowFetch) {
       const page = 0;
       handlePaginatedFetch(searchedValue, page, pageSize);
       setPage(page);
     }
-  }, [isAsync, searchedValue, selectOptions, pageSize, handlePaginatedFetch]);
+  }, [
+    isAsync,
+    searchedValue,
+    pageSize,
+    handlePaginatedFetch,
+    loadingEnabled,
+    fetchOnlyOnSearch,
+  ]);
 
   useEffect(() => {
     if (isSingleMode) {
@@ -382,6 +459,16 @@ const Select = ({
     return error ? <Error error={error} /> : originNode;
   };
 
+  const SuffixIcon = () => {
+    if (isLoading) {
+      return <StyledSpin size="small" />;
+    }
+    if (shouldShowSearch && isDropdownVisible) {
+      return <SearchOutlined />;
+    }
+    return <DownOutlined />;
+  };
+
   return (
     <StyledContainer>
       {header}
@@ -391,7 +478,7 @@ const Select = ({
         dropdownRender={dropdownRender}
         filterOption={handleFilterOption}
         getPopupContainer={triggerNode => triggerNode.parentNode}
-        loading={isLoading}
+        labelInValue={isAsync}
         maxTagCount={MAX_TAG_COUNT}
         mode={mappedMode}
         onDeselect={handleOnDeselect}
@@ -406,6 +493,7 @@ const Select = ({
         showArrow
         tokenSeparators={TOKEN_SEPARATORS}
         value={selectValue}
+        suffixIcon={<SuffixIcon />}
         menuItemSelectedIcon={
           invertSelection ? (
             <StyledStopOutlined iconSize="m" />
diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DatasetSelect.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DatasetSelect.tsx
index c3d589cda114..1fdadcbd0a70 100644
--- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DatasetSelect.tsx
+++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DatasetSelect.tsx
@@ -36,18 +36,11 @@ const cachedSupersetGet = cacheWrapper(
 );
 
 interface DatasetSelectProps {
-  datasetDetails: Record<string, any> | undefined;
-  datasetId: number;
-  onChange: (value: number) => void;
-  value?: { value: number | undefined };
+  onChange: (value: { label: string; value: number }) => void;
+  value?: { label: string; value: number };
 }
 
-const DatasetSelect = ({
-  datasetDetails,
-  datasetId,
-  onChange,
-  value,
-}: DatasetSelectProps) => {
+const DatasetSelect = ({ onChange, value }: DatasetSelectProps) => {
   const getErrorMessage = useCallback(
     ({ error, message }: ClientErrorObject) => {
       let errorText = message || error || t('An error has occurred');
@@ -84,15 +77,6 @@ const DatasetSelect = ({
           .sort((a: { label: string }, b: { label: string }) =>
             a.label.localeCompare(b.label),
           );
-        if (!search) {
-          const found = data.find(element => element.value === datasetId);
-          if (!found && datasetDetails?.table_name) {
-            data.push({
-              label: datasetDetails.table_name,
-              value: datasetId,
-            });
-          }
-        }
         return {
           data,
           totalCount: response.json.count,
@@ -107,7 +91,7 @@ const DatasetSelect = ({
   return (
     <Select
       ariaLabel={t('Dataset')}
-      value={value?.value}
+      value={value}
       options={loadDatasetOptions}
       onChange={onChange}
     />
diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DefaultValue.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DefaultValue.tsx
index b56f2c221b49..e876006ee9cc 100644
--- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DefaultValue.tsx
+++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/DefaultValue.tsx
@@ -16,7 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
-import React, { FC, useEffect, useState } from 'react';
+import React, { FC } from 'react';
 import {
   Behavior,
   SetDataMaskHook,
@@ -48,17 +48,9 @@ const DefaultValue: FC<DefaultValueProps> = ({
   formData,
   enableNoResults,
 }) => {
-  const [loading, setLoading] = useState(hasDataset);
   const formFilter = (form.getFieldValue('filters') || {})[filterId];
   const queriesData = formFilter?.defaultValueQueriesData;
-
-  useEffect(() => {
-    if (!hasDataset || queriesData !== null) {
-      setLoading(false);
-    } else {
-      setLoading(true);
-    }
-  }, [hasDataset, queriesData]);
+  const loading = hasDataset && queriesData === null;
   const value = formFilter.defaultDataMask?.filterState.value;
   const isMissingRequiredValue =
     hasDefaultValue && (value === null || value === undefined);
diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx
index ae27d16e0970..c7300c626e42 100644
--- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx
+++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/FiltersConfigForm.tsx
@@ -355,8 +355,14 @@ const FiltersConfigForm = (
   const hasDataset = !!nativeFilterItems[formFilter?.filterType]?.value
     ?.datasourceCount;
 
+  const datasetId =
+    formFilter?.dataset?.value ??
+    filterToEdit?.targets[0]?.datasetId ??
+    mostUsedDataset(loadedDatasets, charts);
+
   const { controlItems = {}, mainControlItems = {} } = formFilter
     ? getControlItemsMap({
+        datasetId,
         disabled: false,
         forceUpdate,
         form,
@@ -372,10 +378,9 @@ const FiltersConfigForm = (
   const nativeFilterItem = nativeFilterItems[formFilter?.filterType] ?? {};
   // @ts-ignore
   const enableNoResults = !!nativeFilterItem.value?.enableNoResults;
-  const datasetId = formFilter?.dataset?.value;
 
   useEffect(() => {
-    if (datasetId && hasColumn) {
+    if (datasetId) {
       cachedSupersetGet({
         endpoint: `/api/v1/dataset/${datasetId}`,
       })
@@ -391,7 +396,7 @@ const FiltersConfigForm = (
           addDangerToast(response.message);
         });
     }
-  }, [datasetId, hasColumn]);
+  }, [datasetId]);
 
   useImperativeHandle(ref, () => ({
     changeTab(tab: 'configuration' | 'scoping') {
@@ -491,10 +496,6 @@ const FiltersConfigForm = (
     [filterId, forceUpdate, form, formFilter, hasDataset],
   );
 
-  const initialDatasetId =
-    filterToEdit?.targets[0]?.datasetId ??
-    mostUsedDataset(loadedDatasets, charts);
-
   const newFormData = getFormData({
     datasetId,
     groupby: hasColumn ? formFilter?.column : undefined,
@@ -508,6 +509,8 @@ const FiltersConfigForm = (
     setHasDefaultValue,
   ] = useDefaultValue(formFilter, filterToEdit);
 
+  const showDataset = !datasetId || datasetDetails;
+
   useEffect(() => {
     if (hasDataset && hasFilledDataset && hasDefaultValue && isDataDirty) {
       refreshHandler();
@@ -519,6 +522,7 @@ const FiltersConfigForm = (
     formFilter,
     isDataDirty,
     refreshHandler,
+    showDataset,
   ]);
 
   const updateFormValues = useCallback(
@@ -713,24 +717,29 @@ const FiltersConfigForm = (
         </StyledContainer>
         {hasDataset && (
           <StyledRowContainer>
-            <StyledFormItem
-              name={['filters', filterId, 'dataset']}
-              initialValue={{ value: initialDatasetId }}
-              label={<StyledLabel>{t('Dataset')}</StyledLabel>}
-              rules={[
-                { required: !removed, message: t('Dataset is required') },
-              ]}
-              {...getFiltersConfigModalTestId('datasource-input')}
-            >
-              {!datasetId || !hasColumn || datasetDetails ? (
+            {showDataset ? (
+              <StyledFormItem
+                name={['filters', filterId, 'dataset']}
+                label={<StyledLabel>{t('Dataset')}</StyledLabel>}
+                initialValue={
+                  datasetDetails
+                    ? {
+                        label: datasetDetails.table_name,
+                        value: datasetDetails.id,
+                      }
+                    : undefined
+                }
+                rules={[
+                  { required: !removed, message: t('Dataset is required') },
+                ]}
+                {...getFiltersConfigModalTestId('datasource-input')}
+              >
                 <DatasetSelect
-                  datasetDetails={datasetDetails}
-                  datasetId={initialDatasetId}
-                  onChange={(value: number) => {
+                  onChange={(value: { label: string; value: number }) => {
                     // We need to reset the column when the dataset has changed
-                    if (value !== datasetId) {
+                    if (value.value !== datasetId) {
                       setNativeFilterFieldValues(form, filterId, {
-                        dataset: { value },
+                        dataset: value,
                         defaultDataMask: null,
                         column: null,
                       });
@@ -738,10 +747,12 @@ const FiltersConfigForm = (
                     forceUpdate();
                   }}
                 />
-              ) : (
+              </StyledFormItem>
+            ) : (
+              <StyledFormItem label={<StyledLabel>{t('Dataset')}</StyledLabel>}>
                 <Loading position="inline-centered" />
-              )}
-            </StyledFormItem>
+              </StyledFormItem>
+            )}
             {hasDataset &&
               Object.keys(mainControlItems).map(
                 key => mainControlItems[key].element,
@@ -784,11 +795,8 @@ const FiltersConfigForm = (
                   required={hasDefaultValue}
                   rules={[
                     {
-                      validator: (rule, value) => {
-                        const hasValue =
-                          value?.filterState?.value !== null &&
-                          value?.filterState?.value !== undefined;
-                        if (hasValue) {
+                      validator: () => {
+                        if (formFilter?.defaultDataMask?.filterState?.value) {
                           return Promise.resolve();
                         }
                         return Promise.reject(
diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.test.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.test.tsx
index 7c33f394496f..f2f62983f360 100644
--- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.test.tsx
+++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.test.tsx
@@ -63,6 +63,7 @@ const filterMock: Filter = {
 };
 
 const createProps: () => ControlItemsProps = () => ({
+  datasetId: 1,
   disabled: false,
   forceUpdate: jest.fn(),
   form: formMock,
diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.tsx
index 24a938cb5f63..65ac4fbca259 100644
--- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.tsx
+++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigForm/getControlItemsMap.tsx
@@ -41,6 +41,7 @@ import { Filter } from '../../types';
 import { ColumnSelect } from './ColumnSelect';
 
 export interface ControlItemsProps {
+  datasetId: number;
   disabled: boolean;
   forceUpdate: Function;
   form: FormInstance<NativeFiltersForm>;
@@ -56,6 +57,7 @@ const CleanFormItem = styled(FormItem)`
 `;
 
 export default function getControlItemsMap({
+  datasetId,
   disabled,
   forceUpdate,
   form,
@@ -87,7 +89,6 @@ export default function getControlItemsMap({
         filterToEdit?.controlValues?.[mainControlItem.name] ??
         mainControlItem?.config?.default;
       const initColumn = filterToEdit?.targets[0]?.column?.name;
-      const datasetId = formFilter?.dataset?.value;
 
       const element = (
         <>
diff --git a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx
index e3b04c254f0f..a5497d00d8cb 100644
--- a/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx
+++ b/superset-frontend/src/dashboard/components/nativeFilters/FiltersConfigModal/FiltersConfigModal.test.tsx
@@ -28,7 +28,7 @@ import {
   TimeGrainFilterPlugin,
 } from 'src/filters/components';
 import { render, screen, waitFor } from 'spec/helpers/testing-library';
-import mockDatasource, { datasourceId } from 'spec/fixtures/mockDatasource';
+import mockDatasource, { id, datasourceId } from 'spec/fixtures/mockDatasource';
 import chartQueries from 'spec/fixtures/mockChartQueries';
 import {
   FiltersConfigModal,
@@ -71,9 +71,9 @@ const noTemporalColumnsState = () => {
   };
 };
 
-fetchMock.get('glob:*/api/v1/dataset/1', {
+const datasetResult = (id: number) => ({
   description_columns: {},
-  id: 1,
+  id,
   label_columns: {
     columns: 'Columns',
     table_name: 'Table Name',
@@ -87,11 +87,14 @@ fetchMock.get('glob:*/api/v1/dataset/1', {
       },
     ],
     table_name: 'birth_names',
-    id: 1,
+    id,
   },
   show_columns: ['id', 'table_name'],
 });
 
+fetchMock.get('glob:*/api/v1/dataset/1', datasetResult(1));
+fetchMock.get(`glob:*/api/v1/dataset/${id}`, datasetResult(id));
+
 fetchMock.post('glob:*/api/v1/chart/data', {
   result: [
     {
@@ -320,6 +323,11 @@ test('validates the pre-filter value', async () => {
 
 test("doesn't render time range pre-filter if there are no temporal columns in datasource", async () => {
   defaultRender(noTemporalColumnsState());
+  userEvent.click(screen.getByText(DATASET_REGEX));
+  await waitFor(() => {
+    expect(screen.queryByLabelText('Loading')).not.toBeInTheDocument();
+    userEvent.click(screen.getByText('birth_names'));
+  });
   userEvent.click(screen.getByText(ADVANCED_REGEX));
   userEvent.click(getCheckbox(PRE_FILTER_REGEX));
   await waitFor(() =>