allow custom icon for DATA_STRUCTURE_META_TYPES.FEATURE

Signed-off-by: Joshua Li <[email protected]>
opensearch-project · Aug 28, 2024 · 555c299 · 555c299
1 parent 03dcec7
commit 555c299
Show file tree

Hide file tree

Showing 6 changed files with 33 additions and 22 deletions.
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
@@ -124,5 +124,5 @@ const fetchIndexPatterns = async (client: SavedObjectsClientContract): Promise<D
     }
   );
 
-  return injectMetaToDataStructures(dataStructures);
+  return injectMetaToDataStructures(dataStructures, (dataStructure) => dataStructure.parent?.id);
 };
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
@@ -10,14 +10,20 @@ import { getQueryService } from '../../../../services';
  * Inject {@link DataStructureMeta} to DataStructures based on
  * {@link QueryEditorExtensions}.
  *
- * This function combines the meta fields from QueryEditorExtensions and in
+ * 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[]) => {
+export const injectMetaToDataStructures = async (
+  dataStructures: DataStructure[],
+  selectDataSourceId: (dataStructure: DataStructure) => string | undefined = (
+    dataStructure: DataStructure
+  ) => dataStructure.id
+) => {
   const queryEditorExtensions = Object.values(
     getQueryService().queryString.getLanguageService().getQueryEditorExtensionMap()
   );
@@ -26,7 +32,9 @@ export const injectMetaToDataStructures = async (dataStructures: DataStructure[]
   return Promise.all(
     dataStructures.map(async (dataStructure) => {
       const metaArray = await Promise.allSettled(
-        queryEditorExtensions.map((curr) => curr.getDataStructureMeta?.(dataStructure.id))
+        queryEditorExtensions.map((curr) =>
+          curr.getDataStructureMeta?.(selectDataSourceId(dataStructure))
+        )
       ).then((settledResults) =>
         settledResults
           .filter(
@@ -39,7 +47,10 @@ export const injectMetaToDataStructures = async (dataStructures: DataStructure[]
         (acc, curr) => (acc || curr ? ({ ...acc, ...curr } as DataStructureMeta) : undefined),
         undefined
       );
-      return { meta, ...dataStructure };
+      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/query_enhancements/public/query_assist/utils/create_extension.tsx b/src/plugins/query_enhancements/public/query_assist/utils/create_extension.tsx
@@ -68,7 +68,7 @@ export const createQueryAssistExtension = (
       );
       if (isEnabled) {
         return {
-          type: DATA_STRUCTURE_META_TYPES.TYPE,
+          type: DATA_STRUCTURE_META_TYPES.FEATURE,
           icon: { type: assistantMark },
           tooltip: 'Query assist is available',
         };
-Original file line number
+Diff line change
@@ Expand Up @@
         }
       );
-      return injectMetaToDataStructures(dataStructures);
+      return injectMetaToDataStructures(dataStructures, (dataStructure) => dataStructure.parent?.id);
     };