[MDS] Support Vega Visualizations (#5975) (#6212)

* Add MDS support for Vega Signed-off-by: Huy Nguyen <[email protected]> * Refactor field to data_source_id Signed-off-by: Huy Nguyen <[email protected]> * Add to CHANGELOG.md Signed-off-by: Huy Nguyen <[email protected]> * Added test cases and renamed field to use data_source_name Signed-off-by: Huy Nguyen <[email protected]> * Add prefix datasource name test case and add example in default hjson Signed-off-by: Huy Nguyen <[email protected]> * Move CHANGELOG to appropriate section Signed-off-by: Huy Nguyen <[email protected]> * Increased test coverage of search() method Signed-off-by: Huy Nguyen <[email protected]> --------- Signed-off-by: Huy Nguyen <[email protected]> (cherry picked from commit 1c5ad6c) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> # Conflicts: # CHANGELOG.md Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Mar 19, 2024 · 05ede47 · 05ede47
1 parent 693e0c0
commit 05ede47
Show file tree

Hide file tree

Showing 9 changed files with 268 additions and 25 deletions.
diff --git a/src/plugins/vis_type_vega/opensearch_dashboards.json b/src/plugins/vis_type_vega/opensearch_dashboards.json
@@ -11,7 +11,7 @@
     "inspector",
     "uiActions"
   ],
-  "optionalPlugins": ["home", "usageCollection"],
+  "optionalPlugins": ["home", "usageCollection", "dataSource"],
   "requiredBundles": [
     "opensearchDashboardsUtils",
     "opensearchDashboardsReact",

diff --git a/src/plugins/vis_type_vega/public/data_model/opensearch_query_parser.ts b/src/plugins/vis_type_vega/public/data_model/opensearch_query_parser.ts
@@ -229,7 +229,7 @@ export class OpenSearchQueryParser {
       name: getRequestName(r, index),
     }));
 
-    const data$ = this._searchAPI.search(opensearchSearches);
+    const data$ = await this._searchAPI.search(opensearchSearches);
 
     const results = await data$.toPromise();
 

