Skip to content

Commit 328e08e

Browse files
joshuali925joshuali925
authored
[Discover-next] Support data sources for query assist (#6972)
* disable query assist for non-default datasource Signed-off-by: Joshua Li <joshuali925@gmail.com> * disable query assist input when loading Signed-off-by: Joshua Li <joshuali925@gmail.com> * support MDS for query assist Signed-off-by: Joshua Li <joshuali925@gmail.com> * add unit tests for agents Signed-off-by: Joshua Li <joshuali925@gmail.com> * Revert "add unit tests for agents" This reverts commit 983514e. The test configs are not yet setup in query_enhancements plugins. Signed-off-by: Joshua Li <joshuali925@gmail.com> --------- Signed-off-by: Joshua Li <joshuali925@gmail.com>
1 parent 3ef7479 commit 328e08e

File tree

17 files changed

+121
-33
lines changed

17 files changed

+121
-33
lines changed

‎plugins-extra/query_enhancements/common/query_assist/index.ts

+2
Original file line numberDiff line numberDiff line change
@@ -11,4 +11,6 @@ export interface QueryAssistParameters {
1111
question: string;
1212
index: string;
1313
language: string;
14+
// for MDS
15+
dataSourceId?: string;
1416
}

‎plugins-extra/query_enhancements/opensearch_dashboards.json

+1-1
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,6 @@
55
"server": true,
66
"ui": true,
77
"requiredPlugins": ["data"],
8-
"optionalPlugins": [],
8+
"optionalPlugins": ["dataSource", "dataSourceManagement"],
99
"requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact"]
1010
}

‎plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx

