Fix UI and detection of external data source in query assist

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

@@ -39,7 +39,6 @@ monaco.languages.register({ id: LANGUAGE_ID_KUERY });
 
 export interface QueryEditorProps {
   indexPatterns: Array<IIndexPattern | string>;
-  dataSource?: DataSource;
   query: Query;
   dataSetContainerRef?: React.RefCallback<HTMLDivElement>;
   settings: Settings;
@@ -124,6 +123,10 @@ export default class QueryEditorUI extends Component<Props, State> {
     return toUser(this.props.query.query);
   };
 
+  private setIsCollapsed = (isCollapsed: boolean) => {
+    this.setState({ isCollapsed });
+  };
+
   private renderQueryEditorExtensions() {
     if (
       !(
@@ -139,11 +142,12 @@ export default class QueryEditorUI extends Component<Props, State> {
     return (
       <QueryEditorExtensions
         language={this.props.queryLanguage}
+        onSelectLanguage={this.onSelectLanguage}
+        isCollapsed={this.state.isCollapsed}
+        setIsCollapsed={this.setIsCollapsed}
         configMap={this.extensionMap}
         componentContainer={this.headerRef.current}
         bannerContainer={this.bannerRef.current}
-        indexPatterns={this.props.indexPatterns}
-        dataSource={this.props.dataSource}
       />
     );
   }
@@ -375,11 +379,11 @@ export default class QueryEditorUI extends Component<Props, State> {
       // eslint-disable-next-line no-unsanitized/property
       style.innerHTML = `
       .${containerId} .monaco-editor .view-lines {
-        padding-left: 15px; 
+        padding-left: 15px;
       }
       .${containerId} .monaco-editor .cursor {
         height: ${customCursorHeight}px !important;
-        margin-top: ${(38 - customCursorHeight) / 2}px !important; 
+        margin-top: ${(38 - customCursorHeight) / 2}px !important;
       }
     `;
 

src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension.test.tsx

@@ -6,7 +6,6 @@
 import { render, waitFor } from '@testing-library/react';
 import React, { ComponentProps } from 'react';
 import { of } from 'rxjs';
-import { IIndexPattern } from '../../../../common';
 import { QueryEditorExtension } from './query_editor_extension';
 
 jest.mock('react-dom', () => ({
@@ -16,21 +15,6 @@ jest.mock('react-dom', () => ({
 
 type QueryEditorExtensionProps = ComponentProps<typeof QueryEditorExtension>;
 
-const mockIndexPattern = {
-  id: '1234',
-  title: 'logstash-*',
-  fields: [
-    {
-      name: 'response',
-      type: 'number',
-      esTypes: ['integer'],
-      aggregatable: true,
-      filterable: true,
-      searchable: true,
-    },
-  ],
-} as IIndexPattern;
-
 describe('QueryEditorExtension', () => {
   const getComponentMock = jest.fn();
   const getBannerMock = jest.fn();
@@ -45,8 +29,10 @@ describe('QueryEditorExtension', () => {
       getBanner: getBannerMock,
     },
     dependencies: {
-      indexPatterns: [mockIndexPattern],
       language: 'Test',
+      onSelectLanguage: jest.fn(),
+      isCollapsed: false,
+      setIsCollapsed: jest.fn(),
     },
     componentContainer: document.createElement('div'),
     bannerContainer: document.createElement('div'),

src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extension.tsx

@@ -7,8 +7,6 @@ import { EuiErrorBoundary } from '@elastic/eui';
 import React, { useEffect, useMemo, useRef, useState } from 'react';
 import ReactDOM from 'react-dom';
 import { Observable } from 'rxjs';
-import { IIndexPattern } from '../../../../common';
-import { DataSource } from '../../../data_sources/datasource';
 
 interface QueryEditorExtensionProps {
   config: QueryEditorExtensionConfig;
@@ -19,17 +17,21 @@ interface QueryEditorExtensionProps {
 
 export interface QueryEditorExtensionDependencies {
   /**
-   * Currently selected index patterns.
+   * Currently selected query language.
    */
-  indexPatterns?: Array<IIndexPattern | string>;
+  language: string;
   /**
-   * Currently selected data source.
+   * Change the selected query language.
    */
-  dataSource?: DataSource;
+  onSelectLanguage: (language: string) => void;
   /**
-   * Currently selected query language.
+   * Whether the query editor is collapsed.
    */
-  language: string;
+  isCollapsed: boolean;
+  /**
+   * Set whether the query editor is collapsed.
+   */
+  setIsCollapsed: (isCollapsed: boolean) => void;
 }
 
 export interface QueryEditorExtensionConfig {

src/plugins/data/public/ui/query_editor/query_editor_extensions/query_editor_extensions.test.tsx

@@ -14,33 +14,19 @@ type QueryEditorExtensionsProps = ComponentProps<typeof QueryEditorExtensions>;
 jest.mock('./query_editor_extension', () => ({
   QueryEditorExtension: jest.fn(({ config, dependencies }: QueryEditorExtensionProps) => (
     <div>
-      Mocked QueryEditorExtension {config.id} with{' '}
-      {dependencies.indexPatterns?.map((i) => (typeof i === 'string' ? i : i.title)).join(', ')}
+      Mocked QueryEditorExtension {config.id} with {dependencies.language}
     </div>
   )),
 }));
 
 describe('QueryEditorExtensions', () => {
   const defaultProps: QueryEditorExtensionsProps = {
-    indexPatterns: [
-      {
-        id: '1234',
-        title: 'logstash-*',
-        fields: [
-          {
-            name: 'response',
-            type: 'number',
-            esTypes: ['integer'],
-            aggregatable: true,
-            filterable: true,
-            searchable: true,
-          },
-        ],
-      },
-    ],
     componentContainer: document.createElement('div'),
     bannerContainer: document.createElement('div'),
-    language: 'Test',
+    language: 'Test-lang',
+    onSelectLanguage: jest.fn(),
+    isCollapsed: false,
+    setIsCollapsed: jest.fn(),
   };
 
   beforeEach(() => {
@@ -59,8 +45,8 @@ describe('QueryEditorExtensions', () => {
 
   it('correctly orders configurations based on order property', () => {
     const configMap = {
-      '1': { id: '1', order: 2, isEnabled: jest.fn(), getComponent: jest.fn() },
-      '2': { id: '2', order: 1, isEnabled: jest.fn(), getComponent: jest.fn() },
+      '1': { id: '1', order: 2, isEnabled$: jest.fn(), getComponent: jest.fn() },
+      '2': { id: '2', order: 1, isEnabled$: jest.fn(), getComponent: jest.fn() },
     };
 
     const { getAllByText } = render(
@@ -75,18 +61,23 @@ describe('QueryEditorExtensions', () => {
 
   it('passes dependencies correctly to QueryEditorExtension', async () => {
     const configMap = {
-      '1': { id: '1', order: 1, isEnabled: jest.fn(), getComponent: jest.fn() },
+      '1': { id: '1', order: 1, isEnabled$: jest.fn(), getComponent: jest.fn() },
     };
 
     const { getByText } = render(<QueryEditorExtensions {...defaultProps} configMap={configMap} />);
 
     await waitFor(() => {
-      expect(getByText(/logstash-\*/)).toBeInTheDocument();
+      expect(getByText(/Test-lang/)).toBeInTheDocument();
     });
 
     expect(QueryEditorExtension).toHaveBeenCalledWith(
       expect.objectContaining({
-        dependencies: { indexPatterns: defaultProps.indexPatterns, language: 'Test' },
+        dependencies: {
+          language: 'Test-lang',
+          onSelectLanguage: expect.any(Function),
+          isCollapsed: false,
+          setIsCollapsed: expect.any(Function),
+        },
       }),
       expect.anything()
     );

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

@@ -48,7 +48,6 @@ export interface QueryEditorTopRowProps {
   disableAutoFocus?: boolean;
   screenTitle?: string;
   indexPatterns?: Array<IIndexPattern | string>;
-  dataSource?: DataSource;
   isLoading?: boolean;
   prepend?: React.ComponentProps<typeof EuiCompressedFieldText>['prepend'];
   showQueryEditor?: boolean;
@@ -220,7 +219,6 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
         <QueryEditor
           disableAutoFocus={props.disableAutoFocus}
           indexPatterns={props.indexPatterns!}
-          dataSource={props.dataSource}
           prepend={props.prepend}
           query={parsedQuery}
           dataSetContainerRef={props.dataSetContainerRef}

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

@@ -210,7 +210,6 @@ 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

@@ -60,7 +60,6 @@ interface SearchBarInjectedDeps {
 
 export interface SearchBarOwnProps {
   indexPatterns?: IIndexPattern[];
-  dataSource?: DataSource;
   isLoading?: boolean;
   customSubmitButton?: React.ReactNode;
   screenTitle?: string;
@@ -498,7 +497,6 @@ 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/discover/public/application/view_components/canvas/top_nav.tsx

@@ -126,10 +126,6 @@ export const TopNav = ({ opts, showSaveQuery, isEnhancementsEnabled }: TopNavPro
         useDefaultBehaviors
         setMenuMountPoint={opts.setHeaderActionMenu}
         indexPatterns={indexPattern ? [indexPattern] : indexPatterns}
-        // TODO after
-        // https://github.com/opensearch-project/OpenSearch-Dashboards/pull/6833
-        // is ported to main, pass dataSource to TopNavMenu by picking
-        // commit 328e08e688c again.
         onQuerySubmit={opts.onQuerySubmit}
         savedQueryId={state.savedQuery}
         onSavedQueryIdChange={updateSavedQueryId}

src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.test.tsx

@@ -24,6 +24,12 @@ const renderQueryAssistBanner = (overrideProps: Partial<QueryAssistBannerProps>
   >(
     {
       languages: ['test-lang1', 'test-lang2'],
+      dependencies: {
+        language: 'default',
+        onSelectLanguage: jest.fn(),
+        isCollapsed: true,
+        setIsCollapsed: jest.fn(),
+      },
     },
     overrideProps
   );
@@ -47,4 +53,12 @@ describe('<QueryAssistBanner /> spec', () => {
       component.queryByText('Natural Language Query Generation for test-lang1, test-lang2')
     ).toBeNull();
   });
+
+  it('should change language', async () => {
+    const { props, component } = renderQueryAssistBanner();
+
+    fireEvent.click(component.getByTestId('queryAssist-banner-changeLanguage'));
+    expect(props.dependencies.onSelectLanguage).toBeCalledWith('test-lang1');
+    expect(props.dependencies.setIsCollapsed).toBeCalledWith(false);
+  });
 });

src/plugins/query_enhancements/public/query_assist/components/query_assist_banner.tsx

@@ -14,12 +14,14 @@ import {
 } from '@elastic/eui';
 import { FormattedMessage } from '@osd/i18n/react';
 import React, { useState } from 'react';
+import { QueryEditorExtensionDependencies } from '../../../../data/public';
 import assistantMark from '../../assets/query_assist_mark.svg';
 import { getStorage } from '../../services';
 
 const BANNER_STORAGE_KEY = 'queryAssist:banner:show';
 
 interface QueryAssistBannerProps {
+  dependencies: QueryEditorExtensionDependencies;
   languages: string[];
 }
 
@@ -33,7 +35,8 @@ export const QueryAssistBanner: React.FC<QueryAssistBannerProps> = (props) => {
     _setShowCallOut(show);
   };
 
-  if (!showCallOut || storage.get(BANNER_STORAGE_KEY) === false) return null;
+  if (!showCallOut || storage.get(BANNER_STORAGE_KEY) === false || props.languages.length === 0)
+    return null;
 
   return (
     <EuiCallOut
@@ -54,7 +57,13 @@ export const QueryAssistBanner: React.FC<QueryAssistBannerProps> = (props) => {
                 id="queryAssist.banner.title.prefix"
                 defaultMessage="Use natural language to explore your data with "
               />
-              <EuiLink>
+              <EuiLink
+                data-test-subj="queryAssist-banner-changeLanguage"
+                onClick={() => {
+                  props.dependencies.onSelectLanguage(props.languages[0]);
+                  if (props.dependencies.isCollapsed) props.dependencies.setIsCollapsed(false);
+                }}
+              >
                 <FormattedMessage
                   id="queryAssist.banner.title.suffix"
                   defaultMessage="Natural Language Query Generation for {languages}"

src/plugins/query_enhancements/public/query_assist/components/query_assist_bar.tsx

@@ -82,6 +82,8 @@ export const QueryAssistBar: React.FC<QueryAssistInputProps> = (props) => {
     }
   };
 
+  if (props.dependencies.isCollapsed) return null;
+
   return (
     <EuiForm component="form" onSubmit={onSubmit}>
       <EuiFormRow fullWidth>

src/plugins/query_enhancements/public/query_assist/utils/create_extension.test.tsx

@@ -9,7 +9,7 @@ import React from 'react';
 import { of } from 'rxjs';
 import { coreMock } from '../../../../../core/public/mocks';
 import { SimpleDataSet } from '../../../../data/common';
-import { IIndexPattern } from '../../../../data/public';
+import { QueryEditorExtensionDependencies } from '../../../../data/public';
 import { dataPluginMock } from '../../../../data/public/mocks';
 import { DataSetContract } from '../../../../data/public/query';
 import { ConfigSchema } from '../../../common/config';
@@ -42,6 +42,12 @@ jest.mock('../components', () => ({
 }));
 
 describe('CreateExtension', () => {
+  const dependencies: QueryEditorExtensionDependencies = {
+    language: 'PPL',
+    onSelectLanguage: jest.fn(),
+    isCollapsed: false,
+    setIsCollapsed: jest.fn(),
+  };
   afterEach(() => {
     jest.clearAllMocks();
   });
@@ -53,7 +59,7 @@ describe('CreateExtension', () => {
   it('should be enabled if at least one language is configured', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
     const extension = createQueryAssistExtension(httpMock, dataMock, config);
-    const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' }));
+    const isEnabled = await firstValueFrom(extension.isEnabled$(dependencies));
     expect(isEnabled).toBeTruthy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
       query: { dataSourceId: 'mock-data-source-id' },
@@ -63,7 +69,7 @@ describe('CreateExtension', () => {
   it('should be disabled for unsupported language', async () => {
     httpMock.get.mockRejectedValueOnce(new Error('network failure'));
     const extension = createQueryAssistExtension(httpMock, dataMock, config);
-    const isEnabled = await firstValueFrom(extension.isEnabled$({ language: 'PPL' }));
+    const isEnabled = await firstValueFrom(extension.isEnabled$(dependencies));
     expect(isEnabled).toBeFalsy();
     expect(httpMock.get).toBeCalledWith('/api/enhancements/assist/languages', {
       query: { dataSourceId: 'mock-data-source-id' },
@@ -73,10 +79,7 @@ describe('CreateExtension', () => {
   it('should render the component if language is supported', async () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
     const extension = createQueryAssistExtension(httpMock, dataMock, config);
-    const component = extension.getComponent?.({
-      language: 'PPL',
-      indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],
-    });
+    const component = extension.getComponent?.(dependencies);
 
     if (!component) throw new Error('QueryEditorExtensions Component is undefined');
 
@@ -91,8 +94,8 @@ describe('CreateExtension', () => {
     httpMock.get.mockResolvedValueOnce({ configuredLanguages: ['PPL'] });
     const extension = createQueryAssistExtension(httpMock, dataMock, config);
     const banner = extension.getBanner?.({
+      ...dependencies,
       language: 'DQL',
-      indexPatterns: [{ id: 'test-pattern' }] as IIndexPattern[],
     });
 
     if (!banner) throw new Error('QueryEditorExtensions Banner is undefined');