diff --git a/src/plugins/vis_type_vega/public/data_model/search_api.test.ts b/src/plugins/vis_type_vega/public/data_model/search_api.test.ts
@@ -0,0 +1,159 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { SavedObjectsClientContract, SavedObjectsFindOptions } from 'opensearch-dashboards/public';
+import { SearchAPI, SearchAPIDependencies } from './search_api';
+import { ISearchStart } from 'src/plugins/data/public';
+import { IUiSettingsClient } from 'opensearch-dashboards/public';
+
+jest.mock('rxjs', () => ({
+  combineLatest: jest.fn().mockImplementation((obj) => obj),
+}));
+
+jest.mock('../../../data/public', () => ({
+  getSearchParamsFromRequest: jest.fn().mockImplementation((obj, _) => obj),
+}));
+
+interface MockSearch {
+  params?: Record<string, unknown>;
+  dataSourceId?: string;
+  pipe: () => {};
+}
+
+describe('SearchAPI.search', () => {
+  // This will only test that searchApiParams were correctly set. As such, every other function can be mocked
+  const getSearchAPI = (dataSourceEnabled: boolean) => {
+    const savedObjectsClient = {} as SavedObjectsClientContract;
+
+    const searchStartMock = {} as ISearchStart;
+    searchStartMock.search = jest.fn().mockImplementation((obj, _) => {
+      const mockedSearchResults = {} as MockSearch;
+      mockedSearchResults.params = obj;
+      mockedSearchResults.pipe = jest.fn().mockReturnValue(mockedSearchResults.params);
+      return mockedSearchResults;
+    });
+
+    const uiSettings = {} as IUiSettingsClient;
+    uiSettings.get = jest.fn().mockReturnValue(0);
+    uiSettings.get.bind = jest.fn().mockReturnValue(0);
+
+    const dependencies = {
+      savedObjectsClient,
+      dataSourceEnabled,
+      search: searchStartMock,
+      uiSettings,
+    } as SearchAPIDependencies;
+    const searchAPI = new SearchAPI(dependencies);
+    searchAPI.findDataSourceIdbyName = jest.fn().mockImplementation((name) => {
+      if (!dataSourceEnabled) {
+        throw new Error();
+      }
+      if (name === 'exampleName') {
+        return Promise.resolve('some-id');
+      }
+    });
+
+    return searchAPI;
+  };
+
+  test('If MDS is disabled and there is no datasource, return params without datasource id', async () => {
+    const searchAPI = getSearchAPI(false);
+    const requests = [{ name: 'example-id' }];
+    const fetchParams = ((await searchAPI.search(requests)) as unknown) as MockSearch[];
+    expect(fetchParams[0].params).toBe(requests[0]);
+    expect(fetchParams[0].hasOwnProperty('dataSourceId')).toBe(false);
+  });
+
+  test('If MDS is disabled and there is a datasource, it should throw an errorr', () => {
+    const searchAPI = getSearchAPI(false);
+    const requests = [{ name: 'example-id', data_source_name: 'non-existent-datasource' }];
+    expect(searchAPI.search(requests)).rejects.toThrowError();
+  });
+
+  test('If MDS is enabled and there is no datasource, return params without datasource id', async () => {
+    const searchAPI = getSearchAPI(true);
+    const requests = [{ name: 'example-id' }];
+    const fetchParams = ((await searchAPI.search(requests)) as unknown) as MockSearch[];
+    expect(fetchParams[0].params).toBe(requests[0]);
+    expect(fetchParams[0].hasOwnProperty('dataSourceId')).toBe(false);
+  });
+
+  test('If MDS is enabled and there is a datasource, return params with datasource id', async () => {
+    const searchAPI = getSearchAPI(true);
+    const requests = [{ name: 'example-id', data_source_name: 'exampleName' }];
+    const fetchParams = ((await searchAPI.search(requests)) as unknown) as MockSearch[];
+    expect(fetchParams[0].hasOwnProperty('params')).toBe(true);
+    expect(fetchParams[0].dataSourceId).toBe('some-id');
+  });
+});
+
+describe('SearchAPI.findDataSourceIdbyName', () => {
+  const savedObjectsClient = {} as SavedObjectsClientContract;
+  savedObjectsClient.find = jest.fn().mockImplementation((query: SavedObjectsFindOptions) => {
+    if (query.search === `"uniqueDataSource"`) {
+      return Promise.resolve({
+        total: 1,
+        savedObjects: [{ id: 'some-datasource-id', attributes: { title: 'uniqueDataSource' } }],
+      });
+    } else if (query.search === `"duplicateDataSource"`) {
+      return Promise.resolve({
+        total: 2,
+        savedObjects: [
+          { id: 'some-datasource-id', attributes: { title: 'duplicateDataSource' } },
+          { id: 'some-other-datasource-id', attributes: { title: 'duplicateDataSource' } },
+        ],
+      });
+    } else if (query.search === `"DataSource"`) {
+      return Promise.resolve({
+        total: 2,
+        savedObjects: [
+          { id: 'some-datasource-id', attributes: { title: 'DataSource' } },
+          { id: 'some-other-datasource-id', attributes: { title: 'DataSource Copy' } },
+        ],
+      });
+    } else {
+      return Promise.resolve({
+        total: 0,
+        savedObjects: [],
+      });
+    }
+  });
+
+  const getSearchAPI = (dataSourceEnabled: boolean) => {
+    const dependencies = { savedObjectsClient, dataSourceEnabled } as SearchAPIDependencies;
+    return new SearchAPI(dependencies);
+  };
+
+  test('If dataSource is disabled, throw error', () => {
+    const searchAPI = getSearchAPI(false);
+    expect(searchAPI.findDataSourceIdbyName('nonexistentDataSource')).rejects.toThrowError(
+      'data_source_name cannot be used because data_source.enabled is false'
+    );
+  });
+
+  test('If dataSource is enabled but no matching dataSourceName, then throw error', () => {
+    const searchAPI = getSearchAPI(true);
+    expect(searchAPI.findDataSourceIdbyName('nonexistentDataSource')).rejects.toThrowError(
+      'Expected exactly 1 result for data_source_name "nonexistentDataSource" but got 0 results'
+    );
+  });
+
+  test('If dataSource is enabled but multiple dataSourceNames, then throw error', () => {
+    const searchAPI = getSearchAPI(true);
+    expect(searchAPI.findDataSourceIdbyName('duplicateDataSource')).rejects.toThrowError(
+      'Expected exactly 1 result for data_source_name "duplicateDataSource" but got 2 results'
+    );
+  });
+
+  test('If dataSource is enabled but only one dataSourceName, then return id', async () => {
+    const searchAPI = getSearchAPI(true);
+    expect(await searchAPI.findDataSourceIdbyName('uniqueDataSource')).toBe('some-datasource-id');
+  });
+
+  test('If dataSource is enabled and the dataSourceName is a prefix of another, ensure the prefix is only returned', async () => {
+    const searchAPI = getSearchAPI(true);
+    expect(await searchAPI.findDataSourceIdbyName('DataSource')).toBe('some-datasource-id');
+  });
+});
diff --git a/src/plugins/vis_type_vega/public/data_model/search_api.ts b/src/plugins/vis_type_vega/public/data_model/search_api.ts
@@ -31,6 +31,8 @@
 import { combineLatest } from 'rxjs';
 import { map, tap } from 'rxjs/operators';
 import { CoreStart, IUiSettingsClient } from 'opensearch-dashboards/public';
