[Explorer] Data selector enhancement and refactoring adoptions (#1759)

* support new changes in data selector

Signed-off-by: Eric <menwe@amazon.com>

* add a missing dependency

Signed-off-by: Eric <menwe@amazon.com>

* changes addressing reviews

Signed-off-by: Eric <menwe@amazon.com>

* fix delete s3 data source issue

Signed-off-by: Eric <menwe@amazon.com>

* error handling

Signed-off-by: Eric <menwe@amazon.com>

* code clean up

Signed-off-by: Eric <menwe@amazon.com>

---------

Signed-off-by: Eric <menwe@amazon.com>
(cherry picked from commit 552d720)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

common/constants/data_sources.ts

@@ -90,3 +90,7 @@ export const ACCELERATION_AGGREGRATION_FUNCTIONS = [
 ];
 
 export const SPARK_PARTITION_INFO = `# Partition Information`;
+export const OBS_DEFAULT_CLUSTER = 'observability-default'; // prefix key for generating data source id for default cluster in data selector
+export const OBS_S3_DATA_SOURCE = 'observability-s3'; // prefix key for generating data source id for s3 data sources in data selector
+export const S3_DATA_SOURCE_GROUP_DISPLAY_NAME = 'Amazon S3'; // display group name for Amazon-managed-s3 data sources in data selector
+export const S3_DATA_SOURCE_GROUP_SPARK_DISPLAY_NAME = 'Spark'; // display group name for OpenSearch-spark-s3 data sources in data selector

common/constants/shared.ts

@@ -256,7 +256,7 @@ export const VISUALIZATION_ERROR = {
   NO_METRIC: 'Invalid Metric MetaData',
 };
 
-export const S3_DATASOURCE_TYPE = 'S3_DATASOURCE';
+export const S3_DATA_SOURCE_TYPE = 's3glue';
 
 export const ASYNC_QUERY_SESSION_ID = 'async-query-session-id';
 export const ASYNC_QUERY_DATASOURCE_CACHE = 'async-query-catalog-cache';

public/components/event_analytics/explorer/datasources/datasources_selection.tsx

@@ -5,8 +5,10 @@
 
 import React, { useCallback, useContext, useEffect, useMemo, useState } from 'react';
 import { batch, useDispatch, useSelector } from 'react-redux';
+import { htmlIdGenerator } from '@elastic/eui';
 import { LogExplorerRouterContext } from '../..';
 import {
+  DataSource,
   DataSourceGroup,
   DataSourceSelectable,
   DataSourceType,
@@ -20,6 +22,7 @@ import {
   DEFAULT_DATA_SOURCE_TYPE,
   DEFAULT_DATA_SOURCE_TYPE_NAME,
   INDEX_URL_PARAM_KEY,
+  OBS_DEFAULT_CLUSTER,
   OLLY_QUESTION_URL_PARAM_KEY,
   QUERY_LANGUAGE,
 } from '../../../../../common/constants/data_sources';
@@ -90,6 +93,21 @@ const removeDataSourceFromURLParams = (currURL: string) => {
   }
 };
 
+const getMatchedOption = (
+  dataSourceList: DataSourceGroup[],
+  dataSourceName: string,
+  dataSourceType: string
+) => {
+  if (!dataSourceName || !dataSourceType) return [];
+  for (const dsGroup of dataSourceList) {
+    const matchedOption = dsGroup.options.find(
+      (item) => item.type === dataSourceType && item.name === dataSourceName
+    );
+    if (matchedOption !== undefined) return [matchedOption];
+  }
+  return [];
+};
+
 export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
   const { dataSources, http } = coreRefs;
   const sqlService = new SQLService(http!);
@@ -99,7 +117,11 @@ export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
   const [activeDataSources, setActiveDataSources] = useState<DataSourceType[]>([]);
   const [dataSourceOptionList, setDataSourceOptionList] = useState<DataSourceGroup[]>([]);
   const [selectedSources, setSelectedSources] = useState<SelectedDataSource[]>(
-    getDataSourceState(explorerSearchMetadata.datasources)
+    getMatchedOption(
+      dataSourceOptionList,
+      explorerSearchMetadata.datasources?.[0]?.name || '',
+      explorerSearchMetadata.datasources?.[0]?.type || ''
+    )
   );
 
   /**
@@ -149,8 +171,14 @@ export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
   };
 
   useEffect(() => {
-    setSelectedSources(getDataSourceState(explorerSearchMetadata.datasources));
-  }, [explorerSearchMetadata.datasources]);
+    setSelectedSources(
+      getMatchedOption(
+        memorizedDataSourceOptionList,
+        explorerSearchMetadata.datasources?.[0]?.name || '',
+        explorerSearchMetadata.datasources?.[0]?.type || ''
+      )
+    );
+  }, [explorerSearchMetadata.datasources, dataSourceOptionList]);
 
   const handleDataSetFetchError = useCallback(() => {
     return (error: Error) => {
@@ -162,22 +190,33 @@ export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
    * Subscribe to data source updates and manage the active data sources state.
    */
   useEffect(() => {
-    const subscription = dataSources.dataSourceService.dataSources$.subscribe(
-      (currentDataSources) => {
+    const subscription = dataSources.dataSourceService
+      .getDataSources$()
+      .subscribe((currentDataSources: DataSource[]) => {
         // temporary solution for 2.11 to render OpenSearch / default cluster for observability
         // local indices and index patterns, while keep listing all index patterns for data explorer
         // it filters the registered index pattern data sources in data plugin, and attach default cluster
         // for all indices
         setActiveDataSources([
           new ObservabilityDefaultDataSource({
+            id: htmlIdGenerator(OBS_DEFAULT_CLUSTER)(DEFAULT_DATA_SOURCE_TYPE),
             name: DEFAULT_DATA_SOURCE_NAME,
             type: DEFAULT_DATA_SOURCE_TYPE,
-            metadata: null,
+            metadata: {
+              ui: {
+                label: DEFAULT_DATA_SOURCE_OBSERVABILITY_DISPLAY_NAME,
+                groupType: DEFAULT_DATA_SOURCE_OBSERVABILITY_DISPLAY_NAME,
+                selector: {
+                  displayDatasetsAsSource: false, // when true, selector UI will render data sets with source by calling getDataSets()
+                },
+              },
+            },
           }),
-          ...Object.values(currentDataSources).filter((ds) => ds.type !== DEFAULT_DATA_SOURCE_TYPE),
+          ...Object.values(currentDataSources).filter(
+            (ds) => ds.getType() !== DEFAULT_DATA_SOURCE_TYPE
+          ),
         ]);
-      }
-    );
+      });
 
     return () => subscription.unsubscribe();
   }, []);
@@ -259,6 +298,10 @@ export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
     });
   }, [dataSourceOptionList]);
 
