[discover-next][bug] add back data set navigator to control state

src/plugins/data/common/data_frames/utils.ts

@@ -183,18 +183,18 @@
   };
 
   if (data && data.fields && data.fields.length > 0) {
     for (let index = 0; index < data.size; index++) {
       const hit: { [key: string]: any } = {};
       data.fields.forEach((field) => {
         hit[field.name] = field.values[index];
       });
       hits.push({
         _index: data.name,
         _source: hit,
       });
     }
   }
   searchResponse.hits.hits = hits;
 
   if (data.hasOwnProperty('aggs')) {
     const dataWithAggs = data as IDataFrameWithAggs;
@@ -291,14 +291,14 @@
   aggConfig?: DataFrameAggConfig
 ): Partial<IFieldType> | undefined => {
   if (queryConfig?.timeFieldName) {
     return {
       name: queryConfig.timeFieldName,
       type: 'date',
     };
   }
   const fields = data.schema || data.fields;
   if (!fields) {
     throw Error('Invalid time field');
   }
   return aggConfig && aggConfig.date_histogram && aggConfig.date_histogram.field
     ? fields.find((field) => field.name === aggConfig?.date_histogram?.field)
@@ -484,6 +484,11 @@
     id: id ?? DATA_FRAME_TYPES.DEFAULT,
     title: dataFrame.name,
     timeFieldName: getTimeField(dataFrame, dataFrame.meta?.queryConfig)?.name,
+    dataSourceRef: {
+      id: dataFrame.meta?.queryConfig?.dataSourceId,
+      name: dataFrame.meta?.queryConfig?.dataSourceName,
+      type: dataFrame.meta?.queryConfig?.dataSourceType,
+    },
     type: !id ? DATA_FRAME_TYPES.DEFAULT : undefined,
     fields: fields.reduce(flattenFields, {} as IndexPatternFieldMap),
   };

src/plugins/data/common/search/opensearch_search/types.ts

