[Search] Share async strategy utils (#128358)

src/plugins/data/server/search/strategies/common/async_utils.test.ts

@@ -0,0 +1,110 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { getCommonDefaultAsyncSubmitParams, getCommonDefaultAsyncGetParams } from './async_utils';
+import moment from 'moment';
+import { SearchSessionsConfigSchema } from '../../../../config';
+
+const getMockSearchSessionsConfig = ({
+  enabled = true,
+  defaultExpiration = moment.duration(7, 'd'),
+} = {}) =>
+  ({
+    enabled,
+    defaultExpiration,
+  } as SearchSessionsConfigSchema);
+
+describe('request utils', () => {
+  describe('getCommonDefaultAsyncSubmitParams', () => {
+    test('Uses `keep_alive` from default params if no `sessionId` is provided', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+      });
+      const params = getCommonDefaultAsyncSubmitParams(mockConfig, {});
+      expect(params).toHaveProperty('keep_alive', '1m');
+    });
+
+    test('Uses `keep_alive` from config if enabled', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+      });
+      const params = getCommonDefaultAsyncSubmitParams(mockConfig, {
+        sessionId: 'foo',
+      });
+      expect(params).toHaveProperty('keep_alive', '259200000ms');
+    });
+
+    test('Uses `keepAlive` of `1m` if disabled', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+        enabled: false,
+      });
+      const params = getCommonDefaultAsyncSubmitParams(mockConfig, {
+        sessionId: 'foo',
+      });
+      expect(params).toHaveProperty('keep_alive', '1m');
+    });
+
+    test('Uses `keep_on_completion` if enabled', async () => {
+      const mockConfig = getMockSearchSessionsConfig({});
+      const params = getCommonDefaultAsyncSubmitParams(mockConfig, {
+        sessionId: 'foo',
+      });
+      expect(params).toHaveProperty('keep_on_completion', true);
+    });
+
+    test('Does not use `keep_on_completion` if disabled', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+        enabled: false,
+      });
+      const params = getCommonDefaultAsyncSubmitParams(mockConfig, {
+        sessionId: 'foo',
+      });
+      expect(params).toHaveProperty('keep_on_completion', false);
+    });
+  });
+
+  describe('getCommonDefaultAsyncGetParams', () => {
+    test('Uses `wait_for_completion_timeout`', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+        enabled: true,
+      });
+      const params = getCommonDefaultAsyncGetParams(mockConfig, {});
+      expect(params).toHaveProperty('wait_for_completion_timeout');
+    });
+
+    test('Uses `keep_alive` if `sessionId` is not provided', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+        enabled: true,
+      });
+      const params = getCommonDefaultAsyncGetParams(mockConfig, {});
+      expect(params).toHaveProperty('keep_alive', '1m');
+    });
+
+    test('Has no `keep_alive` if `sessionId` is provided', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+        enabled: true,
+      });
+      const params = getCommonDefaultAsyncGetParams(mockConfig, { sessionId: 'foo' });
+      expect(params).not.toHaveProperty('keep_alive');
+    });
+
+    test('Uses `keep_alive` if `sessionId` is provided but sessions disabled', async () => {
+      const mockConfig = getMockSearchSessionsConfig({
+        defaultExpiration: moment.duration(3, 'd'),
+        enabled: false,
+      });
+      const params = getCommonDefaultAsyncGetParams(mockConfig, { sessionId: 'foo' });
+      expect(params).toHaveProperty('keep_alive', '1m');
+    });
+  });
+});

src/plugins/data/server/search/strategies/common/async_utils.ts