+import { SavedObjectsClientContract } from 'src/core/public';
+import { DataSourceAttributes } from 'src/plugins/data_source/common/data_sources';
 import {
   getSearchParamsFromRequest,
   SearchRequest,
@@ -45,6 +47,8 @@ export interface SearchAPIDependencies {
   uiSettings: IUiSettingsClient;
   injectedMetadata: CoreStart['injectedMetadata'];
   search: DataPublicPluginStart['search'];
+  dataSourceEnabled: boolean;
+  savedObjectsClient: SavedObjectsClientContract;
 }
 
 export class SearchAPI {
@@ -54,31 +58,75 @@ export class SearchAPI {
     public readonly inspectorAdapters?: VegaInspectorAdapters
   ) {}
 
-  search(searchRequests: SearchRequest[]) {
+  async search(searchRequests: SearchRequest[]) {
     const { search } = this.dependencies.search;
     const requestResponders: any = {};
 
     return combineLatest(
-      searchRequests.map((request) => {
-        const requestId = request.name;
-        const params = getSearchParamsFromRequest(request, {
-          getConfig: this.dependencies.uiSettings.get.bind(this.dependencies.uiSettings),
-        });
-
-        if (this.inspectorAdapters) {
-          requestResponders[requestId] = this.inspectorAdapters.requests.start(requestId, request);
-          requestResponders[requestId].json(params.body);
-        }
-
-        return search({ params }, { abortSignal: this.abortSignal }).pipe(
-          tap((data) => this.inspectSearchResult(data, requestResponders[requestId])),
-          map((data) => ({
-            name: requestId,
-            rawResponse: data.rawResponse,
-          }))
-        );
-      })
+      await Promise.all(
+        searchRequests.map(async (request) => {
+          const requestId = request.name;
+          const dataSourceId = !!request.data_source_name
+            ? await this.findDataSourceIdbyName(request.data_source_name)
+            : undefined;
+
+          const params = getSearchParamsFromRequest(request, {
+            getConfig: this.dependencies.uiSettings.get.bind(this.dependencies.uiSettings),
+          });
+
+          if (this.inspectorAdapters) {
+            requestResponders[requestId] = this.inspectorAdapters.requests.start(
+              requestId,
+              request
+            );
+            requestResponders[requestId].json(params.body);
+          }
+
+          const searchApiParams =
+            dataSourceId && this.dependencies.dataSourceEnabled
+              ? { params, dataSourceId }
+              : { params };
+
+          return search(searchApiParams, { abortSignal: this.abortSignal }).pipe(
+            tap((data) => this.inspectSearchResult(data, requestResponders[requestId])),
+            map((data) => ({
+              name: requestId,
+              rawResponse: data.rawResponse,
+            }))
+          );
+        })
+      )
+    );
+  }
+
+  async findDataSourceIdbyName(dataSourceName: string) {
+    if (!this.dependencies.dataSourceEnabled) {
+      throw new Error('data_source_name cannot be used because data_source.enabled is false');
+    }
+    const dataSources = await this.dataSourceFindQuery(dataSourceName);
+
+    // In the case that data_source_name is a prefix of another name, match exact data_source_name
+    const possibleDataSourceIds = dataSources.savedObjects.filter(
+      (obj) => obj.attributes.title === dataSourceName
     );
+
+    if (possibleDataSourceIds.length !== 1) {
+      throw new Error(
+        `Expected exactly 1 result for data_source_name "${dataSourceName}" but got ${possibleDataSourceIds.length} results`
+      );
+    }
+
+    return possibleDataSourceIds.pop()?.id;
+  }
+
+  async dataSourceFindQuery(dataSourceName: string) {
+    return await this.dependencies.savedObjectsClient.find<DataSourceAttributes>({
+      type: 'data-source',
+      perPage: 10,
+      search: `"${dataSourceName}"`,
+      searchFields: ['title'],
+      fields: ['id', 'title'],
+    });
   }
 
   public resetSearchStats() {

diff --git a/src/plugins/vis_type_vega/public/data_model/types.ts b/src/plugins/vis_type_vega/public/data_model/types.ts
@@ -186,6 +186,7 @@ export interface UrlObject {
   [CONSTANTS.TYPE]?: string;
   name?: string;
   index?: string;
+  data_source_name?: string;
   body?: Body;
   size?: number;
   timeout?: string;

diff --git a/src/plugins/vis_type_vega/public/default.spec.hjson b/src/plugins/vis_type_vega/public/default.spec.hjson
@@ -29,6 +29,10 @@
 
       // Which index to search
       index: _all
+
+      // If "data_source.enabled: true", optionally set the data source name to query from (omit field if querying from local cluster)
+      // data_source_name: Example US Cluster
+
       // Aggregate data by the time field into time buckets, counting the number of documents in each bucket.
       body: {
         aggs: {

diff --git a/src/plugins/vis_type_vega/public/plugin.ts b/src/plugins/vis_type_vega/public/plugin.ts
@@ -28,6 +28,7 @@
  * under the License.
  */
 
+import { DataSourcePluginSetup } from 'src/plugins/data_source/public';
 import { PluginInitializerContext, CoreSetup, CoreStart, Plugin } from '../../../core/public';
 import { Plugin as ExpressionsPublicPlugin } from '../../expressions/public';
 import { DataPublicPluginSetup, DataPublicPluginStart } from '../../data/public';
@@ -41,6 +42,8 @@ import {
   setUISettings,
   setMapsLegacyConfig,
   setInjectedMetadata,
+  setDataSourceEnabled,
+  setSavedObjectsClient,
 } from './services';
 
 import { createVegaFn } from './expressions/vega_fn';
@@ -69,6 +72,7 @@ export interface VegaPluginSetupDependencies {
   visualizations: VisualizationsSetup;
   inspector: InspectorSetup;
   data: DataPublicPluginSetup;
+  dataSource?: DataSourcePluginSetup;
   mapsLegacy: any;
 }
 
@@ -88,14 +92,22 @@ export class VegaPlugin implements Plugin<Promise<void>, void> {
 
   public async setup(
     core: CoreSetup,
-    { inspector, data, expressions, visualizations, mapsLegacy }: VegaPluginSetupDependencies
+    {
+      inspector,
+      data,
+      expressions,
+      visualizations,
+      mapsLegacy,
+      dataSource,
+    }: VegaPluginSetupDependencies
   ) {
     setInjectedVars({
       enableExternalUrls: this.initializerContext.config.get().enableExternalUrls,
       emsTileLayerId: core.injectedMetadata.getInjectedVar('emsTileLayerId', true),
     });
     setUISettings(core.uiSettings);
     setMapsLegacyConfig(mapsLegacy.config);
+    setDataSourceEnabled({ enabled: !!dataSource });
 
     const visualizationDependencies: Readonly<VegaVisualizationDependencies> = {
       core,
@@ -116,6 +128,7 @@ export class VegaPlugin implements Plugin<Promise<void>, void> {
   public start(core: CoreStart, { data, uiActions }: VegaPluginStartDependencies) {
     setNotifications(core.notifications);
     setData(data);
+    setSavedObjectsClient(core.savedObjects);
     setUiActions(uiActions);
     setInjectedMetadata(core.injectedMetadata);
   }

diff --git a/src/plugins/vis_type_vega/public/services.ts b/src/plugins/vis_type_vega/public/services.ts
@@ -28,14 +28,25 @@
  * under the License.
  */
 
-import { CoreStart, NotificationsStart, IUiSettingsClient } from 'src/core/public';
+import {
+  CoreStart,
+  NotificationsStart,
+  IUiSettingsClient,
+  SavedObjectsStart,
+} from 'src/core/public';
 
 import { DataPublicPluginStart } from '../../data/public';
 import { createGetterSetter } from '../../opensearch_dashboards_utils/public';
 import { MapsLegacyConfig } from '../../maps_legacy/config';
 import { UiActionsStart } from '../../ui_actions/public';
 
 export const [getData, setData] = createGetterSetter<DataPublicPluginStart>('Data');
+export const [getDataSourceEnabled, setDataSourceEnabled] = createGetterSetter<{
+  enabled: boolean;
+}>('DataSource');
+export const [getSavedObjectsClient, setSavedObjectsClient] = createGetterSetter<SavedObjectsStart>(
+  'SavedObjects'
+);
 
 export const [getNotifications, setNotifications] = createGetterSetter<NotificationsStart>(
   'Notifications'