[Discover-next] Support data sources for query assist (opensearch-pro…

…ject#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>
(cherry picked from commit 328e08e)

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

@@ -254,9 +254,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}
       />
     );
   }
@@ -293,9 +294,7 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
   }
 
   function shouldRenderSearchBarExtensions(): boolean {
-    return Boolean(
-      queryLanguage && props.queryEnhancements?.get(queryLanguage)?.searchBar?.extensions?.length
-    );
+    return Boolean(queryUiEnhancement?.extensions?.length);
   }
 
   function renderUpdateButton() {

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

@@ -210,6 +210,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;
@@ -494,6 +496,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 {

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<DiscoverViewServices>();
+  const dataSource = useDataSource(services);
   const { inspectorAdapters, savedSearch, indexPattern } = useDiscoverContext();
   const [indexPatterns, setIndexPatterns] = useState<IndexPattern[] | undefined>(undefined);
 
@@ -89,6 +91,7 @@ export const TopNav = ({ opts }: TopNavProps) => {
       useDefaultBehaviors
       setMenuMountPoint={opts.setHeaderActionMenu}
       indexPatterns={indexPattern ? [indexPattern] : indexPatterns}
+      dataSource={dataSource}
       onQuerySubmit={opts.onQuerySubmit}
     />
   );