@@ -0,0 +1,62 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import {
+  AsyncSearchSubmitRequest,
+  AsyncSearchGetRequest,
+} from '@elastic/elasticsearch/lib/api/types';
+import { SearchSessionsConfigSchema } from '../../../../config';
+import { ISearchOptions } from '../../../../common';
+
+/**
+ @internal
+ */
+export function getCommonDefaultAsyncSubmitParams(
+  searchSessionsConfig: SearchSessionsConfigSchema | null,
+  options: ISearchOptions
+): Pick<
+  AsyncSearchSubmitRequest,
+  'keep_alive' | 'wait_for_completion_timeout' | 'keep_on_completion'
+> {
+  const useSearchSessions = searchSessionsConfig?.enabled && !!options.sessionId;
+
+  const keepAlive = useSearchSessions
+    ? `${searchSessionsConfig!.defaultExpiration.asMilliseconds()}ms`
+    : '1m';
+
+  return {
+    // Wait up to 100ms for the response to return
+    wait_for_completion_timeout: '100ms',
+    // If search sessions are used, store and get an async ID even for short running requests.
+    keep_on_completion: useSearchSessions,
+    // The initial keepalive is as defined in defaultExpiration if search sessions are used or 1m otherwise.
+    keep_alive: keepAlive,
+  };
+}
+
+/**
+ @internal
+ */
+export function getCommonDefaultAsyncGetParams(
+  searchSessionsConfig: SearchSessionsConfigSchema | null,
+  options: ISearchOptions
+): Pick<AsyncSearchGetRequest, 'keep_alive' | 'wait_for_completion_timeout'> {
+  const useSearchSessions = searchSessionsConfig?.enabled && !!options.sessionId;
+
+  return {
+    // Wait up to 100ms for the response to return
+    wait_for_completion_timeout: '100ms',
+    ...(useSearchSessions
+      ? // Don't change the expiration of search requests that are tracked in a search session
+        undefined
+      : {
+          // We still need to do polling for searches not within the context of a search session or when search session disabled
+          keep_alive: '1m',
+        }),
+  };
+}

src/plugins/data/server/search/strategies/eql_search/eql_search_strategy.ts

@@ -19,7 +19,8 @@ import { toEqlKibanaSearchResponse } from './response_utils';
 import { EqlSearchResponse } from './types';
 import { ISearchStrategy } from '../../types';
 import { getDefaultSearchParams } from '../es_search';
