Fix overriding of keep_alive and wait_for_completion_timeout

src/plugins/data/common/search/strategies/es_search/types.ts

@@ -15,7 +15,8 @@ export type ISearchRequestParams = {
   trackTotalHits?: boolean;
 } & estypes.SearchRequest;
 
-export interface IEsSearchRequest extends IKibanaSearchRequest<ISearchRequestParams> {
+export interface IEsSearchRequest<T extends ISearchRequestParams = ISearchRequestParams>
+  extends IKibanaSearchRequest<T> {
   indexType?: string;
 }
 

src/plugins/data/server/search/strategies/ese_search/ese_search_strategy.test.ts

@@ -136,6 +136,30 @@ describe('ES search strategy', () => {
         expect(request).toHaveProperty('keep_alive', '60000ms');
       });
 
+      it('allows overriding keep_alive and wait_for_completion_timeout', async () => {
+        mockGetCaller.mockResolvedValueOnce(mockAsyncResponse);
+
+        const params = {
+          index: 'logstash-*',
+          body: { query: {} },
+          wait_for_completion_timeout: '10s',
+          keep_alive: '5m',
+        };
+        const esSearch = await enhancedEsSearchStrategyProvider(
+          mockLegacyConfig$,
+          mockSearchConfig,
+          mockLogger
+        );
+
+        await esSearch.search({ id: 'foo', params }, {}, mockDeps).toPromise();
+
+        expect(mockGetCaller).toBeCalled();
+        const request = mockGetCaller.mock.calls[0][0];
+        expect(request.id).toEqual('foo');
+        expect(request).toHaveProperty('wait_for_completion_timeout', '10s');
+        expect(request).toHaveProperty('keep_alive', '5m');
+      });
+
       it('sets transport options on POST requests', async () => {
         const transportOptions = { maxRetries: 1 };
         mockSubmitCaller.mockResolvedValueOnce(mockAsyncResponse);

src/plugins/data/server/search/strategies/ese_search/ese_search_strategy.ts

@@ -12,6 +12,7 @@ import { catchError, tap } from 'rxjs/operators';
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
 import { firstValueFrom, from } from 'rxjs';
 import { getKbnServerError } from '@kbn/kibana-utils-plugin/server';
+import { IAsyncSearchRequestParams } from '../..';
 import { getKbnSearchError, KbnSearchError } from '../../report_search_error';
 import type { ISearchStrategy, SearchStrategyDependencies } from '../../types';
 import type {
@@ -43,14 +44,14 @@ export const enhancedEsSearchStrategyProvider = (
   logger: Logger,
   usage?: SearchUsage,
   useInternalUser: boolean = false
-): ISearchStrategy => {
+): ISearchStrategy<IEsSearchRequest<IAsyncSearchRequestParams>> => {
   function cancelAsyncSearch(id: string, esClient: IScopedClusterClient) {
     const client = useInternalUser ? esClient.asInternalUser : esClient.asCurrentUser;
     return client.asyncSearch.delete({ id });
   }
 
   function asyncSearch(
-    { id, ...request }: IEsSearchRequest,
+    { id, ...request }: IEsSearchRequest<IAsyncSearchRequestParams>,
     options: IAsyncSearchOptions,
     { esClient, uiSettingsClient }: SearchStrategyDependencies
   ) {
@@ -60,7 +61,10 @@ export const enhancedEsSearchStrategyProvider = (
       const params = id
         ? {
             ...getDefaultAsyncGetParams(searchConfig, options),
-            ...request.params,
+            ...(request.params?.keep_alive ? { keep_alive: request.params.keep_alive } : {}),
+            ...(request.params?.wait_for_completion_timeout
+              ? { wait_for_completion_timeout: request.params.wait_for_completion_timeout }
+              : {}),
           }
         : {
             ...(await getDefaultAsyncSubmitParams(uiSettingsClient, searchConfig, options)),

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

@@ -6,17 +6,27 @@
  * Side Public License, v 1.
  */
 
-import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import type {
+  AsyncSearchGetRequest,
+  SearchResponse,
+  ShardStatistics,
+} from '@elastic/elasticsearch/lib/api/types';
+import { ISearchRequestParams } from '../../../../common';
+
+export interface IAsyncSearchRequestParams extends ISearchRequestParams {
+  keep_alive?: AsyncSearchGetRequest['keep_alive'];
+  wait_for_completion_timeout?: AsyncSearchGetRequest['wait_for_completion_timeout'];
+}
 
 export interface AsyncSearchResponse<T = unknown> {
   id?: string;
-  response: estypes.SearchResponse<T>;
+  response: SearchResponse<T>;
   start_time_in_millis: number;
   expiration_time_in_millis: number;
   is_partial: boolean;
   is_running: boolean;
 }
 export interface AsyncSearchStatusResponse extends Omit<AsyncSearchResponse, 'response'> {
   completion_status: number;
-  _shards: estypes.ShardStatistics;
+  _shards: ShardStatistics;
 }