Working query enhance

Signed-off-by: Kawika Avilla <kavilla414@gmail.com>

config/opensearch_dashboards.yml

@@ -274,4 +274,6 @@
 # opensearchDashboards.survey.url: "https://survey.opensearch.org"
 
 # Set the value of this setting to true to enable plugin augmentation on Dashboard
-# vis_augmenter.pluginAugmentationEnabled: true
+# vis_augmenter.pluginAugmentationEnabled: true
+opensearch_alerting.enabled: false
+opensearch_security.enabled: false

src/plugins/data/common/opensearch_query/opensearch_query/build_opensearch_query.ts

@@ -73,10 +73,15 @@ export function buildOpenSearchQuery(
     return buildQueryFromSql(sqlQueries, config.dateFormatTZ);
   }
 
-  const validQueries = queries
-    .filter((query) => query.language !== 'SQL')
-    .filter((query) => has(query, 'query'));
+  const validQueries = queries.filter((query) => has(query, 'query'));
   const queriesByLanguage = groupBy(validQueries, 'language');
+  const unsupportedQueries = Object.keys(queriesByLanguage).filter(
+    (language) => language !== 'kuery' && language !== 'lucene'
+  );
+  if (unsupportedQueries.length > 0) {
+    return queries;
+  }
+
   const kueryQuery = buildQueryFromKuery(
     indexPattern,
     queriesByLanguage.kuery,

src/plugins/data/public/plugin.ts

@@ -46,9 +46,9 @@ import {
 } from './types';
 import { AutocompleteService } from './autocomplete';
 import { SearchService } from './search/search_service';
+import { UiService } from './ui/ui_service';
 import { FieldFormatsService } from './field_formats';
 import { QueryService } from './query';
-import { createIndexPatternSelect } from './ui/index_pattern_select';
 import {
   IndexPatternsService,
   onRedirectNoIndexPattern,
@@ -63,9 +63,9 @@ import {
   setOverlays,
   setQueryService,
   setSearchService,
+  setUiService,
   setUiSettings,
 } from './services';
-import { createSearchBar } from './ui/search_bar/create_search_bar';
 import { opensearchaggs } from './search/expressions';
 import {
   SELECT_RANGE_TRIGGER,
@@ -110,12 +110,14 @@ export class DataPublicPlugin
     > {
   private readonly autocomplete: AutocompleteService;
   private readonly searchService: SearchService;
+  private readonly uiService: UiService;
   private readonly fieldFormatsService: FieldFormatsService;
   private readonly queryService: QueryService;
   private readonly storage: IStorageWrapper;
 
   constructor(initializerContext: PluginInitializerContext<ConfigSchema>) {
     this.searchService = new SearchService(initializerContext);
+    this.uiService = new UiService(initializerContext);
     this.queryService = new QueryService();
     this.fieldFormatsService = new FieldFormatsService();
     this.autocomplete = new AutocompleteService(initializerContext);
@@ -159,13 +161,16 @@ export class DataPublicPlugin
       expressions,
     });
 
+    const uiService = this.uiService.setup(core, {});
+
     return {
       autocomplete: this.autocomplete.setup(core),
       search: searchService,
       fieldFormats: this.fieldFormatsService.setup(core),
       query: queryService,
       __enhance: (enhancements: DataPublicPluginEnhancements) => {
-        searchService.__enhance(enhancements.search);
+        if (enhancements.search) searchService.__enhance(enhancements.search);
+        if (enhancements.ui) uiService.__enhance(enhancements.ui);
       },
     };
   }
@@ -234,27 +239,21 @@ export class DataPublicPlugin
         dataSourceFactory,
       },
     };
-
     registerDefaultDatasource(dataServices);
 
-    const SearchBar = createSearchBar({
-      core,
-      data: dataServices,
-      storage: this.storage,
-    });
+    const uiService = this.uiService.start(core, { dataServices, storage: this.storage });
+    setUiService(uiService);
 
     return {
       ...dataServices,
-      ui: {
-        IndexPatternSelect: createIndexPatternSelect(core.savedObjects.client),
-        SearchBar,
-      },
+      ui: uiService,
     };
   }
 
   public stop() {
     this.autocomplete.clearProviders();
     this.queryService.stop();
     this.searchService.stop();
+    this.uiService.stop();
   }
 }

src/plugins/data/public/search/search_service.ts

@@ -74,6 +74,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
   private readonly aggsService = new AggsService();
   private readonly searchSourceService = new SearchSourceService();
   private searchInterceptor!: ISearchInterceptor;