@@ -48,10 +48,6 @@ export interface ISearchOptions {
    * Use this option to enable support for long numerals.
    */
   withLongNumeralsSupport?: boolean;
-  /**
-   * Use this option to enable support for async.
-   */
-  isAsync?: boolean;
 }
 
 export type ISearchRequestParams<T = Record<string, any>> = {

src/plugins/data/public/antlr/opensearch_sql/code_completion.ts

@@ -21,7 +21,6 @@
 import { SqlErrorListener } from './sql_error_listerner';
 import { findCursorTokenIndex } from '../shared/cursor';
 import { openSearchSqlAutocompleteData } from './opensearch_sql_autocomplete';
-import { getUiSettings } from '../../services';
 import { SQL_SYMBOLS } from './constants';
 import { QuerySuggestion, QuerySuggestionGetFnArgs } from '../../autocomplete';
 import { fetchColumnValues, fetchTableSchemas } from '../shared/utils';
@@ -46,8 +45,8 @@
   query,
   services,
 }: QuerySuggestionGetFnArgs): Promise<QuerySuggestion[]> => {
   const { api } = services.uiSettings;
   const dataSetManager = services.data.query.dataSet;
   const suggestions = getOpenSearchSqlAutoCompleteSuggestions(query, {
     line: position?.lineNumber || selectionStart,
     column: position?.column || selectionEnd,
@@ -59,12 +58,12 @@
     // Fetch columns and values
     if ('suggestColumns' in suggestions && (suggestions.suggestColumns?.tables?.length ?? 0) > 0) {
       const tableNames = suggestions.suggestColumns?.tables?.map((table) => table.name) ?? [];
       const schemas = await fetchTableSchemas(tableNames, api, services);
 
       schemas.forEach((schema) => {
         if (schema.body?.fields?.length > 0) {
           const columns = schema.body.fields.find((col: any) => col.name === 'COLUMN_NAME');
           const fieldTypes = schema.body.fields.find((col: any) => col.name === 'DATA_TYPE');
           if (columns && fieldTypes) {
             finalSuggestions.push(
               ...columns.values.map((col: string, index: number) => ({
@@ -86,7 +85,7 @@
           tableNames,
           suggestions.suggestValuesForColumn as string,
           api,
-          connectionService
+          services
         );
         values.forEach((value) => {
           if (value.body?.fields?.length > 0) {

src/plugins/data/public/index.ts

@@ -446,6 +446,9 @@ export {
   // for BWC, keeping the old name
   IUiStart as DataPublicPluginStartUi,
   DataSetNavigator,
+  setAsyncSessionId,
+  getAsyncSessionId,
+  setAsyncSessionIdByObj,
 } from './ui';
 
 /**

src/plugins/data/public/query/dataset_manager/dataset_manager.test.ts

@@ -6,12 +6,13 @@
 import { DataSetManager } from './dataset_manager';
 import { coreMock } from '../../../../../core/public/mocks';
 import { SimpleDataSet } from '../../../common';
-
 describe('DataSetManager', () => {
   let service: DataSetManager;
 
   beforeEach(() => {
-    service = new DataSetManager(coreMock.createSetup().uiSettings);
+    const uiSettingsMock = coreMock.createSetup().uiSettings;
+    uiSettingsMock.get.mockReturnValue(true);
+    service = new DataSetManager(uiSettingsMock);
   });
 
   test('getUpdates$ emits initially and after data set changes', () => {
@@ -20,15 +21,15 @@ describe('DataSetManager', () => {
     obs$.subscribe((v) => {
       emittedValues.push(v);
     });
-    expect(emittedValues).toHaveLength(1);
-    expect(emittedValues[1]).toEqual(undefined);
+    expect(emittedValues).toHaveLength(0);
+    expect(emittedValues[0]).toEqual(undefined);
 
     const newDataSet: SimpleDataSet = { id: 'test_dataset', title: 'Test Dataset' };
     service.setDataSet(newDataSet);
-    expect(emittedValues).toHaveLength(2);
-    expect(emittedValues[1]).toEqual(newDataSet);
+    expect(emittedValues).toHaveLength(1);
+    expect(emittedValues[0]).toEqual(newDataSet);
 
     service.setDataSet({ ...newDataSet });
-    expect(emittedValues).toHaveLength(3);
+    expect(emittedValues).toHaveLength(2);
   });
 });

src/plugins/data/public/query/dataset_manager/dataset_manager.ts

@@ -5,7 +5,13 @@
 
 import { BehaviorSubject } from 'rxjs';
 import { CoreStart } from 'opensearch-dashboards/public';
-import { SIMPLE_DATA_SET_TYPES, SimpleDataSet, SimpleDataSource } from '../../../common';
+import { skip } from 'rxjs/operators';
+import {
+  SIMPLE_DATA_SET_TYPES,
+  SimpleDataSet,
+  SimpleDataSource,
+  UI_SETTINGS,
+} from '../../../common';
 import { IndexPatternsContract } from '../../index_patterns';
 
 export class DataSetManager {
@@ -24,7 +30,7 @@
   };
 
   public getUpdates$ = () => {
-    return this.dataSet$.asObservable();
+    return this.dataSet$.asObservable().pipe(skip(1));
   };
 
   public getDataSet = () => {
@@ -36,6 +42,7 @@
    * @param {Query} query
    */
   public setDataSet = (dataSet: SimpleDataSet | undefined) => {
+    if (!this.uiSettings.get(UI_SETTINGS.QUERY_ENHANCEMENTS_ENABLED)) return;
     this.dataSet$.next(dataSet);
   };
 
@@ -46,7 +53,7 @@
   public fetchDefaultDataSet = async (): Promise<SimpleDataSet | undefined> => {
     const defaultIndexPatternId = this.uiSettings.get('defaultIndex');
     if (!defaultIndexPatternId) {
       return undefined;
     }
 
     const indexPattern = await this.indexPatterns?.get(defaultIndexPatternId);
@@ -55,10 +62,10 @@
     }
 
     if (!indexPattern.id) {
       return undefined;
     }
 
     return {
       id: indexPattern.id,
       title: indexPattern.title,
       type: SIMPLE_DATA_SET_TYPES.INDEX_PATTERN,

src/plugins/data/public/query/state_sync/connect_to_query_state.ts

@@ -31,13 +31,19 @@
 import { Subscription } from 'rxjs';
 import { filter, map } from 'rxjs/operators';
 import _ from 'lodash';
+import { CoreStart } from 'opensearch-dashboards/public';
 import {
   BaseStateContainer,
   IOsdUrlStateStorage,
 } from '../../../../opensearch_dashboards_utils/public';
 import { QuerySetup, QueryStart } from '../query_service';
 import { QueryState, QueryStateChange } from './types';
-import { FilterStateStore, COMPARE_ALL_OPTIONS, compareFilters } from '../../../common';
+import {
+  FilterStateStore,
+  COMPARE_ALL_OPTIONS,
+  compareFilters,
+  UI_SETTINGS,
+} from '../../../common';
 import { validateTimeRange } from '../timefilter';
 
 /**
@@ -63,7 +69,8 @@
     filters: FilterStateStore;
     query: boolean;
     dataSet?: boolean;
-  }
+  },
+  uiSettings?: CoreStart['uiSettings']
 ) => {
   try {
     const syncKeys: Array<keyof QueryStateChange> = [];
@@ -74,13 +81,16 @@
       syncKeys.push('appFilters');
     }
     if (syncConfig.dataSet) {
       syncKeys.push('dataSet');
     }
 
     const initialStateFromURL: QueryState = OsdUrlStateStorage.get('_q') ?? {
       query: queryString.getDefaultQuery(),
       filters: filterManager.getAppFilters(),
-      dataSet: dataSet.getDefaultDataSet(),
+      ...(uiSettings &&
+        uiSettings.get(UI_SETTINGS.QUERY_ENHANCEMENTS_ENABLED) && {
+          dataSet: dataSet.getDataSet(),
+        }),
     };
 
     // set up initial '_q' flag in the URL to sync query and filter changes
@@ -98,11 +108,11 @@
 
     if (syncConfig.dataSet && !_.isEqual(initialStateFromURL.dataSet, dataSet.getDataSet())) {
       if (initialStateFromURL.dataSet) {
         dataSet.setDataSet(_.cloneDeep(initialStateFromURL.dataSet));
       } else {
         const defaultDataSet = await dataSet.getDefaultDataSet();
         if (defaultDataSet) {
           dataSet.setDataSet(defaultDataSet);
         }
       }
     }
@@ -140,7 +150,7 @@
               }
 
               if (syncConfig.dataSet && changes.dataSet) {
                 newState.dataSet = dataSet.getDataSet();
               }
 
               return newState;
@@ -310,7 +320,7 @@
             }
           }
           if (syncConfig.dataSet && changes.dataSet) {
             newState.dataSet = dataSet.getDataSet();
           }
           return newState;
         })
@@ -375,12 +385,12 @@
         const currentDataSet = dataSet.getDataSet();
         if (!_.isEqual(state.dataSet, currentDataSet)) {
           if (state.dataSet) {
             dataSet.setDataSet(state.dataSet);
           } else {
             const defaultDataSet = await dataSet.getDefaultDataSet();
             if (defaultDataSet) {
               dataSet.setDataSet(defaultDataSet);
               stateContainer.set({ ...stateContainer.get(), dataSet: defaultDataSet });
             }
           }
         }

src/plugins/data/public/query/state_sync/sync_state_with_url.ts

@@ -28,6 +28,7 @@
  * under the License.
  */
 
+import { CoreStart } from 'opensearch-dashboards/public';
 import {
   createStateContainer,
   IOsdUrlStateStorage,
@@ -37,6 +38,7 @@ import { QuerySetup, QueryStart } from '../query_service';
 import { connectToQueryState } from './connect_to_query_state';
 import { QueryState } from './types';
 import { FilterStateStore } from '../../../common/opensearch_query/filters';
+import { UI_SETTINGS } from '../../../common';
 
 const GLOBAL_STATE_STORAGE_KEY = '_g';
 
@@ -50,7 +52,8 @@ export const syncQueryStateWithUrl = (
     QueryStart | QuerySetup,
     'filterManager' | 'timefilter' | 'queryString' | 'dataSet' | 'state$'
   >,
-  osdUrlStateStorage: IOsdUrlStateStorage
+  osdUrlStateStorage: IOsdUrlStateStorage,
+  uiSettings?: CoreStart['uiSettings']
 ) => {
   const {
     timefilter: { timefilter },
@@ -61,7 +64,10 @@ export const syncQueryStateWithUrl = (
     time: timefilter.getTime(),
     refreshInterval: timefilter.getRefreshInterval(),
     filters: filterManager.getGlobalFilters(),
-    dataSet: dataSet.getDataSet(),
+    ...(uiSettings &&
+      uiSettings.get(UI_SETTINGS.QUERY_ENHANCEMENTS_ENABLED) && {
+        dataSet: dataSet.getDataSet(),
+      }),
   };
 
   // retrieve current state from `_g` url

src/plugins/data/public/search/search_service.ts

@@ -135,21 +135,21 @@
     { fieldFormats, indexPatterns }: SearchServiceStartDependencies
   ): ISearchStart {
     const search = ((request, options) => {
       const selectedLanguage = getQueryService().queryString.getQuery().language;
       const uiService = getUiService();
       const enhancement = uiService.Settings.getQueryEnhancements(selectedLanguage);
       uiService.Settings.setUiOverridesByUserQueryLanguage(selectedLanguage);
       const isEnhancedEnabled = uiSettings.get(UI_SETTINGS.QUERY_ENHANCEMENTS_ENABLED);
 
       if (enhancement) {
         if (!isEnhancedEnabled) {
           notifications.toasts.addWarning(
             `Query enhancements are disabled. Please enable to use: ${selectedLanguage}.`
           );
         }
         return enhancement.search.search(request, options);
       }
       return this.defaultSearchInterceptor.search(request, options);
     }) as ISearchGeneric;
 
     const loadingCount$ = new BehaviorSubject(0);
@@ -161,6 +161,7 @@
         if (this.dfCache.get() && this.dfCache.get()?.name !== dataFrame.name) {
           indexPatterns.clearCache(this.dfCache.get()!.name, false);
         }
+
         if (
           dataFrame.meta &&
           dataFrame.meta.queryConfig &&
@@ -173,13 +174,12 @@
         }
         this.dfCache.set(dataFrame);
         const dataSetName = `${dataFrame.meta?.queryConfig?.dataSourceId ?? ''}.${dataFrame.name}`;
         const existingIndexPattern = await indexPatterns.get(dataSetName, true);
         const dataSet = await indexPatterns.create(
           dataFrameToSpec(dataFrame, existingIndexPattern?.id ?? dataSetName),
           !existingIndexPattern?.id
         );
-        await indexPatterns.refreshFields(dataSet, true);
         indexPatterns.saveToCache(dataSetName, dataSet);
       },
       clear: () => {
         if (this.dfCache.get() === undefined) return;

src/plugins/data/public/ui/dataset_navigator/dataset_navigator.tsx

@@ -74,16 +74,16 @@
 }
 
 export const DataSetNavigator = (props: DataSetNavigatorProps) => {
   const { savedObjectsClient, http, dataSet: dataSetManager } = props;
   const searchService = getSearchService();
   const queryService = getQueryService();
   const uiService = getUiService();
   const indexPatternsService = getIndexPatterns();
   const notifications = getNotifications();
 
   const { dataSet } = useDataSetManager({ dataSetManager: dataSetManager! });
 
   const [navigatorState, setNavigatorState] = useState<DataSetNavigatorState>({
     isOpen: false,
     isLoading: false,
     isMounted: false,
@@ -97,37 +97,37 @@
     cachedTables: [],
   });
 
   const [selectedDataSetState, setSelectedDataSetState] = useState<SelectedDataSetState>();
 
   const {
     loadStatus: dataSourcesLoadStatus,
     loadExternalDataSources: startLoadingDataSources,
   } = useLoadExternalDataSourcesToCache(http!, notifications);
   const {
     loadStatus: databasesLoadStatus,
     startLoading: startLoadingDatabases,
   } = useLoadDatabasesToCache(http!, notifications);
   const { loadStatus: tablesLoadStatus, startLoading: startLoadingTables } = useLoadTablesToCache(
     http!,
     notifications
   );
 
   const onClick = () => {
     setNavigatorState((prevState) => ({
       ...prevState,
       isOpen: !prevState.isOpen,
     }));
   };
 
   const isLoading = (loading: boolean) => {
     setNavigatorState((prevState) => ({
       ...prevState,
       isLoading: loading,
     }));
   };
 
   const closePopover = () => {
     setNavigatorState((prevState) => ({
       ...prevState,
       isOpen: false,
       externalDataSources: [],
@@ -138,33 +138,33 @@
     }));
   };
 
   const onRefresh = () => {
     if (!isCatalogCacheFetching(dataSourcesLoadStatus) && navigatorState.dataSources.length > 0) {
       startLoadingDataSources(navigatorState.dataSources.map((dataSource) => dataSource.id));
     }
   };
 
   const handleSelectedDataSet = useCallback(
     async (ds?: SimpleDataSet) => {
       const selectedDataSet = ds ?? navigatorState.currentDataSet;
       if (!selectedDataSet || !selectedDataSet.id) return;
 
       const language = uiService.Settings.getUserQueryLanguage();
       const queryEnhancements = uiService.Settings.getQueryEnhancements(language);
       const initialInput = queryEnhancements?.searchBar?.queryStringInput?.initialValue;
 
       // Update query
       const query = initialInput
         ? initialInput.replace('<data_source>', selectedDataSet.title!)
         : '';
       uiService.Settings.setUserQueryString(query);
       queryService.queryString.setQuery({ query, language });
 
       // Update dataset
       queryService.dataSet.setDataSet(selectedDataSet);
 
       // Add to recent datasets
       CatalogCacheManager.addRecentDataSet({
         id: selectedDataSet.id,
         title: selectedDataSet.title ?? selectedDataSet.id!,
         dataSourceRef: selectedDataSet.dataSourceRef,
@@ -173,7 +173,7 @@
       });
 
       // Update data set manager
       dataSetManager!.setDataSet({
         id: selectedDataSet.id,
         title: selectedDataSet.title,
         ...(selectedDataSet.dataSourceRef && {
@@ -187,9 +187,7 @@
         type: selectedDataSet.type,
       });
 
-      setSelectedDataSetState(selectedDataSet);
-
       closePopover();
     },
     [
       dataSetManager,
@@ -200,9 +198,9 @@
     ]
   );
 
   useEffect(() => {
     setNavigatorState((prevState) => ({ ...prevState, isMounted: true, isLoading: true }));
     Promise.all([
       dataSetManager?.init(indexPatternsService),
       fetchIndexPatterns(savedObjectsClient!, ''),
       fetchDataSources(savedObjectsClient!),
@@ -210,7 +208,7 @@
     ])
       .then(([defaultDataSet, indexPatterns, dataSources, isExternalDataSourcesEnabled]) => {
         if (!navigatorState.isMounted) return;
         setNavigatorState((prevState) => ({
           ...prevState,
           isExternalDataSourcesEnabled,
           indexPatterns,
@@ -220,7 +218,7 @@
         const selectedPattern = dataSet ?? defaultDataSet;
 
         if (selectedPattern) {
           setSelectedDataSetState({
             id: selectedPattern.id,
             title: selectedPattern.title,
             type: selectedPattern.type,
@@ -240,28 +238,27 @@
         }
       })
       .finally(() => {
         isLoading(false);
       });
     return () => {
       setNavigatorState((prevState) => ({ ...prevState, isMounted: false }));
     };
   }, [
-    savedObjectsClient,
-    http,
-    navigatorState.isMounted,
     dataSet,
     dataSetManager,
+    http,
     indexPatternsService,
-    handleSelectedDataSet,
+    navigatorState.isMounted,
+    savedObjectsClient,
   ]);
 
   useEffect(() => {
     const status = dataSourcesLoadStatus.toLowerCase();
     const externalDataSourcesCache = CatalogCacheManager.getExternalDataSourcesCache();
     if (status === DirectQueryLoadingStatus.SUCCESS) {
       setNavigatorState((prevState) => ({
         ...prevState,
         externalDataSources: externalDataSourcesCache.externalDataSources.map((ds) => ({
           id: ds.dataSourceRef,
           name: ds.name,
           type: SIMPLE_DATA_SOURCE_TYPES.EXTERNAL,
@@ -271,22 +268,22 @@
       status === DirectQueryLoadingStatus.CANCELED ||
       status === DirectQueryLoadingStatus.FAILED
     ) {
       setNavigatorState((prevState) => ({ ...prevState, failed: true }));
     }
   }, [dataSourcesLoadStatus]);
 
   useEffect(() => {
     const status = databasesLoadStatus.toLowerCase();
     if (
       selectedDataSetState?.dataSourceRef &&
       selectedDataSetState.dataSourceRef.type === SIMPLE_DATA_SOURCE_TYPES.EXTERNAL
     ) {
       const dataSourceCache = CatalogCacheManager.getOrCreateDataSource(
         selectedDataSetState.dataSourceRef.name,
         selectedDataSetState.dataSourceRef.id
       );
       if (status === DirectQueryLoadingStatus.SUCCESS) {
         setNavigatorState((prevState) => ({
           ...prevState,
           cachedDatabases: dataSourceCache.databases,
         }));
@@ -294,15 +291,15 @@
         status === DirectQueryLoadingStatus.CANCELED ||
         status === DirectQueryLoadingStatus.FAILED
       ) {
         setNavigatorState((prevState) => ({ ...prevState, failed: true }));
       }
     }
   }, [databasesLoadStatus, selectedDataSetState?.dataSourceRef]);
 
   const handleSelectExternalDataSource = useCallback(
     async (dataSource) => {
       if (dataSource && dataSource.type === SIMPLE_DATA_SOURCE_TYPES.EXTERNAL) {
         const dataSourceCache = CatalogCacheManager.getOrCreateDataSource(
           dataSource.name,
           dataSource.id
         );
@@ -311,17 +308,17 @@
             dataSourceCache.status === CachedDataSourceStatus.Failed) &&
           !isCatalogCacheFetching(databasesLoadStatus)
         ) {
           await startLoadingDatabases({
             dataSourceName: dataSource.name,
             dataSourceMDSId: dataSource.id,
           });
         } else if (dataSourceCache.status === CachedDataSourceStatus.Updated) {
           setNavigatorState((prevState) => ({
             ...prevState,
             cachedDatabases: dataSourceCache.databases,
           }));
         }
         setSelectedDataSetState((prevState) => ({
           ...prevState,
           dataSourceRef: dataSource,
           isExternal: true,
@@ -332,33 +329,36 @@
   );
 
   // Start loading tables for selected database
   const handleSelectExternalDatabase = useCallback(
     (externalDatabase: SimpleDataSource) => {
       if (selectedDataSetState?.dataSourceRef && externalDatabase) {
         let databaseCache: CachedDatabase;
         try {
           databaseCache = CatalogCacheManager.getDatabase(
             selectedDataSetState.dataSourceRef.name,
             externalDatabase.name,
             selectedDataSetState.dataSourceRef.id
           );
         } catch (error) {
           return;
         }
         if (
           databaseCache.status === CachedDataSourceStatus.Empty ||
           (databaseCache.status === CachedDataSourceStatus.Failed &&
             !isCatalogCacheFetching(tablesLoadStatus))
         ) {
           startLoadingTables({
             dataSourceName: selectedDataSetState.dataSourceRef.name,
             databaseName: externalDatabase.name,
             dataSourceMDSId: selectedDataSetState.dataSourceRef.id,
           });
         } else if (databaseCache.status === CachedDataSourceStatus.Updated) {
           setNavigatorState((prevState) => ({
             ...prevState,
-            cachedTables: databaseCache.tables,
+            cachedTables: databaseCache.tables.map((table) => ({
+              id: table.name,
+              title: table.name,
+            })),
           }));
         }
       }
@@ -367,33 +367,36 @@
   );
 
   // Retrieve tables from cache upon success
   useEffect(() => {
     if (
       selectedDataSetState?.dataSourceRef &&
       selectedDataSetState.dataSourceRef?.type === SIMPLE_DATA_SOURCE_TYPES.EXTERNAL &&
       selectedDataSetState.database
     ) {
       const tablesStatus = tablesLoadStatus.toLowerCase();
       let databaseCache: CachedDatabase;
       try {
         databaseCache = CatalogCacheManager.getDatabase(
           selectedDataSetState.dataSourceRef.name,
           selectedDataSetState.database,
           selectedDataSetState.dataSourceRef.id
         );
       } catch (error) {
         return;
       }
       if (tablesStatus === DirectQueryLoadingStatus.SUCCESS) {
         setNavigatorState((prevState) => ({
           ...prevState,
-          cachedTables: databaseCache.tables,
+          cachedTables: databaseCache.tables.map((table) => ({
+            id: table.name,
+            title: table.name,
+          })),
         }));
       } else if (
         tablesStatus === DirectQueryLoadingStatus.CANCELED ||
         tablesStatus === DirectQueryLoadingStatus.FAILED
       ) {
         notifications.toasts.addWarning('Error loading tables');
       }
     }
   }, [
@@ -403,16 +406,16 @@
     notifications.toasts,
   ]);
 
   const handleSelectedDataSource = useCallback(
     async (source: SimpleDataSource) => {
       if (source) {
         isLoading(true);
         const indices = await fetchIndices(searchService, source.id);
         setNavigatorState((prevState) => ({
           ...prevState,
           currentDataSourceRef: {
             ...source,
             indices: indices.map((indexName: string) => ({
               id: indexName,
               title: indexName,
               dataSourceRef: {
@@ -423,24 +426,24 @@
             })),
           },
         }));
         isLoading(false);
       }
     },
     [searchService]
   );
 
   const handleSelectedObject = useCallback(
     async (object) => {
       if (object) {
         isLoading(true);
         const fields = await indexPatternsService.getFieldsForWildcard({
           pattern: object.title,
           dataSourceId: object.dataSourceRef?.id,
         });
 
         const timeFields = fields.filter((field: any) => field.type === 'date');
         const timeFieldName = timeFields?.length > 0 ? timeFields[0].name : undefined;
         setNavigatorState((prevState) => ({
           ...prevState,
           currentDataSet: {
             id: `${object.dataSourceRef ? object.dataSourceRef.id : ''}.${object.id}`,
@@ -448,69 +451,70 @@
             fields,
             timeFields,
             timeFieldName,
+            dataSourceRef: object.dataSourceRef,
             type: SIMPLE_DATA_SET_TYPES.TEMPORARY,
           },
         }));
         isLoading(false);
       }
     },
     [indexPatternsService]
   );
 
   const indexPatternsLabel = i18n.translate('data.query.dataSetNavigator.indexPatternsName', {
     defaultMessage: 'Index patterns',
   });
   const indicesLabel = i18n.translate('data.query.dataSetNavigator.indicesName', {
     defaultMessage: 'Indexes',
   });
   const S3DataSourcesLabel = i18n.translate('data.query.dataSetNavigator.S3DataSourcesLabel', {
     defaultMessage: 'S3',
   });
 
   const createRefreshButton = () => (
     <EuiButtonEmpty
       iconType="refresh"
       onClick={onRefresh}
       isLoading={isCatalogCacheFetching(databasesLoadStatus, tablesLoadStatus)}
     />
   );
 
   const createLoadingSpinner = () => (
     <EuiPanel hasShadow={false} hasBorder={false}>
       <EuiLoadingSpinner className="dataSetNavigator__loading" size="m" />
     </EuiPanel>
   );
 
   const createIndexPatternsPanel = () => ({
     id: 1,
     title: indexPatternsLabel,
     items: navigatorState.indexPatterns.map((indexPattern) => ({
       name: indexPattern.title,
       onClick: async () => await handleSelectedDataSet(indexPattern),
     })),
     content: <div>{navigatorState.indexPatterns.length === 0 && createLoadingSpinner()}</div>,
   });
 
   const createIndexesPanel = () => ({
     id: 2,
     title: indicesLabel,
     items: [
       ...navigatorState.dataSources.map((dataSource) => ({
         name: dataSource.name,
         panel: 3,
         onClick: async () => await handleSelectedDataSource(dataSource),
       })),
     ],
     content: <div>{navigatorState.isLoading && createLoadingSpinner()}</div>,
   });
 
   const createDataSourcesPanel = () => ({
     id: 3,
     title: navigatorState.currentDataSourceRef?.name ?? indicesLabel,
     items: navigatorState.currentDataSourceRef?.indices?.map((object) => ({
       name: object.title,
       panel: 7,
       onClick: async () => await handleSelectedObject(object),
     })),
     content: (
       <div>
@@ -519,7 +523,7 @@
     ),
   });
 
   const createS3DataSourcesPanel = () => ({
     id: 4,
     title: (
       <div>
@@ -529,29 +533,29 @@
       </div>
     ),
     items: [
       ...navigatorState.externalDataSources.map((dataSource) => ({
         name: dataSource.name,
         onClick: async () => await handleSelectExternalDataSource(dataSource),
         panel: 5,
       })),
     ],
     content: <div>{dataSourcesLoadStatus && createLoadingSpinner()}</div>,
   });
 
   const createDatabasesPanel = () => ({
     id: 5,
     title: selectedDataSetState?.dataSourceRef?.name
       ? selectedDataSetState.dataSourceRef?.name
       : 'Databases',
     items: [
       ...navigatorState.cachedDatabases.map((db) => ({
         name: db.name,
         onClick: async () => {
           setSelectedDataSetState((prevState) => ({
             ...prevState,
             database: db,
           }));
           await handleSelectExternalDatabase(db);
         },
         panel: 6,
       })),
@@ -559,7 +563,7 @@
     content: <div>{isCatalogCacheFetching(databasesLoadStatus) && createLoadingSpinner()}</div>,
   });
 
   return (
     <EuiPopover
       button={
         <EuiButtonEmpty
@@ -613,23 +617,23 @@
                       name: S3DataSourcesLabel,
                       panel: 4,
                       onClick: async () => {
                         const externalDataSourcesCache = CatalogCacheManager.getExternalDataSourcesCache();
                         if (
                           (externalDataSourcesCache.status === CachedDataSourceStatus.Empty ||
                             externalDataSourcesCache.status === CachedDataSourceStatus.Failed) &&
                           !isCatalogCacheFetching(dataSourcesLoadStatus) &&
                           navigatorState.dataSources.length > 0
                         ) {
                           startLoadingDataSources(
                             navigatorState.dataSources.map((dataSource) => dataSource.id)
                           );
                         } else if (
                           externalDataSourcesCache.status === CachedDataSourceStatus.Updated
                         ) {
                           setNavigatorState((prevState) => ({
                             ...prevState,
                             externalDataSources: externalDataSourcesCache.externalDataSources.map(
                               (ds) => ({
                                 id: ds.dataSourceRef,
                                 name: ds.name,
                                 type: SIMPLE_DATA_SOURCE_TYPES.EXTERNAL,
@@ -652,10 +656,10 @@
             id: 6,
             title: selectedDataSetState?.database ? selectedDataSetState.database.name : 'Tables',
             items: [
               ...navigatorState.cachedTables.map((table) => ({
                 name: table.name,
                 onClick: async () => {
                   const tableObject = {
                     ...selectedDataSetState,
                     id: `${selectedDataSetState?.dataSourceRef!.name}.${
                       selectedDataSetState?.database.name
@@ -670,8 +674,7 @@
                     },
                     type: SIMPLE_DATA_SET_TYPES.TEMPORARY_ASYNC,
                   };
-                  setSelectedDataSetState(tableObject);
-                  handleSelectedDataSet(tableObject);
+                  await handleSelectedDataSet(tableObject);
                 },
               })),
             ],
@@ -683,9 +686,7 @@
             id: 7,
             title: navigatorState.currentDataSet?.title,
             content:
-              navigatorState.isLoading ||
-              !navigatorState.currentDataSet ||
-              !navigatorState.currentDataSet?.title ? (
+              !navigatorState.currentDataSet || !navigatorState.currentDataSet?.title ? (
                 <div>{createLoadingSpinner()}</div>
               ) : (
                 <EuiForm className="dataSetNavigatorFormWrapper">
@@ -700,7 +701,7 @@
                         ...(navigatorState.currentDataSet!.timeFields &&
                         navigatorState.currentDataSet!.timeFields.length > 0
                           ? [
                               ...navigatorState.currentDataSet!.timeFields!.map((field: any) => ({
                                 value: field.name,
                                 text: field.name,
                               })),
@@ -709,7 +710,7 @@
                         { value: 'no-time-filter', text: "I don't want to use a time filter" },
                       ]}
                       onChange={(event) => {
                         setNavigatorState((prevState) => ({
                           ...prevState,
                           currentDataSet: {
                             ...prevState.currentDataSet!,
@@ -727,7 +728,7 @@
                     size="s"
                     fullWidth
                     onClick={async () => {
                       await handleSelectedDataSet();
                     }}
                   >
                     Select
@@ -738,17 +739,17 @@
           {
             id: 8,
             title: 'Recently Used',
             items: CatalogCacheManager.getRecentDataSets().map((ds) => ({
               name: ds.title,
               onClick: async () => {
                 setSelectedDataSetState({
                   id: ds.id ?? ds.title,
                   title: ds.title,
                   dataSourceRef: ds.dataSourceRef,
                   database: undefined,
                   timeFieldName: ds.timeFieldName,
                 });
                 await handleSelectedDataSet();
               },
             })),
           },

src/plugins/data/public/ui/dataset_navigator/index.tsx

@@ -5,3 +5,4 @@
 
 export { DataSetNavigator, DataSetNavigatorProps } from './dataset_navigator';
 export { createDataSetNavigator } from './create_dataset_navigator';
+export { setAsyncSessionId, getAsyncSessionId, setAsyncSessionIdByObj } from './lib';