+9-2
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,10 @@ import React, { SyntheticEvent, useMemo, useRef, useState } from 'react';
33
import { IDataPluginServices, PersistedLog } from '../../../../../src/plugins/data/public';
44
import { SearchBarExtensionDependencies } from '../../../../../src/plugins/data/public/ui/search_bar_extensions/search_bar_extension';
55
import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
6+
import { QueryAssistParameters } from '../../../common/query_assist';
67
import { getStorage } from '../../services';
78
import { useGenerateQuery } from '../hooks';
8-
import { getPersistedLog, ProhibitedQueryError } from '../utils';
9+
import { getMdsDataSourceId, getPersistedLog, ProhibitedQueryError } from '../utils';
910
import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs';
1011
import { QueryAssistInput } from './query_assist_input';
1112
import { QueryAssistSubmitButton } from './submit_button';
@@ -45,10 +46,15 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
4546
dismissCallout();
4647
previousQuestionRef.current = inputRef.current.value;
4748
persistedLog.add(inputRef.current.value);
48-
const params = {
49+
const dataSourceId = await getMdsDataSourceId(
50+
services.data.indexPatterns,
51+
selectedIndexPattern
52+
);
53+
const params: QueryAssistParameters = {
4954
question: inputRef.current.value,
5055
index: selectedIndex,
5156
language: props.language,
57+
dataSourceId,
5258
};
5359
const { response, error } = await generateQuery(params);
5460
if (error) {
@@ -75,6 +81,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
7581
<QueryAssistInput
7682
inputRef={inputRef}
7783
persistedLog={persistedLog}
84+
isDisabled={loading}
7885
selectedIndex={selectedIndex}
7986
previousQuestion={previousQuestionRef.current}
8087
/>

‎plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx

+2
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { getData } from '../../services';
77
interface QueryAssistInputProps {
88
inputRef: React.RefObject<HTMLInputElement>;
99
persistedLog: PersistedLog;
10+
isDisabled: boolean;
1011
initialValue?: string;
1112
selectedIndex?: string;
1213
previousQuestion?: string;
@@ -70,6 +71,7 @@ export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
7071
<EuiFieldText
7172
inputRef={props.inputRef}
7273
value={value}
74+
disabled={props.isDisabled}
7375
onClick={() => setIsSuggestionsVisible(true)}
7476
onChange={(e) => setValue(e.target.value)}
7577
onKeyDown={() => setIsSuggestionsVisible(true)}

‎plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx

+25-11
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,9 @@
1-
import React from 'react';
21
import { HttpSetup } from 'opensearch-dashboards/public';
3-
import { QueryAssistBar } from '../components';
2+
import React from 'react';
3+
import { getMdsDataSourceId } from '.';
44
import { SearchBarExtensionConfig } from '../../../../../src/plugins/data/public/ui/search_bar_extensions';
5+
import { getData } from '../../services';
6+
import { QueryAssistBar } from '../components';
57

68
export const createQueryAssistExtension = (
79
http: HttpSetup,
@@ -11,15 +13,27 @@ export const createQueryAssistExtension = (
1113
id: 'query-assist',
1214
order: 1000,
1315
isEnabled: (() => {
14-
let agentConfigured: boolean;
15-
return async () => {
16-
if (agentConfigured === undefined) {
17-
agentConfigured = await http
18-
.get<{ configured: boolean }>(`/api/ql/query_assist/configured/${language}`)
19-
.then((response) => response.configured)
20-
.catch(() => false);
21-
}
22-
return agentConfigured;
16+
const agentConfiguredMap: Map<string | undefined, boolean> = new Map();
17+
return async (dependencies) => {
18+
// currently query assist tool relies on opensearch API to get index
19+
// mappings, other data sources are not supported
20+
if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default')
21+
return false;
22+
23+
const dataSourceId = await getMdsDataSourceId(
24+
getData().indexPatterns,
25+
dependencies.indexPatterns?.at(0)
26+
);
27+
const cached = agentConfiguredMap.get(dataSourceId);
28+
if (cached !== undefined) return cached;
29+
const configured = await http
30+
.get<{ configured: boolean }>(`/api/ql/query_assist/configured/${language}`, {
31+
query: { dataSourceId },
32+
})
33+
.then((response) => response.configured)
34+
.catch(() => false);
35+
agentConfiguredMap.set(dataSourceId, configured);
36+
return configured;
2337
};
2438
})(),
2539
getComponent: (dependencies) => (

‎plugins-extra/query_enhancements/public/query_assist/utils/get_mds_id.ts

+11
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
import { IIndexPattern, IndexPatternsContract } from '../../../../../src/plugins/data/public';
2+
3+
export const getMdsDataSourceId = async (
4+
indexPatterns: IndexPatternsContract,
5+
indexPattern: IIndexPattern | string | undefined
6+
): Promise<string | undefined> => {
7+
if (!indexPattern || typeof indexPattern !== 'object' || !indexPattern.id) return undefined;
8+
return indexPatterns
9+
.get(indexPattern.id)
10+
.then((indexPatternEntity) => indexPatternEntity.dataSourceRef?.id);
11+
};

‎plugins-extra/query_enhancements/public/query_assist/utils/index.ts

+1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
export * from './create_extension';
22
export * from './errors';
3+
export * from './get_mds_id';
34
export * from './get_persisted_log';

‎plugins-extra/query_enhancements/server/plugin.ts

+15-11
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,24 @@
11
import { Observable } from 'rxjs';
22
import {
3-
PluginInitializerContext,
43
CoreSetup,
54
CoreStart,
6-
Plugin,
75
Logger,
6+
Plugin,
7+
PluginInitializerContext,
88
SharedGlobalConfig,
99
} from '../../../src/core/server';
10-
10+
import { PPL_SEARCH_STRATEGY, SQL_ASYNC_SEARCH_STRATEGY, SQL_SEARCH_STRATEGY } from '../common';
11+
import { defineRoutes } from './routes';
12+
import { EnginePlugin } from './search/engine_plugin';
13+
import { PPLPlugin } from './search/ppl/ppl_plugin';
14+
import { pplSearchStrategyProvider } from './search/ppl/ppl_search_strategy';
15+
import { sqlAsyncSearchStrategyProvider } from './search/sql/sql_async_search_strategy';
16+
import { sqlSearchStrategyProvider } from './search/sql/sql_search_strategy';
1117
import {
1218
QueryEnhancementsPluginSetup,
1319
QueryEnhancementsPluginSetupDependencies,
1420
QueryEnhancementsPluginStart,
1521
} from './types';
16-
import { defineRoutes } from './routes';
17-
import { PPLPlugin } from './search/ppl/ppl_plugin';
18-
import { EnginePlugin } from './search/engine_plugin';
19-
import { PPL_SEARCH_STRATEGY, SQL_SEARCH_STRATEGY, SQL_ASYNC_SEARCH_STRATEGY } from '../common';
20-
import { pplSearchStrategyProvider } from './search/ppl/ppl_search_strategy';
21-
import { sqlSearchStrategyProvider } from './search/sql/sql_search_strategy';
22-
import { sqlAsyncSearchStrategyProvider } from './search/sql/sql_async_search_strategy';
2322
import { uiSettings } from './ui_settings';
2423

2524
export class QueryEnhancementsPlugin
@@ -31,7 +30,7 @@ export class QueryEnhancementsPlugin
3130
this.config$ = initializerContext.config.legacy.globalConfig$;
3231
}
3332

34-
public setup(core: CoreSetup, { data }: QueryEnhancementsPluginSetupDependencies) {
33+
public setup(core: CoreSetup, { data, dataSource }: QueryEnhancementsPluginSetupDependencies) {
3534
this.logger.debug('queryEnhancements: Setup');
3635
const router = core.http.createRouter();
3736
// Register server side APIs
@@ -53,6 +52,11 @@ export class QueryEnhancementsPlugin
5352
data.search.registerSearchStrategy(SQL_SEARCH_STRATEGY, sqlSearchStrategy);
5453
data.search.registerSearchStrategy(SQL_ASYNC_SEARCH_STRATEGY, sqlAsyncSearchStrategy);
5554

55+
core.http.registerRouteHandlerContext('query_assist', () => ({
56+
logger: this.logger,
57+
dataSourceEnabled: !!dataSource,
58+
}));
59+
5660
defineRoutes(this.logger, router, {
5761
ppl: pplSearchStrategy,
5862
sql: sqlSearchStrategy,

‎plugins-extra/query_enhancements/server/routes/index.ts

+2-1
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ export function defineRoutes(
2020
ISearchStrategy<IOpenSearchDashboardsSearchRequest, IDataFrameResponse>
2121
>
2222
) {
23-
registerQueryAssistRoutes(logger, router);
23+
registerQueryAssistRoutes(router);
24+
2425
router.post(
2526
{
2627
path: `/api/pplql/search`,

‎plugins-extra/query_enhancements/server/routes/query_assist/agents.ts

+6-2
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,13 @@ export const requestAgentByConfig = async (options: {
4747
context: RequestHandlerContext;
4848
configName: string;
4949
body: RequestBody;
50+
dataSourceId?: string;
5051
}): Promise<AgentResponse> => {
51-
const { context, configName, body } = options;
52-
const client = context.core.opensearch.client.asCurrentUser;
52+
const { context, configName, body, dataSourceId } = options;
53+
const client =
54+
context.query_assist.dataSourceEnabled && dataSourceId
55+
? await context.dataSource.opensearch.getClient(dataSourceId)
56+
: context.core.opensearch.client.asCurrentUser;
5357
const agentId = await getAgentIdByConfig(client, configName);
5458
return client.transport.request(
5559
{

‎plugins-extra/query_enhancements/server/routes/query_assist/routes.ts

+11-3
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
import { schema, Type } from '@osd/config-schema';
2-
import { IRouter, Logger } from 'opensearch-dashboards/server';
2+
import { IRouter } from 'opensearch-dashboards/server';
33
import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors';
44
import { ERROR_DETAILS } from '../../../common/query_assist';
55
import { getAgentIdByConfig, requestAgentByConfig } from './agents';
66
import { AGENT_CONFIG_NAME_MAP } from './index';
77
import { createPPLResponseBody } from './ppl/create_response';
88

9-
export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
9+
export function registerQueryAssistRoutes(router: IRouter) {
1010
const languageSchema = schema.oneOf(
1111
Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'PPL'>]
1212
);
@@ -18,10 +18,16 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
1818
params: schema.object({
1919
language: languageSchema,
2020
}),
21+
query: schema.object({
22+
dataSourceId: schema.maybe(schema.string()),
23+
}),
2124
},
2225
},
2326
async (context, request, response) => {
24-
const client = context.core.opensearch.client.asCurrentUser;
27+
const client =
28+
context.query_assist.dataSourceEnabled && request.query.dataSourceId
29+
? await context.dataSource.opensearch.getClient(request.query.dataSourceId)
30+
: context.core.opensearch.client.asCurrentUser;
2531
try {
2632
// if the call does not throw any error, then the agent is properly configured
2733
await getAgentIdByConfig(client, AGENT_CONFIG_NAME_MAP[request.params.language]);
@@ -40,6 +46,7 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
4046
index: schema.string(),
4147
question: schema.string(),
4248
language: languageSchema,
49+
dataSourceId: schema.maybe(schema.string()),
4350
}),
4451
},
4552
},
@@ -56,6 +63,7 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
5663
question: request.body.question,
5764
},
5865
},
66+
dataSourceId: request.body.dataSourceId,
5967
});
6068
const responseBody = createPPLResponseBody(agentResponse);
6169
return response.ok({ body: responseBody });

‎plugins-extra/query_enhancements/server/types.ts

+13
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,17 @@
44
*/
55

66
import { DataPluginSetup } from 'src/plugins/data/server/plugin';
7+
import { HomeServerPluginSetup } from 'src/plugins/home/server/plugin';
8+
import { Logger } from '../../../src/core/server';
9+
import { DataSourcePluginStart } from '../../../src/plugins/data_source/server';
710

811
// eslint-disable-next-line @typescript-eslint/no-empty-interface
912
export interface QueryEnhancementsPluginSetup {}
1013
// eslint-disable-next-line @typescript-eslint/no-empty-interface
1114
export interface QueryEnhancementsPluginStart {}
1215
export interface QueryEnhancementsPluginSetupDependencies {
1316
data: DataPluginSetup;
17+
dataSource?: DataSourcePluginStart;
1418
}
1519

1620
export interface ISchema {
@@ -31,3 +35,12 @@ export interface IPPLEventsDataSource {
3135
datarows: any[];
3236
jsonData?: any[];
3337
}
38+
39+
declare module '../../../src/core/server' {
40+
interface RequestHandlerContext {
41+
query_assist: {
42+
logger: Logger;
43+
dataSourceEnabled: boolean;
44+
};
45+
}
46+
}

‎src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx

+11-2
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,14 @@ import {
1616
// @ts-ignore
1717
import { EuiSuperUpdateButton, OnRefreshProps } from '@elastic/eui';
1818
import { isEqual, compact } from 'lodash';
19-
import { IDataPluginServices, IIndexPattern, TimeRange, TimeHistoryContract, Query } from '../..';
19+
import {
20+
IDataPluginServices,
21+
IIndexPattern,
22+
TimeRange,
23+
TimeHistoryContract,
24+
Query,
25+
DataSource,
26+
} from '../..';
2027
import {
2128
useOpenSearchDashboards,
2229
withOpenSearchDashboards,
@@ -44,6 +51,7 @@ export interface QueryEditorTopRowProps {
4451
disableAutoFocus?: boolean;
4552
screenTitle?: string;
4653
indexPatterns?: Array<IIndexPattern | string>;
54+
dataSource?: DataSource;
4755
isLoading?: boolean;
4856
prepend?: React.ComponentProps<typeof EuiFieldText>['prepend'];
4957
showQueryEditor?: boolean;
@@ -270,9 +278,10 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
270278
if (!shouldRenderSearchBarExtensions() || !queryEditorHeaderRef.current) return;
271279
return (
272280
<SearchBarExtensions
273-
configs={props.queryEnhancements?.get(queryLanguage!)?.searchBar?.extensions}
281+
configs={queryUiEnhancement?.extensions}
274282
portalContainer={queryEditorHeaderRef.current}
275283
indexPatterns={props.indexPatterns}
284+
dataSource={props.dataSource}
276285
/>
277286
);
278287
}

‎src/plugins/data/public/ui/search_bar/create_search_bar.tsx

+1
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,7 @@ export function createSearchBar({
212212
showSaveQuery={props.showSaveQuery}
213213
screenTitle={props.screenTitle}
214214
indexPatterns={props.indexPatterns}
215+
dataSource={props.dataSource}
215216
indicateNoData={props.indicateNoData}
216217
timeHistory={data.query.timefilter.history}
217218
dateRangeFrom={timeRange.from}

‎src/plugins/data/public/ui/search_bar/search_bar.tsx

+3
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import classNames from 'classnames';
3333
import { compact, get, isEqual } from 'lodash';
3434
import React, { Component } from 'react';
3535
import ResizeObserver from 'resize-observer-polyfill';
36+
import { DataSource } from '../..';
3637
import {
3738
OpenSearchDashboardsReactContextValue,
3839
withOpenSearchDashboards,
@@ -59,6 +60,7 @@ interface SearchBarInjectedDeps {
5960

6061
export interface SearchBarOwnProps {
6162
indexPatterns?: IIndexPattern[];
63+
dataSource?: DataSource;
6264
isLoading?: boolean;
6365
customSubmitButton?: React.ReactNode;
6466
screenTitle?: string;
@@ -501,6 +503,7 @@ class SearchBarUI extends Component<SearchBarProps, State> {
501503
screenTitle={this.props.screenTitle}
502504
onSubmit={this.onQueryBarSubmit}
503505
indexPatterns={this.props.indexPatterns}
506+
dataSource={this.props.dataSource}
504507
isLoading={this.props.isLoading}
505508
prepend={this.props.showFilterBar ? savedQueryManagement : undefined}
506509
showDatePicker={this.props.showDatePicker}

‎src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx

+5
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { EuiErrorBoundary } from '@elastic/eui';
77
import React, { useEffect, useMemo, useRef, useState } from 'react';
88
import ReactDOM from 'react-dom';
99
import { IIndexPattern } from '../../../common';
10+
import { DataSource } from '../../data_sources/datasource';
1011

1112
interface SearchBarExtensionProps {
1213
config: SearchBarExtensionConfig;
@@ -19,6 +20,10 @@ export interface SearchBarExtensionDependencies {
1920
* Currently selected index patterns.
2021
*/
2122
indexPatterns?: Array<IIndexPattern | string>;
23+
/**
24+
* Currently selected data source.
25+
*/
26+
dataSource?: DataSource;
2227
}
2328

2429
export interface SearchBarExtensionConfig {

0 commit comments

Comments
 (0)