[data.search.session] Add extend functionality to server-side search service (#86195) (#87320)

* [data.search.session] Store search strategy in saved object

* [data.search.session] Add extend functionality

* Update docs

* Update unit test to check strategy

* Throw kbnServerError instead of error

* Fix test

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

docs/development/plugins/data/server/kibana-plugin-plugins-data-server.isearchstrategy.extend.md

@@ -0,0 +1,11 @@
+<!-- Do not edit this file. It is automatically generated by API Documenter. -->
+
+[Home](./index.md) &gt; [kibana-plugin-plugins-data-server](./kibana-plugin-plugins-data-server.md) &gt; [ISearchStrategy](./kibana-plugin-plugins-data-server.isearchstrategy.md) &gt; [extend](./kibana-plugin-plugins-data-server.isearchstrategy.extend.md)
+
+## ISearchStrategy.extend property
+
+<b>Signature:</b>
+
+```typescript
+extend?: (id: string, keepAlive: string, options: ISearchOptions, deps: SearchStrategyDependencies) => Promise<void>;
+```

docs/development/plugins/data/server/kibana-plugin-plugins-data-server.isearchstrategy.md

@@ -17,5 +17,6 @@ export interface ISearchStrategy<SearchStrategyRequest extends IKibanaSearchRequ
 |  Property | Type | Description |
 |  --- | --- | --- |
 |  [cancel](./kibana-plugin-plugins-data-server.isearchstrategy.cancel.md) | <code>(id: string, options: ISearchOptions, deps: SearchStrategyDependencies) =&gt; Promise&lt;void&gt;</code> |  |
+|  [extend](./kibana-plugin-plugins-data-server.isearchstrategy.extend.md) | <code>(id: string, keepAlive: string, options: ISearchOptions, deps: SearchStrategyDependencies) =&gt; Promise&lt;void&gt;</code> |  |
 |  [search](./kibana-plugin-plugins-data-server.isearchstrategy.search.md) | <code>(request: SearchStrategyRequest, options: ISearchOptions, deps: SearchStrategyDependencies) =&gt; Observable&lt;SearchStrategyResponse&gt;</code> |  |
 

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

@@ -29,10 +29,22 @@ export type ISearchGeneric = <
 ) => Observable<SearchStrategyResponse>;
 
 export type ISearchCancelGeneric = (id: string, options?: ISearchOptions) => Promise<void>;
+export type ISearchExtendGeneric = (
+  id: string,
+  keepAlive: string,
+  options?: ISearchOptions
+) => Promise<void>;
 
 export interface ISearchClient {
   search: ISearchGeneric;
+  /**
+   * Used to cancel an in-progress search request.
+   */
   cancel: ISearchCancelGeneric;
+  /**
+   * Used to extend the TTL of an in-progress search request.
+   */
+  extend: ISearchExtendGeneric;
 }
 
 export interface IKibanaSearchResponse<RawResponse = any> {

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

@@ -315,6 +315,19 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     return strategy.cancel(id, options, deps);
   };
 
+  private extend = (
+    id: string,
+    keepAlive: string,
+    options: ISearchOptions,
+    deps: SearchStrategyDependencies
+  ) => {
+    const strategy = this.getSearchStrategy(options.strategy);
+    if (!strategy.extend) {
+      throw new KbnServerError(`Search strategy ${options.strategy} does not support extend`, 400);
+    }
+    return strategy.extend(id, keepAlive, options, deps);
+  };
+
   private getSearchStrategy = <
     SearchStrategyRequest extends IKibanaSearchRequest = IEsSearchRequest,
     SearchStrategyResponse extends IKibanaSearchResponse = IEsSearchResponse
@@ -344,6 +357,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
         search: (searchRequest, options = {}) =>
           this.search(scopedSession, searchRequest, options, deps),
         cancel: (id, options = {}) => this.cancel(id, options, deps),
+        extend: (id, keepAlive, options = {}) => this.extend(id, keepAlive, options, deps),
       };
     };
   };

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

@@ -86,6 +86,12 @@ export interface ISearchStrategy<
     deps: SearchStrategyDependencies
   ) => Observable<SearchStrategyResponse>;
   cancel?: (id: string, options: ISearchOptions, deps: SearchStrategyDependencies) => Promise<void>;
+  extend?: (
+    id: string,
+    keepAlive: string,
+    options: ISearchOptions,
+    deps: SearchStrategyDependencies
+  ) => Promise<void>;
 }
 
 export interface ISearchStart<

src/plugins/data/server/server.api.md

@@ -956,6 +956,8 @@ export interface ISearchStrategy<SearchStrategyRequest extends IKibanaSearchRequ
     // (undocumented)
     cancel?: (id: string, options: ISearchOptions, deps: SearchStrategyDependencies) => Promise<void>;
     // (undocumented)
