From 1e874f853564f0864343d9abc530efc9c992e5b2 Mon Sep 17 00:00:00 2001 From: Zhongnan Su Date: Thu, 18 Aug 2022 18:53:11 -0700 Subject: [PATCH 1/3] Refactor search strategy to conditionally user datasource client Signed-off-by: Zhongnan Su --- .../common/search/opensearch_search/types.ts | 1 + .../search/search_source/search_source.ts | 15 ++++++++++++--- src/plugins/data/opensearch_dashboards.json | 7 ++----- .../search/opensearch_search/decide_client.ts | 19 +++++++++++++++++++ .../opensearch_search_strategy.ts | 8 ++++---- 5 files changed, 38 insertions(+), 12 deletions(-) create mode 100644 src/plugins/data/server/search/opensearch_search/decide_client.ts diff --git a/src/plugins/data/common/search/opensearch_search/types.ts b/src/plugins/data/common/search/opensearch_search/types.ts index 5ff0fb689f80..a45d4d0660ce 100644 --- a/src/plugins/data/common/search/opensearch_search/types.ts +++ b/src/plugins/data/common/search/opensearch_search/types.ts @@ -52,6 +52,7 @@ export type ISearchRequestParams> = { export interface IOpenSearchSearchRequest extends IOpenSearchDashboardsSearchRequest { indexType?: string; + dataSourceId?: string; } export type IOpenSearchSearchResponse = IOpenSearchDashboardsSearchResponse< diff --git a/src/plugins/data/common/search/search_source/search_source.ts b/src/plugins/data/common/search/search_source/search_source.ts index 229357abffe4..848ec8d3328f 100644 --- a/src/plugins/data/common/search/search_source/search_source.ts +++ b/src/plugins/data/common/search/search_source/search_source.ts @@ -280,8 +280,16 @@ export class SearchSource { let response; if (getConfig(UI_SETTINGS.COURIER_BATCH_SEARCHES)) { + /** + * batch search will issue `_msearch` requests to OpenSearch, which is not likely to work with + * multiple data sources. maybe deprecate this, or simply disable it when multiple data source is supported? + * zengyan NOTE: Kibana is going to remove it: https://github.com/elastic/kibana/issues/55140 + */ 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 +343,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)); } /** diff --git a/src/plugins/data/opensearch_dashboards.json b/src/plugins/data/opensearch_dashboards.json index b75610115176..1193def7cb91 100644 --- a/src/plugins/data/opensearch_dashboards.json +++ b/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" ] -} \ No newline at end of file +} diff --git a/src/plugins/data/server/search/opensearch_search/decide_client.ts b/src/plugins/data/server/search/opensearch_search/decide_client.ts new file mode 100644 index 000000000000..5f38b82ffb4b --- /dev/null +++ b/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 => { + // 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; +}; diff --git a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts b/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts index a29885a0a9f2..4ccc8db05728 100644 --- a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.ts +++ b/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, @@ -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>; if (usage) usage.trackSuccess(rawResponse.took); From 46842b59304e0799ba2ec9dab7671f1daf24af06 Mon Sep 17 00:00:00 2001 From: Zhongnan Su Date: Thu, 18 Aug 2022 21:06:58 -0700 Subject: [PATCH 2/3] update tests Signed-off-by: Zhongnan Su --- .../search_source/search_source.test.ts | 20 ++++++ .../search/search_source/search_source.ts | 2 +- .../opensearch_search_strategy.test.ts | 64 ++++++++++++++++--- 3 files changed, 77 insertions(+), 9 deletions(-) diff --git a/src/plugins/data/common/search/search_source/search_source.test.ts b/src/plugins/data/common/search/search_source/search_source.test.ts index 40193e1a4a30..92cc0682a136 100644 --- a/src/plugins/data/common/search/search_source/search_source.test.ts +++ b/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', () => { diff --git a/src/plugins/data/common/search/search_source/search_source.ts b/src/plugins/data/common/search/search_source/search_source.ts index 848ec8d3328f..ba9ea7910a6d 100644 --- a/src/plugins/data/common/search/search_source/search_source.ts +++ b/src/plugins/data/common/search/search_source/search_source.ts @@ -475,7 +475,7 @@ export class SearchSource { } } - private flatten() { + flatten() { const searchRequest = this.mergeProps(); searchRequest.body = searchRequest.body || {}; diff --git a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts b/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts index 6207273c95c9..fe95e3d7d4eb 100644 --- a/src/plugins/data/server/search/opensearch_search/opensearch_search_strategy.test.ts +++ b/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,7 +45,19 @@ 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: { @@ -53,13 +65,18 @@ describe('OpenSearch search strategy', () => { get: () => {}, }, }, - opensearch: { client: { asCurrentUser: { search: mockApiCaller } } }, + opensearch: { client: { asCurrentUser: { search: mockOpenSearchApiCaller } } }, }, }; + const mockDataSourceEnabledContext = { + ...mockContext, + ...mockDataSourceContext, + }; const mockConfig$ = pluginInitializerContextConfigMock({}).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(); + }); }); From a5c046e7d5a087a80a545454cd4d9db4f9133722 Mon Sep 17 00:00:00 2001 From: Zhongnan Su Date: Fri, 19 Aug 2022 10:37:40 -0700 Subject: [PATCH 3/3] address comments Signed-off-by: Zhongnan Su --- .../data/common/search/search_source/search_source.ts | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/plugins/data/common/search/search_source/search_source.ts b/src/plugins/data/common/search/search_source/search_source.ts index ba9ea7910a6d..9a85ec85ce87 100644 --- a/src/plugins/data/common/search/search_source/search_source.ts +++ b/src/plugins/data/common/search/search_source/search_source.ts @@ -280,11 +280,6 @@ export class SearchSource { let response; if (getConfig(UI_SETTINGS.COURIER_BATCH_SEARCHES)) { - /** - * batch search will issue `_msearch` requests to OpenSearch, which is not likely to work with - * multiple data sources. maybe deprecate this, or simply disable it when multiple data source is supported? - * zengyan NOTE: Kibana is going to remove it: https://github.com/elastic/kibana/issues/55140 - */ response = await this.legacyFetch(searchRequest, options); } else { const indexPattern = this.getField('index');