[discover] fix too many queries sent (#8347)

Fixes the issue where too many queries were being fired. A number of things were being called because the dependency array was triggering updates and re-calls. Forcing a re-render of elements that had the possibility of updating the query manager. Includes some tests. Signed-off-by: Kawika Avilla <[email protected]> --------- Signed-off-by: Kawika Avilla <[email protected]> Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>
opensearch-project · Sep 26, 2024 · 0882435 · 0882435
1 parent d1caa92
commit 0882435
Show file tree

Hide file tree

Showing 8 changed files with 243 additions and 57 deletions.
diff --git a/changelogs/fragments/8347.yml b/changelogs/fragments/8347.yml
@@ -0,0 +1,2 @@
+fix:
+- Prevent too many queries sent from dataset selector ([#8347](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8347))
diff --git a/src/plugins/data/common/search/search_source/search_source.ts b/src/plugins/data/common/search/search_source/search_source.ts
@@ -431,7 +431,6 @@ export class SearchSource {
 
     const params = getExternalSearchParamsFromRequest(searchRequest, {
       getConfig,
-      getDataFrame: this.getDataFrame.bind(this),
     });
 
     return search({ params }, options).then(async (response: any) => {

diff --git a/src/plugins/data/public/query/query_string/query_string_manager.ts b/src/plugins/data/public/query/query_string/query_string_manager.ts
@@ -60,7 +60,7 @@ export class QueryStringManager {
     return this.storage.get('userQueryString') || '';
   }
 
-  public getDefaultQuery() {
+  public getDefaultQuery(): Query {
     const defaultLanguageId = this.getDefaultLanguage();
     const defaultQuery = this.getDefaultQueryString();
     const defaultDataset = this.datasetService?.getDefault();

diff --git a/src/plugins/data/public/ui/dataset_selector/dataset_selector.test.tsx b/src/plugins/data/public/ui/dataset_selector/dataset_selector.test.tsx
@@ -0,0 +1,116 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { mount } from 'enzyme';
+import { DatasetSelector } from './dataset_selector';
+import { getQueryService } from '../../services';
+import { CoreStart } from 'opensearch-dashboards/public';
+import { DataStorage, Dataset } from '../../../common';
+import { IDataPluginServices, DataPublicPluginStart } from '../../types';
+import { DatasetTypeConfig } from '../..';
+
+jest.mock('../../services', () => ({
+  getQueryService: jest.fn(),
+}));
+
+describe('DatasetSelector', () => {
+  const nextTick = () => new Promise((res) => process.nextTick(res));
+
+  const fetchMock = jest.fn().mockResolvedValue({ children: [] });
+  const setDatasetsMock = jest.fn();
+  const mockDatasetTypeConfig: DatasetTypeConfig = {
+    id: 'mockType',
+    title: 'Mock Type',
+    meta: { icon: { type: 'mockIcon' }, tooltip: 'Mock Tooltip' },
+    fetch: fetchMock,
+    toDataset: jest.fn(),
+    fetchFields: jest.fn(),
+    supportedLanguages: jest.fn().mockReturnValue(['mockLanguage']),
+  };
+
+  const getTypeMock = jest.fn().mockReturnValue(mockDatasetTypeConfig);
+
+  const mockServices: IDataPluginServices = {
+    appName: 'testApp',
+    uiSettings: {} as CoreStart['uiSettings'],
+    savedObjects: {} as CoreStart['savedObjects'],
+    notifications: ({
+      toasts: {
+        addSuccess: jest.fn(),
+        addError: jest.fn(),
+      },
+    } as unknown) as CoreStart['notifications'],
+    http: {} as CoreStart['http'],
+    storage: {} as DataStorage,
+    data: {} as DataPublicPluginStart,
+    overlays: ({
+      openModal: jest.fn(),
+    } as unknown) as CoreStart['overlays'],
+  };
+
+  beforeEach(() => {
+    (getQueryService as jest.Mock).mockReturnValue({
+      queryString: {
+        getDatasetService: jest.fn().mockReturnValue({
+          getType: getTypeMock,
+        }),
+      },
+    });
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('should fetch datasets once on mount', async () => {
+    const props = {
+      selectedDataset: undefined as Dataset | undefined,
+      setSelectedDataset: jest.fn(),
+      services: mockServices,
+    };
+
+    const wrapper = mount(<DatasetSelector {...props} />);
+
+    await nextTick();
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should not fetch datasets on re-render', async () => {
+    const props = {
+      selectedDataset: undefined as Dataset | undefined,
+      setSelectedDataset: jest.fn(),
+      services: mockServices,
+    };
+
+    const wrapper = mount(<DatasetSelector {...props} />);
+    await nextTick();
+
+    // Simulate a re-render
+    wrapper.setProps({});
+    await nextTick();
+
+    expect(fetchMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should not update datasets state on re-render', async () => {
+    const props = {
+      selectedDataset: undefined as Dataset | undefined,
+      setSelectedDataset: jest.fn(),
+      services: mockServices,
+    };
+
+    const wrapper = mount(<DatasetSelector {...props} />);
+    await nextTick();
+
+    setDatasetsMock.mockClear();
+
+    // Simulate a re-render
+    wrapper.setProps({});
+    await nextTick();
+
+    expect(setDatasetsMock).not.toHaveBeenCalled();
+  });
+});
diff --git a/src/plugins/data/public/ui/dataset_selector/dataset_selector.tsx b/src/plugins/data/public/ui/dataset_selector/dataset_selector.tsx
@@ -27,58 +27,63 @@ interface DatasetSelectorProps {
   services: IDataPluginServices;
 }
 
+/**
+ * This component provides a dropdown selector for datasets and an advanced selector modal.
+ * It fetches datasets once on mount to populate the selector options.
+ *
+ * @remarks
+ * The component uses several optimizations to prevent unnecessary re-renders:
+ * 1. It uses `useRef` and `useEffect` to ensure datasets are fetched only once on mount.
+ * 2. It uses `useMemo` and `useCallback` to memoize computed values and functions.
+ * 3. It intentionally omits some dependencies from the `useEffect` dependency array to prevent re-fetching.
+ */
 export const DatasetSelector = ({
   selectedDataset,
   setSelectedDataset,
   services,
 }: DatasetSelectorProps) => {
+  const isMounted = useRef(false);
   const [isOpen, setIsOpen] = useState(false);
   const [datasets, setDatasets] = useState<Dataset[]>([]);
   const { overlays } = services;
-
-  const isMounted = useRef(true);
-
-  useEffect(() => {
-    return () => {
-      isMounted.current = false;
-    };
-  }, []);
-
   const datasetService = getQueryService().queryString.getDatasetService();
-
   const datasetIcon =
     datasetService.getType(selectedDataset?.type || '')?.meta.icon.type || 'database';
 
-  const fetchDatasets = useCallback(async () => {
-    const typeConfig = datasetService.getType(DEFAULT_DATA.SET_TYPES.INDEX_PATTERN);
-    if (!typeConfig) return;
+  useEffect(() => {
+    isMounted.current = true;
+    const fetchDatasets = async () => {
+      if (!isMounted.current) return;
 
-    const fetchedIndexPatternDataStructures = await typeConfig.fetch(services, []);
+      const typeConfig = datasetService.getType(DEFAULT_DATA.SET_TYPES.INDEX_PATTERN);
+      if (!typeConfig) return;
 
-    if (!isMounted.current) return;
+      const fetchedIndexPatternDataStructures = await typeConfig.fetch(services, []);
 
-    const fetchedDatasets =
-      fetchedIndexPatternDataStructures.children?.map((pattern) =>
-        typeConfig.toDataset([pattern])
-      ) ?? [];
-    setDatasets(fetchedDatasets);
+      const fetchedDatasets =
+        fetchedIndexPatternDataStructures.children?.map((pattern) =>
+          typeConfig.toDataset([pattern])
+        ) ?? [];
+      setDatasets(fetchedDatasets);
 
-    // If no dataset is selected, select the first one
-    if (!selectedDataset && fetchedDatasets.length > 0) {
-      setSelectedDataset(fetchedDatasets[0]);
-    }
-  }, [datasetService, selectedDataset, services, setSelectedDataset]);
+      // If no dataset is selected, select the first one
+      if (!selectedDataset && fetchedDatasets.length > 0) {
+        setSelectedDataset(fetchedDatasets[0]);
+      }
+    };
 
-  useEffect(() => {
     fetchDatasets();
-  }, [fetchDatasets]);
+
+    return () => {
+      isMounted.current = false;
+    };
+    // NOTE: Intentionally omitting dependencies which can cause unnecessary re-renders
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [datasetService]);
 
   const togglePopover = useCallback(async () => {
-    if (!isOpen) {
-      await fetchDatasets();
-    }
     setIsOpen(!isOpen);
-  }, [isOpen, fetchDatasets]);
+  }, [isOpen]);
 
   const closePopover = useCallback(() => setIsOpen(false), []);
 

diff --git a/src/plugins/data/public/ui/dataset_selector/index.test.tsx b/src/plugins/data/public/ui/dataset_selector/index.test.tsx
@@ -0,0 +1,72 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import { mount } from 'enzyme';
+import { DatasetSelector as ConnectedDatasetSelector } from './index';
+import { DatasetSelector } from './dataset_selector';
+import { useOpenSearchDashboards } from '../../../../opensearch_dashboards_react/public';
+import { Dataset } from '../../../common';
+
+jest.mock('../../../../opensearch_dashboards_react/public', () => ({
+  useOpenSearchDashboards: jest.fn(),
+}));
+
+jest.mock('./dataset_selector', () => ({
+  DatasetSelector: jest.fn(() => null),
+}));
+
+describe('ConnectedDatasetSelector', () => {
+  const mockQueryString = {
+    getQuery: jest.fn().mockReturnValue({}),
+    getDefaultQuery: jest.fn().mockReturnValue({}),
+    getInitialQueryByDataset: jest.fn().mockReturnValue({}),
+    setQuery: jest.fn(),
+  };
+  const mockOnSubmit = jest.fn();
+  const mockServices = {
+    data: {
+      query: {
+        // @ts-ignore
+        queryString: mockQueryString,
+      },
+    },
+  };
+
+  beforeEach(() => {
+    (useOpenSearchDashboards as jest.Mock).mockReturnValue({ services: mockServices });
+  });
+
+  it('should render DatasetSelector with correct props', () => {
+    const wrapper = mount(<ConnectedDatasetSelector onSubmit={mockOnSubmit} />);
+    expect(wrapper.find(DatasetSelector).props()).toEqual({
+      selectedDataset: undefined,
+      setSelectedDataset: expect.any(Function),
+      services: mockServices,
+    });
+  });
+
+  it('should initialize selectedDataset correctly', () => {
+    const mockDataset: Dataset = { id: 'initial', title: 'Initial Dataset', type: 'test' };
+    mockQueryString.getQuery.mockReturnValueOnce({ dataset: mockDataset });
+
+    const wrapper = mount(<ConnectedDatasetSelector onSubmit={mockOnSubmit} />);
+    expect(wrapper.find(DatasetSelector).prop('selectedDataset')).toEqual(mockDataset);
+  });
+
+  it('should call handleDatasetChange only once when dataset changes', () => {
+    const wrapper = mount(<ConnectedDatasetSelector onSubmit={mockOnSubmit} />);
+    const setSelectedDataset = wrapper.find(DatasetSelector).prop('setSelectedDataset') as (
+      dataset?: Dataset
+    ) => void;
+
+    const newDataset: Dataset = { id: 'test', title: 'Test Dataset', type: 'test' };
+    setSelectedDataset(newDataset);
+
+    expect(mockQueryString.getInitialQueryByDataset).toHaveBeenCalledTimes(1);
+    expect(mockQueryString.setQuery).toHaveBeenCalledTimes(1);
+    expect(mockOnSubmit).toHaveBeenCalledTimes(1);
+  });
+});
diff --git a/src/plugins/data/public/ui/dataset_selector/index.tsx b/src/plugins/data/public/ui/dataset_selector/index.tsx
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { useEffect, useState } from 'react';
+import { useCallback, useState } from 'react';
 import React from 'react';
 import { Dataset, Query, TimeRange } from '../../../common';
 import { DatasetSelector } from './dataset_selector';
@@ -17,25 +17,21 @@ interface ConnectedDatasetSelectorProps {
 const ConnectedDatasetSelector = ({ onSubmit }: ConnectedDatasetSelectorProps) => {
   const { services } = useOpenSearchDashboards<IDataPluginServices>();
   const queryString = services.data.query.queryString;
-  const initialDataset = queryString.getQuery().dataset || queryString.getDefaultQuery().dataset;
-  const [selectedDataset, setSelectedDataset] = useState<Dataset | undefined>(initialDataset);
-
-  useEffect(() => {
-    const subscription = queryString.getUpdates$().subscribe((query) => {
-      setSelectedDataset(query.dataset);
-    });
-
-    return () => subscription.unsubscribe();
-  }, [queryString]);
+  const [selectedDataset, setSelectedDataset] = useState<Dataset | undefined>(
+    () => queryString.getQuery().dataset || queryString.getDefaultQuery().dataset
+  );
 
-  const handleDatasetChange = (dataset?: Dataset) => {
-    setSelectedDataset(dataset);
-    if (dataset) {
-      const query = queryString.getInitialQueryByDataset(dataset);
-      queryString.setQuery(query);
-      onSubmit!(queryString.getQuery());
-    }
-  };
+  const handleDatasetChange = useCallback(
+    (dataset?: Dataset) => {
+      setSelectedDataset(dataset);
+      if (dataset) {
+        const query = queryString.getInitialQueryByDataset(dataset);
+        queryString.setQuery(query);
+        onSubmit!(queryString.getQuery());
+      }
+    },
+    [onSubmit, queryString]
+  );
 
   return (
     <DatasetSelector

diff --git a/src/plugins/discover/public/application/view_components/utils/update_search_source.ts b/src/plugins/discover/public/application/view_components/utils/update_search_source.ts
@@ -30,12 +30,8 @@ export const updateSearchSource = async ({
   histogramConfigs,
 }: Props) => {
   const { uiSettings, data } = services;
-  const queryDataset = data.query.queryString.getQuery().dataset;
 
-  const dataset =
-    indexPattern.id === queryDataset?.id
-      ? await data.indexPatterns.get(queryDataset?.id!)
-      : indexPattern;
+  const dataset = indexPattern;
 
   const sortForSearchSource = getSortForSearchSource(
     sort,