[data.search.session] Store search strategy in saved object for background sessions (#86057)

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

* Update unit test to check strategy

* Fix test

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

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

@@ -15,4 +15,5 @@ export {
   BackgroundSessionSavedObjectAttributes,
   BackgroundSessionFindOptions,
   BackgroundSessionStatus,
+  BackgroundSessionSearchInfo,
 } from './search';

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

@@ -19,7 +19,12 @@ export interface BackgroundSessionSavedObjectAttributes {
   urlGeneratorId: string;
   initialState: Record<string, unknown>;
   restoreState: Record<string, unknown>;
-  idMapping: Record<string, string>;
+  idMapping: Record<string, BackgroundSessionSearchInfo>;
+}
+
+export interface BackgroundSessionSearchInfo {
+  id: string; // ID of the async search request
+  strategy: string; // Search strategy used to submit the search request
 }
 
 export interface BackgroundSessionFindOptions {

x-pack/plugins/data_enhanced/server/search/session/session_service.test.ts

@@ -30,6 +30,7 @@ describe('BackgroundSessionService', () => {
 
   const MOCK_SESSION_ID = 'session-id-mock';
   const MOCK_ASYNC_ID = '123456';
+  const MOCK_STRATEGY = 'ese';
   const MOCK_KEY_HASH = '608de49a4600dbb5b173492759792e4a';
 
   const createMockInternalSavedObjectClient = (
@@ -47,7 +48,10 @@ describe('BackgroundSessionService', () => {
                   attributes: {
                     sessionId: MOCK_SESSION_ID,
                     idMapping: {
-                      'another-key': 'another-async-id',
+                      'another-key': {
+                        id: 'another-async-id',
+                        strategy: 'another-strategy',
+                      },
                     },
                   },
                   id: MOCK_SESSION_ID,
@@ -283,7 +287,7 @@ describe('BackgroundSessionService', () => {
       await service.trackId(
         searchRequest,
         searchId,
-        { sessionId, isStored },
+        { sessionId, isStored, strategy: MOCK_STRATEGY },
         { savedObjectsClient }
       );
 
@@ -313,7 +317,8 @@ describe('BackgroundSessionService', () => {
       );
 
       const [setSessionId, setParams] = setSpy.mock.calls[0];
-      expect(setParams.ids.get(requestHash)).toBe(searchId);
+      expect(setParams.ids.get(requestHash).id).toBe(searchId);
+      expect(setParams.ids.get(requestHash).strategy).toBe(MOCK_STRATEGY);
       expect(setSessionId).toBe(sessionId);
     });
 
@@ -326,12 +331,17 @@ describe('BackgroundSessionService', () => {
       await service.trackId(
         searchRequest,
         searchId,
-        { sessionId, isStored },
+        { sessionId, isStored, strategy: MOCK_STRATEGY },
         { savedObjectsClient }
       );
 
       expect(savedObjectsClient.update).toHaveBeenCalledWith(BACKGROUND_SESSION_TYPE, sessionId, {
-        idMapping: { [requestHash]: searchId },
+        idMapping: {
+          [requestHash]: {
+            id: searchId,
+            strategy: MOCK_STRATEGY,
+          },
+        },
       });
     });
   });
@@ -380,7 +390,12 @@ describe('BackgroundSessionService', () => {
           name: 'my_name',
           appId: 'my_app_id',
           urlGeneratorId: 'my_url_generator_id',
-          idMapping: { [requestHash]: searchId },
+          idMapping: {
+            [requestHash]: {
+              id: searchId,
+              strategy: MOCK_STRATEGY,
+            },
+          },
         },
         references: [],
       };
@@ -419,7 +434,10 @@ describe('BackgroundSessionService', () => {
       const findSpy = jest.fn().mockResolvedValue({ saved_objects: [] });
       createMockInternalSavedObjectClient(findSpy);
 
-      const mockIdMapping = createMockIdMapping([[MOCK_KEY_HASH, MOCK_ASYNC_ID]], moment());
+      const mockIdMapping = createMockIdMapping(
+        [[MOCK_KEY_HASH, { id: MOCK_ASYNC_ID, strategy: MOCK_STRATEGY }]],
+        moment()
+      );
 
       Object.defineProperty(service, 'sessionSearchMap', {
         get: () => mockIdMapping,
@@ -438,7 +456,7 @@ describe('BackgroundSessionService', () => {
       createMockInternalSavedObjectClient(findSpy);
 
       const mockIdMapping = createMockIdMapping(
-        [[MOCK_KEY_HASH, MOCK_ASYNC_ID]],
+        [[MOCK_KEY_HASH, { id: MOCK_ASYNC_ID, strategy: MOCK_STRATEGY }]],
         moment().subtract(2, 'm')
       );
 
@@ -459,7 +477,7 @@ describe('BackgroundSessionService', () => {
       createMockInternalSavedObjectClient(findSpy);
 
       const mockIdMapping = createMockIdMapping(
-        [[MOCK_KEY_HASH, MOCK_ASYNC_ID]],
+        [[MOCK_KEY_HASH, { id: MOCK_ASYNC_ID, strategy: MOCK_STRATEGY }]],
         moment(),
         MAX_UPDATE_RETRIES
       );
@@ -528,7 +546,10 @@ describe('BackgroundSessionService', () => {
             attributes: {
               idMapping: {
                 b: 'c',
-                [MOCK_KEY_HASH]: MOCK_ASYNC_ID,
+                [MOCK_KEY_HASH]: {
+                  id: MOCK_ASYNC_ID,
+                  strategy: MOCK_STRATEGY,
+                },
               },
             },
           },
@@ -566,7 +587,10 @@ describe('BackgroundSessionService', () => {
             id: MOCK_SESSION_ID,
             attributes: {
               idMapping: {
-                b: 'c',
+                b: {
+                  id: 'c',
+                  strategy: MOCK_STRATEGY,
+                },
               },
             },
           },

x-pack/plugins/data_enhanced/server/search/session/session_service.ts

@@ -32,6 +32,7 @@ import {
 import {
   BackgroundSessionSavedObjectAttributes,
   BackgroundSessionFindOptions,
+  BackgroundSessionSearchInfo,
   BackgroundSessionStatus,
 } from '../../../common';
 import { BACKGROUND_SESSION_TYPE } from '../../saved_objects';
@@ -51,7 +52,7 @@ export interface BackgroundSessionDependencies {
 export interface SessionInfo {
   insertTime: Moment;
   retryCount: number;
-  ids: Map<string, string>;
+  ids: Map<string, BackgroundSessionSearchInfo>;
 }
 
 export class BackgroundSessionService implements ISessionService {
@@ -316,25 +317,31 @@ export class BackgroundSessionService implements ISessionService {
   public trackId = async (
     searchRequest: IKibanaSearchRequest,
     searchId: string,
-    { sessionId, isStored }: ISearchOptions,
+    { sessionId, isStored, strategy }: ISearchOptions,
     deps: BackgroundSessionDependencies
   ) => {
     if (!sessionId || !searchId) return;
     this.logger.debug(`trackId | ${sessionId} | ${searchId}`);
     const requestHash = createRequestHash(searchRequest.params);
+    const searchInfo = {
+      id: searchId,
+      strategy: strategy!,
+    };
 
     // If there is already a saved object for this session, update it to include this request/ID.
     // Otherwise, just update the in-memory mapping for this session for when the session is saved.
     if (isStored) {
-      const attributes = { idMapping: { [requestHash]: searchId } };
+      const attributes = {
+        idMapping: { [requestHash]: searchInfo },
+      };
       await this.update(sessionId, attributes, deps);
     } else {
       const map = this.sessionSearchMap.get(sessionId) ?? {
         insertTime: moment(),
         retryCount: 0,
-        ids: new Map<string, string>(),
+        ids: new Map<string, BackgroundSessionSearchInfo>(),
       };
-      map.ids.set(requestHash, searchId);
+      map.ids.set(requestHash, searchInfo);
       this.sessionSearchMap.set(sessionId, map);
     }
   };
@@ -363,7 +370,7 @@ export class BackgroundSessionService implements ISessionService {
       throw new Error('No search ID in this session matching the given search request');
     }
 
-    return session.attributes.idMapping[requestHash];
+    return session.attributes.idMapping[requestHash].id;
   };
 
   public asScopedProvider = ({ savedObjects }: CoreStart) => {

x-pack/test/api_integration/apis/search/session.ts

@@ -119,8 +119,9 @@ export default function ({ getService }: FtrProviderContext) {
 
         const { idMapping } = resp.body.attributes;
 
-        expect(Object.values(idMapping)).to.contain(id1);
+        const idMappings = Object.values(idMapping).map((value: any) => value.id);
-        expect(Object.values(idMapping)).to.contain(id2);
+        expect(idMappings).to.contain(id1);
+        expect(idMappings).to.contain(id2);
       });
     });
   });