Skip to content

Commit

Permalink
[Discover-next] Support data sources for query assist (#6972)
Browse files Browse the repository at this point in the history
* disable query assist for non-default datasource

Signed-off-by: Joshua Li <[email protected]>

* disable query assist input when loading

Signed-off-by: Joshua Li <[email protected]>

* support MDS for query assist

Signed-off-by: Joshua Li <[email protected]>

* add unit tests for agents

Signed-off-by: Joshua Li <[email protected]>

* 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 <[email protected]>

---------

Signed-off-by: Joshua Li <[email protected]>
  • Loading branch information
joshuali925 authored Jun 7, 2024
1 parent 3ef7479 commit 328e08e
Show file tree
Hide file tree
Showing 17 changed files with 121 additions and 33 deletions.
2 changes: 2 additions & 0 deletions plugins-extra/query_enhancements/common/query_assist/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,4 +11,6 @@ export interface QueryAssistParameters {
question: string;
index: string;
language: string;
// for MDS
dataSourceId?: string;
}
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,6 @@
"server": true,
"ui": true,
"requiredPlugins": ["data"],
"optionalPlugins": [],
"optionalPlugins": ["dataSource", "dataSourceManagement"],
"requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact"]
}
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,10 @@ import React, { SyntheticEvent, useMemo, useRef, useState } from 'react';
import { IDataPluginServices, PersistedLog } from '../../../../../src/plugins/data/public';
import { SearchBarExtensionDependencies } from '../../../../../src/plugins/data/public/ui/search_bar_extensions/search_bar_extension';
import { useOpenSearchDashboards } from '../../../../../src/plugins/opensearch_dashboards_react/public';
import { QueryAssistParameters } from '../../../common/query_assist';
import { getStorage } from '../../services';
import { useGenerateQuery } from '../hooks';
import { getPersistedLog, ProhibitedQueryError } from '../utils';
import { getMdsDataSourceId, getPersistedLog, ProhibitedQueryError } from '../utils';
import { QueryAssistCallOut, QueryAssistCallOutType } from './call_outs';
import { QueryAssistInput } from './query_assist_input';
import { QueryAssistSubmitButton } from './submit_button';
Expand Down Expand Up @@ -45,10 +46,15 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
dismissCallout();
previousQuestionRef.current = inputRef.current.value;
persistedLog.add(inputRef.current.value);
const params = {
const dataSourceId = await getMdsDataSourceId(
services.data.indexPatterns,
selectedIndexPattern
);
const params: QueryAssistParameters = {
question: inputRef.current.value,
index: selectedIndex,
language: props.language,
dataSourceId,
};
const { response, error } = await generateQuery(params);
if (error) {
Expand All @@ -75,6 +81,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
<QueryAssistInput
inputRef={inputRef}
persistedLog={persistedLog}
isDisabled={loading}
selectedIndex={selectedIndex}
previousQuestion={previousQuestionRef.current}
/>
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import { getData } from '../../services';
interface QueryAssistInputProps {
inputRef: React.RefObject<HTMLInputElement>;
persistedLog: PersistedLog;
isDisabled: boolean;
initialValue?: string;
selectedIndex?: string;
previousQuestion?: string;
Expand Down Expand Up @@ -70,6 +71,7 @@ export const QueryAssistInput: React.FC<QueryAssistInputProps> = (props) => {
<EuiFieldText
inputRef={props.inputRef}
value={value}
disabled={props.isDisabled}
onClick={() => setIsSuggestionsVisible(true)}
onChange={(e) => setValue(e.target.value)}
onKeyDown={() => setIsSuggestionsVisible(true)}
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,9 @@
import React from 'react';
import { HttpSetup } from 'opensearch-dashboards/public';
import { QueryAssistBar } from '../components';
import React from 'react';
import { getMdsDataSourceId } from '.';
import { SearchBarExtensionConfig } from '../../../../../src/plugins/data/public/ui/search_bar_extensions';
import { getData } from '../../services';
import { QueryAssistBar } from '../components';

export const createQueryAssistExtension = (
http: HttpSetup,
Expand All @@ -11,15 +13,27 @@ export const createQueryAssistExtension = (
id: 'query-assist',
order: 1000,
isEnabled: (() => {
let agentConfigured: boolean;
return async () => {
if (agentConfigured === undefined) {
agentConfigured = await http
.get<{ configured: boolean }>(`/api/ql/query_assist/configured/${language}`)
.then((response) => response.configured)
.catch(() => false);
}
return agentConfigured;
const agentConfiguredMap: Map<string | undefined, boolean> = new Map();
return async (dependencies) => {
// currently query assist tool relies on opensearch API to get index
// mappings, other data sources are not supported
if (dependencies.dataSource && dependencies.dataSource?.getType() !== 'default')
return false;

const dataSourceId = await getMdsDataSourceId(
getData().indexPatterns,
dependencies.indexPatterns?.at(0)
);
const cached = agentConfiguredMap.get(dataSourceId);
if (cached !== undefined) return cached;
const configured = await http
.get<{ configured: boolean }>(`/api/ql/query_assist/configured/${language}`, {
query: { dataSourceId },
})
.then((response) => response.configured)
.catch(() => false);
agentConfiguredMap.set(dataSourceId, configured);
return configured;
};
})(),
getComponent: (dependencies) => (
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,11 @@
import { IIndexPattern, IndexPatternsContract } from '../../../../../src/plugins/data/public';

export const getMdsDataSourceId = async (
indexPatterns: IndexPatternsContract,
indexPattern: IIndexPattern | string | undefined
): Promise<string | undefined> => {
if (!indexPattern || typeof indexPattern !== 'object' || !indexPattern.id) return undefined;
return indexPatterns
.get(indexPattern.id)
.then((indexPatternEntity) => indexPatternEntity.dataSourceRef?.id);
};
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
export * from './create_extension';
export * from './errors';
export * from './get_mds_id';
export * from './get_persisted_log';
26 changes: 15 additions & 11 deletions plugins-extra/query_enhancements/server/plugin.ts
Original file line number Diff line number Diff line change
@@ -1,25 +1,24 @@
import { Observable } from 'rxjs';
import {
PluginInitializerContext,
CoreSetup,
CoreStart,
Plugin,
Logger,
Plugin,
PluginInitializerContext,
SharedGlobalConfig,
} from '../../../src/core/server';

import { PPL_SEARCH_STRATEGY, SQL_ASYNC_SEARCH_STRATEGY, SQL_SEARCH_STRATEGY } from '../common';
import { defineRoutes } from './routes';
import { EnginePlugin } from './search/engine_plugin';
import { PPLPlugin } from './search/ppl/ppl_plugin';
import { pplSearchStrategyProvider } from './search/ppl/ppl_search_strategy';
import { sqlAsyncSearchStrategyProvider } from './search/sql/sql_async_search_strategy';
import { sqlSearchStrategyProvider } from './search/sql/sql_search_strategy';
import {
QueryEnhancementsPluginSetup,
QueryEnhancementsPluginSetupDependencies,
QueryEnhancementsPluginStart,
} from './types';
import { defineRoutes } from './routes';
import { PPLPlugin } from './search/ppl/ppl_plugin';
import { EnginePlugin } from './search/engine_plugin';
import { PPL_SEARCH_STRATEGY, SQL_SEARCH_STRATEGY, SQL_ASYNC_SEARCH_STRATEGY } from '../common';
import { pplSearchStrategyProvider } from './search/ppl/ppl_search_strategy';
import { sqlSearchStrategyProvider } from './search/sql/sql_search_strategy';
import { sqlAsyncSearchStrategyProvider } from './search/sql/sql_async_search_strategy';
import { uiSettings } from './ui_settings';

export class QueryEnhancementsPlugin
Expand All @@ -31,7 +30,7 @@ export class QueryEnhancementsPlugin
this.config$ = initializerContext.config.legacy.globalConfig$;
}

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

core.http.registerRouteHandlerContext('query_assist', () => ({
logger: this.logger,
dataSourceEnabled: !!dataSource,
}));

defineRoutes(this.logger, router, {
ppl: pplSearchStrategy,
sql: sqlSearchStrategy,
Expand Down
3 changes: 2 additions & 1 deletion plugins-extra/query_enhancements/server/routes/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,8 @@ export function defineRoutes(
ISearchStrategy<IOpenSearchDashboardsSearchRequest, IDataFrameResponse>
>
) {
registerQueryAssistRoutes(logger, router);
registerQueryAssistRoutes(router);

router.post(
{
path: `/api/pplql/search`,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,13 @@ export const requestAgentByConfig = async (options: {
context: RequestHandlerContext;
configName: string;
body: RequestBody;
dataSourceId?: string;
}): Promise<AgentResponse> => {
const { context, configName, body } = options;
const client = context.core.opensearch.client.asCurrentUser;
const { context, configName, body, dataSourceId } = options;
const client =
context.query_assist.dataSourceEnabled && dataSourceId
? await context.dataSource.opensearch.getClient(dataSourceId)
: context.core.opensearch.client.asCurrentUser;
const agentId = await getAgentIdByConfig(client, configName);
return client.transport.request(
{
Expand Down
Original file line number Diff line number Diff line change
@@ -1,12 +1,12 @@
import { schema, Type } from '@osd/config-schema';
import { IRouter, Logger } from 'opensearch-dashboards/server';
import { IRouter } from 'opensearch-dashboards/server';
import { isResponseError } from '../../../../../src/core/server/opensearch/client/errors';
import { ERROR_DETAILS } from '../../../common/query_assist';
import { getAgentIdByConfig, requestAgentByConfig } from './agents';
import { AGENT_CONFIG_NAME_MAP } from './index';
import { createPPLResponseBody } from './ppl/create_response';

export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
export function registerQueryAssistRoutes(router: IRouter) {
const languageSchema = schema.oneOf(
Object.keys(AGENT_CONFIG_NAME_MAP).map(schema.literal) as [Type<'PPL'>]
);
Expand All @@ -18,10 +18,16 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
params: schema.object({
language: languageSchema,
}),
query: schema.object({
dataSourceId: schema.maybe(schema.string()),
}),
},
},
async (context, request, response) => {
const client = context.core.opensearch.client.asCurrentUser;
const client =
context.query_assist.dataSourceEnabled && request.query.dataSourceId
? await context.dataSource.opensearch.getClient(request.query.dataSourceId)
: context.core.opensearch.client.asCurrentUser;
try {
// if the call does not throw any error, then the agent is properly configured
await getAgentIdByConfig(client, AGENT_CONFIG_NAME_MAP[request.params.language]);
Expand All @@ -40,6 +46,7 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
index: schema.string(),
question: schema.string(),
language: languageSchema,
dataSourceId: schema.maybe(schema.string()),
}),
},
},
Expand All @@ -56,6 +63,7 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) {
question: request.body.question,
},
},
dataSourceId: request.body.dataSourceId,
});
const responseBody = createPPLResponseBody(agentResponse);
return response.ok({ body: responseBody });
Expand Down
13 changes: 13 additions & 0 deletions plugins-extra/query_enhancements/server/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,17 @@
*/

import { DataPluginSetup } from 'src/plugins/data/server/plugin';
import { HomeServerPluginSetup } from 'src/plugins/home/server/plugin';
import { Logger } from '../../../src/core/server';
import { DataSourcePluginStart } from '../../../src/plugins/data_source/server';

// eslint-disable-next-line @typescript-eslint/no-empty-interface
export interface QueryEnhancementsPluginSetup {}
// eslint-disable-next-line @typescript-eslint/no-empty-interface
export interface QueryEnhancementsPluginStart {}
export interface QueryEnhancementsPluginSetupDependencies {
data: DataPluginSetup;
dataSource?: DataSourcePluginStart;
}

export interface ISchema {
Expand All @@ -31,3 +35,12 @@ export interface IPPLEventsDataSource {
datarows: any[];
jsonData?: any[];
}

declare module '../../../src/core/server' {
interface RequestHandlerContext {
query_assist: {
logger: Logger;
dataSourceEnabled: boolean;
};
}
}
13 changes: 11 additions & 2 deletions src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,14 @@ import {
// @ts-ignore
import { EuiSuperUpdateButton, OnRefreshProps } from '@elastic/eui';
import { isEqual, compact } from 'lodash';
import { IDataPluginServices, IIndexPattern, TimeRange, TimeHistoryContract, Query } from '../..';
import {
IDataPluginServices,
IIndexPattern,
TimeRange,
TimeHistoryContract,
Query,
DataSource,
} from '../..';
import {
useOpenSearchDashboards,
withOpenSearchDashboards,
Expand Down Expand Up @@ -44,6 +51,7 @@ export interface QueryEditorTopRowProps {
disableAutoFocus?: boolean;
screenTitle?: string;
indexPatterns?: Array<IIndexPattern | string>;
dataSource?: DataSource;
isLoading?: boolean;
prepend?: React.ComponentProps<typeof EuiFieldText>['prepend'];
showQueryEditor?: boolean;
Expand Down Expand Up @@ -270,9 +278,10 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
if (!shouldRenderSearchBarExtensions() || !queryEditorHeaderRef.current) return;
return (
<SearchBarExtensions
configs={props.queryEnhancements?.get(queryLanguage!)?.searchBar?.extensions}
configs={queryUiEnhancement?.extensions}
portalContainer={queryEditorHeaderRef.current}
indexPatterns={props.indexPatterns}
dataSource={props.dataSource}
/>
);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -212,6 +212,7 @@ export function createSearchBar({
showSaveQuery={props.showSaveQuery}
screenTitle={props.screenTitle}
indexPatterns={props.indexPatterns}
dataSource={props.dataSource}
indicateNoData={props.indicateNoData}
timeHistory={data.query.timefilter.history}
dateRangeFrom={timeRange.from}
Expand Down
3 changes: 3 additions & 0 deletions src/plugins/data/public/ui/search_bar/search_bar.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ import classNames from 'classnames';
import { compact, get, isEqual } from 'lodash';
import React, { Component } from 'react';
import ResizeObserver from 'resize-observer-polyfill';
import { DataSource } from '../..';
import {
OpenSearchDashboardsReactContextValue,
withOpenSearchDashboards,
Expand All @@ -59,6 +60,7 @@ interface SearchBarInjectedDeps {

export interface SearchBarOwnProps {
indexPatterns?: IIndexPattern[];
dataSource?: DataSource;
isLoading?: boolean;
customSubmitButton?: React.ReactNode;
screenTitle?: string;
Expand Down Expand Up @@ -501,6 +503,7 @@ class SearchBarUI extends Component<SearchBarProps, State> {
screenTitle={this.props.screenTitle}
onSubmit={this.onQueryBarSubmit}
indexPatterns={this.props.indexPatterns}
dataSource={this.props.dataSource}
isLoading={this.props.isLoading}
prepend={this.props.showFilterBar ? savedQueryManagement : undefined}
showDatePicker={this.props.showDatePicker}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import { EuiErrorBoundary } from '@elastic/eui';
import React, { useEffect, useMemo, useRef, useState } from 'react';
import ReactDOM from 'react-dom';
import { IIndexPattern } from '../../../common';
import { DataSource } from '../../data_sources/datasource';

interface SearchBarExtensionProps {
config: SearchBarExtensionConfig;
Expand All @@ -19,6 +20,10 @@ export interface SearchBarExtensionDependencies {
* Currently selected index patterns.
*/
indexPatterns?: Array<IIndexPattern | string>;
/**
* Currently selected data source.
*/
dataSource?: DataSource;
}

export interface SearchBarExtensionConfig {
Expand Down
Loading

0 comments on commit 328e08e

Please sign in to comment.