Added framework to get default query string using dataset and languag…

…e combination (opensearch-project#8896)

* added framework to get default query using dataset

Signed-off-by: Amardeepsingh Siglani <[email protected]>

* Changeset file for PR opensearch-project#8896 created/updated

* deduped code

Signed-off-by: Amardeepsingh Siglani <[email protected]>

* added UTs; minor refactor

Signed-off-by: Amardeepsingh Siglani <[email protected]>

---------

Signed-off-by: Amardeepsingh Siglani <[email protected]>
Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com>

changelogs/fragments/8896.yml

@@ -0,0 +1,2 @@
+feat:
+- Added framework to get default query string using dataset and language combination ([#8896](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8896))

src/plugins/data/public/query/query_string/dataset_service/types.ts

@@ -8,6 +8,7 @@ import {
   DatasetField,
   DatasetSearchOptions,
   DataStructure,
+  Query,
   SavedObject,
 } from '../../../../common';
 import { IDataPluginServices } from '../../../types';
@@ -104,4 +105,8 @@ export interface DatasetTypeConfig {
    * Service used for indexedViews related operations
    */
   indexedViewsService?: DatasetIndexedViewsService;
+  /**
+   * Returns the initial query that is added to the query editor when a dataset is selected.
+   */
+  getInitialQueryString?: (query: Query) => string | void;
 }

src/plugins/data/public/query/query_string/query_string_manager.test.ts

@@ -308,5 +308,37 @@ describe('QueryStringManager', () => {
       expect(result.dataset).toEqual(currentDataset);
       expect(result.query).toBeDefined();
     });
+
+    test('getInitialQueryByLanguage returns the initial query from the dataset config if present', () => {
+      service.getDatasetService().getType = jest.fn().mockReturnValue({
+        supportedLanguages: jest.fn(),
+        getInitialQueryString: jest.fn().mockImplementation(({ language }) => {
+          switch (language) {
+            case 'sql':
+              return 'default sql dataset query';
+            case 'ppl':
+              return 'default ppl dataset query';
+          }
+        }),
+      });
+
+      const sqlQuery = service.getInitialQueryByLanguage('sql');
+      expect(sqlQuery).toHaveProperty('query', 'default sql dataset query');
+
+      const pplQuery = service.getInitialQueryByLanguage('ppl');
+      expect(pplQuery).toHaveProperty('query', 'default ppl dataset query');
+    });
+
+    test('getInitialQueryByLanguage returns the initial query from the language config if dataset does not provide one', () => {
+      service.getDatasetService().getType = jest.fn().mockReturnValue({
+        supportedLanguages: jest.fn(),
+      });
+      service.getLanguageService().getLanguage = jest.fn().mockReturnValue({
+        getQueryString: jest.fn().mockReturnValue('default-language-service-query'),
+      });
+
+      const sqlQuery = service.getInitialQueryByLanguage('sql');
+      expect(sqlQuery).toHaveProperty('query', 'default-language-service-query');
+    });
   });
 });

src/plugins/data/public/query/query_string/query_string_manager.ts

@@ -63,6 +63,21 @@ export class QueryStringManager {
     return this.storage.get('userQueryString') || '';
   }
 
+  private getInitialDatasetQueryString(query: Query) {
+    const { language, dataset } = query;
+
+    const languageConfig = this.languageService.getLanguage(language);
+    let typeConfig;
+
+    if (dataset) {
+      typeConfig = this.datasetService.getType(dataset.type);
+    }
+
+    return (
+      typeConfig?.getInitialQueryString?.(query) ?? (languageConfig?.getQueryString(query) || '')
+    );
+  }
+
   public getDefaultQuery(): Query {
     const defaultLanguageId = this.getDefaultLanguage();
     const defaultQuery = this.getDefaultQueryString();
@@ -79,13 +94,11 @@ export class QueryStringManager {
       defaultDataset &&
       this.languageService
     ) {
-      const language = this.languageService.getLanguage(defaultLanguageId);
       const newQuery = { ...query, dataset: defaultDataset };
-      const newQueryString = language?.getQueryString(newQuery) || '';
 
       return {
         ...newQuery,
-        query: newQueryString,
+        query: this.getInitialDatasetQueryString(newQuery),
       };
     }
 
@@ -244,13 +257,12 @@ export class QueryStringManager {
 
     // Both language and dataset provided - generate fresh query
     if (language && dataset) {
-      const languageService = this.languageService.getLanguage(language);
       const newQuery = {
         language,
         dataset,
         query: '',
       };
-      newQuery.query = languageService?.getQueryString(newQuery) || '';
+      newQuery.query = this.getInitialDatasetQueryString(newQuery);
       return newQuery;
     }
 
@@ -274,12 +286,12 @@ export class QueryStringManager {
    */
   public getInitialQueryByLanguage = (languageId: string) => {
     const curQuery = this.query$.getValue();
-    const language = this.languageService.getLanguage(languageId);
     const newQuery = {
       ...curQuery,
       language: languageId,
     };
-    const queryString = language?.getQueryString(newQuery) || '';
+
+    const queryString = this.getInitialDatasetQueryString(newQuery);
     this.languageService.setUserQueryString(queryString);
 
     return {
@@ -296,17 +308,15 @@ export class QueryStringManager {
     const curQuery = this.query$.getValue();
     // Use dataset's preferred language or fallback to current language
     const languageId = newDataset.language || curQuery.language;
-    const language = this.languageService.getLanguage(languageId);
     const newQuery = {
       ...curQuery,
       language: languageId,
       dataset: newDataset,
     };
-    const queryString = language?.getQueryString(newQuery) || '';
 
     return {
       ...newQuery,
-      query: queryString,
+      query: this.getInitialDatasetQueryString(newQuery, newDataset),
     };
   };