[data.search] Move search method inside session service and add tests (#84715) (#84953)

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

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

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-import { BehaviorSubject, from, Observable } from 'rxjs';
+import { BehaviorSubject, Observable } from 'rxjs';
 import { pick } from 'lodash';
 import {
   CoreSetup,
@@ -29,7 +29,7 @@ import {
   SharedGlobalConfig,
   StartServicesAccessor,
 } from 'src/core/server';
-import { catchError, first, map, switchMap } from 'rxjs/operators';
+import { catchError, first, map } from 'rxjs/operators';
 import { BfetchServerSetup } from 'src/plugins/bfetch/server';
 import { ExpressionsServerSetup } from 'src/plugins/expressions/server';
 import {
@@ -50,7 +50,11 @@ import { DataPluginStart } from '../plugin';
 import { UsageCollectionSetup } from '../../../usage_collection/server';
 import { registerUsageCollector } from './collectors/register';
 import { usageProvider } from './collectors/usage';
-import { BACKGROUND_SESSION_TYPE, searchTelemetry } from '../saved_objects';
+import {
+  BACKGROUND_SESSION_TYPE,
+  backgroundSessionMapping,
+  searchTelemetry,
+} from '../saved_objects';
 import {
   IEsSearchRequest,
   IEsSearchResponse,
@@ -74,8 +78,6 @@ import { aggShardDelay } from '../../common/search/aggs/buckets/shard_delay_fn';
 import { ConfigSchema } from '../../config';
 import { BackgroundSessionService, ISearchSessionClient } from './session';
 import { registerSessionRoutes } from './routes/session';
-import { backgroundSessionMapping } from '../saved_objects';
-import { tapFirst } from '../../common/utils';
 
 declare module 'src/core/server' {
   interface RequestHandlerContext {
@@ -297,7 +299,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     SearchStrategyRequest extends IKibanaSearchRequest = IEsSearchRequest,
     SearchStrategyResponse extends IKibanaSearchResponse = IEsSearchResponse
   >(
-    searchRequest: SearchStrategyRequest,
+    request: SearchStrategyRequest,
     options: ISearchOptions,
     deps: SearchStrategyDependencies
   ) => {
@@ -305,24 +307,9 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
       options.strategy
     );
 
-    // If this is a restored background search session, look up the ID using the provided sessionId
-    const getSearchRequest = async () =>
-      !options.isRestore || searchRequest.id
-        ? searchRequest
-        : {
-            ...searchRequest,
-            id: await this.sessionService.getId(searchRequest, options, deps),
-          };
-
-    return from(getSearchRequest()).pipe(
-      switchMap((request) => strategy.search(request, options, deps)),
-      tapFirst((response) => {
-        if (searchRequest.id || !options.sessionId || !response.id || options.isRestore) return;
-        this.sessionService.trackId(searchRequest, response.id, options, {
-          savedObjectsClient: deps.savedObjectsClient,
-        });
-      })
-    );
+    return options.sessionId
+      ? this.sessionService.search(strategy, request, options, deps)
+      : strategy.search(request, options, deps);
   };
 
   private cancel = (id: string, options: ISearchOptions, deps: SearchStrategyDependencies) => {

src/plugins/data/server/search/session/session_service.test.ts

@@ -17,7 +17,9 @@
  * under the License.
  */
 
+import { of } from 'rxjs';
 import type { SavedObject, SavedObjectsClientContract } from 'kibana/server';
+import type { SearchStrategyDependencies } from '../types';
 import { savedObjectsClientMock } from '../../../../../core/server/mocks';
 import { BackgroundSessionStatus } from '../../../common';
 import { BACKGROUND_SESSION_TYPE } from '../../saved_objects';
@@ -28,6 +30,7 @@ describe('BackgroundSessionService', () => {
   let savedObjectsClient: jest.Mocked<SavedObjectsClientContract>;
   let service: BackgroundSessionService;
 
+  const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
   const mockSavedObject: SavedObject = {
     id: 'd7170a35-7e2c-48d6-8dec-9a056721b489',
     type: BACKGROUND_SESSION_TYPE,
@@ -45,9 +48,13 @@ describe('BackgroundSessionService', () => {
     service = new BackgroundSessionService();
   });
 
-  it('save throws if `name` is not provided', () => {
-    const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
+  it('search throws if `name` is not provided', () => {
+    expect(() => service.save(sessionId, {}, { savedObjectsClient })).rejects.toMatchInlineSnapshot(
+      `[Error: Name is required]`
+    );
+  });
 
+  it('save throws if `name` is not provided', () => {
     expect(() => service.save(sessionId, {}, { savedObjectsClient })).rejects.toMatchInlineSnapshot(
       `[Error: Name is required]`
     );
@@ -56,7 +63,6 @@ describe('BackgroundSessionService', () => {
   it('get calls saved objects client', async () => {
     savedObjectsClient.get.mockResolvedValue(mockSavedObject);
 
-    const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
     const response = await service.get(sessionId, { savedObjectsClient });
 
     expect(response).toBe(mockSavedObject);
@@ -93,7 +99,6 @@ describe('BackgroundSessionService', () => {
     };
     savedObjectsClient.update.mockResolvedValue(mockUpdateSavedObject);
 
-    const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
     const attributes = { name: 'new_name' };
     const response = await service.update(sessionId, attributes, { savedObjectsClient });
 
@@ -108,19 +113,87 @@ describe('BackgroundSessionService', () => {
   it('delete calls saved objects client', async () => {
     savedObjectsClient.delete.mockResolvedValue({});
 
-    const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
     const response = await service.delete(sessionId, { savedObjectsClient });
 
     expect(response).toEqual({});
     expect(savedObjectsClient.delete).toHaveBeenCalledWith(BACKGROUND_SESSION_TYPE, sessionId);
   });
 
+  describe('search', () => {
+    const mockSearch = jest.fn().mockReturnValue(of({}));
+    const mockStrategy = { search: mockSearch };
+    const mockDeps = {} as SearchStrategyDependencies;
+
+    beforeEach(() => {
+      mockSearch.mockClear();
+    });
+
+    it('searches using the original request if not restoring', async () => {
+      const searchRequest = { params: {} };
+      const options = { sessionId, isStored: false, isRestore: false };
+
+      await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();
+
+      expect(mockSearch).toBeCalledWith(searchRequest, options, mockDeps);
+    });
+
+    it('searches using the original request if `id` is provided', async () => {
+      const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
+      const searchRequest = { id: searchId, params: {} };
+      const options = { sessionId, isStored: true, isRestore: true };
+
+      await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();
+
+      expect(mockSearch).toBeCalledWith(searchRequest, options, mockDeps);
+    });
+
+    it('searches by looking up an `id` if restoring and `id` is not provided', async () => {
+      const searchRequest = { params: {} };
+      const options = { sessionId, isStored: true, isRestore: true };
+      const spyGetId = jest.spyOn(service, 'getId').mockResolvedValueOnce('my_id');
+
+      await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();
+
+      expect(mockSearch).toBeCalledWith({ ...searchRequest, id: 'my_id' }, options, mockDeps);
+
+      spyGetId.mockRestore();
+    });
+
+    it('calls `trackId` once if the response contains an `id` and not restoring', async () => {
+      const searchRequest = { params: {} };
+      const options = { sessionId, isStored: false, isRestore: false };
+      const spyTrackId = jest.spyOn(service, 'trackId').mockResolvedValue();
+      mockSearch.mockReturnValueOnce(of({ id: 'my_id' }, { id: 'my_id' }));
+
+      await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();
+
+      expect(spyTrackId).toBeCalledTimes(1);
+      expect(spyTrackId).toBeCalledWith(searchRequest, 'my_id', options, mockDeps);
+
+      spyTrackId.mockRestore();
+    });
+
+    it('does not call `trackId` if restoring', async () => {
+      const searchRequest = { params: {} };
+      const options = { sessionId, isStored: true, isRestore: true };
+      const spyGetId = jest.spyOn(service, 'getId').mockResolvedValueOnce('my_id');
+      const spyTrackId = jest.spyOn(service, 'trackId').mockResolvedValue();
+      mockSearch.mockReturnValueOnce(of({ id: 'my_id' }));
+
+      await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();
+
+      expect(spyTrackId).not.toBeCalled();
+
+      spyGetId.mockRestore();
+      spyTrackId.mockRestore();
+    });
+  });
+
   describe('trackId', () => {
     it('stores hash in memory when `isStored` is `false` for when `save` is called', async () => {
       const searchRequest = { params: {} };
       const requestHash = createRequestHash(searchRequest.params);
       const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
-      const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
       const isStored = false;
       const name = 'my saved background search session';
       const appId = 'my_app_id';
@@ -164,7 +237,6 @@ describe('BackgroundSessionService', () => {
       const searchRequest = { params: {} };
       const requestHash = createRequestHash(searchRequest.params);
       const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
-      const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
       const isStored = true;
 
       await service.trackId(
@@ -191,7 +263,6 @@ describe('BackgroundSessionService', () => {
 
     it('throws if there is not a saved object', () => {
       const searchRequest = { params: {} };
-      const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
 
       expect(() =>
         service.getId(searchRequest, { sessionId, isStored: false }, { savedObjectsClient })
@@ -202,7 +273,6 @@ describe('BackgroundSessionService', () => {
 
     it('throws if not restoring a saved session', () => {
       const searchRequest = { params: {} };
-      const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
 
       expect(() =>
         service.getId(
@@ -219,7 +289,6 @@ describe('BackgroundSessionService', () => {
       const searchRequest = { params: {} };
       const requestHash = createRequestHash(searchRequest.params);
       const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
-      const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
       const mockSession = {
         id: 'd7170a35-7e2c-48d6-8dec-9a056721b489',
         type: BACKGROUND_SESSION_TYPE,

src/plugins/data/server/search/session/session_service.ts

@@ -18,14 +18,19 @@
  */
 
 import { CoreStart, KibanaRequest, SavedObjectsClientContract } from 'kibana/server';
+import { from } from 'rxjs';
+import { switchMap } from 'rxjs/operators';
 import {
   BackgroundSessionSavedObjectAttributes,
+  BackgroundSessionStatus,
   IKibanaSearchRequest,
+  IKibanaSearchResponse,
   ISearchOptions,
   SearchSessionFindOptions,
-  BackgroundSessionStatus,
+  tapFirst,
 } from '../../../common';
 import { BACKGROUND_SESSION_TYPE } from '../../saved_objects';
+import { ISearchStrategy, SearchStrategyDependencies } from '../types';
 import { createRequestHash } from './utils';
 
 const DEFAULT_EXPIRATION = 7 * 24 * 60 * 60 * 1000;
@@ -59,6 +64,32 @@ export class BackgroundSessionService {
     this.sessionSearchMap.clear();
   };
 
+  public search = <Request extends IKibanaSearchRequest, Response extends IKibanaSearchResponse>(
+    strategy: ISearchStrategy<Request, Response>,
+    searchRequest: Request,
+    options: ISearchOptions,
+    deps: SearchStrategyDependencies
+  ) => {
+    // If this is a restored background search session, look up the ID using the provided sessionId
+    const getSearchRequest = async () =>
+      !options.isRestore || searchRequest.id
+        ? searchRequest
+        : {
+            ...searchRequest,
+            id: await this.getId(searchRequest, options, deps),
+          };
+
+    return from(getSearchRequest()).pipe(
+      switchMap((request) => strategy.search(request, options, deps)),
+      tapFirst((response) => {
+        if (searchRequest.id || !options.sessionId || !response.id || options.isRestore) return;
+        this.trackId(searchRequest, response.id, options, {
+          savedObjectsClient: deps.savedObjectsClient,
+        });
+      })
+    );
+  };
+
   // TODO: Generate the `userId` from the realm type/realm name/username
   public save = async (
     sessionId: string,
@@ -208,10 +239,6 @@ export class BackgroundSessionService {
         update: (sessionId: string, attributes: Partial<BackgroundSessionSavedObjectAttributes>) =>
           this.update(sessionId, attributes, deps),
         delete: (sessionId: string) => this.delete(sessionId, deps),
-        trackId: (searchRequest: IKibanaSearchRequest, searchId: string, options: ISearchOptions) =>
-          this.trackId(searchRequest, searchId, options, deps),
-        getId: (searchRequest: IKibanaSearchRequest, options: ISearchOptions) =>
-          this.getId(searchRequest, options, deps),
       };
     };
   };