Skip to content

Commit

Permalink
[MD] Refactor search strategy to conditionally user datasource client (
Browse files Browse the repository at this point in the history
  • Loading branch information
zhongnansu authored and kristenTian committed Sep 14, 2022
1 parent fbcdd94 commit ce21de8
Show file tree
Hide file tree
Showing 7 changed files with 110 additions and 21 deletions.
1 change: 1 addition & 0 deletions src/plugins/data/common/search/opensearch_search/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<
Expand Down
20 changes: 20 additions & 0 deletions src/plugins/data/common/search/search_source/search_source.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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', () => {
Expand Down
12 changes: 8 additions & 4 deletions src/plugins/data/common/search/search_source/search_source.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}

Expand Down Expand Up @@ -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));
}

/**
Expand Down Expand Up @@ -466,7 +470,7 @@ export class SearchSource {
}
}

private flatten() {
flatten() {
const searchRequest = this.mergeProps();

searchRequest.body = searchRequest.body || {};
Expand Down
7 changes: 2 additions & 5 deletions src/plugins/data/opensearch_dashboards.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": [
Expand All @@ -15,4 +12,4 @@
"opensearchDashboardsReact",
"inspector"
]
}
}
19 changes: 19 additions & 0 deletions src/plugins/data/server/search/opensearch_search/decide_client.ts
Original file line number Diff line number Diff line change
@@ -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;
};
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@ describe('OpenSearch search strategy', () => {
const mockLogger: any = {
debug: () => {},
};
const mockApiCaller = jest.fn().mockResolvedValue({
const body = {
body: {
_shards: {
total: 10,
Expand All @@ -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 () => {
Expand All @@ -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,
Expand All @@ -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,
});
Expand All @@ -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();
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ import {
getShardTimeout,
shimAbortSignal,
} from '..';
import { decideClient } from './decide_client';

export const opensearchSearchStrategyProvider = (
config$: Observable<SharedGlobalConfig>,
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit ce21de8

Please sign in to comment.