+    extend?: (id: string, keepAlive: string, options: ISearchOptions, deps: SearchStrategyDependencies) => Promise<void>;
+    // (undocumented)
     search: (request: SearchStrategyRequest, options: ISearchOptions, deps: SearchStrategyDependencies) => Observable<SearchStrategyResponse>;
 }
 
@@ -1434,7 +1436,7 @@ export function usageProvider(core: CoreSetup_2): SearchUsage;
 // src/plugins/data/server/index.ts:279:1 - (ae-forgotten-export) The symbol "calcAutoIntervalLessThan" needs to be exported by the entry point index.d.ts
 // src/plugins/data/server/index_patterns/index_patterns_service.ts:70:14 - (ae-forgotten-export) The symbol "IndexPatternsService" needs to be exported by the entry point index.d.ts
 // src/plugins/data/server/plugin.ts:90:74 - (ae-forgotten-export) The symbol "DataEnhancements" needs to be exported by the entry point index.d.ts
-// src/plugins/data/server/search/types.ts:106:5 - (ae-forgotten-export) The symbol "ISearchStartSearchSource" needs to be exported by the entry point index.d.ts
+// src/plugins/data/server/search/types.ts:112:5 - (ae-forgotten-export) The symbol "ISearchStartSearchSource" needs to be exported by the entry point index.d.ts
 
 // (No @packageDocumentation comment for this package)
 

x-pack/plugins/data_enhanced/server/search/eql_search_strategy.ts

@@ -56,5 +56,10 @@ export const eqlSearchStrategyProvider = (
 
       return pollSearch(search, options).pipe(tap((response) => (id = response.id)));
     },
+
+    extend: async (id, keepAlive, options, { esClient }) => {
+      logger.debug(`_eql/extend ${id} by ${keepAlive}`);
+      await esClient.asCurrentUser.eql.get({ id, keep_alive: keepAlive });
+    },
   };
 };

x-pack/plugins/data_enhanced/server/search/es_search_strategy.test.ts

