add query assist icons to clusters if available

Signed-off-by: Joshua Li <[email protected]>

src/plugins/data/public/query/query_string/dataset_service/lib/index_pattern_type.ts

@@ -15,6 +15,7 @@ import {
 } from '../../../../../common';
 import { DatasetTypeConfig } from '../types';
 import { getIndexPatterns } from '../../../../services';
+import { injectMetaToDataStructures } from './utils';
 
 export const indexPatternTypeConfig: DatasetTypeConfig = {
   id: DEFAULT_DATA.SET_TYPES.INDEX_PATTERN,
@@ -97,7 +98,7 @@ const fetchIndexPatterns = async (client: SavedObjectsClientContract): Promise<D
     });
   }
 
-  return resp.savedObjects.map(
+  const dataStructures = resp.savedObjects.map(
     (savedObject): DataStructure => {
       const dataSourceId = savedObject.references.find((ref) => ref.type === 'data-source')?.id;
       const dataSource = dataSourceId ? dataSourceMap[dataSourceId] : undefined;
@@ -122,4 +123,6 @@ const fetchIndexPatterns = async (client: SavedObjectsClientContract): Promise<D
       return indexPatternDataStructure;
     }
   );
+
+  return injectMetaToDataStructures(dataStructures);
 };

src/plugins/data/public/query/query_string/dataset_service/lib/index_type.ts

@@ -8,6 +8,7 @@ import { map } from 'rxjs/operators';
 import { DEFAULT_DATA, DataStructure, Dataset } from '../../../../../common';
 import { DatasetTypeConfig } from '../types';
 import { getSearchService, getIndexPatterns } from '../../../../services';
