Support injecting DataStructureMeta from QueryEditorExtensions fo…

…r Query Assist (opensearch-project#7871) * fix query editor extensions enhance key Signed-off-by: Joshua Li <joshuali925@gmail.com> * add query assist icons to clusters if available Signed-off-by: Joshua Li <joshuali925@gmail.com> * allow custom icon for `DATA_STRUCTURE_META_TYPES.FEATURE` Signed-off-by: Joshua Li <joshuali925@gmail.com> * Changeset file for PR opensearch-project#7871 created/updated --------- Signed-off-by: Joshua Li <joshuali925@gmail.com> Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com> Co-authored-by: Ashwin P Chandran <ashwinpc@amazon.com>
sumukhswamy · Aug 28, 2024 · 9bc2433 · 9bc2433
1 parent 5e0ce2b
commit 9bc2433
Show file tree

Hide file tree

Showing 10 changed files with 129 additions and 60 deletions.
diff --git a/changelogs/fragments/7871.yml b/changelogs/fragments/7871.yml
@@ -0,0 +1,2 @@
+feat:
+- Support injecting `DataStructureMeta` from `QueryEditorExtensions` for Query Assist ([#7871](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7871))
diff --git a/src/plugins/data/common/datasets/types.ts b/src/plugins/data/common/datasets/types.ts
@@ -139,7 +139,7 @@ export enum DATA_STRUCTURE_META_TYPES {
  */
 export interface DataStructureFeatureMeta {
   type: DATA_STRUCTURE_META_TYPES.FEATURE;
-  icon?: string;
+  icon?: EuiIconProps;
   tooltip?: string;
 }
 
@@ -157,6 +157,8 @@ export interface DataStructureDataTypeMeta {
  */
 export interface DataStructureCustomMeta {
   type: DATA_STRUCTURE_META_TYPES.CUSTOM;
+  icon?: EuiIconProps;
+  tooltip?: string;
   [key: string]: any;
 }
 

diff --git a/src/plugins/data/public/query/query_string/dataset_service/lib/index_pattern_type.ts b/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, (dataStructure) => dataStructure.parent?.id);
 };
diff --git a/src/plugins/data/public/query/query_string/dataset_service/lib/index_type.ts b/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[]> => {

diff --git a/src/plugins/data/public/query/query_string/dataset_service/lib/utils.ts b/src/plugins/data/public/query/query_string/dataset_service/lib/utils.ts
@@ -0,0 +1,56 @@
+/*
+ * 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
+ * provided data structures. Lower extension order is higher priority, and
+ * existing meta fields have highest priority.
+ *
+ * @param dataStructures - {@link DataStructure}
+ * @param selectDataSourceId - function to get data source id given a data structure
+ * @returns data structures with meta
+ */
+export const injectMetaToDataStructures = async (
+  dataStructures: DataStructure[],
+  selectDataSourceId: (dataStructure: DataStructure) => string | undefined = (
+    dataStructure: DataStructure
+  ) => dataStructure.id
+) => {
+  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?.(selectDataSourceId(dataStructure))
+        )
+      ).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
+      );
+      if (meta || dataStructure.meta) {
+        dataStructure.meta = { ...meta, ...dataStructure.meta } as DataStructureMeta;
+      }
+      return dataStructure;
+    })
+  );
+};
diff --git a/src/plugins/data/public/ui/dataset_selector/advanced_selector.tsx b/src/plugins/data/public/ui/dataset_selector/advanced_selector.tsx
@@ -37,11 +37,11 @@ export const AdvancedSelector = ({
         .getTypes()
         .map((type) => {
           return {
-            id: type!.id,
-            title: type!.title,
-            type: type!.id,
+            id: type.id,
+            title: type.title,
+            type: type.id,
             meta: {
-              ...type!.meta,
+              ...type.meta,
               type: DATA_STRUCTURE_META_TYPES.TYPE,
             },
           } as DataStructure;

diff --git a/src/plugins/data/public/ui/dataset_selector/dataset_explorer.tsx b/src/plugins/data/public/ui/dataset_selector/dataset_explorer.tsx
@@ -190,26 +190,24 @@ const LoadingEmptyColumn = ({ isLoading }: { isLoading: boolean }) =>
     <EmptyColumn />
   );
 const appendIcon = (item: DataStructure) => {
-  if (item.meta?.type === DATA_STRUCTURE_META_TYPES.FEATURE) {
+  if (item.meta?.type === DATA_STRUCTURE_META_TYPES.TYPE) {
+    return (
+      <EuiToolTip content={item.meta.tooltip}>
+        <EuiIcon type="iInCircle" />
+      </EuiToolTip>
+    );
+  } else {
     if (item.meta?.icon && item.meta?.tooltip) {
       return (
         <EuiToolTip content={item.meta.tooltip}>
-          <EuiIcon type={item.meta.icon} />
+          <EuiIcon {...item.meta.icon} />
         </EuiToolTip>
       );
     } else if (item.meta?.icon) {
-      return <EuiIcon type={item.meta.icon} />;
+      return <EuiIcon {...item.meta.icon} />;
     }
   }
 
-  if (item.meta?.type === DATA_STRUCTURE_META_TYPES.TYPE) {
-    return (
-      <EuiToolTip content={item.meta.tooltip}>
-        <EuiIcon type="iInCircle" />
-      </EuiToolTip>
-    );
-  }
-
   return null;
 };
 

diff --git a/src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension.tsx b/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.

diff --git a/src/plugins/query_enhancements/public/plugin.tsx b/src/plugins/query_enhancements/public/plugin.tsx
@@ -95,7 +95,7 @@ export class QueryEnhancementsPlugin
     queryString.getLanguageService().registerLanguage(sqlLanguageConfig);
 
     data.__enhance({
-      ui: {
+      editor: {
         queryEditorExtension: createQueryAssistExtension(core.http, data, this.config.queryAssist),
       },
     });

diff --git a/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx b/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.FEATURE,
+          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);
     });