[Discover-next] Add query editor extensions

changelogs/fragments/7034.yml

@@ -0,0 +1,2 @@
+feat:
+- Add search bar extensions ([#7034](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7034))

src/plugins/data/public/index.ts

@@ -526,3 +526,5 @@ export {
   DataSourceGroup,
   DataSourceOption,
 } from './data_sources/datasource_selector';
+
+export { PersistedLog } from './query';

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

@@ -43,6 +43,8 @@ export interface QueryEditorProps {
   isInvalid?: boolean;
   queryEditorHeaderRef: React.RefObject<HTMLDivElement>;
   queryEditorHeaderClassName?: string;
+  queryEditorBannerRef: React.RefObject<HTMLDivElement>;
+  queryEditorBannerClassName?: string;
 }
 
 interface Props extends QueryEditorProps {
@@ -253,8 +255,14 @@ export default class QueryEditorUI extends Component<Props, State> {
       this.props.queryEditorHeaderClassName
     );
 
+    const queryEditorBannerClassName = classNames(
+      'osdQueryEditorBanner',
+      this.props.queryEditorBannerClassName
+    );
+
     return (
       <div className={className}>
+        <div ref={this.props.queryEditorBannerRef} className={queryEditorBannerClassName} />
         <EuiFlexGroup gutterSize="xs" direction="column">
           <EuiFlexItem grow={false}>
             <EuiFlexGroup gutterSize="xs">

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

@@ -2,38 +2,37 @@
  * Copyright OpenSearch Contributors
  * SPDX-License-Identifier: Apache-2.0
  */
-
 import dateMath from '@elastic/datemath';
-import classNames from 'classnames';
-import React, { useRef, useState } from 'react';
-
 import {
+  EuiFieldText,
   EuiFlexGroup,
   EuiFlexItem,
   EuiSuperDatePicker,
-  EuiFieldText,
+  EuiSuperUpdateButton,
+  OnRefreshProps,
   prettyDuration,
 } from '@elastic/eui';
-// @ts-ignore
-import { EuiSuperUpdateButton, OnRefreshProps } from '@elastic/eui';
-import { isEqual, compact } from 'lodash';
+import classNames from 'classnames';
+import { compact, isEqual } from 'lodash';
+import React, { useRef, useState } from 'react';
 import {
+  DataSource,
   IDataPluginServices,
   IIndexPattern,
-  TimeRange,
-  TimeHistoryContract,
   Query,
-  DataSource,
+  TimeHistoryContract,
+  TimeRange,
 } from '../..';
 import {
   useOpenSearchDashboards,
   withOpenSearchDashboards,
 } from '../../../../opensearch_dashboards_react/public';
-import QueryEditorUI from './query_editor';
 import { UI_SETTINGS } from '../../../common';
-import { PersistedLog, fromUser, getQueryLog } from '../../query';
-import { NoDataPopover } from './no_data_popover';
+import { fromUser, getQueryLog, PersistedLog } from '../../query';
+import { SearchBarExtensions } from '../search_bar_extensions';
 import { Settings } from '../types';
+import { NoDataPopover } from './no_data_popover';
+import QueryEditorUI from './query_editor';
 
 const QueryEditor = withOpenSearchDashboards(QueryEditorUI);
 
@@ -73,6 +72,7 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
   const [isDateRangeInvalid, setIsDateRangeInvalid] = useState(false);
   const [isQueryEditorFocused, setIsQueryEditorFocused] = useState(false);
   const queryEditorHeaderRef = useRef<HTMLDivElement | null>(null);
+  const queryEditorBannerRef = useRef<HTMLDivElement | null>(null);
 
   const opensearchDashboards = useOpenSearchDashboards<IDataPluginServices>();
   const { uiSettings, storage, appName } = opensearchDashboards.services;
@@ -85,6 +85,7 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
       props.settings &&
       props.settings.getQueryEnhancements(queryLanguage)?.searchBar) ||
     null;
+  const searchBarExtensions = props.settings?.getSearchBarExtensions();
   const parsedQuery =
     !queryUiEnhancement || isValidQuery(props.query)
       ? props.query!
@@ -246,11 +247,32 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
           persistedLog={persistedLog}
           dataTestSubj={props.dataTestSubj}
           queryEditorHeaderRef={queryEditorHeaderRef}
+          queryEditorBannerRef={queryEditorBannerRef}
         />
       </EuiFlexItem>
     );
   }
 
+  function renderSearchBarExtensions() {
+    if (
+      !shouldRenderSearchBarExtensions() ||
+      !queryEditorHeaderRef.current ||
+      !queryEditorBannerRef.current ||
+      !queryLanguage
+    )
+      return;
+    return (
+      <SearchBarExtensions
+        language={queryLanguage}
+        configs={searchBarExtensions}
+        componentContainer={queryEditorHeaderRef.current}
+        bannerContainer={queryEditorBannerRef.current}
+        indexPatterns={props.indexPatterns}
+        dataSource={props.dataSource}
+      />
+    );
+  }
+
   function renderSharingMetaFields() {
     const { from, to } = getDateRange();
     const dateRangePretty = prettyDuration(
@@ -282,6 +304,10 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
     );
   }
 
+  function shouldRenderSearchBarExtensions(): boolean {
+    return Boolean(searchBarExtensions?.length);
+  }
+
   function renderUpdateButton() {
     const button = props.customSubmitButton ? (
       React.cloneElement(props.customSubmitButton, { onClick: onClickSubmitButton })
@@ -374,6 +400,7 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
       direction="column"
       justifyContent="flexEnd"
     >
+      {renderSearchBarExtensions()}
       {renderQueryEditor()}
       <EuiFlexItem>
         <EuiFlexGroup responsive={false} gutterSize="none">

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

@@ -28,27 +28,25 @@
  * under the License.
  */
 
-import { compact } from 'lodash';
 import { InjectedIntl, injectI18n } from '@osd/i18n/react';
 import classNames from 'classnames';
+import { compact, get, isEqual } from 'lodash';
 import React, { Component } from 'react';
 import ResizeObserver from 'resize-observer-polyfill';
-import { get, isEqual } from 'lodash';
-
+import { DataSource } from '../..';
 import {
-  withOpenSearchDashboards,
   OpenSearchDashboardsReactContextValue,
+  withOpenSearchDashboards,
 } from '../../../../opensearch_dashboards_react/public';
-
-import QueryBarTopRow from '../query_string_input/query_bar_top_row';
-import { SavedQueryAttributes, TimeHistoryContract, SavedQuery } from '../../query';
+import { Filter, IIndexPattern, Query, TimeRange, UI_SETTINGS } from '../../../common';
+import { SavedQuery, SavedQueryAttributes, TimeHistoryContract } from '../../query';
 import { IDataPluginServices } from '../../types';
-import { TimeRange, Query, Filter, IIndexPattern, UI_SETTINGS } from '../../../common';
 import { FilterBar } from '../filter_bar/filter_bar';
+import { QueryEditorTopRow } from '../query_editor';
+import QueryBarTopRow from '../query_string_input/query_bar_top_row';
 import { SavedQueryMeta, SaveQueryForm } from '../saved_query_form';
 import { SavedQueryManagementComponent } from '../saved_query_management';
 import { Settings } from '../types';
-import { QueryEditorTopRow } from '../query_editor';
 
 interface SearchBarInjectedDeps {
   opensearchDashboards: OpenSearchDashboardsReactContextValue<IDataPluginServices>;
@@ -62,6 +60,7 @@ interface SearchBarInjectedDeps {
 
 export interface SearchBarOwnProps {
   indexPatterns?: IIndexPattern[];
+  dataSource?: DataSource;
   isLoading?: boolean;
   customSubmitButton?: React.ReactNode;
   screenTitle?: string;
@@ -497,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/index.tsx

@@ -0,0 +1,17 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React, { ComponentProps } from 'react';
+
+const Fallback = () => <div />;
+
+const LazySearchBarExtensions = React.lazy(() => import('./search_bar_extensions'));
+export const SearchBarExtensions = (props: ComponentProps<typeof LazySearchBarExtensions>) => (
+  <React.Suspense fallback={<Fallback />}>
+    <LazySearchBarExtensions {...props} />
+  </React.Suspense>
+);
+
+export { SearchBarExtensionConfig } from './search_bar_extension';

src/plugins/data/public/ui/search_bar_extensions/search_bar_extension.test.tsx

@@ -0,0 +1,86 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { render, waitFor } from '@testing-library/react';
+import React, { ComponentProps } from 'react';
+import { IIndexPattern } from '../../../common';
+import { SearchBarExtension } from './search_bar_extension';
+
+jest.mock('react-dom', () => ({
+  ...jest.requireActual('react-dom'),
+  createPortal: jest.fn((element) => element),
+}));
+
+type SearchBarExtensionProps = ComponentProps<typeof SearchBarExtension>;
+
+const mockIndexPattern = {
+  id: '1234',
+  title: 'logstash-*',
+  fields: [
+    {
+      name: 'response',
+      type: 'number',
+      esTypes: ['integer'],
+      aggregatable: true,
+      filterable: true,
+      searchable: true,
+    },
+  ],
+} as IIndexPattern;
+
+describe('SearchBarExtension', () => {
+  const getComponentMock = jest.fn();
+  const getBannerMock = jest.fn();
+  const isEnabledMock = jest.fn();
+
+  const defaultProps: SearchBarExtensionProps = {
+    config: {
+      id: 'test-extension',
+      order: 1,
+      isEnabled: isEnabledMock,
+      getComponent: getComponentMock,
+      getBanner: getBannerMock,
+    },
+    dependencies: {
+      indexPatterns: [mockIndexPattern],
+      language: 'Test',
+    },
+    componentContainer: document.createElement('div'),
+    bannerContainer: document.createElement('div'),
+  };
+
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it('renders correctly when isEnabled is true', async () => {
+    isEnabledMock.mockResolvedValue(true);
+    getComponentMock.mockReturnValue(<div>Test Component</div>);
+    getBannerMock.mockReturnValue(<div>Test Banner</div>);
+
+    const { getByText } = render(<SearchBarExtension {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(getByText('Test Component')).toBeInTheDocument();
+      expect(getByText('Test Banner')).toBeInTheDocument();
+    });
+
+    expect(isEnabledMock).toHaveBeenCalled();
+    expect(getComponentMock).toHaveBeenCalledWith(defaultProps.dependencies);
+  });
+
+  it('does not render when isEnabled is false', async () => {
+    isEnabledMock.mockResolvedValue(false);
+    getComponentMock.mockReturnValue(<div>Test Component</div>);
+
+    const { queryByText } = render(<SearchBarExtension {...defaultProps} />);
+
+    await waitFor(() => {
+      expect(queryByText('Test Component')).toBeNull();
+    });
+
+    expect(isEnabledMock).toHaveBeenCalled();
+  });
+});