-import { getDefaultAsyncGetParams, getIgnoreThrottled } from '../ese_search/request_utils';
+import { getIgnoreThrottled } from '../ese_search/request_utils';
+import { getCommonDefaultAsyncGetParams } from '../common/async_utils';
 
 export const eqlSearchStrategyProvider = (
   logger: Logger
@@ -45,11 +46,11 @@ export const eqlSearchStrategyProvider = (
           uiSettingsClient
         );
         const params = id
-          ? getDefaultAsyncGetParams(null, options)
+          ? getCommonDefaultAsyncGetParams(null, options)
           : {
               ...(await getIgnoreThrottled(uiSettingsClient)),
               ...defaultParams,
-              ...getDefaultAsyncGetParams(null, options),
+              ...getCommonDefaultAsyncGetParams(null, options),
               ...request.params,
             };
         const response = id

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

@@ -12,6 +12,10 @@ import { AsyncSearchSubmitRequest } from '@elastic/elasticsearch/lib/api/types';
 import { ISearchOptions, UI_SETTINGS } from '../../../../common';
 import { getDefaultSearchParams } from '../es_search';
 import { SearchSessionsConfigSchema } from '../../../../config';
+import {
+  getCommonDefaultAsyncGetParams,
+  getCommonDefaultAsyncSubmitParams,
+} from '../common/async_utils';
 
 /**
  * @internal
@@ -43,23 +47,10 @@ export async function getDefaultAsyncSubmitParams(
     | 'keep_on_completion'
   >
 > {
-  const useSearchSessions = searchSessionsConfig?.enabled && !!options.sessionId;
-
-  // TODO: searchSessionsConfig could be "null" if we are running without x-pack which happens only in tests.
-  // This can be cleaned up when we completely stop separating basic and oss
-  const keepAlive = useSearchSessions
-    ? `${searchSessionsConfig!.defaultExpiration.asMilliseconds()}ms`
-    : '1m';
-
   return {
     // TODO: adjust for partial results
     batched_reduce_size: 64,
-    // Wait up to 100ms for the response to return
-    wait_for_completion_timeout: '100ms',
-    // If search sessions are used, store and get an async ID even for short running requests.
-    keep_on_completion: useSearchSessions,
-    // The initial keepalive is as defined in defaultExpiration if search sessions are used or 1m otherwise.
-    keep_alive: keepAlive,
+    ...getCommonDefaultAsyncSubmitParams(searchSessionsConfig, options),
     ...(await getIgnoreThrottled(uiSettingsClient)),
     ...(await getDefaultSearchParams(uiSettingsClient)),
     // If search sessions are used, set the initial expiration time.
@@ -73,17 +64,7 @@ export function getDefaultAsyncGetParams(
   searchSessionsConfig: SearchSessionsConfigSchema | null,
   options: ISearchOptions
 ): Pick<AsyncSearchGetRequest, 'keep_alive' | 'wait_for_completion_timeout'> {
-  const useSearchSessions = searchSessionsConfig?.enabled && !!options.sessionId;
-
   return {
-    // Wait up to 100ms for the response to return
-    wait_for_completion_timeout: '100ms',
-    ...(useSearchSessions
-      ? // Don't change the expiration of search requests that are tracked in a search session
-        undefined
-      : {
-          // We still need to do polling for searches not within the context of a search session or when search session disabled
-          keep_alive: '1m',
-        }),
+    ...getCommonDefaultAsyncGetParams(searchSessionsConfig, options),
   };
 }

src/plugins/data/server/search/strategies/sql_search/request_utils.ts

@@ -9,6 +9,10 @@
 import { SqlGetAsyncRequest, SqlQueryRequest } from '@elastic/elasticsearch/lib/api/types';
 import { ISearchOptions } from '../../../../common';
 import { SearchSessionsConfigSchema } from '../../../../config';
+import {
+  getCommonDefaultAsyncGetParams,
+  getCommonDefaultAsyncSubmitParams,
+} from '../common/async_utils';
 
 /**
  @internal
@@ -17,19 +21,8 @@ export function getDefaultAsyncSubmitParams(
   searchSessionsConfig: SearchSessionsConfigSchema | null,
   options: ISearchOptions
 ): Pick<SqlQueryRequest, 'keep_alive' | 'wait_for_completion_timeout' | 'keep_on_completion'> {
-  const useSearchSessions = searchSessionsConfig?.enabled && !!options.sessionId;
-
-  const keepAlive = useSearchSessions
-    ? `${searchSessionsConfig!.defaultExpiration.asMilliseconds()}ms`
-    : '1m';
-
   return {
-    // Wait up to 100ms for the response to return
-    wait_for_completion_timeout: '100ms',
-    // If search sessions are used, store and get an async ID even for short running requests.
-    keep_on_completion: useSearchSessions,
-    // The initial keepalive is as defined in defaultExpiration if search sessions are used or 1m otherwise.
-    keep_alive: keepAlive,
+    ...getCommonDefaultAsyncSubmitParams(searchSessionsConfig, options),
   };
 }
 
@@ -40,17 +33,7 @@ export function getDefaultAsyncGetParams(
   searchSessionsConfig: SearchSessionsConfigSchema | null,
   options: ISearchOptions
 ): Pick<SqlGetAsyncRequest, 'keep_alive' | 'wait_for_completion_timeout'> {
-  const useSearchSessions = searchSessionsConfig?.enabled && !!options.sessionId;
-
   return {
-    // Wait up to 100ms for the response to return
-    wait_for_completion_timeout: '100ms',
-    ...(useSearchSessions
-      ? // Don't change the expiration of search requests that are tracked in a search session
-        undefined
-      : {
-          // We still need to do polling for searches not within the context of a search session or when search session disabled
-          keep_alive: '1m',
-        }),
+    ...getCommonDefaultAsyncGetParams(searchSessionsConfig, options),
   };
 }