+  const onRefresh = useCallback(() => {
+    dataSources.dataSourceService.reload();
+  }, [dataSources.dataSourceService]);
+
   return (
     <DataSourceSelectable
       className="dsc-selector"
@@ -267,9 +310,10 @@ export const DataSourceSelection = ({ tabId }: { tabId: string }) => {
       setDataSourceOptionList={setDataSourceOptionList}
       selectedSources={selectedSources}
       onDataSourceSelect={handleSourceChange}
-      onFetchDataSetError={handleDataSetFetchError}
       singleSelection={{ asPlainText: true }}
       dataSourceSelectorConfigs={DATA_SOURCE_SELECTOR_CONFIGS}
+      onGetDataSetError={handleDataSetFetchError}
+      onRefresh={onRefresh}
     />
   );
 };

public/components/event_analytics/explorer/log_explorer.tsx

@@ -5,7 +5,7 @@
 /* eslint-disable react-hooks/exhaustive-deps */
 import { isEmpty } from 'lodash';
 import React, { useContext, useEffect, useRef, useState } from 'react';
-import { batch, useSelector, useDispatch } from 'react-redux';
+import { useSelector } from 'react-redux';
 import { useHistory } from 'react-router-dom';
 import { LogExplorerRouterContext } from '..';
 import {
@@ -26,14 +26,6 @@ import {
 } from '../../../../common/constants/shared';
 import { coreRefs } from '../../../../public/framework/core_refs';
 
-import { init as initFields } from '../../event_analytics/redux/slices/field_slice';
-import { init as initPatterns } from '../../event_analytics/redux/slices/patterns_slice';
-import { init as initQueryResult } from '../../event_analytics/redux/slices/query_result_slice';
-import { init as initQuery } from '../../event_analytics/redux/slices/query_slice';
-import { init as initVisualizationConfig } from '../../event_analytics/redux/slices/viualization_config_slice';
-import { resetSummary as initQueryAssistSummary } from '../../event_analytics/redux/slices/query_assistant_summarization_slice';
-import { init as initSearchMetaData } from '../../event_analytics/redux/slices/search_meta_data_slice';
-
 const searchBarConfigs = {
   [TAB_EVENT_ID]: {
     showSaveButton: true,
@@ -60,7 +52,6 @@ export const LogExplorer = ({
   dataSourcePluggables,
 }: ILogExplorerProps) => {
   const history = useHistory();
-  const dispatch = useDispatch();
   const routerContext = useContext(LogExplorerRouterContext);
   const tabIds = useSelector(selectQueryTabs).queryTabIds.filter(
     (tabid: string) => !tabid.match(APP_ANALYTICS_TAB_ID_REGEX)
@@ -97,20 +88,8 @@ export const LogExplorer = ({
   useEffect(() => {
     if (!isEmpty(savedObjectId)) {
       dispatchSavedObjectId();
-    } else {
-      // below piece of code was added to simulate creating a new tab if saved obj isn't being loaded,
-      // since tabs being visually removed means 'new tab' cannot be created any other way
-      const tabId = tabIds[0];
-      batch(() => {
-        dispatch(initQuery({ tabId }));
-        dispatch(initQueryResult({ tabId }));
-        dispatch(initFields({ tabId }));
-        dispatch(initVisualizationConfig({ tabId }));
-        dispatch(initPatterns({ tabId }));
-        dispatch(initQueryAssistSummary({ tabId }));
-        dispatch(initSearchMetaData({ tabId }));
-      });
     }
+
     if (routerContext && routerContext.searchParams.has(CREATE_TAB_PARAM_KEY)) {
       // need to wait for current redux event loop to finish
       setImmediate(() => {

public/framework/datasources/obs_opensearch_datasource.ts

@@ -6,25 +6,26 @@
 import { DataSource } from '../../../../../src/plugins/data/public';
 
 interface DataSourceConfig {
+  id: string;
   name: string;
   type: string;
   metadata: any;
 }
 
 export class ObservabilityDefaultDataSource extends DataSource<any, any, any, any, any> {
-  constructor({ name, type, metadata }: DataSourceConfig) {
-    super(name, type, metadata);
+  constructor({ id, name, type, metadata }: DataSourceConfig) {
+    super({ id, name, type, metadata });
   }
 
-  async getDataSet(dataSetParams?: any) {
+  async getDataSet() {
     return ['Default data source'];
   }
 
   async testConnection(): Promise<boolean> {
     return true;
   }
 
-  async runQuery(queryParams: any) {
+  async runQuery() {
     return null;
   }
 }

public/framework/datasources/s3_datasource.ts

@@ -9,22 +9,23 @@ interface DataSourceConfig {
   name: string;
   type: string;
   metadata: any;
+  id: string;
 }
 
-export class S3DataSource extends DataSource<any, any, any, any, any> {
-  constructor({ name, type, metadata }: DataSourceConfig) {
-    super(name, type, metadata);
+export class S3DataSource extends DataSource {
+  constructor({ id, name, type, metadata }: DataSourceConfig) {
+    super({ id, name, type, metadata });
   }
 
-  async getDataSet(dataSetParams?: any) {
-    return [this.getName()];
+  async getDataSet() {
+    return { dataSets: [this.getName()] };
   }
 
-  async testConnection(): Promise<void> {
-    throw new Error('This operation is not supported for this class.');
+  async testConnection(): Promise<boolean> {
+    return true;
   }
 
-  async runQuery(queryParams: any) {
-    return null;
+  async runQuery() {
+    return { data: {} };
   }
 }

public/plugin.tsx

@@ -3,8 +3,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { i18n } from '@osd/i18n';
 import React from 'react';
+import { i18n } from '@osd/i18n';
+import { htmlIdGenerator } from '@elastic/eui';
 import {
   AppCategory,
   AppMountParameters,
@@ -46,7 +47,7 @@ import {
   observabilityTracesID,
   observabilityTracesPluginOrder,
   observabilityTracesTitle,
-  S3_DATASOURCE_TYPE,
+  S3_DATA_SOURCE_TYPE,
 } from '../common/constants/shared';
 import { QueryManager } from '../common/query_manager';
 import { AssociatedObject, CachedAcceleration } from '../common/types/data_connections';
@@ -95,6 +96,12 @@ import {
   ObservabilityStart,
   SetupDependencies,
 } from './types';
+import {
+  DATA_SOURCE_TYPES,
+  OBS_S3_DATA_SOURCE,
+  S3_DATA_SOURCE_GROUP_DISPLAY_NAME,
+  S3_DATA_SOURCE_GROUP_SPARK_DISPLAY_NAME,
+} from '../common/constants/data_sources';
 
 interface PublicConfig {
   query_assist: {
@@ -387,38 +394,63 @@ export class ObservabilityPlugin
     coreRefs.overlays = core.overlays;
 
     const { dataSourceService, dataSourceFactory } = startDeps.data.dataSources;
+    dataSourceFactory.registerDataSourceType(S3_DATA_SOURCE_TYPE, S3DataSource);
+
+    const getDataSourceTypeLabel = (type: string) => {
+      if (type === DATA_SOURCE_TYPES.S3Glue) return S3_DATA_SOURCE_GROUP_DISPLAY_NAME;
+      if (type === DATA_SOURCE_TYPES.SPARK) return S3_DATA_SOURCE_GROUP_SPARK_DISPLAY_NAME;
+      return `${type.charAt(0).toUpperCase()}${type.slice(1)}`;
+    };
 
     // register all s3 datasources
-    const registerS3Datasource = () => {
-      dataSourceFactory.registerDataSourceType(S3_DATASOURCE_TYPE, S3DataSource);
-      core.http.get(`${DATACONNECTIONS_BASE}`).then((s3DataSources) => {
-        s3DataSources.map((s3ds) => {
-          dataSourceService.registerDataSource(
-            dataSourceFactory.getDataSourceInstance(S3_DATASOURCE_TYPE, {
-              name: s3ds.name,
-              type: s3ds.connector.toLowerCase(),
-              metadata: s3ds,
-            })
-          );
+    const registerDataSources = () => {
+      try {
+        core.http.get(`${DATACONNECTIONS_BASE}`).then((s3DataSources) => {
+          s3DataSources.map((s3ds) => {
+            dataSourceService.registerDataSource(
+              dataSourceFactory.getDataSourceInstance(S3_DATA_SOURCE_TYPE, {
+                id: htmlIdGenerator(OBS_S3_DATA_SOURCE)(),
+                name: s3ds.name,
+                type: s3ds.connector.toLowerCase(),
+                metadata: {
+                  ...s3ds.properties,
+                  ui: {
+                    label: s3ds.name,
+                    typeLabel: getDataSourceTypeLabel(s3ds.connector.toLowerCase()),
+                    groupType: s3ds.connector.toLowerCase(),
+                    selector: {
+                      displayDatasetsAsSource: false,
+                    },
+                  },
+                },
+              })
+            );
+          });
         });
-      });
+      } catch (error) {
+        console.error('Error registering S3 datasources', error);
+      }
     };
 
+    dataSourceService.registerDataSourceFetchers([
+      { type: S3_DATA_SOURCE_TYPE, registerDataSources },
+    ]);
+
     if (startDeps.securityDashboards) {
       core.http
         .get(SECURITY_PLUGIN_ACCOUNT_API)
         .then(() => {
-          registerS3Datasource();
+          registerDataSources();
         })
         .catch((e) => {
           if (e?.response?.status !== 401) {
             // accounts api should not return any error status other than 401 if security installed,
             // this datasource register is included just in case
-            registerS3Datasource();
+            registerDataSources();
           }
         });
     } else {
-      registerS3Datasource();
+      registerDataSources();
     }
 
     core.http.intercept({