From 335978d6da0473b3a433cc932b016ae008ee16bb Mon Sep 17 00:00:00 2001
From: Kawika Avilla <kavilla414@gmail.com>
Date: Mon, 29 Jul 2024 21:05:53 +0000
Subject: [PATCH] fix initial load

Signed-off-by: Kawika Avilla <kavilla414@gmail.com>
---
 .../dataset_navigator/dataset_navigator.tsx   | 48 +++++++++++--------
 1 file changed, 27 insertions(+), 21 deletions(-)

diff --git a/src/plugins/data/public/ui/dataset_navigator/dataset_navigator.tsx b/src/plugins/data/public/ui/dataset_navigator/dataset_navigator.tsx
index ee4b976670f5..cc0269c741ee 100644
--- a/src/plugins/data/public/ui/dataset_navigator/dataset_navigator.tsx
+++ b/src/plugins/data/public/ui/dataset_navigator/dataset_navigator.tsx
@@ -328,27 +328,32 @@ export const DataSetNavigator: React.FC<DataSetNavigatorProps> = ({
     async (source: SimpleDataSource) => {
       if (source) {
         setIsLoading(true);
-        const indices = await fetchIndices(searchService, source.id);
-        setSelectedDataSetState((prevState) => {
-          if (!prevState) return;
-
-          return {
-            ...prevState,
-            dataSourceRef: {
-              ...source,
-              indices: indices.map((indexName: string) => ({
-                id: indexName,
-                title: indexName,
-                dataSourceRef: {
-                  id: source.id,
-                  name: source.name,
-                  type: source.type,
-                },
-              })),
-            },
+        try {
+          const indices = await fetchIndices(searchService, source.id);
+          const updatedSource = {
+            ...source,
+            indices: indices.map((indexName: string) => ({
+              id: indexName,
+              title: indexName,
+              dataSourceRef: {
+                id: source.id,
+                name: source.name,
+                type: source.type,
+              },
+            })),
           };
-        });
-        setIsLoading(false);
+
+          setDataSources((prevDataSources) =>
+            prevDataSources.map((ds) => (ds.id === source.id ? updatedSource : ds))
+          );
+
+          setSelectedDataSetState((prevState) => ({
+            ...prevState!,
+            dataSourceRef: updatedSource,
+          }));
+        } finally {
+          setIsLoading(false);
+        }
       }
     },
     [searchService]
@@ -485,7 +490,8 @@ export const DataSetNavigator: React.FC<DataSetNavigatorProps> = ({
           panel: 7,
           onClick: () => handleSelectedObject(object),
         })),
-        content: isLoading && !selectedDataSetState?.dataSourceRef && createLoadingSpinner(),
+        content:
+          isLoading && !selectedDataSetState?.dataSourceRef?.indices && createLoadingSpinner(),
       },
       {
         id: 4,