Use Search API in Vega

src/legacy/core_plugins/kibana/public/__tests__/vis_type_vega/vega_visualization.js

@@ -44,7 +44,7 @@ import vegaMapImage256 from './vega_map_image_256.png';
 // eslint-disable-next-line @kbn/eslint/no-restricted-paths
 import { VegaParser } from '../../../../../../plugins/vis_type_vega/public/data_model/vega_parser';
 // eslint-disable-next-line @kbn/eslint/no-restricted-paths
-import { SearchCache } from '../../../../../../plugins/vis_type_vega/public/data_model/search_cache';
+import { SearchAPI } from '../../../../../../plugins/vis_type_vega/public/data_model/search_api';
 
 // eslint-disable-next-line @kbn/eslint/no-restricted-paths
 import { createVegaTypeDefinition } from '../../../../../../plugins/vis_type_vega/public/vega_type';
@@ -205,7 +205,12 @@ describe('VegaVisualizations', () => {
       try {
         vegaVis = new VegaVisualization(domNode, vis);
 
-        const vegaParser = new VegaParser(vegaliteGraph, new SearchCache());
+        const vegaParser = new VegaParser(
+          vegaliteGraph,
+          new SearchAPI({
+            search: npStart.plugins.data.search,
+          })
+        );
         await vegaParser.parseAsync();
 
         await vegaVis.render(vegaParser, vis.params, { data: true });
@@ -227,7 +232,12 @@ describe('VegaVisualizations', () => {
       let vegaVis;
       try {
         vegaVis = new VegaVisualization(domNode, vis);
-        const vegaParser = new VegaParser(vegaGraph, new SearchCache());
+        const vegaParser = new VegaParser(
+          vegaGraph,
+          new SearchAPI({
+            search: npStart.plugins.data.search,
+          })
+        );
         await vegaParser.parseAsync();
 
         await vegaVis.render(vegaParser, vis.params, { data: true });
@@ -243,7 +253,12 @@ describe('VegaVisualizations', () => {
       let vegaVis;
       try {
         vegaVis = new VegaVisualization(domNode, vis);
-        const vegaParser = new VegaParser(vegaTooltipGraph, new SearchCache());
+        const vegaParser = new VegaParser(
+          vegaTooltipGraph,
+          new SearchAPI({
+            search: npStart.plugins.data.search,
+          })
+        );
         await vegaParser.parseAsync();
         await vegaVis.render(vegaParser, vis.params, { data: true });
 
@@ -285,7 +300,12 @@ describe('VegaVisualizations', () => {
       let vegaVis;
       try {
         vegaVis = new VegaVisualization(domNode, vis);
-        const vegaParser = new VegaParser(vegaMapGraph, new SearchCache());
+        const vegaParser = new VegaParser(
+          vegaMapGraph,
+          new SearchAPI({
+            search: npStart.plugins.data.search,
+          })
+        );
         await vegaParser.parseAsync();
 
         domNode.style.width = '256px';
@@ -324,7 +344,9 @@ describe('VegaVisualizations', () => {
               }
             ]
           }`,
-          new SearchCache()
+          new SearchAPI({
+            search: npStart.plugins.data.search,
+          })
         );
         await vegaParser.parseAsync();
 

src/plugins/data/public/search/fetch/get_search_params.ts

@@ -17,7 +17,8 @@
  * under the License.
  */
 
-import { IUiSettingsClient } from 'kibana/public';
+import { IUiSettingsClient, CoreStart } from 'kibana/public';
+import { SearchRequest } from './types';
 
 const sessionId = Date.now();
 
@@ -52,3 +53,18 @@ export function getPreference(config: IUiSettingsClient) {
 export function getTimeout(esShardTimeout: number) {
   return esShardTimeout > 0 ? `${esShardTimeout}ms` : undefined;
 }
+
+export function getSearchParamsFromRequest(
+  searchRequest: SearchRequest,
+  dependencies: { injectedMetadata: CoreStart['injectedMetadata']; uiSettings: IUiSettingsClient }
+) {
+  const { injectedMetadata, uiSettings } = dependencies;
+  const esShardTimeout = injectedMetadata.getInjectedVar('esShardTimeout') as number;
+  const searchParams = getSearchParams(uiSettings, esShardTimeout);
+
+  return {
+    index: searchRequest.index.title || searchRequest.index,
+    body: searchRequest.body,
+    ...searchParams,
+  };
+}

src/plugins/data/public/search/fetch/index.ts

@@ -20,6 +20,7 @@
 export * from './types';
 export {
   getSearchParams,
+  getSearchParamsFromRequest,
   getPreference,
   getTimeout,
   getIgnoreThrottled,

src/plugins/data/public/search/mocks.ts

@@ -31,6 +31,9 @@ const searchStartMock: jest.Mocked<ISearchStart> = {
   aggs: searchAggsStartMock(),
   setInterceptor: jest.fn(),
   search: jest.fn(),
+  searchParams: {
+    createFromRequest: jest.fn(),
+  },
   searchSource: searchSourceMock,
   __LEGACY: {
     esClient: {

src/plugins/data/public/search/search_service.ts

@@ -40,6 +40,7 @@ import {
 } from './aggs';
 import { FieldFormatsStart } from '../field_formats';
 import { ISearchGeneric } from './i_search';
+import { getSearchParamsFromRequest, SearchRequest } from './fetch';
 
 interface SearchServiceSetupDependencies {
   expressions: ExpressionsSetup;
@@ -171,6 +172,13 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
           return new SearchSource({}, searchSourceDependencies);
         },
       },
+      searchParams: {
+        createFromRequest: (searchRequest: SearchRequest) =>
+          getSearchParamsFromRequest(searchRequest, {
+            injectedMetadata: core.injectedMetadata,
+            uiSettings: core.uiSettings,
+          }),
+      },
       setInterceptor: (searchInterceptor: SearchInterceptor) => {
         // TODO: should an intercepror have a destroy method?
         this.searchInterceptor = searchInterceptor;

src/plugins/data/public/search/search_source/search_source.ts

@@ -78,7 +78,7 @@ import { fieldWildcardFilter } from '../../../../kibana_utils/public';
 import { META_FIELDS_SETTING, DOC_HIGHLIGHT_SETTING } from '../../../common';
 import { IIndexPattern, ISearchGeneric, SearchRequest } from '../..';
 import { SearchSourceOptions, SearchSourceFields } from './types';
-import { FetchOptions, RequestFailure, getSearchParams, handleResponse } from '../fetch';
+import { FetchOptions, RequestFailure, handleResponse, getSearchParamsFromRequest } from '../fetch';
 
 import { getEsQueryConfig, buildEsQuery, Filter } from '../../../common';
 import { getHighlightRequest } from '../../../common/field_formats';
@@ -205,13 +205,12 @@ export class SearchSource {
    */
   private fetch$(searchRequest: SearchRequest, signal?: AbortSignal) {
     const { search, injectedMetadata, uiSettings } = this.dependencies;
-    const esShardTimeout = injectedMetadata.getInjectedVar('esShardTimeout') as number;
-    const searchParams = getSearchParams(uiSettings, esShardTimeout);
-    const params = {
-      index: searchRequest.index.title || searchRequest.index,
-      body: searchRequest.body,
-      ...searchParams,
-    };
+
+    const params = getSearchParamsFromRequest(searchRequest, {
+      injectedMetadata,
+      uiSettings,
+    });
+
     return search({ params, indexType: searchRequest.indexType }, { signal }).pipe(
       map(({ rawResponse }) => handleResponse(searchRequest, rawResponse))
     );

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

@@ -18,12 +18,14 @@
  */
 
 import { CoreStart } from 'kibana/public';
+import { SearchParams } from 'elasticsearch';
 import { SearchAggsSetup, SearchAggsStart } from './aggs';
 import { ISearch, ISearchGeneric } from './i_search';
 import { TStrategyTypes } from './strategy_types';
 import { LegacyApiCaller } from './legacy/es_client';
 import { SearchInterceptor } from './search_interceptor';
 import { ISearchSource, SearchSourceFields } from './search_source';
+import { SearchRequest } from './fetch';
 
 export interface ISearchContext {
   core: CoreStart;
@@ -85,5 +87,8 @@ export interface ISearchStart {
     create: (fields?: SearchSourceFields) => Promise<ISearchSource>;
     createEmpty: () => ISearchSource;
   };
+  searchParams: {
+    createFromRequest: (searchRequest: SearchRequest) => SearchParams;
+  };
   __LEGACY: ISearchStartLegacy;
 }

src/plugins/vis_type_vega/public/__mocks__/services.ts

@@ -40,16 +40,13 @@ export const [getSavedObjects, setSavedObjects] = createGetterSetter<SavedObject
 setSavedObjects(coreMock.createStart().savedObjects);
 
 export const [getInjectedVars, setInjectedVars] = createGetterSetter<{
-  esShardTimeout: number;
   enableExternalUrls: boolean;
   emsTileLayerId: unknown;
 }>('InjectedVars');
 setInjectedVars({
   emsTileLayerId: {},
   enableExternalUrls: true,
-  esShardTimeout: 10000,
 });
 
-export const getEsShardTimeout = () => getInjectedVars().esShardTimeout;
 export const getEnableExternalUrls = () => getInjectedVars().enableExternalUrls;
 export const getEmsTileLayerId = () => getInjectedVars().emsTileLayerId;

src/plugins/vis_type_vega/public/data_model/es_query_parser.js

@@ -17,11 +17,9 @@
  * under the License.
  */
 
-import _ from 'lodash';
 import moment from 'moment';
 import { i18n } from '@kbn/i18n';
-
-import { getEsShardTimeout } from '../services';
+import { isPlainObject, cloneDeep } from 'lodash';
 
 const TIMEFILTER = '%timefilter%';
 const AUTOINTERVAL = '%autointerval%';
@@ -37,12 +35,11 @@ const TIMEFIELD = '%timefield%';
  * This class parses ES requests specified in the data.url objects.
  */
 export class EsQueryParser {
-  constructor(timeCache, searchCache, filters, onWarning) {
+  constructor(timeCache, searchAPI, filters, onWarning) {
     this._timeCache = timeCache;
-    this._searchCache = searchCache;
+    this._searchAPI = searchAPI;
     this._filters = filters;
     this._onWarning = onWarning;
-    this._esShardTimeout = getEsShardTimeout();
   }
 
   // noinspection JSMethodCanBeStatic
@@ -59,7 +56,7 @@ export class EsQueryParser {
 
     if (body === undefined) {
       url.body = body = {};
-    } else if (!_.isPlainObject(body)) {
+    } else if (!isPlainObject(body)) {
       throw new Error(
         i18n.translate('visTypeVega.esQueryParser.urlBodyValueTypeErrorMessage', {
           defaultMessage: '{configName} must be an object',
@@ -167,7 +164,7 @@ export class EsQueryParser {
 
       if (context) {
         // Use dashboard context
-        const newQuery = _.cloneDeep(this._filters);
+        const newQuery = cloneDeep(this._filters);
         if (timefield) {
           newQuery.bool.must.push(body.query);
         }
@@ -179,34 +176,20 @@ export class EsQueryParser {
     return { dataObject, url };
   }
 
-  mapRequest = (request) => {
-    const esRequest = request.url;
-    if (this._esShardTimeout) {
-      // remove possible timeout query param to prevent two conflicting timeout parameters
-      const { body = {}, timeout, ...rest } = esRequest; //eslint-disable-line no-unused-vars
-      body.timeout = `${this._esShardTimeout}ms`;
-      return {
-        body,
-        ...rest,
-      };
-    } else {
-      return esRequest;
-    }
-  };
-
   /**
    * Process items generated by parseUrl()
    * @param {object[]} requests each object is generated by parseUrl()
    * @returns {Promise<void>}
    */
   async populateData(requests) {
-    const esSearches = requests.map(this.mapRequest);
+    const esSearches = requests.map((r) => r.url);
+    const data$ = this._searchAPI.search(esSearches);
 
-    const results = await this._searchCache.search(esSearches);
+    const results = await data$.toPromise();
 
-    for (let i = 0; i < requests.length; i++) {
-      requests[i].dataObject.values = results[i];
-    }
+    results.forEach((data) => {
+      requests[data.id].dataObject.values = data.rawResponse;
+    });
   }
 
   /**
@@ -222,7 +205,7 @@ export class EsQueryParser {
           const item = obj[pos];
           if (isQuery && (item === MUST_CLAUSE || item === MUST_NOT_CLAUSE)) {
             const ctxTag = item === MUST_CLAUSE ? 'must' : 'must_not';
-            const ctx = _.cloneDeep(this._filters);
+            const ctx = cloneDeep(this._filters);
             if (ctx && ctx.bool && ctx.bool[ctxTag]) {
               if (Array.isArray(ctx.bool[ctxTag])) {
                 // replace one value with an array of values

src/plugins/vis_type_vega/public/data_model/es_query_parser.test.js

@@ -94,28 +94,36 @@ describe(`EsQueryParser time`, () => {
 });
 
 describe('EsQueryParser.populateData', () => {
-  let searchStub;
+  let searchApiStub;
+  let data;
   let parser;
 
   beforeEach(() => {
-    searchStub = jest.fn(() => Promise.resolve([{}, {}]));
-    parser = new EsQueryParser({}, { search: searchStub }, undefined, undefined);
+    searchApiStub = {
+      search: jest.fn(() => ({
+        toPromise: jest.fn(() => Promise.resolve(data)),
+      })),
+    };
+    parser = new EsQueryParser({}, searchApiStub, undefined, undefined);
   });
 
   test('should set the timeout for each request', async () => {
+    data = [
+      { id: 0, rawResponse: {} },
+      { id: 1, rawResponse: {} },
+    ];
     await parser.populateData([
       { url: { body: {} }, dataObject: {} },
       { url: { body: {} }, dataObject: {} },
     ]);
-    expect(searchStub.mock.calls[0][0][0].body.timeout).toBe.defined;
+
+    expect(searchApiStub.search.mock.calls[0][0][0].body).toBeDefined();
   });
 
   test('should remove possible timeout parameters on a request', async () => {
-    await parser.populateData([
-      { url: { timeout: '500h', body: { timeout: '500h' } }, dataObject: {} },
-    ]);
-    expect(searchStub.mock.calls[0][0][0].body.timeout).toBe.defined;
-    expect(searchStub.mock.calls[0][0][0].timeout).toBe(undefined);
+    data = [{ id: 0, rawResponse: {} }];
+    await parser.populateData([{ url: { body: { timeout: '500h' } }, dataObject: {} }]);
+    expect(searchApiStub.search.mock.calls[0][0][0].body.timeout).toBeDefined();
   });
 });