[MD] Refactor search strategy to conditionally user datasource client (…

…opensearch-project#2171)

Signed-off-by: Kristen Tian <[email protected]>

src/plugins/data/common/search/opensearch_search/types.ts

@@ -52,6 +52,7 @@ export type ISearchRequestParams<T = Record<string, any>> = {
 export interface IOpenSearchSearchRequest
   extends IOpenSearchDashboardsSearchRequest<ISearchRequestParams> {
   indexType?: string;
+  dataSourceId?: string;
 }
 
 export type IOpenSearchSearchResponse<Source = any> = IOpenSearchDashboardsSearchResponse<

src/plugins/data/common/search/search_source/search_source.test.ts

@@ -59,6 +59,16 @@ const indexPattern2 = ({
   getSourceFiltering: () => mockSource2,
 } as unknown) as IndexPattern;
 
+const dataSourceId = 'dataSourceId';
+const indexPattern3 = ({
+  dataSourceRef: {
+    id: dataSourceId,
+    type: 'dataSource',
+  },
+  getComputedFields,
+  getSourceFiltering: () => mockSource,
+} as unknown) as IndexPattern;
+
 describe('SearchSource', () => {
   let mockSearchMethod: any;
   let searchSourceDependencies: SearchSourceDependencies;
@@ -209,6 +219,16 @@ describe('SearchSource', () => {
       await searchSource.fetch(options);
       expect(mockSearchMethod).toBeCalledTimes(1);
     });
+
+    test('index pattern with dataSourceId will generate request with dataSourceId', async () => {
+      const searchSource = new SearchSource({}, searchSourceDependencies);
+      searchSource.setField('index', indexPattern3);
+      const options = {};
+      await searchSource.fetch(options);
+      const request = searchSource.history[0];
+      expect(mockSearchMethod).toBeCalledTimes(1);
+      expect(request.dataSourceId).toEqual(dataSourceId);
+    });
   });
 
   describe('#serialize', () => {

src/plugins/data/common/search/search_source/search_source.ts

@@ -282,6 +282,9 @@ export class SearchSource {
     if (getConfig(UI_SETTINGS.COURIER_BATCH_SEARCHES)) {
       response = await this.legacyFetch(searchRequest, options);
     } else {
+      const indexPattern = this.getField('index');
+      searchRequest.dataSourceId = indexPattern?.dataSourceRef?.id;
+
       response = await this.fetchSearch(searchRequest, options);
     }
 
@@ -335,9 +338,10 @@ export class SearchSource {
       getConfig,
     });
 
-    return search({ params, indexType: searchRequest.indexType }, options).then(({ rawResponse }) =>
-      onResponse(searchRequest, rawResponse)
-    );
+    return search(
+      { params, indexType: searchRequest.indexType, dataSourceId: searchRequest.dataSourceId },
+      options
+    ).then(({ rawResponse }) => onResponse(searchRequest, rawResponse));
   }
 
   /**
@@ -466,7 +470,7 @@ export class SearchSource {
     }
   }
 
-  private flatten() {
+  flatten() {
     const searchRequest = this.mergeProps();
 
     searchRequest.body = searchRequest.body || {};

src/plugins/data/opensearch_dashboards.json

@@ -3,10 +3,7 @@
   "version": "opensearchDashboards",
   "server": true,
   "ui": true,
-  "requiredPlugins": [
-    "expressions",
-    "uiActions"
-  ],
+  "requiredPlugins": ["expressions", "uiActions"],
   "optionalPlugins": ["usageCollection", "dataSource"],
   "extraPublicDirs": ["common", "common/utils/abort_utils"],
   "requiredBundles": [
@@ -15,4 +12,4 @@
     "opensearchDashboardsReact",
     "inspector"
   ]
-}
+}

src/plugins/data/server/search/opensearch_search/decide_client.ts

@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { OpenSearchClient, RequestHandlerContext } from 'src/core/server';
+import { IOpenSearchSearchRequest } from '..';
+
+export const decideClient = async (
+  context: RequestHandlerContext,
+  request: IOpenSearchSearchRequest
+): Promise<OpenSearchClient> => {
+  // if data source feature is disabled, return default opensearch client of current user
+  const client =
+    request.dataSourceId && context.dataSource
+      ? await context.dataSource.opensearch.getClient(request.dataSourceId)
+      : context.core.opensearch.client.asCurrentUser;
+  return client;
+};

src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts

@@ -36,7 +36,7 @@ describe('OpenSearch search strategy', () => {
   const mockLogger: any = {
     debug: () => {},
   };
-  const mockApiCaller = jest.fn().mockResolvedValue({
+  const body = {
     body: {
       _shards: {
         total: 10,
@@ -45,21 +45,38 @@ describe('OpenSearch search strategy', () => {
         successful: 7,
       },
     },
-  });
+  };
+  const mockOpenSearchApiCaller = jest.fn().mockResolvedValue(body);
+  const mockDataSourceApiCaller = jest.fn().mockResolvedValue(body);
+  const dataSourceId = 'test-data-source-id';
+  const mockDataSourceContext = {
+    dataSource: {
+      opensearch: {
+        getClient: () => {
+          return { search: mockDataSourceApiCaller };
+        },
+      },
+    },
+  };
   const mockContext = {
     core: {
       uiSettings: {
         client: {
           get: () => {},
         },
       },
-      opensearch: { client: { asCurrentUser: { search: mockApiCaller } } },
+      opensearch: { client: { asCurrentUser: { search: mockOpenSearchApiCaller } } },
     },
   };
+  const mockDataSourceEnabledContext = {
+    ...mockContext,
+    ...mockDataSourceContext,
+  };
   const mockConfig$ = pluginInitializerContextConfigMock<any>({}).legacy.globalConfig$;
 
   beforeEach(() => {
-    mockApiCaller.mockClear();
+    mockOpenSearchApiCaller.mockClear();
+    mockDataSourceApiCaller.mockClear();
   });
 
   it('returns a strategy with `search`', async () => {
@@ -74,8 +91,8 @@ describe('OpenSearch search strategy', () => {
 
     await opensearchSearch.search((mockContext as unknown) as RequestHandlerContext, { params });
 
-    expect(mockApiCaller).toBeCalled();
-    expect(mockApiCaller.mock.calls[0][0]).toEqual({
+    expect(mockOpenSearchApiCaller).toBeCalled();
+    expect(mockOpenSearchApiCaller.mock.calls[0][0]).toEqual({
       ...params,
       ignore_unavailable: true,
       track_total_hits: true,
@@ -88,8 +105,8 @@ describe('OpenSearch search strategy', () => {
 
     await opensearchSearch.search((mockContext as unknown) as RequestHandlerContext, { params });
 
-    expect(mockApiCaller).toBeCalled();
-    expect(mockApiCaller.mock.calls[0][0]).toEqual({
+    expect(mockOpenSearchApiCaller).toBeCalled();
+    expect(mockOpenSearchApiCaller.mock.calls[0][0]).toEqual({
       ...params,
       track_total_hits: true,
     });
@@ -111,4 +128,35 @@ describe('OpenSearch search strategy', () => {
     expect(response).toHaveProperty('loaded');
     expect(response).toHaveProperty('rawResponse');
   });
+
+  it('dataSource enabled, send request with dataSourceId get data source client', async () => {
+    const opensearchSearch = await opensearchSearchStrategyProvider(mockConfig$, mockLogger);
+
+    await opensearchSearch.search(
+      (mockDataSourceEnabledContext as unknown) as RequestHandlerContext,
+      {
+        dataSourceId,
+      }
+    );
+    expect(mockDataSourceApiCaller).toBeCalled();
+    expect(mockOpenSearchApiCaller).not.toBeCalled();
+  });
+
+  it('dataSource disabled, send request with dataSourceId get default client', async () => {
+    const opensearchSearch = await opensearchSearchStrategyProvider(mockConfig$, mockLogger);
+
+    await opensearchSearch.search((mockContext as unknown) as RequestHandlerContext, {
+      dataSourceId,
+    });
+    expect(mockOpenSearchApiCaller).toBeCalled();
+    expect(mockDataSourceApiCaller).not.toBeCalled();
+  });
+
+  it('dataSource enabled, send request without dataSourceId get default client', async () => {
+    const opensearchSearch = await opensearchSearchStrategyProvider(mockConfig$, mockLogger);
+
+    await opensearchSearch.search((mockContext as unknown) as RequestHandlerContext, {});
+    expect(mockOpenSearchApiCaller).toBeCalled();
+    expect(mockDataSourceApiCaller).not.toBeCalled();
+  });
 });

src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts

@@ -42,6 +42,7 @@ import {
   getShardTimeout,
   shimAbortSignal,
 } from '..';
+import { decideClient } from './decide_client';
 
 export const opensearchSearchStrategyProvider = (
   config$: Observable<SharedGlobalConfig>,
@@ -70,10 +71,9 @@ export const opensearchSearchStrategyProvider = (
       });
 
       try {
-        const promise = shimAbortSignal(
-          context.core.opensearch.client.asCurrentUser.search(params),
-          options?.abortSignal
-        );
+        const client = await decideClient(context, request);
+        const promise = shimAbortSignal(client.search(params), options?.abortSignal);
+
         const { body: rawResponse } = (await promise) as ApiResponse<SearchResponse<any>>;
 
         if (usage) usage.trackSuccess(rawResponse.took);