+  private defaultSearchInterceptor!: ISearchInterceptor;
   private usageCollector?: SearchUsageCollector;
 
   constructor(private initializerContext: PluginInitializerContext<ConfigSchema>) {}
@@ -95,6 +96,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
       startServices: getStartServices(),
       usageCollector: this.usageCollector!,
     });
+    this.defaultSearchInterceptor = this.searchInterceptor;
 
     expressions.registerFunction(opensearchdsl);
     expressions.registerType(opensearchRawResponse);
@@ -154,6 +156,10 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
         this.searchInterceptor.showError(e);
       },
       searchSource: this.searchSourceService.start(indexPatterns, searchSourceDependencies),
+      __enhance: (enhancements: SearchEnhancements) => {
+        this.searchInterceptor = enhancements.searchInterceptor;
+      },
+      getDefaultSearchInterceptor: () => this.defaultSearchInterceptor,
     };
   }
 

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

@@ -78,6 +78,8 @@ export interface ISearchStart {
    * {@link ISearchStartSearchSource}
    */
   searchSource: ISearchStartSearchSource;
+  __enhance: (enhancements: SearchEnhancements) => void;
+  getDefaultSearchInterceptor: () => ISearchInterceptor;
 }
 
 export { SEARCH_EVENT_TYPE } from './collectors';

src/plugins/data/public/services.ts

@@ -59,3 +59,5 @@ export const [getQueryService, setQueryService] = createGetterSetter<
 export const [getSearchService, setSearchService] = createGetterSetter<
   DataPublicPluginStart['search']
 >('Search');
+
+export const [getUiService, setUiService] = createGetterSetter<DataPublicPluginStart['ui']>('Ui');

src/plugins/data/public/types.ts

@@ -28,7 +28,6 @@
  * under the License.
  */
 
-import React from 'react';
 import { CoreStart } from 'src/core/public';
 import { IStorageWrapper } from 'src/plugins/opensearch_dashboards_utils/public';
 import { ExpressionsSetup } from 'src/plugins/expressions/public';
@@ -39,12 +38,14 @@ import { createFiltersFromRangeSelectAction, createFiltersFromValueClickAction }
 import { ISearchSetup, ISearchStart, SearchEnhancements } from './search';
 import { QuerySetup, QueryStart } from './query';
 import { IndexPatternsContract } from './index_patterns';
-import { IndexPatternSelectProps, StatefulSearchBarProps } from './ui';
 import { UsageCollectionSetup } from '../../usage_collection/public';
 import { DataSourceStart } from './data_sources/datasource_services/types';
