From 896a9735e503f8baa4b47594b93dd3ce47555f13 Mon Sep 17 00:00:00 2001
From: Sean Li <lnse@amazon.com>
Date: Tue, 17 Sep 2024 03:11:15 -0700
Subject: [PATCH] moving recent datasets to fetched datasets

Signed-off-by: Sean Li <lnse@amazon.com>
---
 .../ui/dataset_selector/dataset_selector.tsx  | 98 +++++++++++--------
 1 file changed, 57 insertions(+), 41 deletions(-)

diff --git a/src/plugins/data/public/ui/dataset_selector/dataset_selector.tsx b/src/plugins/data/public/ui/dataset_selector/dataset_selector.tsx
index 63ac10c04fd5..ea27f22da150 100644
--- a/src/plugins/data/public/ui/dataset_selector/dataset_selector.tsx
+++ b/src/plugins/data/public/ui/dataset_selector/dataset_selector.tsx
@@ -20,6 +20,7 @@ import { Dataset, DEFAULT_DATA } from '../../../common';
 import { IDataPluginServices } from '../../types';
 import { AdvancedSelector } from './advanced_selector';
 import { getQueryService } from '../../services';
+import { DatasetServiceContract } from '../../query';
 
 interface DatasetSelectorProps {
   selectedDataset?: Dataset;
@@ -33,7 +34,7 @@ export const DatasetSelector = ({
   services,
 }: DatasetSelectorProps) => {
   const [isOpen, setIsOpen] = useState(false);
-  const [datasets, setDatasets] = useState<Dataset[]>([]);
+  const [datasets, setDatasets] = useState<Array<{ dataset: Dataset; recent: boolean }>>([]);
   const { overlays } = services;
 
   const isMounted = useRef(true);
@@ -57,15 +58,26 @@ export const DatasetSelector = ({
 
     if (!isMounted.current) return;
 
+    const recentDatasets = datasetService.getRecentDatasets();
     const fetchedDatasets =
       fetchedIndexPatternDataStructures.children?.map((pattern) =>
         typeConfig.toDataset([pattern])
       ) ?? [];
-    setDatasets(fetchedDatasets);
+    const filteredFetchedDatasets = fetchedDatasets.filter(
+      (dataset) => !recentDatasets.some((ds) => ds.id === dataset.id)
+    );
+    const allDatasets = [
+      ...recentDatasets.map((ds) => ({ dataset: ds, recent: true })),
+      ...filteredFetchedDatasets.map((ds) => ({ dataset: ds, recent: false })),
+    ];
+    setDatasets(allDatasets);
 
     // If no dataset is selected, select the first one
     if (!selectedDataset && fetchedDatasets.length > 0) {
-      setSelectedDataset(fetchedDatasets[0]);
+      setSelectedDataset(
+        allDatasets.filter((ds) => ds.recent)[0].dataset ??
+          allDatasets.filter((ds) => !ds.recent)[0].dataset
+      );
     }
   }, [datasetService, selectedDataset, services, setSelectedDataset]);
 
@@ -83,57 +95,61 @@ export const DatasetSelector = ({
   const closePopover = useCallback(() => setIsOpen(false), []);
 
   const options = useMemo(() => {
-    const newOptions: EuiSelectableOption[] = [
-      {
-        label: 'Index patterns',
-        isGroupLabel: true,
-      },
-    ];
-
-    datasets.forEach(({ id, title, type, dataSource }) => {
-      const label = dataSource ? `${dataSource.title}::${title}` : title;
-      newOptions.push({
-        label,
-        checked: id === selectedDataset?.id ? 'on' : undefined,
-        key: id,
-        prepend: <EuiIcon type={datasetService.getType(type)!.meta.icon.type} />,
-      });
-    });
-    const recentlySelected: EuiSelectableOption[] = [
-      {
-        label: 'Recently selected data',
-        isGroupLabel: true,
-      },
-    ];
-    datasetService.getRecentDatasets().forEach(({ id, title, type, dataSource }) => {
-      const label = dataSource ? `${dataSource.title}::${title}` : title;
-      recentlySelected.push({
-        label,
-        checked: id === selectedDataset?.id ? 'on' : undefined,
-        key: id,
-        prepend: <EuiIcon type={datasetService.getType(type)!.meta.icon.type} />,
+    const buildDatasetOptions = (
+      groupLabel: string,
+      ds: Dataset[],
+      selectedDatasetId: string | undefined,
+      dsService: DatasetServiceContract
+    ): EuiSelectableOption[] => {
+      const datasetOptions: EuiSelectableOption[] = [
+        {
+          label: groupLabel,
+          isGroupLabel: true,
+        },
+      ];
+      ds.forEach(({ id, title, type, dataSource }) => {
+        const label = dataSource ? `${dataSource.title}::${title}` : title;
+        datasetOptions.push({
+          label,
+          checked: id === selectedDatasetId ? 'on' : undefined,
+          key: id,
+          prepend: <EuiIcon type={datasetService.getType(type)!.meta.icon.type} />,
+        });
       });
-    });
-    if (recentlySelected.length > 1) {
-      newOptions.unshift(...recentlySelected);
-    }
-    return newOptions;
+      return datasetOptions;
+    };
+
+    const recentDatasetOptions = buildDatasetOptions(
+      'Recently selected data',
+      datasets.filter((ds) => ds.recent).map((ds) => ds.dataset),
+      selectedDataset?.id,
+      datasetService
+    );
+    const indexPatternOptions = buildDatasetOptions(
+      'Index patterns',
+      datasets.filter((ds) => !ds.recent).map((ds) => ds.dataset),
+      selectedDataset?.id,
+      datasetService
+    );
+    return [...recentDatasetOptions, ...indexPatternOptions];
   }, [datasets, selectedDataset?.id, datasetService]);
 
   const handleOptionChange = useCallback(
     (newOptions: EuiSelectableOption[]) => {
       const selectedOption = newOptions.find((option) => option.checked === 'on');
       if (selectedOption) {
-        const foundDataset =
-          datasets.find((dataset) => dataset.id === selectedOption.key) ??
-          datasetService.getRecentDataset(selectedOption.key);
+        const recentDatasets = datasets.filter((ds) => ds.recent).map((ds) => ds.dataset);
+        const indexPatterns = datasets.filter((ds) => !ds.recent).map((ds) => ds.dataset);
+        const foundDataset = [...recentDatasets, ...indexPatterns].find(
+          (dataset) => dataset.id === selectedOption.key
+        );
         if (foundDataset) {
           closePopover();
           setSelectedDataset(foundDataset);
         }
       }
     },
-    [datasets, setSelectedDataset, closePopover, datasetService]
+    [datasets, setSelectedDataset, closePopover]
   );
 
   const datasetTitle = useMemo(() => {