[Discover-next] Support data sources for query assist

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;
 }

plugins-extra/query_enhancements/opensearch_dashboards.json

@@ -5,6 +5,6 @@
   "server": true,
   "ui": true,
   "requiredPlugins": ["data"],
-  "optionalPlugins": [],
+  "optionalPlugins": ["dataSource", "dataSourceManagement"],
   "requiredBundles": ["opensearchDashboardsUtils", "opensearchDashboardsReact"]
 }

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<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) {
@@ -75,6 +81,7 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
             <QueryAssistInput
               inputRef={inputRef}
               persistedLog={persistedLog}
+              isDisabled={loading}
               selectedIndex={selectedIndex}
               previousQuestion={previousQuestionRef.current}
             />

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<HTMLInputElement>;
   persistedLog: PersistedLog;
+  isDisabled: boolean;
   initialValue?: string;
   selectedIndex?: string;
   previousQuestion?: string;
@@ -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)}

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<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) => (

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<string | undefined> => {
+  if (!indexPattern || typeof indexPattern !== 'object' || !indexPattern.id) return undefined;
+  return indexPatterns
+    .get(indexPattern.id)
+    .then((indexPatternEntity) => indexPatternEntity.dataSourceRef?.id);
+};

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';

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,

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

@@ -20,7 +20,8 @@ export function defineRoutes(
     ISearchStrategy<IOpenSearchDashboardsSearchRequest, IDataFrameResponse>
   >
 ) {
-  registerQueryAssistRoutes(logger, router);
+  registerQueryAssistRoutes(router);
+
   router.post(
     {
       path: `/api/pplql/search`,

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<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(
     {

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 });

plugins-extra/query_enhancements/server/types.ts

@@ -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 {
@@ -31,3 +35,12 @@ export interface IPPLEventsDataSource {
   datarows: any[];
   jsonData?: any[];
 }
+
+declare module '../../../src/core/server' {
+  interface RequestHandlerContext {
+    query_assist: {
+      logger: Logger;
+      dataSourceEnabled: boolean;
+    };
+  }
+}

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<IIndexPattern | string>;
+  dataSource?: DataSource;
   isLoading?: boolean;
   prepend?: React.ComponentProps<typeof EuiFieldText>['prepend'];
   showQueryEditor?: boolean;
@@ -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}
       />
     );
   }

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}

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<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}

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<IIndexPattern | string>;
+  /**
+   * Currently selected data source.
+   */
+  dataSource?: DataSource;
 }
 
 export interface SearchBarExtensionConfig {