From 328e08e688c39de1f47fee1c357e9928c0576390 Mon Sep 17 00:00:00 2001 From: Joshua Li Date: Fri, 7 Jun 2024 15:08:18 -0700 Subject: [PATCH] [Discover-next] Support data sources for query assist (#6972) * disable query assist for non-default datasource Signed-off-by: Joshua Li * disable query assist input when loading Signed-off-by: Joshua Li * support MDS for query assist Signed-off-by: Joshua Li * add unit tests for agents Signed-off-by: Joshua Li * Revert "add unit tests for agents" This reverts commit 983514ee11362c5efe4cdb59802b3ff402b61ef2. The test configs are not yet setup in query_enhancements plugins. Signed-off-by: Joshua Li --------- Signed-off-by: Joshua Li --- .../common/query_assist/index.ts | 2 ++ .../opensearch_dashboards.json | 2 +- .../components/query_assist_bar.tsx | 11 ++++-- .../components/query_assist_input.tsx | 2 ++ .../query_assist/utils/create_extension.tsx | 36 +++++++++++++------ .../public/query_assist/utils/get_mds_id.ts | 11 ++++++ .../public/query_assist/utils/index.ts | 1 + .../query_enhancements/server/plugin.ts | 26 ++++++++------ .../query_enhancements/server/routes/index.ts | 3 +- .../server/routes/query_assist/agents.ts | 8 +++-- .../server/routes/query_assist/routes.ts | 14 ++++++-- .../query_enhancements/server/types.ts | 13 +++++++ .../ui/query_editor/query_editor_top_row.tsx | 13 +++++-- .../ui/search_bar/create_search_bar.tsx | 1 + .../data/public/ui/search_bar/search_bar.tsx | 3 ++ .../search_bar_extension.tsx | 5 +++ .../view_components/canvas/top_nav.tsx | 3 ++ 17 files changed, 121 insertions(+), 33 deletions(-) create mode 100644 plugins-extra/query_enhancements/public/query_assist/utils/get_mds_id.ts diff --git a/plugins-extra/query_enhancements/common/query_assist/index.ts b/plugins-extra/query_enhancements/common/query_assist/index.ts index d4173b12f8c..2834fbd0b09 100644 --- a/plugins-extra/query_enhancements/common/query_assist/index.ts +++ b/plugins-extra/query_enhancements/common/query_assist/index.ts @@ -11,4 +11,6 @@ export interface QueryAssistParameters { question: string; index: string; language: string; + // for MDS + dataSourceId?: string; } diff --git a/plugins-extra/query_enhancements/opensearch_dashboards.json b/plugins-extra/query_enhancements/opensearch_dashboards.json index 41c8af395ed..5e345624837 100644 --- a/plugins-extra/query_enhancements/opensearch_dashboards.json +++ b/plugins-extra/query_enhancements/opensearch_dashboards.json @@ -5,6 +5,6 @@ "server": true, "ui": true, "requiredPlugins": ["data"], - "optionalPlugins": [], + "optionalPlugins": ["dataSource", "dataSourceManagement"], "requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact"] } diff --git a/plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx b/plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx index 904382c1899..db1a24c3e0e 100644 --- a/plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx +++ b/plugins-extra/query_enhancements/public/query_assist/components/query_assist_bar.tsx @@ -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'; @@ -45,10 +46,15 @@ export const QueryAssistBar: React.FC = (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) { @@ -75,6 +81,7 @@ export const QueryAssistBar: React.FC = (props) => { diff --git a/plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx b/plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx index 476e2d1b7f5..248dc5454fd 100644 --- a/plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx +++ b/plugins-extra/query_enhancements/public/query_assist/components/query_assist_input.tsx @@ -7,6 +7,7 @@ import { getData } from '../../services'; interface QueryAssistInputProps { inputRef: React.RefObject; persistedLog: PersistedLog; + isDisabled: boolean; initialValue?: string; selectedIndex?: string; previousQuestion?: string; @@ -70,6 +71,7 @@ export const QueryAssistInput: React.FC = (props) => { setIsSuggestionsVisible(true)} onChange={(e) => setValue(e.target.value)} onKeyDown={() => setIsSuggestionsVisible(true)} diff --git a/plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx b/plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx index ef3bcb7fc33..0993d9fe238 100644 --- a/plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx +++ b/plugins-extra/query_enhancements/public/query_assist/utils/create_extension.tsx @@ -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, @@ -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 = 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) => ( diff --git a/plugins-extra/query_enhancements/public/query_assist/utils/get_mds_id.ts b/plugins-extra/query_enhancements/public/query_assist/utils/get_mds_id.ts new file mode 100644 index 00000000000..6dab73ee9cb --- /dev/null +++ b/plugins-extra/query_enhancements/public/query_assist/utils/get_mds_id.ts @@ -0,0 +1,11 @@ +import { IIndexPattern, IndexPatternsContract } from '../../../../../src/plugins/data/public'; + +export const getMdsDataSourceId = async ( + indexPatterns: IndexPatternsContract, + indexPattern: IIndexPattern | string | undefined +): Promise => { + if (!indexPattern || typeof indexPattern !== 'object' || !indexPattern.id) return undefined; + return indexPatterns + .get(indexPattern.id) + .then((indexPatternEntity) => indexPatternEntity.dataSourceRef?.id); +}; diff --git a/plugins-extra/query_enhancements/public/query_assist/utils/index.ts b/plugins-extra/query_enhancements/public/query_assist/utils/index.ts index 5384401204c..421893f700f 100644 --- a/plugins-extra/query_enhancements/public/query_assist/utils/index.ts +++ b/plugins-extra/query_enhancements/public/query_assist/utils/index.ts @@ -1,3 +1,4 @@ export * from './create_extension'; export * from './errors'; +export * from './get_mds_id'; export * from './get_persisted_log'; diff --git a/plugins-extra/query_enhancements/server/plugin.ts b/plugins-extra/query_enhancements/server/plugin.ts index 9691f1d269c..839b58df12c 100644 --- a/plugins-extra/query_enhancements/server/plugin.ts +++ b/plugins-extra/query_enhancements/server/plugin.ts @@ -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 @@ -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 @@ -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, diff --git a/plugins-extra/query_enhancements/server/routes/index.ts b/plugins-extra/query_enhancements/server/routes/index.ts index 526c780eb17..c34246b2511 100644 --- a/plugins-extra/query_enhancements/server/routes/index.ts +++ b/plugins-extra/query_enhancements/server/routes/index.ts @@ -20,7 +20,8 @@ export function defineRoutes( ISearchStrategy > ) { - registerQueryAssistRoutes(logger, router); + registerQueryAssistRoutes(router); + router.post( { path: `/api/pplql/search`, diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts b/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts index f0a8dcfe382..5f0abd4072d 100644 --- a/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts +++ b/plugins-extra/query_enhancements/server/routes/query_assist/agents.ts @@ -47,9 +47,13 @@ export const requestAgentByConfig = async (options: { context: RequestHandlerContext; configName: string; body: RequestBody; + dataSourceId?: string; }): Promise => { - 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( { diff --git a/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts b/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts index 0b8c5c89323..8e9e1510a4e 100644 --- a/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts +++ b/plugins-extra/query_enhancements/server/routes/query_assist/routes.ts @@ -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'>] ); @@ -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]); @@ -40,6 +46,7 @@ export function registerQueryAssistRoutes(logger: Logger, router: IRouter) { index: schema.string(), question: schema.string(), language: languageSchema, + dataSourceId: schema.maybe(schema.string()), }), }, }, @@ -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 }); diff --git a/plugins-extra/query_enhancements/server/types.ts b/plugins-extra/query_enhancements/server/types.ts index b4f3bd12c18..762f386ca07 100644 --- a/plugins-extra/query_enhancements/server/types.ts +++ b/plugins-extra/query_enhancements/server/types.ts @@ -4,6 +4,9 @@ */ 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 {} @@ -11,6 +14,7 @@ export interface QueryEnhancementsPluginSetup {} export interface QueryEnhancementsPluginStart {} export interface QueryEnhancementsPluginSetupDependencies { data: DataPluginSetup; + dataSource?: DataSourcePluginStart; } export interface ISchema { @@ -31,3 +35,12 @@ export interface IPPLEventsDataSource { datarows: any[]; jsonData?: any[]; } + +declare module '../../../src/core/server' { + interface RequestHandlerContext { + query_assist: { + logger: Logger; + dataSourceEnabled: boolean; + }; + } +} diff --git a/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx b/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx index d453d5f6c84..4437d3866c1 100644 --- a/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx +++ b/src/plugins/data/public/ui/query_editor/query_editor_top_row.tsx @@ -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, @@ -44,6 +51,7 @@ export interface QueryEditorTopRowProps { disableAutoFocus?: boolean; screenTitle?: string; indexPatterns?: Array; + dataSource?: DataSource; isLoading?: boolean; prepend?: React.ComponentProps['prepend']; showQueryEditor?: boolean; @@ -270,9 +278,10 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) { if (!shouldRenderSearchBarExtensions() || !queryEditorHeaderRef.current) return; return ( ); } diff --git a/src/plugins/data/public/ui/search_bar/create_search_bar.tsx b/src/plugins/data/public/ui/search_bar/create_search_bar.tsx index 05538bbd465..ae6884fed2f 100644 --- a/src/plugins/data/public/ui/search_bar/create_search_bar.tsx +++ b/src/plugins/data/public/ui/search_bar/create_search_bar.tsx @@ -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} diff --git a/src/plugins/data/public/ui/search_bar/search_bar.tsx b/src/plugins/data/public/ui/search_bar/search_bar.tsx index d98679c7c1d..398eac51fe4 100644 --- a/src/plugins/data/public/ui/search_bar/search_bar.tsx +++ b/src/plugins/data/public/ui/search_bar/search_bar.tsx @@ -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, @@ -59,6 +60,7 @@ interface SearchBarInjectedDeps { export interface SearchBarOwnProps { indexPatterns?: IIndexPattern[]; + dataSource?: DataSource; isLoading?: boolean; customSubmitButton?: React.ReactNode; screenTitle?: string; @@ -501,6 +503,7 @@ class SearchBarUI extends Component { 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} diff --git a/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx b/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx index 505846c66b0..bbcfc30b331 100644 --- a/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx +++ b/src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.tsx @@ -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; @@ -19,6 +20,10 @@ export interface SearchBarExtensionDependencies { * Currently selected index patterns. */ indexPatterns?: Array; + /** + * Currently selected data source. + */ + dataSource?: DataSource; } export interface SearchBarExtensionConfig { diff --git a/src/plugins/discover/public/application/view_components/canvas/top_nav.tsx b/src/plugins/discover/public/application/view_components/canvas/top_nav.tsx index feb7b91e7c5..134362c7b7e 100644 --- a/src/plugins/discover/public/application/view_components/canvas/top_nav.tsx +++ b/src/plugins/discover/public/application/view_components/canvas/top_nav.tsx @@ -14,6 +14,7 @@ import { getTopNavLinks } from '../../components/top_nav/get_top_nav_links'; import { useDiscoverContext } from '../context'; import { getRootBreadcrumbs } from '../../helpers/breadcrumbs'; import { opensearchFilters, connectStorageToQueryState } from '../../../../../data/public'; +import { useDataSource } from '../utils/use_datasource'; export interface TopNavProps { opts: { @@ -24,6 +25,7 @@ export interface TopNavProps { export const TopNav = ({ opts }: TopNavProps) => { const { services } = useOpenSearchDashboards(); + const dataSource = useDataSource(services); const { inspectorAdapters, savedSearch, indexPattern } = useDiscoverContext(); const [indexPatterns, setIndexPatterns] = useState(undefined); @@ -89,6 +91,7 @@ export const TopNav = ({ opts }: TopNavProps) => { useDefaultBehaviors setMenuMountPoint={opts.setHeaderActionMenu} indexPatterns={indexPattern ? [indexPattern] : indexPatterns} + dataSource={dataSource} onQuerySubmit={opts.onQuerySubmit} /> );