Extended Dataset type config to provide search options like strategy (#…

…8113) (#8276) * added function to provide search strategy; updated interceptors to use the strategy * minor refactor * refactored code --------- (cherry picked from commit b438c2e) Signed-off-by: Amardeepsingh Siglani <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Sep 21, 2024 · a513a8c · a513a8c
1 parent cb366fa
commit a513a8c
Show file tree

Hide file tree

Showing 4 changed files with 34 additions and 5 deletions.
diff --git a/src/plugins/data/common/datasets/types.ts b/src/plugins/data/common/datasets/types.ts
@@ -251,3 +251,7 @@ export interface DatasetField {
   displayName?: string;
   // TODO:  osdFieldType?
 }
+
+export interface DatasetSearchOptions {
+  strategy?: string;
+}
diff --git a/src/plugins/data/public/query/query_string/dataset_service/types.ts b/src/plugins/data/public/query/query_string/dataset_service/types.ts
@@ -3,7 +3,7 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 import { EuiIconProps } from '@elastic/eui';
-import { Dataset, DatasetField, DataStructure } from '../../../../common';
+import { Dataset, DatasetField, DatasetSearchOptions, DataStructure } from '../../../../common';
 import { IDataPluginServices } from '../../../types';
 
 /**
@@ -44,4 +44,9 @@ export interface DatasetTypeConfig {
    * @returns {Promise<string[]>} A promise that resolves to an array of supported language ids.
    */
   supportedLanguages: (dataset: Dataset) => string[];
+  /**
+   * Retrieves the search options to be used for running the query on the data connection associated
+   * with this Dataset
+   */
+  getSearchOptions?: () => DatasetSearchOptions;
 }
diff --git a/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts b/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts
@@ -45,7 +45,7 @@ export class PPLSearchInterceptor extends SearchInterceptor {
     const { id, ...searchRequest } = request;
     const context: EnhancedFetchContext = {
       http: this.deps.http,
-      path: trimEnd(API.PPL_SEARCH),
+      path: trimEnd(`${API.SEARCH}/${strategy}`),
       signal,
     };
 
@@ -55,7 +55,18 @@ export class PPLSearchInterceptor extends SearchInterceptor {
   }
 
   public search(request: IOpenSearchDashboardsSearchRequest, options: ISearchOptions) {
-    return this.runSearch(request, options.abortSignal, SEARCH_STRATEGY.PPL);
+    const dataset = this.queryService.queryString.getQuery().dataset;
+    const datasetType = dataset?.type;
+    let strategy = SEARCH_STRATEGY.PPL;
+
+    if (datasetType) {
+      const datasetTypeConfig = this.queryService.queryString
+        .getDatasetService()
+        .getType(datasetType);
+      strategy = datasetTypeConfig?.getSearchOptions?.().strategy ?? strategy;
+    }
+
+    return this.runSearch(request, options.abortSignal, strategy);
   }
 
   private buildQuery() {

diff --git a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts
@@ -39,7 +39,7 @@ export class SQLSearchInterceptor extends SearchInterceptor {
     const isAsync = strategy === SEARCH_STRATEGY.SQL_ASYNC;
     const context: EnhancedFetchContext = {
       http: this.deps.http,
-      path: trimEnd(isAsync ? API.SQL_ASYNC_SEARCH : API.SQL_SEARCH),
+      path: trimEnd(`${API.SEARCH}/${strategy}`),
       signal,
     };
 
@@ -54,7 +54,16 @@ export class SQLSearchInterceptor extends SearchInterceptor {
 
   public search(request: IOpenSearchDashboardsSearchRequest, options: ISearchOptions) {
     const dataset = this.queryService.queryString.getQuery().dataset;
-    const strategy = dataset?.type === DATASET.S3 ? SEARCH_STRATEGY.SQL_ASYNC : SEARCH_STRATEGY.SQL;
+    const datasetType = dataset?.type;
+    let strategy = datasetType === DATASET.S3 ? SEARCH_STRATEGY.SQL_ASYNC : SEARCH_STRATEGY.SQL;
+
+    if (datasetType) {
+      const datasetTypeConfig = this.queryService.queryString
+        .getDatasetService()
+        .getType(datasetType);
+      strategy = datasetTypeConfig?.getSearchOptions?.().strategy ?? strategy;
+    }
+
     return this.runSearch(request, options.abortSignal, strategy);
   }
 }