ppl query

Signed-off-by: Kawika Avilla <kavilla414@gmail.com>
kavilla · Aug 25, 2024 · 8fce6b1 · 8fce6b1
1 parent 906d4a2
commit 8fce6b1
Show file tree

Hide file tree

Showing 13 changed files with 49 additions and 60 deletions.
diff --git a/src/plugins/data/public/antlr/shared/utils.test.ts b/src/plugins/data/public/antlr/shared/utils.test.ts
@@ -5,7 +5,7 @@
 
 import { of } from 'rxjs';
 import { fetchData } from './utils';
-import { DatasetManager } from '../../query';
+import { QueryStringManager } from '../../query';
 
 describe('fetchData', () => {
   it('should fetch data using the dataSourceRequestHandler', async () => {
@@ -26,18 +26,20 @@ describe('fetchData', () => {
         fetch: jest.fn().mockResolvedValue('fetchedData'),
       },
     };
-    const mockDatasetManager: Partial<DatasetManager> = {
+    const mockQueryString: Partial<QueryStringManager> = {
       getUpdates$: jest
         .fn()
         .mockReturnValue(of({ dataSourceRef: { id: 'testId', name: 'testTitle' } })),
-      getDataset: jest.fn().mockReturnValue({ dataSourceRef: { id: 'testId', name: 'testTitle' } }),
+      getDatasetService: jest
+        .fn()
+        .mockReturnValue({ dataSourceRef: { id: 'testId', name: 'testTitle' } }),
     };
 
     const result = await fetchData(
       mockTables,
       mockQueryFormatter,
       mockApi,
-      mockDatasetManager as DatasetManager
+      mockQueryString as QueryStringManager
     );
     expect(result).toEqual(['fetchedData', 'fetchedData']);
     expect(mockQueryFormatter).toHaveBeenCalledWith('table1', 'testId', 'testTitle');
@@ -59,16 +61,16 @@ describe('fetchData', () => {
         fetch: jest.fn().mockResolvedValue('fetchedData'),
       },
     };
-    const mockDatasetManager: Partial<DatasetManager> = {
+    const mockQueryString: Partial<QueryStringManager> = {
       getUpdates$: jest.fn().mockReturnValue(of(undefined)),
-      getDataset: jest.fn().mockReturnValue(undefined),
+      getDatasetService: jest.fn().mockReturnValue(undefined),
     };
 
     const result = await fetchData(
       mockTables,
       mockQueryFormatter,
       mockApi,
-      mockDatasetManager as DatasetManager
+      mockQueryString as QueryStringManager
     );
     expect(result).toEqual(['fetchedData', 'fetchedData']);
     expect(mockQueryFormatter).toHaveBeenCalledWith('table1');

diff --git a/src/plugins/data/public/query/query_string/dataset_service/lib/index_type.ts b/src/plugins/data/public/query/query_string/dataset_service/lib/index_type.ts
@@ -87,7 +87,7 @@ export const indexTypeConfig: DatasetTypeConfig = {
   },
 
   supportedLanguages: (): string[] => {
-    return ['SQL', 'PPL', 'kuery', 'lucene'];
+    return ['SQL', 'PPL', 'DQL', 'Lucene'];
   },
 };
 

diff --git a/src/plugins/data/public/query/query_string/query_string_manager.mock.ts b/src/plugins/data/public/query/query_string/query_string_manager.mock.ts
@@ -43,6 +43,7 @@ const createSetupContractMock = () => {
     clearQueryHistory: jest.fn(),
     changeQueryHistory: jest.fn(),
     getInitialQuery: jest.fn(),
+    getInitialQueryByLanguage: jest.fn(),
     getDatasetService: jest.fn(),
     getLanguageService: jest.fn(),
   };

diff --git a/src/plugins/data/public/query/query_string/query_string_manager.ts b/src/plugins/data/public/query/query_string/query_string_manager.ts
@@ -142,14 +142,18 @@ export class QueryStringManager {
   };
 
   public getInitialQuery = () => {
+    return this.getInitialQueryByLanguage(this.query$.getValue().language);
+  };
+
+  public getInitialQueryByLanguage = (languageId: string) => {
     const curQuery = this.query$.getValue();
-    const language = this.languageService.getLanguage(curQuery.language);
+    const language = this.languageService.getLanguage(languageId);
     const dataset = curQuery.dataset;
     const input = language?.searchBar?.queryStringInput?.initialValue || '';
 
     return {
       query: input || input.replace('<data_source>', dataset?.title ?? ''),
-      language: curQuery.language,
+      language: languageId,
       dataset: curQuery.dataset,
     };
   };

diff --git a/src/plugins/data/public/ui/query_editor/language_selector.tsx b/src/plugins/data/public/ui/query_editor/language_selector.tsx
@@ -41,14 +41,13 @@ export const QueryLanguageSelector = (props: QueryLanguageSelectorProps) => {
     const subscription = queryString.getUpdates$().subscribe((query: Query) => {
       if (query.language !== currentLanguage) {
         setCurrentLanguage(query.language);
-        props.onSelectLanguage(query.language);
       }
     });
 
     return () => {
       subscription.unsubscribe();
     };
-  }, [queryString, currentLanguage, props.onSelectLanguage, props]);
+  }, [queryString, currentLanguage, props]);
 
   const onButtonClick = () => {
     setPopover(!isPopoverOpen);
@@ -75,24 +74,6 @@ export const QueryLanguageSelector = (props: QueryLanguageSelectorProps) => {
   const handleLanguageChange = (newLanguage: string) => {
     setCurrentLanguage(newLanguage);
     props.onSelectLanguage(newLanguage);
-    uiService.Settings.setUserQueryLanguage(newLanguage);
-
-    // Update the query in the QueryStringManager
-    const currentQuery = queryString.getQuery();
-    const input = languageService.getLanguage(newLanguage)?.searchBar?.queryStringInput
-      ?.initialValue;
-
-    if (!input) return '';
-    const newQuery = input?.replace(
-      '<data_source>',
-      currentQuery.dataset?.title ?? currentQuery.dataset?.title ?? ''
-    );
-
-    queryString.setQuery({
-      query: newQuery,
-      language: newLanguage,
-      dataset: currentQuery.dataset,
-    });
   };
 
   uiService.Settings.setUserQueryLanguage(currentLanguage);

diff --git a/src/plugins/data/public/ui/query_editor/query_editor.tsx b/src/plugins/data/public/ui/query_editor/query_editor.tsx
@@ -181,21 +181,18 @@ export default class QueryEditorUI extends Component<Props, State> {
       body: JSON.stringify({ opt_in: languageId === 'kuery' }),
     });
 
-    const language = this.queryString.getLanguageService().getLanguage(this.props.query.language);
-    const newQuery = {
-      query: this.props.getQueryStringInitialValue?.(language!) ?? '',
-      language: language?.id || languageId,
-    };
+    const languageConfig = this.queryString.getLanguageService().getLanguage(languageId);
+    const newQuery = this.queryString.getInitialQueryByLanguage(languageId);
 
-    const fields = language?.fields;
+    const fields = languageConfig?.fields;
     const newSettings: DataSettings = {
       userQueryLanguage: newQuery.language,
       userQueryString: newQuery.query,
       ...(fields && { uiOverrides: { fields } }),
     };
     this.props.settings?.updateSettings(newSettings);
 
-    const dateRangeEnhancement = language?.searchBar?.dateRange;
+    const dateRangeEnhancement = languageConfig?.searchBar?.dateRange;
     const dateRange = dateRangeEnhancement
       ? {
           from: dateRangeEnhancement.initialFrom!,

diff --git a/src/plugins/data/public/ui/search_bar/search_bar.tsx b/src/plugins/data/public/ui/search_bar/search_bar.tsx
@@ -237,8 +237,8 @@ class SearchBarUI extends Component<SearchBarProps, State> {
       (!this.useNewHeader || this.props.filters.length > 0) &&
       this.props.indexPatterns &&
       compact(this.props.indexPatterns).length > 0 &&
-      (this.queryService.queryString.getLanguage(this.state.query?.language!)?.searchBar
-        ?.showFilterBar ??
+      (this.queryService.queryString.getLanguageService().getLanguage(this.state.query?.language!)
+        ?.searchBar?.showFilterBar ??
         true)
     );
   }

diff --git a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts
@@ -7,7 +7,7 @@ import { trimEnd } from 'lodash';
 import { Observable, throwError } from 'rxjs';
 import { i18n } from '@osd/i18n';
 import { concatMap, map } from 'rxjs/operators';
-import { DATA_FRAME_TYPES, getRawDataFrame, getRawQueryString } from '../../../data/common';
+import { DATA_FRAME_TYPES, getRawDataFrame } from '../../../data/common';
 import {
   DataPublicPluginStart,
   IOpenSearchDashboardsSearchRequest,

diff --git a/src/plugins/query_enhancements/server/routes/index.ts b/src/plugins/query_enhancements/server/routes/index.ts
@@ -32,7 +32,9 @@ function defineRoute(
       validate: {
         body: schema.object({
           query: schema.object({
-            query: schema.object({}, { unknowns: 'allow' }),
+            query: schema.string(),
+            language: schema.string(),
+            dataset: schema.nullable(schema.object({}, { unknowns: 'allow' })),
             format: schema.string(),
           }),
           df: schema.nullable(schema.object({}, { unknowns: 'allow' })),

diff --git a/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts b/src/plugins/query_enhancements/server/search/ppl_search_strategy.ts
@@ -3,7 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { first } from 'rxjs/operators';
 import { SharedGlobalConfig, Logger, ILegacyClusterClient } from 'opensearch-dashboards/server';
 import { Observable } from 'rxjs';
 import { ISearchStrategy, getDefaultSearchParams, SearchUsage } from '../../../data/server';
@@ -13,6 +12,7 @@ import {
   IDataFrameResponse,
   IDataFrameWithAggs,
   IOpenSearchDashboardsSearchRequest,
+  Query,
   createDataFrame,
 } from '../../../data/common';
 import { getFields } from '../../common/utils';
@@ -45,8 +45,6 @@ export const pplSearchStrategyProvider = (
 
     const source = pipeMap.get('source');
 
-    const searchQuery = query;
-
     const filters = pipeMap.get('where');
 
     const stats = pipeMap.get('stats');
@@ -55,25 +53,24 @@ export const pplSearchStrategyProvider = (
       : undefined;
 
     return {
-      map: pipeMap,
-      search: searchQuery,
       aggs: aggsQuery,
     };
   };
 
   return {
     search: async (context, request: any, options) => {
-      const config = await config$.pipe(first()).toPromise();
       const uiSettingsClient = await context.core.uiSettings.client;
 
       const { dataFrameHydrationStrategy, ...defaultParams } = await getDefaultSearchParams(
         uiSettingsClient
       );
 
       try {
-        const { meta } = request.body.df;
+        const query: Query = request.body.query;
+        const {
+          df: { meta },
+        } = request.body;
 
-        request.body.query = request.query.query;
         const rawResponse: any = await pplFacet.describeQuery(context, request);
 
         if (!rawResponse.success) {
@@ -85,7 +82,7 @@ export const pplSearchStrategyProvider = (
         }
 
         const dataFrame = createDataFrame({
-          name: request.query.dataset?.id,
+          name: query.dataset?.id,
           schema: rawResponse.data.schema,
           meta,
           fields: getFields(rawResponse),
@@ -98,8 +95,7 @@ export const pplSearchStrategyProvider = (
         if (dataFrame.meta?.aggsQs) {
           for (const [key, aggQueryString] of Object.entries(dataFrame.meta.aggsQs)) {
             const aggRequest = parseRequest(aggQueryString as string);
-            const query = aggRequest.aggs;
-            request.body.query = query;
+            request.body.query = aggRequest.aggs;
             const rawAggs: any = await pplFacet.describeQuery(context, request);
             (dataFrame as IDataFrameWithAggs).aggs = {};
             (dataFrame as IDataFrameWithAggs).aggs[key] = rawAggs.data.datarows?.map((hit: any) => {

diff --git a/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts b/src/plugins/query_enhancements/server/search/sql_async_search_strategy.ts
@@ -12,6 +12,7 @@ import {
   IDataFrameResponse,
   IOpenSearchDashboardsSearchRequest,
   PartialDataFrame,
+  Query,
   createDataFrame,
 } from '../../../data/common';
 import { Facet } from '../utils';
@@ -38,12 +39,13 @@ export const sqlAsyncSearchStrategyProvider = (
   return {
     search: async (context, request: any, options) => {
       try {
+        const query: Query = request?.body?.query;
         // Create job: this should return a queryId and sessionId
-        if (request?.body?.query?.query) {
+        if (query) {
           const df = request.body?.df;
           request.body = {
-            query: request.body.query.query,
-            datasource: request.body.query.dataset?.dataSource?.title,
+            query: query.query,
+            datasource: query.dataset?.dataSource?.title,
             lang: SEARCH_STRATEGY.SQL,
             sessionId: df?.meta?.sessionId,
           };
@@ -60,13 +62,13 @@ export const sqlAsyncSearchStrategyProvider = (
           const sessionId = rawResponse.data?.sessionId;
 
           const partial: PartialDataFrame = {
-            ...request.body.df,
+            ...df,
             fields: rawResponse?.data?.schema || [],
           };
           const dataFrame = createDataFrame(partial);
           dataFrame.meta = {
             ...dataFrame.meta,
-            query: request.body.query,
+            query: query.query,
             queryId,
             sessionId,
           };

diff --git a/src/plugins/query_enhancements/server/search/sql_search_strategy.ts b/src/plugins/query_enhancements/server/search/sql_search_strategy.ts
@@ -12,6 +12,7 @@ import {
   IDataFrameResponse,
   IOpenSearchDashboardsSearchRequest,
   PartialDataFrame,
+  Query,
   createDataFrame,
 } from '../../../data/common';
 import { Facet } from '../utils';
@@ -27,7 +28,8 @@ export const sqlSearchStrategyProvider = (
   return {
     search: async (context, request: any, _options) => {
       try {
-        request.body.query = request.body.query.query;
+        const query: Query = request.body.query;
+        request.body.query = query.query;
         const rawResponse: any = await sqlFacet.describeQuery(context, request);
 
         if (!rawResponse.success) {

diff --git a/src/plugins/query_enhancements/server/utils/facet.ts b/src/plugins/query_enhancements/server/utils/facet.ts
@@ -6,6 +6,7 @@
 import { Logger } from 'opensearch-dashboards/server';
 import { FacetResponse, IPPLEventsDataSource, IPPLVisualizationDataSource } from '../types';
 import { shimSchemaRow, shimStats } from '.';
+import { Query } from '../../../data/common';
 
 export interface FacetProps {
   client: any;
@@ -36,12 +37,13 @@ export class Facet {
     endpoint: string
   ): Promise<FacetResponse> => {
     try {
-      const { format, df, dataSourceId, ...query } = request.body;
+      const query: Query = request.body.query;
+      const { format, df } = request.body;
       const params = {
         body: { ...query },
         ...(format !== 'jdbc' && { format }),
       };
-      const clientId = dataSourceId ?? df?.meta?.queryConfig?.dataSourceId;
+      const clientId = query.dataset?.dataSource?.id ?? df?.meta?.queryConfig?.dataSourceId;
       const client = clientId
         ? context.dataSource.opensearch.legacy.getClient(clientId).callAPI
         : this.defaultClient.asScoped(request).callAsCurrentUser;