types and datasource prepend

Signed-off-by: Kawika Avilla <[email protected]>

src/core/public/saved_objects/saved_objects_client.ts

@@ -457,14 +457,14 @@ export class SavedObjectsClient {
    *   { id: 'foo', type: 'index-pattern' }
    * ])
    */
-  public bulkGet = (objects: Array<{ id: string; type: string }> = []) => {
+  public bulkGet = <T = unknown>(objects: Array<{ id: string; type: string }> = []) => {
     const filteredObjects = objects.map((obj) => pick(obj, ['id', 'type']));
     return this.performBulkGet(filteredObjects).then((resp) => {
       resp.saved_objects = resp.saved_objects.map((d) => this.createSavedObject(d));
       return renameKeys<
         PromiseType<ReturnType<SavedObjectsApi['bulkGet']>>,
-        SavedObjectsBatchResponse
-      >({ saved_objects: 'savedObjects' }, resp) as SavedObjectsBatchResponse;
+        SavedObjectsBatchResponse<T>
+      >({ saved_objects: 'savedObjects' }, resp) as SavedObjectsBatchResponse<T>;
     });
   };
 

src/plugins/data/common/datasets/types.ts

@@ -119,8 +119,6 @@ export interface DataStructure {
   parent?: DataStructure;
   /** Optional array of child data structures */
   children?: DataStructure[];
-  /** Optional array of data structures of ancestors */
-  path?: DataStructure[];
   hasNext?: boolean;
   columnHeader?: string;
   /** Optional metadata for the data structure */
@@ -212,12 +210,16 @@ export interface BaseDataset {
  *
  * @example
  * Example of a Dataset for an OpenSearch index pattern
- * Index patterns have a reference to the data source and their title is pre appended with the data source they belong to so we dont need to append the data source to the dataset object
  * const logsIndexDataset: Dataset = {
  *   id: "2e1b1b80-9c4d-11ee-8c90-0242ac120001",
- *   title: "Cluster1::logs-*",
+ *   title: "logs-*",
  *   type: "INDEX_PATTERN",
  *   timeFieldName: "@timestamp",
+ *   dataSource: {
+ *     id: "2e1b1b80-9c4d-11ee-8c90-0242ac120001",
+ *     title: "Cluster1",
+ *     type: "OpenSearch"
+ *   }
  * };
  *
  * @example

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

@@ -98,9 +98,9 @@ export class DatasetService {
       return undefined;
     }
 
-    const handler = this.typesRegistry.get(DEFAULT_DATA.SET_TYPES.INDEX_PATTERN);
-    if (handler) {
-      const dataset = handler.toDataset([
+    const dataType = this.typesRegistry.get(DEFAULT_DATA.SET_TYPES.INDEX_PATTERN);
+    if (dataType) {
+      const dataset = dataType.toDataset([
         {
           id: indexPattern.id,
           title: indexPattern.title,

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

@@ -4,12 +4,14 @@
  */
 
 import { SavedObjectsClientContract } from 'opensearch-dashboards/public';
+import { DataSourceAttributes } from '../../../../../../data_source/common/data_sources';
 import {
   DEFAULT_DATA,
   DataStructure,
   DatasetField,
   Dataset,
   IIndexPattern,
+  DATA_STRUCTURE_META_TYPES,
 } from '../../../../../common';
 import { DatasetTypeConfig } from '../types';
 import { getIndexPatterns } from '../../../../services';
@@ -22,7 +24,7 @@ export const indexPatternTypeConfig: DatasetTypeConfig = {
     tooltip: 'OpenSearch Index Patterns',
   },
 
-  toDataset: (path: DataStructure[]): Dataset => {
+  toDataset: (path) => {
     const pattern = path[path.length - 1];
     return {
       id: pattern.id,
@@ -35,13 +37,10 @@ export const indexPatternTypeConfig: DatasetTypeConfig = {
             type: pattern.parent.type,
           }
         : undefined,
-    };
+    } as Dataset;
   },
 
-  fetch: async (
-    savedObjects: SavedObjectsClientContract,
-    path: DataStructure[]
-  ): Promise<DataStructure> => {
+  fetch: async (savedObjects, path) => {
     const dataStructure = path[path.length - 1];
     const indexPatterns = await fetchIndexPatterns(savedObjects);
     return {
@@ -60,8 +59,11 @@ export const indexPatternTypeConfig: DatasetTypeConfig = {
     }));
   },
 
-  supportedLanguages: (): string[] => {
-    return ['SQL', 'PPL', 'DQL', 'Lucene'];
+  supportedLanguages: (dataset): string[] => {
+    if (dataset.dataSource?.type === 'OpenSearch Serverless') {
+      return ['DQL', 'Lucene'];
+    }
+    return ['DQL', 'Lucene', 'PPL', 'SQL'];
   },
 };
 
@@ -73,19 +75,51 @@ const fetchIndexPatterns = async (client: SavedObjectsClientContract): Promise<D
     searchFields: ['title'],
     perPage: 100,
   });
-  return resp.savedObjects.map((savedObject) => ({
-    id: savedObject.id,
-    title: savedObject.attributes.title,
-    timeFieldName: savedObject.attributes.timeFieldName,
-    type: DEFAULT_DATA.SET_TYPES.INDEX_PATTERN,
-    ...(savedObject.references[0]
-      ? {
-          dataSource: {
-            id: savedObject.references[0]?.id,
-            name: savedObject.references[0]?.name,
-            type: 'data-source',
-          },
-        }
-      : {}),
-  }));
+
+  // Get all unique data source ids
+  const datasourceIds = Array.from(
+    new Set(
+      resp.savedObjects
+        .filter((savedObject) => savedObject.references.length > 0)
+        .map((savedObject) => savedObject.references.find((ref) => ref.type === 'data-source')?.id)
+        .filter(Boolean)
+    )
+  ) as string[];
+
+  const dataSourceMap: Record<string, DataSourceAttributes> = {};
+  if (datasourceIds.length > 0) {
+    const dataSourceResp = await client.bulkGet<DataSourceAttributes>(
+      datasourceIds.map((id) => ({ id, type: 'data-source' }))
+    );
+
+    dataSourceResp.savedObjects.forEach((savedObject) => {
+      dataSourceMap[savedObject.id] = savedObject.attributes;
+    });
+  }
+
+  return resp.savedObjects.map(
+    (savedObject): DataStructure => {
+      const dataSourceId = savedObject.references.find((ref) => ref.type === 'data-source')?.id;
+      const dataSource = dataSourceId ? dataSourceMap[dataSourceId] : undefined;
+
+      const indexPatternDataStructure: DataStructure = {
+        id: savedObject.id,
+        title: savedObject.attributes.title,
+        type: DEFAULT_DATA.SET_TYPES.INDEX_PATTERN,
+        meta: {
+          type: DATA_STRUCTURE_META_TYPES.CUSTOM,
+          timeFieldName: savedObject.attributes.timeFieldName,
+        },
+      };
+
+      if (dataSource) {
+        indexPatternDataStructure.parent = {
+          id: dataSourceId!, // Since we know it exists
+          title: dataSource.title,
+          type: dataSource.dataSourceEngineType ?? 'OpenSearch',
+        };
+      }
+      return indexPatternDataStructure;
+    }
+  );
 };

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

@@ -5,7 +5,7 @@
 
 import { SavedObjectsClientContract } from 'opensearch-dashboards/public';
 import { map } from 'rxjs/operators';
-import { DEFAULT_DATA, DataStructure, DatasetField, Dataset } from '../../../../../common';
+import { DEFAULT_DATA, DataStructure, Dataset } from '../../../../../common';
 import { DatasetTypeConfig } from '../types';
 import { getSearchService, getIndexPatterns } from '../../../../services';
 
@@ -25,7 +25,7 @@ export const indexTypeConfig: DatasetTypeConfig = {
     tooltip: 'OpenSearch Indexes',
   },
 
-  toDataset: (path: DataStructure[]): Dataset => {
+  toDataset: (path) => {
     const index = path[path.length - 1];
     const dataSource = path.find((ds) => ds.type === 'DATA_SOURCE');
 
@@ -43,10 +43,7 @@ export const indexTypeConfig: DatasetTypeConfig = {
     };
   },
 
-  fetch: async (
-    savedObjects: SavedObjectsClientContract,
-    path: DataStructure[]
-  ): Promise<DataStructure> => {
+  fetch: async (savedObjects, path) => {
     const dataStructure = path[path.length - 1];
     switch (dataStructure.type) {
       case 'DATA_SOURCE': {
@@ -75,7 +72,7 @@ export const indexTypeConfig: DatasetTypeConfig = {
     }
   },
 
-  fetchFields: async (dataset: Dataset): Promise<DatasetField[]> => {
+  fetchFields: async (dataset) => {
     const fields = await getIndexPatterns().getFieldsForWildcard({
       pattern: dataset.title,
       dataSourceId: dataset.dataSource?.id,
@@ -86,7 +83,7 @@ export const indexTypeConfig: DatasetTypeConfig = {
     }));
   },
 
-  supportedLanguages: (): string[] => {
+  supportedLanguages: (dataset: Dataset): string[] => {
     return ['SQL', 'PPL', 'DQL', 'Lucene'];
   },
 };

src/plugins/data/public/query/query_string/dataset_service/types.ts

@@ -43,5 +43,5 @@ export interface DatasetTypeConfig {
    * Retrieves the supported query languages for this dataset type.
    * @returns {Promise<string[]>} A promise that resolves to an array of supported language names.
    */
-  supportedLanguages: () => string[];
+  supportedLanguages: (dataset: Dataset) => string[];
 }

src/plugins/data/public/query/query_string/query_string_manager.mock.ts

@@ -46,6 +46,7 @@ const createSetupContractMock = () => {
     getInitialQueryByLanguage: jest.fn(),
     getDatasetService: jest.fn(),
     getLanguageService: jest.fn(),
+    getInitialQueryByDataset: jest.fn(),
   };
   return queryStringManagerMock;
 };