From 1cc5c0fcd78eb947ace44b3bdf250be39aa94533 Mon Sep 17 00:00:00 2001 From: Anan Zhuang Date: Thu, 31 Oct 2024 09:54:09 -0700 Subject: [PATCH] [Manual backport 2.x] Add support for async ppl to discover (#8706) (#8769) Backport PR: https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8706 From original PR: * add support for async ppl to discover * Changeset file for PR #8706 created/updated * update s3_type test to add PPL as supported lang * fix lint error --------- Signed-off-by: Shenoy Pratik Co-authored-by: Shenoy Pratik Co-authored-by: opensearch-changeset-bot[bot] <154024398+opensearch-changeset-bot[bot]@users.noreply.github.com> --- changelogs/fragments/8706.yml | 2 + .../query_enhancements/common/constants.ts | 2 + .../public/datasets/s3_type.test.ts | 4 +- .../public/datasets/s3_type.ts | 4 +- .../query_enhancements/public/plugin.tsx | 19 ++-- .../public/search/ppl_search_interceptor.ts | 3 +- .../public/search/sql_search_interceptor.ts | 2 +- .../query_enhancements/server/plugin.ts | 12 ++- .../routes/data_source_connection/routes.ts | 2 +- .../query_enhancements/server/search/index.ts | 5 +- .../search/ppl_async_search_strategy.ts | 102 ++++++++++++++++++ 11 files changed, 139 insertions(+), 18 deletions(-) create mode 100644 changelogs/fragments/8706.yml create mode 100644 src/plugins/query_enhancements/server/search/ppl_async_search_strategy.ts diff --git a/changelogs/fragments/8706.yml b/changelogs/fragments/8706.yml new file mode 100644 index 000000000000..5d4977a48129 --- /dev/null +++ b/changelogs/fragments/8706.yml @@ -0,0 +1,2 @@ +feat: +- Add support for async ppl to discover ([#8706](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/8706)) \ No newline at end of file diff --git a/src/plugins/query_enhancements/common/constants.ts b/src/plugins/query_enhancements/common/constants.ts index f729c40211aa..f116da2f9724 100644 --- a/src/plugins/query_enhancements/common/constants.ts +++ b/src/plugins/query_enhancements/common/constants.ts @@ -17,6 +17,7 @@ export const SEARCH_STRATEGY = { PPL_RAW: 'pplraw', SQL: 'sql', SQL_ASYNC: 'sqlasync', + PPL_ASYNC: 'pplasync', }; export const API = { @@ -24,6 +25,7 @@ export const API = { PPL_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.PPL}`, SQL_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.SQL}`, SQL_ASYNC_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.SQL_ASYNC}`, + PPL_ASYNC_SEARCH: `${BASE_API}/search/${SEARCH_STRATEGY.PPL_ASYNC}`, QUERY_ASSIST: { LANGUAGES: `${BASE_API}/assist/languages`, GENERATE: `${BASE_API}/assist/generate`, diff --git a/src/plugins/query_enhancements/public/datasets/s3_type.test.ts b/src/plugins/query_enhancements/public/datasets/s3_type.test.ts index 689f996291b8..3f085c135f28 100644 --- a/src/plugins/query_enhancements/public/datasets/s3_type.test.ts +++ b/src/plugins/query_enhancements/public/datasets/s3_type.test.ts @@ -244,9 +244,9 @@ describe('s3TypeConfig', () => { expect(result[3].type).toBe('number'); }); - test('supportedLanguages returns SQL', () => { + test('supportedLanguages returns SQL, PPL', () => { const mockDataset: Dataset = { id: 'table1', title: 'Table 1', type: 'S3' }; - expect(s3TypeConfig.supportedLanguages(mockDataset)).toEqual(['SQL']); + expect(s3TypeConfig.supportedLanguages(mockDataset)).toEqual(['SQL', 'PPL']); }); describe('castS3FieldTypeToOSDFieldType()', () => { diff --git a/src/plugins/query_enhancements/public/datasets/s3_type.ts b/src/plugins/query_enhancements/public/datasets/s3_type.ts index 67da4e85c2a6..a62f2d21cecc 100644 --- a/src/plugins/query_enhancements/public/datasets/s3_type.ts +++ b/src/plugins/query_enhancements/public/datasets/s3_type.ts @@ -118,7 +118,7 @@ export const s3TypeConfig: DatasetTypeConfig = { }, supportedLanguages: (dataset: Dataset): string[] => { - return ['SQL']; + return ['SQL', 'PPL']; }, getSampleQueries: (dataset: Dataset, language: string) => { @@ -129,7 +129,7 @@ export const s3TypeConfig: DatasetTypeConfig = { title: i18n.translate('queryEnhancements.s3Type.sampleQuery.basicPPLQuery', { defaultMessage: 'Sample query for PPL', }), - query: `source = ${dataset.title}`, + query: `source = ${dataset.title} | head 10`, }, ]; case 'SQL': diff --git a/src/plugins/query_enhancements/public/plugin.tsx b/src/plugins/query_enhancements/public/plugin.tsx index ce4203d815c2..4058b5a595db 100644 --- a/src/plugins/query_enhancements/public/plugin.tsx +++ b/src/plugins/query_enhancements/public/plugin.tsx @@ -4,21 +4,26 @@ */ import { i18n } from '@osd/i18n'; import { CoreSetup, CoreStart, Plugin, PluginInitializerContext } from '../../../core/public'; +import { DataStorage } from '../../data/common'; +import { + createEditor, + DefaultInput, + LanguageConfig, + Query, + SingleLineInput, +} from '../../data/public'; import { ConfigSchema } from '../common/config'; -import { setData, setStorage } from './services'; +import { s3TypeConfig } from './datasets'; import { createQueryAssistExtension } from './query_assist'; +import { pplLanguageReference, sqlLanguageReference } from './query_editor_extensions'; import { PPLSearchInterceptor, SQLSearchInterceptor } from './search'; +import { setData, setStorage } from './services'; import { QueryEnhancementsPluginSetup, QueryEnhancementsPluginSetupDependencies, QueryEnhancementsPluginStart, QueryEnhancementsPluginStartDependencies, } from './types'; -import { LanguageConfig, Query } from '../../data/public'; -import { s3TypeConfig } from './datasets'; -import { createEditor, DefaultInput, SingleLineInput } from '../../data/public'; -import { DataStorage } from '../../data/common'; -import { pplLanguageReference, sqlLanguageReference } from './query_editor_extensions'; export class QueryEnhancementsPlugin implements @@ -57,7 +62,7 @@ export class QueryEnhancementsPlugin startServices: core.getStartServices(), usageCollector: data.search.usageCollector, }), - getQueryString: (currentQuery: Query) => `source = ${currentQuery.dataset?.title}`, + getQueryString: (currentQuery: Query) => `source = ${currentQuery.dataset?.title} | head 10`, fields: { filterable: false, visualizable: false }, docLink: { title: i18n.translate('queryEnhancements.pplLanguage.docLink', { diff --git a/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts b/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts index c52bebf49a0a..57152dbe98ea 100644 --- a/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts +++ b/src/plugins/query_enhancements/public/search/ppl_search_interceptor.ts @@ -16,6 +16,7 @@ import { } from '../../../data/public'; import { API, + DATASET, EnhancedFetchContext, fetch, formatDate, @@ -60,7 +61,7 @@ export class PPLSearchInterceptor extends SearchInterceptor { public search(request: IOpenSearchDashboardsSearchRequest, options: ISearchOptions) { const dataset = this.queryService.queryString.getQuery().dataset; const datasetType = dataset?.type; - let strategy = SEARCH_STRATEGY.PPL; + let strategy = datasetType === DATASET.S3 ? SEARCH_STRATEGY.PPL_ASYNC : SEARCH_STRATEGY.PPL; if (datasetType) { const datasetTypeConfig = this.queryService.queryString diff --git a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts index 4e62526653f9..9fe17fc79322 100644 --- a/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts +++ b/src/plugins/query_enhancements/public/search/sql_search_interceptor.ts @@ -4,9 +4,9 @@ */ import { trimEnd } from 'lodash'; +import { CoreStart } from 'opensearch-dashboards/public'; import { Observable, throwError } from 'rxjs'; import { catchError } from 'rxjs/operators'; -import { CoreStart } from 'opensearch-dashboards/public'; import { DataPublicPluginStart, IOpenSearchDashboardsSearchRequest, diff --git a/src/plugins/query_enhancements/server/plugin.ts b/src/plugins/query_enhancements/server/plugin.ts index 4a624bfccb1b..dfaea828209d 100644 --- a/src/plugins/query_enhancements/server/plugin.ts +++ b/src/plugins/query_enhancements/server/plugin.ts @@ -17,10 +17,11 @@ import { SEARCH_STRATEGY } from '../common'; import { ConfigSchema } from '../common/config'; import { defineRoutes, defineSearchStrategyRouteProvider } from './routes'; import { - pplSearchStrategyProvider, + pplAsyncSearchStrategyProvider, pplRawSearchStrategyProvider, - sqlSearchStrategyProvider, + pplSearchStrategyProvider, sqlAsyncSearchStrategyProvider, + sqlSearchStrategyProvider, } from './search'; import { QueryEnhancementsPluginSetup, @@ -58,11 +59,17 @@ export class QueryEnhancementsPlugin this.logger, client ); + const pplAsyncSearchStrategy = pplAsyncSearchStrategyProvider( + this.config$, + this.logger, + client + ); data.search.registerSearchStrategy(SEARCH_STRATEGY.PPL, pplSearchStrategy); data.search.registerSearchStrategy(SEARCH_STRATEGY.PPL_RAW, pplRawSearchStrategy); data.search.registerSearchStrategy(SEARCH_STRATEGY.SQL, sqlSearchStrategy); data.search.registerSearchStrategy(SEARCH_STRATEGY.SQL_ASYNC, sqlAsyncSearchStrategy); + data.search.registerSearchStrategy(SEARCH_STRATEGY.PPL_ASYNC, pplAsyncSearchStrategy); core.http.registerRouteHandlerContext('query_assist', () => ({ logger: this.logger, @@ -86,6 +93,7 @@ export class QueryEnhancementsPlugin ppl: pplSearchStrategy, sql: sqlSearchStrategy, sqlasync: sqlAsyncSearchStrategy, + pplasync: pplAsyncSearchStrategy, }); this.logger.info('queryEnhancements: Setup complete'); diff --git a/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts b/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts index fb887c6df0d5..0581cf0fa896 100644 --- a/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts +++ b/src/plugins/query_enhancements/server/routes/data_source_connection/routes.ts @@ -4,7 +4,7 @@ */ import { schema } from '@osd/config-schema'; -import { IRouter, ILegacyClusterClient } from 'opensearch-dashboards/server'; +import { ILegacyClusterClient, IRouter } from 'opensearch-dashboards/server'; import { API } from '../../../common'; export function registerDataSourceConnectionsRoutes( diff --git a/src/plugins/query_enhancements/server/search/index.ts b/src/plugins/query_enhancements/server/search/index.ts index 129ce971662f..af921cf546f1 100644 --- a/src/plugins/query_enhancements/server/search/index.ts +++ b/src/plugins/query_enhancements/server/search/index.ts @@ -3,7 +3,8 @@ * SPDX-License-Identifier: Apache-2.0 */ -export { pplSearchStrategyProvider } from './ppl_search_strategy'; +export { pplAsyncSearchStrategyProvider } from './ppl_async_search_strategy'; export { pplRawSearchStrategyProvider } from './ppl_raw_search_strategy'; -export { sqlSearchStrategyProvider } from './sql_search_strategy'; +export { pplSearchStrategyProvider } from './ppl_search_strategy'; export { sqlAsyncSearchStrategyProvider } from './sql_async_search_strategy'; +export { sqlSearchStrategyProvider } from './sql_search_strategy'; diff --git a/src/plugins/query_enhancements/server/search/ppl_async_search_strategy.ts b/src/plugins/query_enhancements/server/search/ppl_async_search_strategy.ts new file mode 100644 index 000000000000..9ac010b204c0 --- /dev/null +++ b/src/plugins/query_enhancements/server/search/ppl_async_search_strategy.ts @@ -0,0 +1,102 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ILegacyClusterClient, Logger, SharedGlobalConfig } from 'opensearch-dashboards/server'; +import { Observable } from 'rxjs'; +import { + createDataFrame, + DATA_FRAME_TYPES, + IDataFrameResponse, + IOpenSearchDashboardsSearchRequest, + Query, +} from '../../../data/common'; +import { ISearchStrategy, SearchUsage } from '../../../data/server'; +import { buildQueryStatusConfig, getFields, handleFacetError, SEARCH_STRATEGY } from '../../common'; +import { Facet } from '../utils'; + +export const pplAsyncSearchStrategyProvider = ( + config$: Observable, + logger: Logger, + client: ILegacyClusterClient, + usage?: SearchUsage +): ISearchStrategy => { + const pplAsyncFacet = new Facet({ + client, + logger, + endpoint: 'enhancements.runDirectQuery', + }); + const pplAsyncJobsFacet = new Facet({ + client, + logger, + endpoint: 'enhancements.getJobStatus', + useJobs: true, + }); + + return { + search: async (context, request: any, options) => { + try { + const query: Query = request.body.query; + const pollQueryResultsParams = request.body.pollQueryResultsParams; + const inProgressQueryId = pollQueryResultsParams?.queryId; + + if (!inProgressQueryId) { + request.body = { ...request.body, lang: SEARCH_STRATEGY.PPL }; + const rawResponse: any = await pplAsyncFacet.describeQuery(context, request); + + if (!rawResponse.success) handleFacetError(rawResponse); + + const statusConfig = buildQueryStatusConfig(rawResponse); + + return { + type: DATA_FRAME_TYPES.POLLING, + status: 'started', + body: { + queryStatusConfig: statusConfig, + }, + } as IDataFrameResponse; + } else { + request.params = { queryId: inProgressQueryId }; + const queryStatusResponse: any = await pplAsyncJobsFacet.describeQuery(context, request); + const queryStatus = queryStatusResponse?.data?.status; + logger.info(`pplAsyncSearchStrategy: JOB: ${inProgressQueryId} - STATUS: ${queryStatus}`); + + if (queryStatus?.toUpperCase() === 'SUCCESS') { + const dataFrame = createDataFrame({ + name: query.dataset?.id, + schema: queryStatusResponse.data.schema, + meta: { ...pollQueryResultsParams }, + fields: getFields(queryStatusResponse), + }); + + dataFrame.size = queryStatusResponse.data.datarows.length; + + return { + type: DATA_FRAME_TYPES.POLLING, + status: 'success', + body: dataFrame, + } as IDataFrameResponse; + } else if (queryStatus?.toUpperCase() === 'FAILED') { + return { + type: DATA_FRAME_TYPES.POLLING, + status: 'failed', + body: { + error: `JOB: ${inProgressQueryId} failed: ${queryStatusResponse.data.error}`, + }, + } as IDataFrameResponse; + } + + return { + type: DATA_FRAME_TYPES.POLLING, + status: queryStatus, + } as IDataFrameResponse; + } + } catch (e: any) { + logger.error(`pplAsyncSearchStrategy: ${e.message}`); + if (usage) usage.trackError(); + throw e; + } + }, + }; +};