+import { UiEnhancements } from './ui';
+import { DataPublicPluginStartUi } from './ui/types';
 
 export interface DataPublicPluginEnhancements {
-  search: SearchEnhancements;
+  search?: SearchEnhancements;
+  ui?: UiEnhancements;
 }
 
 export interface DataSetupDependencies {
@@ -71,14 +72,6 @@ export interface DataPublicPluginSetup {
   __enhance: (enhancements: DataPublicPluginEnhancements) => void;
 }
 
-/**
- * Data plugin prewired UI components
- */
-export interface DataPublicPluginStartUi {
-  IndexPatternSelect: React.ComponentType<IndexPatternSelectProps>;
-  SearchBar: React.ComponentType<StatefulSearchBarProps>;
-}
-
 /**
  * utilities to generate filters from action context
  */

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

@@ -28,6 +28,7 @@
  * under the License.
  */
 
+export { UiEnhancements } from './types';
 export { IndexPatternSelectProps } from './index_pattern_select';
 export { FilterLabel } from './filter_bar';
 export { QueryStringInput, QueryStringInputProps } from './query_string_input';

src/plugins/data/public/ui/query_string_input/language_switcher.tsx

@@ -31,13 +31,21 @@
 import { EuiComboBox, EuiComboBoxOptionOption, PopoverAnchorPosition } from '@elastic/eui';
 import { i18n } from '@osd/i18n';
 import React from 'react';
+import { getSearchService, getUiService } from '../../services';
 
 interface Props {
   language: string;
   onSelectLanguage: (newLanguage: string) => void;
   anchorPosition?: PopoverAnchorPosition;
 }
 
+function mapExternalLanguageToOptions(language: string) {
+  return {
+    label: language,
+    value: language,
+  };
+}
+
 export function QueryLanguageSwitcher(props: Props) {
   const luceneLabel = i18n.translate('data.query.queryBar.luceneLanguageName', {
     defaultMessage: 'Lucene',
@@ -66,20 +74,30 @@ export function QueryLanguageSwitcher(props: Props) {
     },
   ];
 
+  const queryEnhancements = getUiService().queryEnhancements;
+  queryEnhancements.forEach((enhancement) =>
+    languageOptions.push(mapExternalLanguageToOptions(enhancement.language))
+  );
+
   const selectedLanguage = {
     label: props.language === 'kuery' ? 'DQL' : props.language,
   };
 
   const handleLanguageChange = (newLanguage: EuiComboBoxOptionOption[]) => {
     const queryLanguage = newLanguage[0].label === 'DQL' ? 'kuery' : newLanguage[0].label;
     props.onSelectLanguage(queryLanguage);
+    const queryEnhancement = queryEnhancements.get(queryLanguage);
+    getSearchService().__enhance({
+      searchInterceptor: queryEnhancement
+        ? queryEnhancement.search
+        : getSearchService().getDefaultSearchInterceptor(),
+    });
   };
 
   return (
     <EuiComboBox
       className="languageSwitcher"
       data-test-subj="languageSelect"
-      // TODO: FIX THIS ANY
       options={languageOptions}
       selectedOptions={[selectedLanguage]}
       onChange={handleLanguageChange}

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

@@ -0,0 +1,32 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Any modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+import { SearchInterceptor } from '../search';
+import { IndexPatternSelectProps } from './index_pattern_select';
+import { StatefulSearchBarProps } from './search_bar';
+
+export interface QueryEnhancement {
+  language: string;
+  search: SearchInterceptor;
+}
+
+export interface UiEnhancements {
+  query: QueryEnhancement;
+}
+
+/**
+ * Data plugin prewired UI components
+ */
+export interface DataPublicPluginStartUi {
+  queryEnhancements: Map<string, QueryEnhancement>;
+  IndexPatternSelect: React.ComponentType<IndexPatternSelectProps>;
+  SearchBar: React.ComponentType<StatefulSearchBarProps>;
+}

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

@@ -0,0 +1,74 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ * Any modifications Copyright OpenSearch Contributors. See
+ * GitHub history for details.
+ */
+
+import { Plugin, CoreSetup, CoreStart, PluginInitializerContext } from 'src/core/public';
+import { DataPublicPluginStartUi, QueryEnhancement, UiEnhancements } from './types';
+
+import { ConfigSchema } from '../../config';
+import { createIndexPatternSelect } from './index_pattern_select';
+import { createSearchBar } from './search_bar/create_search_bar';
+import { DataPublicPluginStart } from '../types';
+import { IStorageWrapper } from '../../../opensearch_dashboards_utils/public';
+
+/** @internal */
+// eslint-disable-next-line @typescript-eslint/no-empty-interface
+export interface UiServiceSetupDependencies {}
+
+/** @internal */
+export interface UiServiceStartDependencies {
+  dataServices: Omit<DataPublicPluginStart, 'ui'>;
+  storage: IStorageWrapper;
+}
+
+export class UiService implements Plugin<any, DataPublicPluginStartUi> {
+  constructor(private _: PluginInitializerContext<ConfigSchema>) {}
+  private queryEnhancements: Map<string, QueryEnhancement> = new Map();
+
+  public setup(
+    { http, getStartServices, notifications, uiSettings }: CoreSetup,
+    {}: UiServiceSetupDependencies
+  ) {
+    /**
+     * A global object that intercepts all searches and provides convenience methods for cancelling
+     * all pending search requests, as well as getting the number of pending search requests.
+     */
+    // this.searchInterceptor = new SearchInterceptor({
+    //   toasts: notifications.toasts,
+    //   http,
+    //   uiSettings,
+    //   startServices: getStartServices(),
+    //   usageCollector: this.usageCollector!,
+    // });
+
+    return {
+      __enhance: (enhancements?: UiEnhancements) => {
+        if (!enhancements) return;
+        this.queryEnhancements.set(enhancements.query.language, enhancements.query);
+      },
+    };
+  }
+
+  public start(core: CoreStart, { dataServices, storage }: UiServiceStartDependencies) {
+    const SearchBar = createSearchBar({
+      core,
+      data: dataServices,
+      storage,
+    });
+
+    return {
+      queryEnhancements: this.queryEnhancements,
+      IndexPatternSelect: createIndexPatternSelect(core.savedObjects.client),
+      SearchBar,
+    };
+  }
+
+  public stop() {}
+}

src/plugins/data/server/search/search_service.ts

@@ -164,7 +164,8 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
       });
 
     return {
-      __enhance: (enhancements: SearchEnhancements) => {
+      __enhance: (enhancements?: SearchEnhancements) => {
+        if (!enhancements) return;
         if (this.searchStrategies.hasOwnProperty(enhancements.defaultStrategy)) {
           this.defaultSearchStrategyName = enhancements.defaultStrategy;
         }