@@ -37,6 +37,7 @@ describe('ES search strategy', () => {
   const mockApiCaller = jest.fn();
   const mockGetCaller = jest.fn();
   const mockSubmitCaller = jest.fn();
+  const mockDeleteCaller = jest.fn();
   const mockLogger: any = {
     debug: () => {},
   };
@@ -49,6 +50,7 @@ describe('ES search strategy', () => {
         asyncSearch: {
           get: mockGetCaller,
           submit: mockSubmitCaller,
+          delete: mockDeleteCaller,
         },
         transport: { request: mockApiCaller },
       },
@@ -66,77 +68,113 @@ describe('ES search strategy', () => {
 
   beforeEach(() => {
     mockApiCaller.mockClear();
+    mockGetCaller.mockClear();
+    mockSubmitCaller.mockClear();
+    mockDeleteCaller.mockClear();
   });
 
-  it('returns a strategy with `search`', async () => {
+  it('returns a strategy with `search and `cancel`', async () => {
     const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
 
     expect(typeof esSearch.search).toBe('function');
   });
 
-  it('makes a POST request to async search with params when no ID is provided', async () => {
-    mockSubmitCaller.mockResolvedValueOnce(mockAsyncResponse);
+  describe('search', () => {
+    it('makes a POST request to async search with params when no ID is provided', async () => {
+      mockSubmitCaller.mockResolvedValueOnce(mockAsyncResponse);
 
-    const params = { index: 'logstash-*', body: { query: {} } };
-    const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
+      const params = { index: 'logstash-*', body: { query: {} } };
+      const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
 
-    await esSearch.search({ params }, {}, mockDeps).toPromise();
+      await esSearch.search({ params }, {}, mockDeps).toPromise();
 
-    expect(mockSubmitCaller).toBeCalled();
-    const request = mockSubmitCaller.mock.calls[0][0];
-    expect(request.index).toEqual(params.index);
-    expect(request.body).toEqual(params.body);
+      expect(mockSubmitCaller).toBeCalled();
+      const request = mockSubmitCaller.mock.calls[0][0];
+      expect(request.index).toEqual(params.index);
+      expect(request.body).toEqual(params.body);
+    });
+
+    it('makes a GET request to async search with ID when ID is provided', async () => {
+      mockGetCaller.mockResolvedValueOnce(mockAsyncResponse);
+
+      const params = { index: 'logstash-*', body: { query: {} } };
+      const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, 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');
+      expect(request).toHaveProperty('keep_alive');
+    });
+
+    it('calls the rollup API if the index is a rollup type', async () => {
+      mockApiCaller.mockResolvedValueOnce(mockRollupResponse);
+
+      const params = { index: 'foo-程', body: {} };
+      const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
+
+      await esSearch
+        .search(
+          {
+            indexType: 'rollup',
+            params,
+          },
+          {},
+          mockDeps
+        )
+        .toPromise();
+
+      expect(mockApiCaller).toBeCalled();
+      const { method, path } = mockApiCaller.mock.calls[0][0];
+      expect(method).toBe('POST');
+      expect(path).toBe('/foo-%E7%A8%8B/_rollup_search');
+    });
+
+    it('sets wait_for_completion_timeout and keep_alive in the request', async () => {
+      mockSubmitCaller.mockResolvedValueOnce(mockAsyncResponse);
+
+      const params = { index: 'foo-*', body: {} };
+      const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
+
+      await esSearch.search({ params }, {}, mockDeps).toPromise();
+
+      expect(mockSubmitCaller).toBeCalled();
+      const request = mockSubmitCaller.mock.calls[0][0];
+      expect(request).toHaveProperty('wait_for_completion_timeout');
+      expect(request).toHaveProperty('keep_alive');
+    });
   });
 
-  it('makes a GET request to async search with ID when ID is provided', async () => {
-    mockGetCaller.mockResolvedValueOnce(mockAsyncResponse);
+  describe('cancel', () => {
+    it('makes a DELETE request to async search with the provided ID', async () => {
+      mockDeleteCaller.mockResolvedValueOnce(200);
 
-    const params = { index: 'logstash-*', body: { query: {} } };
-    const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
+      const id = 'some_id';
+      const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
 
-    await esSearch.search({ id: 'foo', params }, {}, mockDeps).toPromise();
+      await esSearch.cancel!(id, {}, mockDeps);
 
-    expect(mockGetCaller).toBeCalled();
-    const request = mockGetCaller.mock.calls[0][0];
-    expect(request.id).toEqual('foo');
-    expect(request).toHaveProperty('wait_for_completion_timeout');
-    expect(request).toHaveProperty('keep_alive');
+      expect(mockDeleteCaller).toBeCalled();
+      const request = mockDeleteCaller.mock.calls[0][0];
+      expect(request).toEqual({ id });
+    });
   });
 
-  it('calls the rollup API if the index is a rollup type', async () => {
-    mockApiCaller.mockResolvedValueOnce(mockRollupResponse);
+  describe('extend', () => {
+    it('makes a GET request to async search with the provided ID and keepAlive', async () => {
+      mockGetCaller.mockResolvedValueOnce(mockAsyncResponse);
 
-    const params = { index: 'foo-程', body: {} };
-    const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
+      const id = 'some_other_id';
+      const keepAlive = '1d';
+      const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
 
-    await esSearch
-      .search(
-        {
-          indexType: 'rollup',
-          params,
-        },
-        {},
-        mockDeps
-      )
-      .toPromise();
+      await esSearch.extend!(id, keepAlive, {}, mockDeps);
 
-    expect(mockApiCaller).toBeCalled();
-    const { method, path } = mockApiCaller.mock.calls[0][0];
-    expect(method).toBe('POST');
-    expect(path).toBe('/foo-%E7%A8%8B/_rollup_search');
-  });
-
-  it('sets wait_for_completion_timeout and keep_alive in the request', async () => {
-    mockSubmitCaller.mockResolvedValueOnce(mockAsyncResponse);
-
-    const params = { index: 'foo-*', body: {} };
-    const esSearch = await enhancedEsSearchStrategyProvider(mockConfig$, mockLogger);
-
-    await esSearch.search({ params }, {}, mockDeps).toPromise();
-
-    expect(mockSubmitCaller).toBeCalled();
-    const request = mockSubmitCaller.mock.calls[0][0];
-    expect(request).toHaveProperty('wait_for_completion_timeout');
-    expect(request).toHaveProperty('keep_alive');
+      expect(mockGetCaller).toBeCalled();
+      const request = mockGetCaller.mock.calls[0][0];
+      expect(request).toEqual({ id, keep_alive: keepAlive });
+    });
   });
 });

x-pack/plugins/data_enhanced/server/search/es_search_strategy.ts

@@ -111,5 +111,9 @@ export const enhancedEsSearchStrategyProvider = (
       logger.debug(`cancel ${id}`);
       await esClient.asCurrentUser.asyncSearch.delete({ id });
     },
+    extend: async (id, keepAlive, options, { esClient }) => {
+      logger.debug(`extend ${id} by ${keepAlive}`);
+      await esClient.asCurrentUser.asyncSearch.get({ id, keep_alive: keepAlive });
+    },
   };
 };