+import { injectMetaToDataStructures } from './utils';
 
 const INDEX_INFO = {
   LOCAL_DATASOURCE: {
@@ -93,14 +94,15 @@ const fetchDataSources = async (client: SavedObjectsClientContract) => {
     type: 'data-source',
     perPage: 10000,
   });
-  const dataSources: DataStructure[] = [INDEX_INFO.LOCAL_DATASOURCE];
-  return dataSources.concat(
+  const dataSources: DataStructure[] = [INDEX_INFO.LOCAL_DATASOURCE].concat(
     resp.savedObjects.map((savedObject) => ({
       id: savedObject.id,
       title: savedObject.attributes.title,
       type: 'DATA_SOURCE',
     }))
   );
+
+  return injectMetaToDataStructures(dataSources);
 };
 
 const fetchIndices = async (dataStructure: DataStructure): Promise<string[]> => {

src/plugins/data/public/query/query_string/dataset_service/lib/utils.ts

@@ -0,0 +1,45 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { DataStructure, DataStructureMeta } from '../../../../../common';
+import { getQueryService } from '../../../../services';
+
+/**
+ * Inject {@link DataStructureMeta} to DataStructures based on
+ * {@link QueryEditorExtensions}.
+ *
+ * This function combines the meta fields from QueryEditorExtensions and in
+ * provided data structures. Lower extension order is higher priority, and
+ * existing meta fields have highest priority.
+ *
+ * @param dataStructures - {@link DataStructure}
+ * @returns data structures with meta
+ */
+export const injectMetaToDataStructures = async (dataStructures: DataStructure[]) => {
+  const queryEditorExtensions = Object.values(
+    getQueryService().queryString.getLanguageService().getQueryEditorExtensionMap()
+  );
+  queryEditorExtensions.sort((a, b) => b.order - a.order);
+
+  return Promise.all(
+    dataStructures.map(async (dataStructure) => {
+      const metaArray = await Promise.allSettled(
+        queryEditorExtensions.map((curr) => curr.getDataStructureMeta?.(dataStructure.id))
+      ).then((settledResults) =>
+        settledResults
+          .filter(
+            <T>(result: PromiseSettledResult<T>): result is PromiseFulfilledResult<T> =>
+              result.status === 'fulfilled'
+          )
+          .map((result) => result.value)
+      );
+      const meta = metaArray.reduce(
+        (acc, curr) => (acc || curr ? ({ ...acc, ...curr } as DataStructureMeta) : undefined),
+        undefined
+      );
+      return { meta, ...dataStructure };
+    })
+  );
+};

src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension.tsx

@@ -7,6 +7,7 @@ import { EuiErrorBoundary } from '@elastic/eui';
 import React, { useEffect, useMemo, useRef, useState } from 'react';
 import ReactDOM from 'react-dom';
 import { Observable } from 'rxjs';
+import { DataStructureMeta } from '../../../../common';
 
 interface QueryEditorExtensionProps {
   config: QueryEditorExtensionConfig;
@@ -48,6 +49,12 @@ export interface QueryEditorExtensionConfig {
    * @returns whether the extension is enabled.
    */
   isEnabled$: (dependencies: QueryEditorExtensionDependencies) => Observable<boolean>;
+  /**
+   * @returns DataStructureMeta for a given data source id.
+   */
+  getDataStructureMeta?: (
+    dataSourceId: string | undefined
+  ) => Promise<DataStructureMeta | undefined>;
   /**
    * A function that returns the query editor extension component. The component
    * will be displayed on top of the query editor in the search bar.

src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx

@@ -6,7 +6,7 @@
 import { HttpSetup } from 'opensearch-dashboards/public';
 import React, { useEffect, useState } from 'react';
 import { distinctUntilChanged, map, startWith, switchMap } from 'rxjs/operators';
-import { DEFAULT_DATA } from '../../../../data/common';
+import { DATA_STRUCTURE_META_TYPES, DEFAULT_DATA } from '../../../../data/common';
 import {
   DataPublicPluginSetup,
   QueryEditorExtensionConfig,
@@ -15,16 +15,32 @@ import {
 import { API } from '../../../common';
 import { ConfigSchema } from '../../../common/config';
 import { QueryAssistBanner, QueryAssistBar } from '../components';
+import assistantMark from '../../assets/query_assist_mark.svg';
+
+/**
+ * @returns list of query assist supported languages for the given data source.
+ */
+const getAvailableLanguagesForDataSource = (() => {
+  const availableLanguagesByDataSource: Map<string | undefined, string[]> = new Map();
+  return async (http: HttpSetup, dataSourceId: string | undefined) => {
+    const cached = availableLanguagesByDataSource.get(dataSourceId);
+    if (cached !== undefined) return cached;
+    const languages = await http
+      .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, {
+        query: { dataSourceId },
+      })
+      .then((response) => response.configuredLanguages)
+      .catch(() => []);
+    availableLanguagesByDataSource.set(dataSourceId, languages);
+    return languages;
+  };
+})();
 
 /**
  * @returns observable list of query assist agent configured languages in the
  * selected data source.
  */
-const getAvailableLanguages$ = (
-  availableLanguagesByDataSource: Map<string | undefined, string[]>,
-  http: HttpSetup,
-  data: DataPublicPluginSetup
-) =>
+const getAvailableLanguages$ = (http: HttpSetup, data: DataPublicPluginSetup) =>
   data.query.queryString.getUpdates$().pipe(
     startWith(data.query.queryString.getQuery()),
     distinctUntilChanged(),
@@ -34,16 +50,7 @@ const getAvailableLanguages$ = (
       if (query.dataset?.dataSource?.type !== DEFAULT_DATA.SOURCE_TYPES.OPENSEARCH) return [];
 
       const dataSourceId = query.dataset?.dataSource?.id;
-      const cached = availableLanguagesByDataSource.get(dataSourceId);
-      if (cached !== undefined) return cached;
-      const languages = await http
-        .get<{ configuredLanguages: string[] }>(API.QUERY_ASSIST.LANGUAGES, {
-          query: { dataSourceId },
-        })
-        .then((response) => response.configuredLanguages)
-        .catch(() => []);
-      availableLanguagesByDataSource.set(dataSourceId, languages);
-      return languages;
+      return getAvailableLanguagesForDataSource(http, dataSourceId);
     })
   );
 
@@ -52,38 +59,35 @@ export const createQueryAssistExtension = (
   data: DataPublicPluginSetup,
   config: ConfigSchema['queryAssist']
 ): QueryEditorExtensionConfig => {
-  const availableLanguagesByDataSource: Map<string | undefined, string[]> = new Map();
-
   return {
     id: 'query-assist',
     order: 1000,
+    getDataStructureMeta: async (dataSourceId) => {
+      const isEnabled = await getAvailableLanguagesForDataSource(http, dataSourceId).then(
+        (languages) => languages.length > 0
+      );
+      if (isEnabled) {
+        return {
+          type: DATA_STRUCTURE_META_TYPES.TYPE,
+          icon: { type: assistantMark },
+          tooltip: 'Query assist is available',
+        };
+      }
+    },
     isEnabled$: () =>
-      getAvailableLanguages$(availableLanguagesByDataSource, http, data).pipe(
-        map((languages) => languages.length > 0)
-      ),
+      getAvailableLanguages$(http, data).pipe(map((languages) => languages.length > 0)),
     getComponent: (dependencies) => {
       // only show the component if user is on a supported language.
       return (
-        <QueryAssistWrapper
-          availableLanguagesByDataSource={availableLanguagesByDataSource}
-          dependencies={dependencies}
-          http={http}
-          data={data}
-        >
+        <QueryAssistWrapper dependencies={dependencies} http={http} data={data}>
           <QueryAssistBar dependencies={dependencies} />
         </QueryAssistWrapper>
       );
     },
     getBanner: (dependencies) => {
       // advertise query assist if user is not on a supported language.
       return (
-        <QueryAssistWrapper
-          availableLanguagesByDataSource={availableLanguagesByDataSource}
-          dependencies={dependencies}
-          http={http}
-          data={data}
-          invert
-        >
+        <QueryAssistWrapper dependencies={dependencies} http={http} data={data} invert>
           <QueryAssistBanner
             dependencies={dependencies}
             languages={config.supportedLanguages.map((conf) => conf.language)}
@@ -95,7 +99,6 @@ export const createQueryAssistExtension = (
 };
 
 interface QueryAssistWrapperProps {
-  availableLanguagesByDataSource: Map<string | undefined, string[]>;
   dependencies: QueryEditorExtensionDependencies;
   http: HttpSetup;
   data: DataPublicPluginSetup;
@@ -108,11 +111,7 @@ const QueryAssistWrapper: React.FC<QueryAssistWrapperProps> = (props) => {
   useEffect(() => {
     let mounted = true;
 
-    const subscription = getAvailableLanguages$(
-      props.availableLanguagesByDataSource,
-      props.http,
-      props.data
-    ).subscribe((languages) => {
+    const subscription = getAvailableLanguages$(props.http, props.data).subscribe((languages) => {
       const available = languages.includes(props.dependencies.language);
       if (mounted) setVisible(props.invert ? !available : available);
     });