[Search service] Add enhanced ES search strategy (elastic#59224)

* Add async search strategy

* Add async search

* Fix async strategy and add tests

* Move types to separate file

* Revert changes to demo search

* Update demo search strategy to use async

* Add async es search strategy

* Return response as rawResponse

* Poll after initial request

* Add cancellation to search strategies

* Add tests

* Simplify async search strategy

* Move loadingCount to search strategy

* Update abort controller library

* Bootstrap

* Abort when the request is aborted

* Add utility and update value suggestions route

* Fix bad merge conflict

* Update tests

* Move to data_enhanced plugin

* Remove bad merge

* Revert switching abort controller libraries

* Revert package.json in lib

* Move to previous abort controller

* Add support for frozen indices

* Fix test to use fake timers to run debounced handlers

* Revert changes to example plugin

* Fix loading bar not going away when cancelling

* Call getSearchStrategy instead of passing  directly

* Add async demo search strategy

* Fix error with setting state

* Update how aborting works

* Fix type checks

* Add test for loading count

* Attempt to fix broken example test

* Revert changes to test

* Fix test

* Update name to camelCase

* Fix failing test

* Don't require data_enhanced in example plugin

* Actually send DELETE request

* Use waitForCompletion parameter

* Use default search params

* Add support for rollups

* Only make changes needed for frozen indices/rollups

* Fix tests/types

* Don't include skipped in loaded/total

Co-authored-by: Elastic Machine <[email protected]>

src/plugins/data/common/index.ts

@@ -24,4 +24,5 @@ export * from './index_patterns';
 export * from './es_query';
 export * from './utils';
 export * from './types';
+export * from './search';
 export * from './constants';

src/plugins/data/public/search/es_search/es_search_strategy.ts

@@ -30,11 +30,10 @@ export const esSearchStrategyProvider: TSearchStrategyProvider<typeof ES_SEARCH_
   const { search } = syncStrategyProvider(context);
   return {
     search: (request, options) => {
-      if (typeof request.params.preference === 'undefined') {
-        const setPreference = context.core.uiSettings.get('courier:setRequestPreference');
-        const customPreference = context.core.uiSettings.get('courier:customRequestPreference');
-        request.params.preference = getEsPreference(setPreference, customPreference);
-      }
+      request.params = {
+        preference: getEsPreference(context.core.uiSettings),
+        ...request.params,
+      };
       return search({ ...request, serverStrategy: ES_SEARCH_STRATEGY }, options) as Observable<
         IEsSearchResponse
       >;

src/plugins/data/public/search/es_search/get_es_preference.test.ts

@@ -18,29 +18,40 @@
  */
 
 import { getEsPreference } from './get_es_preference';
-
-jest.useFakeTimers();
+import { CoreStart } from '../../../../../core/public';
+import { coreMock } from '../../../../../core/public/mocks';
 
 describe('Get ES preference', () => {
+  let mockCoreStart: MockedKeys<CoreStart>;
+
+  beforeEach(() => {
+    mockCoreStart = coreMock.createStart();
+  });
+
   test('returns the session ID if set to sessionId', () => {
-    const setPreference = 'sessionId';
-    const customPreference = 'foobar';
-    const sessionId = 'my_session_id';
-    const preference = getEsPreference(setPreference, customPreference, sessionId);
-    expect(preference).toBe(sessionId);
+    mockCoreStart.uiSettings.get.mockImplementation((key: string) => {
+      if (key === 'courier:setRequestPreference') return 'sessionId';
+      if (key === 'courier:customRequestPreference') return 'foobar';
+    });
+    const preference = getEsPreference(mockCoreStart.uiSettings, 'my_session_id');
+    expect(preference).toBe('my_session_id');
   });
 
   test('returns the custom preference if set to custom', () => {
-    const setPreference = 'custom';
-    const customPreference = 'foobar';
-    const preference = getEsPreference(setPreference, customPreference);
-    expect(preference).toBe(customPreference);
+    mockCoreStart.uiSettings.get.mockImplementation((key: string) => {
+      if (key === 'courier:setRequestPreference') return 'custom';
+      if (key === 'courier:customRequestPreference') return 'foobar';
+    });
+    const preference = getEsPreference(mockCoreStart.uiSettings);
+    expect(preference).toBe('foobar');
   });
 
   test('returns undefined if set to none', () => {
-    const setPreference = 'none';
-    const customPreference = 'foobar';
-    const preference = getEsPreference(setPreference, customPreference);
+    mockCoreStart.uiSettings.get.mockImplementation((key: string) => {
+      if (key === 'courier:setRequestPreference') return 'none';
+      if (key === 'courier:customRequestPreference') return 'foobar';
+    });
+    const preference = getEsPreference(mockCoreStart.uiSettings);
     expect(preference).toBe(undefined);
   });
 });

src/plugins/data/public/search/es_search/get_es_preference.ts

@@ -17,13 +17,13 @@
  * under the License.
  */
 
+import { IUiSettingsClient } from '../../../../../core/public';
+
 const defaultSessionId = `${Date.now()}`;
 
-export function getEsPreference(
-  setRequestPreference: string,
-  customRequestPreference?: string,
-  sessionId: string = defaultSessionId
-) {
-  if (setRequestPreference === 'sessionId') return `${sessionId}`;
-  return setRequestPreference === 'custom' ? customRequestPreference : undefined;
+export function getEsPreference(uiSettings: IUiSettingsClient, sessionId = defaultSessionId) {
+  const setPreference = uiSettings.get('courier:setRequestPreference');
+  if (setPreference === 'sessionId') return `${sessionId}`;
+  const customPreference = uiSettings.get('courier:customRequestPreference');
+  return setPreference === 'custom' ? customPreference : undefined;
 }

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

@@ -0,0 +1,21 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+export { esSearchStrategyProvider } from './es_search_strategy';
+export { getEsPreference } from './get_es_preference';

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

@@ -36,6 +36,7 @@ export {
 export { IEsSearchResponse, IEsSearchRequest, ES_SEARCH_STRATEGY } from '../../common/search';
 
 export { ISyncSearchRequest, SYNC_SEARCH_STRATEGY } from './sync_search_strategy';
+export { esSearchStrategyProvider, getEsPreference } from './es_search';
 
 export { IKibanaSearchResponse, IKibanaSearchRequest } from '../../common/search';
 

src/plugins/data/server/index.ts

@@ -151,8 +151,16 @@ export {
  * Search
  */
 
-export { IRequestTypesMap, IResponseTypesMap } from './search';
-export * from './search';
+export {
+  ISearch,
+  ICancel,
+  ISearchOptions,
+  IRequestTypesMap,
+  IResponseTypesMap,
+  ISearchContext,
+  TSearchStrategyProvider,
+  getDefaultSearchParams,
+} from './search';
 
 /**
  * Types to be shared externally

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

@@ -23,8 +23,6 @@ import { TSearchStrategiesMap } from './i_search_strategy';
 import { IRouteHandlerSearchContext } from './i_route_handler_search_context';
 import { DEFAULT_SEARCH_STRATEGY } from '../../common/search';
 
-// let mockCoreSetup: MockedKeys<CoreSetup>;
-
 const mockDefaultSearch = jest.fn(() => Promise.resolve({ total: 100, loaded: 0 }));
 const mockDefaultSearchStrategyProvider = jest.fn(() =>
   Promise.resolve({
@@ -59,4 +57,15 @@ describe('createApi', () => {
       `"No strategy found for noneByThisName"`
     );
   });
+
+  it('logs the response if `debug` is set to `true`', async () => {
+    const spy = jest.spyOn(console, 'log');
+    await api.search({ params: {} });
+
+    expect(spy).not.toBeCalled();
+
+    await api.search({ debug: true, params: {} });
+
+    expect(spy).toBeCalled();
+  });
 });

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

@@ -31,6 +31,10 @@ export function createApi({
 }) {
   const api: IRouteHandlerSearchContext = {
     search: async (request, options, strategyName) => {
+      if (request.debug) {
+        // eslint-disable-next-line
+        console.log(JSON.stringify(request, null, 2));
+      }
       const name = strategyName ?? DEFAULT_SEARCH_STRATEGY;
       const strategyProvider = searchStrategies[name];
       if (!strategyProvider) {

src/plugins/data/server/search/es_search/es_search_strategy.test.ts

@@ -51,24 +51,6 @@ describe('ES search strategy', () => {
     expect(typeof esSearch.search).toBe('function');
   });
 
-  it('logs the response if `debug` is set to `true`', async () => {
-    const spy = jest.spyOn(console, 'log');
-    const esSearch = esSearchStrategyProvider(
-      {
-        core: mockCoreSetup,
-        config$: mockConfig$,
-      },
-      mockApiCaller,
-      mockSearch
-    );
-
-    expect(spy).not.toBeCalled();
-
-    await esSearch.search({ params: {}, debug: true });
-
-    expect(spy).toBeCalled();
-  });
-
   it('calls the API caller with the params with defaults', async () => {
     const params = { index: 'logstash-*' };
     const esSearch = esSearchStrategyProvider(

src/plugins/data/server/search/es_search/es_search_strategy.ts

@@ -21,7 +21,7 @@ import { APICaller } from 'kibana/server';
 import { SearchResponse } from 'elasticsearch';
 import { ES_SEARCH_STRATEGY } from '../../../common/search';
 import { ISearchStrategy, TSearchStrategyProvider } from '../i_search_strategy';
-import { ISearchContext } from '..';
+import { getDefaultSearchParams, ISearchContext } from '..';
 
 export const esSearchStrategyProvider: TSearchStrategyProvider<typeof ES_SEARCH_STRATEGY> = (
   context: ISearchContext,
@@ -30,28 +30,18 @@ export const esSearchStrategyProvider: TSearchStrategyProvider<typeof ES_SEARCH_
   return {
     search: async (request, options) => {
       const config = await context.config$.pipe(first()).toPromise();
+      const defaultParams = getDefaultSearchParams(config);
       const params = {
-        timeout: `${config.elasticsearch.shardTimeout.asMilliseconds()}ms`,
-        ignoreUnavailable: true, // Don't fail if the index/indices don't exist
-        restTotalHitsAsInt: true, // Get the number of hits as an int rather than a range
+        ...defaultParams,
         ...request.params,
       };
-      if (request.debug) {
-        // eslint-disable-next-line
-        console.log(JSON.stringify(params, null, 2));
-      }
-      const esSearchResponse = (await caller('search', params, options)) as SearchResponse<any>;
+      const rawResponse = (await caller('search', params, options)) as SearchResponse<any>;
 
       // The above query will either complete or timeout and throw an error.
       // There is no progress indication on this api.
-      return {
-        total: esSearchResponse._shards.total,
-        loaded:
-          esSearchResponse._shards.failed +
-          esSearchResponse._shards.skipped +
-          esSearchResponse._shards.successful,
-        rawResponse: esSearchResponse,
-      };
+      const { total, failed, successful } = rawResponse._shards;
+      const loaded = failed + successful;
+      return { total, loaded, rawResponse };
     },
   };
 };

src/plugins/data/server/search/es_search/get_default_search_params.ts

@@ -0,0 +1,28 @@
+/*
+ * Licensed to Elasticsearch B.V. under one or more contributor
+ * license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright
+ * ownership. Elasticsearch B.V. licenses this file to you under
+ * the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import { SharedGlobalConfig } from '../../../../../core/server';
+
+export function getDefaultSearchParams(config: SharedGlobalConfig) {
+  return {
+    timeout: `${config.elasticsearch.shardTimeout.asMilliseconds()}ms`,
+    ignoreUnavailable: true, // Don't fail if the index/indices don't exist
+    restTotalHitsAsInt: true, // Get the number of hits as an int rather than a range
+  };
+}

src/plugins/data/server/search/es_search/index.ts

@@ -17,6 +17,6 @@
  * under the License.
  */
 
-export { esSearchStrategyProvider } from './es_search_strategy';
-
 export { ES_SEARCH_STRATEGY, IEsSearchRequest, IEsSearchResponse } from '../../../common/search';
+export { esSearchStrategyProvider } from './es_search_strategy';
+export { getDefaultSearchParams } from './get_default_search_params';

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

@@ -21,8 +21,10 @@ export { ISearchSetup } from './i_search_setup';
 
 export { ISearchContext } from './i_search_context';
 
-export { IRequestTypesMap, IResponseTypesMap } from './i_search';
+export { ISearch, ICancel, ISearchOptions, IRequestTypesMap, IResponseTypesMap } from './i_search';
 
 export { TStrategyTypes } from './strategy_types';
 
 export { TSearchStrategyProvider } from './i_search_strategy';
+
+export { getDefaultSearchParams } from './es_search';

x-pack/plugins/data_enhanced/common/index.ts

@@ -0,0 +1,7 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+export { EnhancedSearchParams, IEnhancedEsSearchRequest } from './search';

x-pack/plugins/data_enhanced/common/search/index.ts

@@ -0,0 +1,7 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+export { EnhancedSearchParams, IEnhancedEsSearchRequest } from './types';

x-pack/plugins/data_enhanced/common/search/types.ts

@@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+
+import { SearchParams } from 'elasticsearch';
+import { IEsSearchRequest } from '../../../../../src/plugins/data/common';
+
+export interface EnhancedSearchParams extends SearchParams {
+  ignoreThrottled: boolean;
+}
+
+export interface IEnhancedEsSearchRequest extends IEsSearchRequest {
+  /**
+   * Used to determine whether to use the _rollups_search or a regular search endpoint.
+   */
+  isRollup?: boolean;
+}

x-pack/plugins/data_enhanced/kibana.json

@@ -8,6 +8,6 @@
   "requiredPlugins": [
     "data"
   ],
-  "server": false,
+  "server": true,
   "ui": true
 }

x-pack/plugins/data_enhanced/public/plugin.ts

@@ -5,10 +5,18 @@
  */
 
 import { CoreSetup, CoreStart, Plugin } from 'src/core/public';
-import { DataPublicPluginSetup, DataPublicPluginStart } from '../../../../src/plugins/data/public';
+import {
+  DataPublicPluginSetup,
+  DataPublicPluginStart,
+  ES_SEARCH_STRATEGY,
+} from '../../../../src/plugins/data/public';
 import { setAutocompleteService } from './services';
 import { setupKqlQuerySuggestionProvider, KUERY_LANGUAGE_NAME } from './autocomplete';
-import { ASYNC_SEARCH_STRATEGY, asyncSearchStrategyProvider } from './search';
+import {
+  ASYNC_SEARCH_STRATEGY,
+  asyncSearchStrategyProvider,
+  enhancedEsSearchStrategyProvider,
+} from './search';
 
 export interface DataEnhancedSetupDependencies {
   data: DataPublicPluginSetup;
@@ -29,6 +37,10 @@ export class DataEnhancedPlugin implements Plugin {
       setupKqlQuerySuggestionProvider(core)
     );
     data.search.registerSearchStrategyProvider(ASYNC_SEARCH_STRATEGY, asyncSearchStrategyProvider);
+    data.search.registerSearchStrategyProvider(
+      ES_SEARCH_STRATEGY,
+      enhancedEsSearchStrategyProvider
+    );
   }
 
   public start(core: CoreStart, plugins: DataEnhancedStartDependencies) {