[Backport 2.16] feat: support PPL in vega visualization

changelogs/fragments/7285.yml

@@ -0,0 +1,2 @@
+feat:
+- Support PPL in vega visualization ([#7285](https://github.com/opensearch-project/OpenSearch-Dashboards/pull/7285))

src/plugins/data/common/search/search_source/fetch/get_search_params.ts

@@ -94,7 +94,7 @@ export function getSearchParamsFromRequest(
   }
 
   return {
-    index: searchRequest.index.title || searchRequest.index,
+    index: searchRequest.index?.title || searchRequest.index,
     body: searchRequest.body,
     ...searchParams,
   };

src/plugins/data/server/search/routes/search.test.ts

@@ -94,7 +94,10 @@ describe('Search service', () => {
     await handler((mockContext as unknown) as RequestHandlerContext, mockRequest, mockResponse);
 
     expect(mockDataStart.search.search).toBeCalled();
-    expect(mockDataStart.search.search.mock.calls[0][1]).toStrictEqual(mockBody);
+    expect(mockDataStart.search.search.mock.calls[0][1]).toStrictEqual({
+      ...mockBody,
+      rawRequest: mockRequest,
+    });
     expect(mockResponse.ok).toBeCalled();
     expect(mockResponse.ok.mock.calls[0][0]).toEqual({
       body: response,
@@ -125,7 +128,10 @@ describe('Search service', () => {
     await handler((mockContext as unknown) as RequestHandlerContext, mockRequest, mockResponse);
 
     expect(mockDataStart.search.search).toBeCalled();
-    expect(mockDataStart.search.search.mock.calls[0][1]).toStrictEqual(mockBody);
+    expect(mockDataStart.search.search.mock.calls[0][1]).toStrictEqual({
+      ...mockBody,
+      rawRequest: mockRequest,
+    });
     expect(mockResponse.customError).toBeCalled();
     const error: any = mockResponse.customError.mock.calls[0][0];
     expect(error.body.message).toBe('oh no');

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

@@ -62,7 +62,7 @@ export function registerSearchRoute(
       try {
         const { withLongNumeralsSupport, ...response } = await selfStart.search.search(
           context,
-          { ...searchRequest, id },
+          { ...searchRequest, id, rawRequest: request },
           {
             abortSignal,
             strategy,

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

@@ -72,7 +72,9 @@ export interface ISearchSetup {
 }
 
 export interface ISearchStart<
-  SearchStrategyRequest extends IOpenSearchDashboardsSearchRequest = IOpenSearchSearchRequest,
+  SearchStrategyRequest extends IOpenSearchDashboardsSearchRequest = IOpenSearchSearchRequest & {
+    rawRequest?: OpenSearchDashboardsRequest;
+  },
   SearchStrategyResponse extends IOpenSearchDashboardsSearchResponse = IOpenSearchSearchResponse
 > {
   aggs: AggsStart;

src/plugins/query_enhancements/common/constants.ts

@@ -10,6 +10,7 @@ export const BASE_API = '/api/enhancements';
 
 export const SEARCH_STRATEGY = {
   PPL: 'ppl',
+  PPL_RAW: 'pplraw',
   SQL: 'sql',
   SQL_ASYNC: 'sqlasync',
 };

src/plugins/query_enhancements/server/plugin.ts

@@ -25,8 +25,10 @@ import {
   QueryEnhancementsPluginSetup,
   QueryEnhancementsPluginSetupDependencies,
   QueryEnhancementsPluginStart,
+  QueryEnhancementsPluginStartDependencies,
 } from './types';
 import { OpenSearchObservabilityPlugin, OpenSearchPPLPlugin } from './utils';
+import { pplRawSearchStrategyProvider } from './search/ppl_raw_search_strategy';
 
 export class QueryEnhancementsPlugin
   implements Plugin<QueryEnhancementsPluginSetup, QueryEnhancementsPluginStart> {
@@ -51,6 +53,7 @@ export class QueryEnhancementsPlugin
     }
 
     const pplSearchStrategy = pplSearchStrategyProvider(this.config$, this.logger, client);
+    const pplRawSearchStrategy = pplRawSearchStrategyProvider(this.config$, this.logger, client);
     const sqlSearchStrategy = sqlSearchStrategyProvider(this.config$, this.logger, client);
     const sqlAsyncSearchStrategy = sqlAsyncSearchStrategyProvider(
       this.config$,
@@ -59,6 +62,7 @@ export class QueryEnhancementsPlugin
     );
 
     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);
 

src/plugins/query_enhancements/server/search/ppl_raw_search_strategy.ts

@@ -0,0 +1,38 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { SharedGlobalConfig, Logger, ILegacyClusterClient } from 'opensearch-dashboards/server';
+import { Observable } from 'rxjs';
+import { ISearchStrategy, SearchUsage } from '../../../data/server';
+import { shimSchemaRow } from '../utils';
+
+export const pplRawSearchStrategyProvider = (
+  config$: Observable<SharedGlobalConfig>,
+  logger: Logger,
+  client: ILegacyClusterClient,
+  usage?: SearchUsage
+): ISearchStrategy => {
+  return {
+    search: async (context, request: any, options) => {
+      const runSearch = request.dataSourceId
+        ? context.dataSource.opensearch.legacy.getClient(request.dataSourceId).callAPI
+        : client.asScoped(request.rawRequest).callAsCurrentUser;
+
+      try {
+        const rawResponse: any = await runSearch('ppl.pplQuery', { body: request.params.body });
+        const data = shimSchemaRow(rawResponse);
+        rawResponse.jsonData = data.jsonData;
+
+        return {
+          rawResponse,
+        };
+      } catch (e) {
+        logger.error(`pplRawSearchStrategy: ${e.message}`);
+        if (usage) usage.trackError();
+        throw e;
+      }
+    },
+  };
+};

src/plugins/vis_type_vega/public/data_model/ppl_parser.test.ts

@@ -0,0 +1,49 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { PPLQueryParser } from './ppl_parser';
+
+test('it should throw error if with invalid url object', () => {
+  const searchApiMock = {
+    search: jest.fn(() => ({
+      toPromise: jest.fn(() => Promise.resolve({})),
+    })),
+  };
+  const parser = new PPLQueryParser(searchApiMock);
+  expect(() => parser.parseUrl({}, {})).toThrowError();
+  expect(() => parser.parseUrl({}, { body: {} })).toThrowError();
+  expect(() => parser.parseUrl({}, { body: { query: {} } })).toThrowError();
+});
+
+test('it should parse url object', () => {
+  const searchApiMock = {
+    search: jest.fn(() => ({
+      toPromise: jest.fn(() => Promise.resolve({})),
+    })),
+  };
+  const parser = new PPLQueryParser(searchApiMock);
+  const result = parser.parseUrl({}, { body: { query: 'source=test_index' } });
+  expect(result.dataObject).toEqual({});
+  expect(result.url).toEqual({ body: { query: 'source=test_index' } });
+});
+
+it('should populate data to request', async () => {
+  const searchApiMock = {
+    search: jest.fn(() => ({
+      toPromise: jest.fn(() =>
+        Promise.resolve([{ name: 'request name', rawResponse: { jsonData: [{ id: 'id1' }] } }])
+      ),
+    })),
+  };
+  const parser = new PPLQueryParser(searchApiMock);
+  const request = {
+    url: { body: { query: 'source=test_index' } },
+    dataObject: {
+      name: 'request name',
+    },
+  };
+  await parser.populateData([request]);
+  expect(request.dataObject.values).toEqual([{ id: 'id1' }]);
+});

src/plugins/vis_type_vega/public/data_model/ppl_parser.ts

@@ -0,0 +1,57 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { i18n } from '@osd/i18n';
+import { Data, UrlObject, PPLQueryRequest } from './types';
+import { SearchAPI } from './search_api';
+
+const getRequestName = (request: PPLQueryRequest, index: number) =>
+  request.dataObject.name ||
+  i18n.translate('visTypeVega.opensearchQueryParser.unnamedRequest', {
+    defaultMessage: 'Unnamed request #{index}',
+    values: { index },
+  });
+
+export class PPLQueryParser {
+  searchAPI: SearchAPI;
+
+  constructor(searchAPI: SearchAPI) {
+    this.searchAPI = searchAPI;
+  }
+
+  parseUrl(dataObject: Data, url: UrlObject) {
+    // data.url.body.query must be defined
+    if (!url.body || !url.body.query || typeof url.body.query !== 'string') {
+      throw new Error(
+        i18n.translate('visTypeVega.pplQueryParser.dataUrl.PPL.queryCannotBeEmpty', {
+          defaultMessage: '{dataUrlParam} must have query specified',
+          values: {
+            dataUrlParam: '"data.url"',
+          },
+        })
+      );
+    }
+
+    return { dataObject, url };
+  }
+
+  async populateData(requests: PPLQueryRequest[]) {
+    const searchRequests = requests.map((r, index) => ({
+      ...r.url,
+      name: getRequestName(r, index),
+    }));
+
+    const data$ = await this.searchAPI.search(searchRequests, { strategy: 'pplraw' });
+    const results = await data$.toPromise();
+    results.forEach((data, index) => {
+      const requestObject = requests.find((item) => getRequestName(item, index) === data.name);
+
+      if (requestObject) {
+        requestObject.dataObject.url = requestObject.url;
+        requestObject.dataObject.values = (data.rawResponse as any).jsonData;
+      }
+    });
+  }
+}

src/plugins/vis_type_vega/public/data_model/search_api.ts

@@ -28,6 +28,7 @@
  * under the License.
  */
 
+import { i18n } from '@osd/i18n';
 import { combineLatest } from 'rxjs';
 import { map, tap } from 'rxjs/operators';
 import { CoreStart, IUiSettingsClient } from 'opensearch-dashboards/public';
@@ -38,10 +39,21 @@
   SearchRequest,
   DataPublicPluginStart,
   IOpenSearchSearchResponse,
+  IOpenSearchSearchRequest,
 } from '../../../data/public';
 import { search as dataPluginSearch } from '../../../data/public';
 import { VegaInspectorAdapters } from '../vega_inspector';
-import { RequestResponder } from '../../../inspector/public';
+import { RequestResponder, RequestStatistics } from '../../../inspector/public';
+
+interface RawPPLStrategySearchResponse {
+  rawResponse: {
+    datarows: any[];
+    jsonData: any[];
+    schema: any[];
+    size: number;
+    total: number;
+  };
+}
 
 export interface SearchAPIDependencies {
   uiSettings: IUiSettingsClient;
@@ -58,7 +70,7 @@
     public readonly inspectorAdapters?: VegaInspectorAdapters
   ) {}
 
-  async search(searchRequests: SearchRequest[]) {
+  async search(searchRequests: SearchRequest[], options?: { strategy?: string }) {
     const { search } = this.dependencies.search;
     const requestResponders: any = {};
 
@@ -87,7 +99,13 @@
               ? { params, dataSourceId }
               : { params };
 
-          return search(searchApiParams, { abortSignal: this.abortSignal }).pipe(
+          return search<
+            IOpenSearchSearchRequest,
+            IOpenSearchSearchResponse | RawPPLStrategySearchResponse
+          >(searchApiParams, {
+            abortSignal: this.abortSignal,
+            strategy: options?.strategy,
+          }).pipe(
             tap((data) => this.inspectSearchResult(data, requestResponders[requestId])),
             map((data) => ({
               name: requestId,
@@ -135,14 +153,45 @@
     }
   }
 
+  private getPPLRawResponseInspectorStats(response: RawPPLStrategySearchResponse['rawResponse']) {
+    const stats: RequestStatistics = {};
+    stats.hitsTotal = {
+      label: i18n.translate('data.search.searchSource.hitsTotalLabel', {
+        defaultMessage: 'Hits (total)',
+      }),
+      value: `${response.total}`,
+      description: i18n.translate('data.search.searchSource.hitsTotalDescription', {
+        defaultMessage: 'The number of documents that match the query.',
+      }),
+    };
+
+    stats.hits = {
+      label: i18n.translate('data.search.searchSource.hitsLabel', {
+        defaultMessage: 'Hits',
+      }),
+      value: `${response.size}`,
+      description: i18n.translate('data.search.searchSource.hitsDescription', {
+        defaultMessage: 'The number of documents returned by the query.',
+      }),
+    };
+    return stats;
+  }
+
   private inspectSearchResult(
-    response: IOpenSearchSearchResponse,
+    response: IOpenSearchSearchResponse | RawPPLStrategySearchResponse,
     requestResponder: RequestResponder
   ) {
     if (requestResponder) {
-      requestResponder
-        .stats(dataPluginSearch.getResponseInspectorStats(response.rawResponse))
-        .ok({ json: response.rawResponse });
+      // inspect ppl response
+      if ('jsonData' in response.rawResponse) {
+        requestResponder
+          .stats(this.getPPLRawResponseInspectorStats(response.rawResponse))
+          .ok({ json: response.rawResponse });
+      } else {
+        requestResponder
+          .stats(dataPluginSearch.getResponseInspectorStats(response.rawResponse))
+          .ok({ json: response.rawResponse });
+      }
     }
   }
 }

src/plugins/vis_type_vega/public/data_model/types.ts

@@ -40,7 +40,7 @@ import { UrlParser } from './url_parser';
 
 interface Body {
   aggs?: SearchParams['body']['aggs'];
-  query?: Query;
+  query?: Query | string;
   timeout?: string;
 }
 
@@ -214,6 +214,7 @@ export type OpenSearchQueryRequest = Requests;
 export type EmsQueryRequest = Requests & {
   obj: UrlObject;
 };
+export type PPLQueryRequest = Requests;
 
 export interface ContextVarsObject {
   [index: string]: any;

src/plugins/vis_type_vega/public/data_model/vega_parser.ts

@@ -37,6 +37,8 @@ import { euiThemeVars } from '@osd/ui-shared-deps/theme';
 import { i18n } from '@osd/i18n';
 // @ts-ignore
 import { Signal } from 'vega';
+
+import { HttpSetup } from 'opensearch-dashboards/public';
 import { vega, vegaLite } from '../lib/vega';
 import { OpenSearchQueryParser } from './opensearch_query_parser';
 import { Utils } from './utils';
@@ -59,6 +61,7 @@ import {
   ControlsDirection,
   OpenSearchDashboards,
 } from './types';
+import { PPLQueryParser } from './ppl_parser';
 
 // Set default single color to match other OpenSearch Dashboards visualizations
 const defaultColor: string = euiPaletteColorBlind()[0];
@@ -653,6 +656,7 @@ The URL is an identifier only. OpenSearch Dashboards and your browser will never
         opensearch: new OpenSearchQueryParser(this.timeCache, this.searchAPI, this.filters, onWarn),
         emsfile: new EmsFileParser(serviceSettings),
         url: new UrlParser(onWarn),
+        ppl: new PPLQueryParser(this.searchAPI),
       };
     }
     const pending: PendingType = {};

src/plugins/visualizations/public/legacy/build_pipeline.ts

@@ -390,8 +390,8 @@ export const buildVislibDimensions = async (
 
 export const buildPipeline = async (vis: Vis, params: BuildPipelineParams) => {
   const { indexPattern, searchSource } = vis.data;
-  const query = searchSource!.getField('query');
-  const filters = searchSource!.getField('filter');
+  const query = searchSource?.getField('query');
+  const filters = searchSource?.getField('filter');
   const { uiState, title } = vis;
 
   // context