feat: added a new API to allow adding additional editor controls

Signed-off-by: Yulong Ruan <[email protected]>

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

@@ -16,6 +16,7 @@ import {
 } from '@elastic/eui';
 import classNames from 'classnames';
 import { isEqual } from 'lodash';
+import { BehaviorSubject, Subscription } from 'rxjs';
 import React, { Component, createRef, RefObject } from 'react';
 import { monaco } from '@osd/monaco';
 import {
@@ -37,6 +38,7 @@ import { QueryControls } from '../../query/query_string/language_service/get_que
 import { RecentQueriesTable } from '../../query/query_string/language_service/recent_query';
 import { DefaultInputProps } from './editors';
 import { MonacoCompatibleQuerySuggestion } from '../../autocomplete/providers/query_suggestion_provider';
+import { SearchBarControl } from '../types';
 
 export interface QueryEditorProps {
   query: Query;
@@ -61,6 +63,7 @@ export interface QueryEditorProps {
   filterBar?: any;
   prepend?: React.ComponentProps<typeof EuiCompressedFieldText>['prepend'];
   savedQueryManagement?: any;
+  additionalControls$?: BehaviorSubject<SearchBarControl[]>;
 }
 
 interface Props extends QueryEditorProps {
@@ -76,6 +79,7 @@ interface State {
   timeStamp: IFieldType | null;
   lineCount: number | undefined;
   isRecentQueryVisible: boolean;
+  additionalControls?: SearchBarControl[];
 }
 
 // Needed for React.lazy
@@ -91,6 +95,7 @@ export default class QueryEditorUI extends Component<Props, State> {
     timeStamp: null,
     lineCount: undefined,
     isRecentQueryVisible: false,
+    additionalControls: [],
   };
 
   public inputRef: monaco.editor.IStandaloneCodeEditor | null = null;
@@ -104,6 +109,7 @@ export default class QueryEditorUI extends Component<Props, State> {
   private headerRef: RefObject<HTMLDivElement> = createRef();
   private bannerRef: RefObject<HTMLDivElement> = createRef();
   private extensionMap = this.languageManager.getQueryEditorExtensionMap();
+  private additionalControlsSubscription?: Subscription;
 
   private getQueryString = () => {
     return toUser(this.props.query.query);
@@ -228,6 +234,12 @@ export default class QueryEditorUI extends Component<Props, State> {
     }
 
     this.initPersistedLog();
+
+    if (this.props.additionalControls$) {
+      this.additionalControlsSubscription = this.props.additionalControls$.subscribe((controls) => {
+        this.setState({ additionalControls: controls });
+      });
+    }
   }
 
   public componentDidUpdate(prevProps: Props) {
@@ -254,6 +266,7 @@ export default class QueryEditorUI extends Component<Props, State> {
 
   public componentWillUnmount() {
     if (this.abortController) this.abortController.abort();
+    if (this.additionalControlsSubscription) this.additionalControlsSubscription.unsubscribe();
   }
 
   handleOnFocus = () => {
@@ -330,6 +343,18 @@ export default class QueryEditorUI extends Component<Props, State> {
     return <QueryControls queryControls={queryControls} />;
   };
 
+  private renderAdditionalControls = () => {
+    if (!this.state.additionalControls) return null;
+    const sorted = this.state.additionalControls.sort((c1, c2) => c1.order - c2.order);
+    return (
+      <>
+        {sorted.map((control, i) => {
+          return <React.Fragment key={i}>{control.render()}</React.Fragment>;
+        })}
+      </>
+    );
+  };
+
   public render() {
     const className = classNames(this.props.className);
 
@@ -446,6 +471,7 @@ export default class QueryEditorUI extends Component<Props, State> {
           <div className="osdQueryEditor__querycontrols">
             <EuiFlexGroup responsive={false} gutterSize="s" alignItems="center">
               {this.renderQueryControls(languageEditor.TopBar.Controls)}
+              {this.renderAdditionalControls()}
               {!languageEditor.TopBar.Expanded && this.renderToggleIcon()}
               {this.props.savedQueryManagement}
             </EuiFlexGroup>

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

@@ -15,6 +15,7 @@ import {
 import classNames from 'classnames';
 import React, { useState } from 'react';
 import { createPortal } from 'react-dom';
+import { BehaviorSubject } from 'rxjs';
 import { IDataPluginServices, IIndexPattern, Query, TimeHistoryContract, TimeRange } from '../..';
 import {
   useOpenSearchDashboards,
@@ -24,6 +25,7 @@ import { UI_SETTINGS } from '../../../common';
 import { getQueryLog, PersistedLog } from '../../query';
 import { NoDataPopover } from './no_data_popover';
 import QueryEditorUI from './query_editor';
+import { SearchBarControl } from '../types';
 
 const QueryEditor = withOpenSearchDashboards(QueryEditorUI);
 
@@ -54,6 +56,7 @@ export interface QueryEditorTopRowProps {
   indicateNoData?: boolean;
   datePickerRef?: React.RefObject<HTMLDivElement>;
   savedQueryManagement?: any;
+  additionalControls$?: BehaviorSubject<SearchBarControl[]>;
 }
 
 // Needed for React.lazy
@@ -186,6 +189,7 @@ export default function QueryEditorTopRow(props: QueryEditorTopRowProps) {
           dataTestSubj={props.dataTestSubj}
           filterBar={props.filterBar}
           savedQueryManagement={props.savedQueryManagement}
+          additionalControls$={props.additionalControls$}
         />
       </EuiFlexItem>
     );

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

@@ -30,6 +30,7 @@
 
 import _ from 'lodash';
 import React, { useEffect, useRef } from 'react';
+import { BehaviorSubject } from 'rxjs';
 import { CoreStart } from 'src/core/public';
 import { OpenSearchDashboardsContextProvider } from '../../../../opensearch_dashboards_react/public';
 import { QueryStart, SavedQuery } from '../../query';
@@ -40,11 +41,13 @@ import { useSavedQuery } from './lib/use_saved_query';
 import { DataPublicPluginStart } from '../../types';
 import { DataStorage, Filter, Query, TimeRange } from '../../../common';
 import { useQueryStringManager } from './lib/use_query_string_manager';
+import { SearchBarControl } from '../types';
 
 interface StatefulSearchBarDeps {
   core: CoreStart;
   data: Omit<DataPublicPluginStart, 'ui'>;
   storage: DataStorage;
+  searchBarControls$: BehaviorSubject<SearchBarControl[]>;
 }
 
 export type StatefulSearchBarProps = SearchBarOwnProps & {
@@ -129,7 +132,12 @@ const overrideDefaultBehaviors = (props: StatefulSearchBarProps) => {
   return props.useDefaultBehaviors ? {} : props;
 };
 
-export function createSearchBar({ core, storage, data }: StatefulSearchBarDeps) {
+export function createSearchBar({
+  core,
+  storage,
+  data,
+  searchBarControls$,
+}: StatefulSearchBarDeps) {
   // App name should come from the core application service.
   // Until it's available, we'll ask the user to provide it for the pre-wired component.
   return (props: StatefulSearchBarProps) => {
@@ -210,6 +218,7 @@ export function createSearchBar({ core, storage, data }: StatefulSearchBarDeps)
           onSaved={defaultOnSavedQueryUpdated(props, setSavedQuery)}
           datePickerRef={props.datePickerRef}
           isFilterBarPortable={props.isFilterBarPortable}
+          searchBarControls$={searchBarControls$}
           {...overrideDefaultBehaviors(props)}
         />
       </OpenSearchDashboardsContextProvider>

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

@@ -32,6 +32,7 @@ import { InjectedIntl, injectI18n } from '@osd/i18n/react';
 import classNames from 'classnames';
 import { compact, get, isEqual } from 'lodash';
 import React, { Component } from 'react';
+import { BehaviorSubject } from 'rxjs';
 import ResizeObserver from 'resize-observer-polyfill';
 import {
   OpenSearchDashboardsReactContextValue,
@@ -45,6 +46,7 @@ import { QueryEditorTopRow } from '../query_editor';
 import QueryBarTopRow from '../query_string_input/query_bar_top_row';
 import { SavedQueryMeta, SaveQueryForm } from '../saved_query_form';
 import { FilterOptions } from '../filter_bar/filter_options';
+import { SearchBarControl } from '../types';
 
 interface SearchBarInjectedDeps {
   opensearchDashboards: OpenSearchDashboardsReactContextValue<IDataPluginServices>;
@@ -91,6 +93,7 @@ export interface SearchBarOwnProps {
   onClearSavedQuery?: () => void;
 
   onRefresh?: (payload: { dateRange: TimeRange }) => void;
+  searchBarControls$?: BehaviorSubject<SearchBarControl[]>;
   indicateNoData?: boolean;
 }
 
@@ -550,6 +553,7 @@ class SearchBarUI extends Component<SearchBarProps, State> {
           indicateNoData={this.props.indicateNoData}
           datePickerRef={this.props.datePickerRef}
           savedQueryManagement={searchBarMenu(false, true)}
+          additionalControls$={this.props.searchBarControls$}
         />
       );
     }

src/plugins/data/public/ui/types.ts

@@ -3,10 +3,16 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import React from 'react';
 import { IndexPatternSelectProps } from './index_pattern_select';
 import { StatefulSearchBarProps } from './search_bar';
 import { SuggestionsComponentProps } from './typeahead/suggestions_component';
 
+export interface SearchBarControl {
+  order: number;
+  render: () => React.ReactNode;
+}
+
 /**
  * The setup contract exposed by the Search plugin exposes the search strategy extension
  * point.
@@ -21,4 +27,5 @@ export interface IUiStart {
   IndexPatternSelect: React.ComponentType<IndexPatternSelectProps>;
   SearchBar: React.ComponentType<StatefulSearchBarProps>;
   SuggestionsComponent: React.ComponentType<SuggestionsComponentProps>;
+  addSearchBarControl: (control: SearchBarControl) => void;
 }

src/plugins/data/public/ui/ui_service.ts

@@ -3,13 +3,14 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+import { BehaviorSubject } from 'rxjs';
 import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from 'src/core/public';
 import { ConfigSchema } from '../../config';
 import { DataPublicPluginStart } from '../types';
 import { createIndexPatternSelect } from './index_pattern_select';
 import { createSearchBar } from './search_bar/create_search_bar';
 import { SuggestionsComponent } from './typeahead';
-import { IUiSetup, IUiStart } from './types';
+import { IUiSetup, IUiStart, SearchBarControl } from './types';
 import { DataStorage } from '../../common';
 
 /** @internal */
@@ -24,13 +25,20 @@ export interface UiServiceStartDependencies {
 
 export class UiService implements Plugin<IUiSetup, IUiStart> {
   enhancementsConfig: ConfigSchema['enhancements'];
+  searchBarControls$: BehaviorSubject<SearchBarControl[]> = new BehaviorSubject(
+    [] as SearchBarControl[]
+  );
 
   constructor(initializerContext: PluginInitializerContext<ConfigSchema>) {
     const { enhancements } = initializerContext.config.get<ConfigSchema>();
 
     this.enhancementsConfig = enhancements;
   }
 
+  addSearchBarControl = (control: SearchBarControl) => {
+    this.searchBarControls$.next([...this.searchBarControls$.value, control]);
+  };
+
   public setup(core: CoreSetup, {}: UiServiceSetupDependencies): IUiSetup {
     return {};
   }
@@ -40,12 +48,14 @@ export class UiService implements Plugin<IUiSetup, IUiStart> {
       core,
       data: dataServices,
       storage,
+      searchBarControls$: this.searchBarControls$,
     });
 
     return {
       IndexPatternSelect: createIndexPatternSelect(core.savedObjects.client),
       SearchBar,
       SuggestionsComponent,
+      addSearchBarControl: this.addSearchBarControl,
     };
   }