[Search Session] Fix integration in vega, timelion and TSVB (#87862)

src/plugins/data/public/public.api.md

@@ -2624,7 +2624,7 @@ export const UI_SETTINGS: {
 // src/plugins/data/public/index.ts:433:1 - (ae-forgotten-export) The symbol "propFilter" needs to be exported by the entry point index.d.ts
 // src/plugins/data/public/index.ts:436:1 - (ae-forgotten-export) The symbol "toAbsoluteDates" needs to be exported by the entry point index.d.ts
 // src/plugins/data/public/query/state_sync/connect_to_query_state.ts:45:5 - (ae-forgotten-export) The symbol "FilterStateStore" needs to be exported by the entry point index.d.ts
-// src/plugins/data/public/search/session/session_service.ts:51:5 - (ae-forgotten-export) The symbol "UrlGeneratorStateMapping" needs to be exported by the entry point index.d.ts
+// src/plugins/data/public/search/session/session_service.ts:52:5 - (ae-forgotten-export) The symbol "UrlGeneratorStateMapping" needs to be exported by the entry point index.d.ts
 
 // (No @packageDocumentation comment for this package)
 

src/plugins/data/public/search/search_interceptor.test.ts

@@ -115,12 +115,14 @@ describe('SearchInterceptor', () => {
       }: {
         isRestore?: boolean;
         isStored?: boolean;
-        sessionId?: string;
+        sessionId: string;
       }) => {
         const sessionServiceMock = searchMock.session as jest.Mocked<ISessionService>;
-        sessionServiceMock.getSessionId.mockImplementation(() => sessionId);
+        sessionServiceMock.getSearchOptions.mockImplementation(() => ({
-        sessionServiceMock.isRestore.mockImplementation(() => isRestore);
+          sessionId,
-        sessionServiceMock.isStored.mockImplementation(() => isStored);
+          isRestore,
+          isStored,
+        }));
         fetchMock.mockResolvedValue({ result: 200 });
       };
 
@@ -130,30 +132,14 @@ describe('SearchInterceptor', () => {
 
       afterEach(() => {
         const sessionServiceMock = searchMock.session as jest.Mocked<ISessionService>;
-        sessionServiceMock.getSessionId.mockReset();
+        sessionServiceMock.getSearchOptions.mockReset();
-        sessionServiceMock.isRestore.mockReset();
-        sessionServiceMock.isStored.mockReset();
         fetchMock.mockReset();
       });
 
-      test('infers isRestore from session service state', async () => {
+      test('gets session search options from session service', async () => {
         const sessionId = 'sid';
         setup({
           isRestore: true,
-          sessionId,
-        });
-
-        await searchInterceptor.search(mockRequest, { sessionId }).toPromise();
-        expect(fetchMock.mock.calls[0][0]).toEqual(
-          expect.objectContaining({
-            options: { sessionId: 'sid', isStored: false, isRestore: true },
-          })
-        );
-      });
-
-      test('infers isStored from session service state', async () => {
-        const sessionId = 'sid';
-        setup({
           isStored: true,
           sessionId,
         });
@@ -161,41 +147,13 @@ describe('SearchInterceptor', () => {
         await searchInterceptor.search(mockRequest, { sessionId }).toPromise();
         expect(fetchMock.mock.calls[0][0]).toEqual(
           expect.objectContaining({
-            options: { sessionId: 'sid', isStored: true, isRestore: false },
+            options: { sessionId, isStored: true, isRestore: true },
           })
         );
-      });
 
-      test('skips isRestore & isStore in case not a current session Id', async () => {
+        expect(
-        setup({
+          (searchMock.session as jest.Mocked<ISessionService>).getSearchOptions
-          isStored: true,
+        ).toHaveBeenCalledWith(sessionId);
-          isRestore: true,
-          sessionId: 'session id',
-        });
-
-        await searchInterceptor
-          .search(mockRequest, { sessionId: 'different session id' })
-          .toPromise();
-        expect(fetchMock.mock.calls[0][0]).toEqual(
-          expect.objectContaining({
-            options: { sessionId: 'different session id', isStored: false, isRestore: false },
-          })
-        );
-      });
-
-      test('skips isRestore & isStore in case no session Id', async () => {
-        setup({
-          isStored: true,
-          isRestore: true,
-          sessionId: undefined,
-        });
-
-        await searchInterceptor.search(mockRequest, { sessionId: 'sessionId' }).toPromise();
-        expect(fetchMock.mock.calls[0][0]).toEqual(
-          expect.objectContaining({
-            options: { sessionId: 'sessionId', isStored: false, isRestore: false },
-          })
-        );
       });
     });
 

src/plugins/data/public/search/search_interceptor.ts

@@ -130,16 +130,12 @@ export class SearchInterceptor {
   ): Promise<IKibanaSearchResponse> {
     const { abortSignal, ...requestOptions } = options || {};
 
-    const isCurrentSession =
-      options?.sessionId && this.deps.session.getSessionId() === options.sessionId;
-
     return this.batchedFetch(
       {
         request,
         options: {
           ...requestOptions,
-          isStored: isCurrentSession ? this.deps.session.isStored() : false,
+          ...(options?.sessionId && this.deps.session.getSearchOptions(options.sessionId)),
-          isRestore: isCurrentSession ? this.deps.session.isRestore() : false,
         },
       },
       abortSignal

src/plugins/data/public/search/session/mocks.ts

@@ -49,5 +49,7 @@ export function getSessionServiceMock(): jest.Mocked<ISessionService> {
     isStored: jest.fn(),
     isRestore: jest.fn(),
     save: jest.fn(),
+    isCurrentSession: jest.fn(),
+    getSearchOptions: jest.fn(),
   };
 }

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

@@ -33,10 +33,18 @@ describe('Session service', () => {
 
   beforeEach(() => {
     const initializerContext = coreMock.createPluginInitializerContext();
+    const startService = coreMock.createSetup().getStartServices;
     nowProvider = createNowProviderMock();
     sessionService = new SessionService(
       initializerContext,
-      coreMock.createSetup().getStartServices,
+      () =>
+        startService().then(([coreStart, ...rest]) => [
+          {
+            ...coreStart,
+            application: { ...coreStart.application, currentAppId$: new BehaviorSubject('app') },
+          },
+          ...rest,
+        ]),
       getSessionsClientMock(),
       nowProvider,
       { freezeState: false } // needed to use mocks inside state container
@@ -100,4 +108,63 @@ describe('Session service', () => {
       expect(abort).toBeCalledTimes(3);
     });
   });
+
+  test('getSearchOptions infers isRestore & isStored from state', async () => {
+    const sessionId = sessionService.start();
+    const someOtherId = 'some-other-id';
+
+    expect(sessionService.getSearchOptions(someOtherId)).toEqual({
+      isStored: false,
+      isRestore: false,
+      sessionId: someOtherId,
+    });
+    expect(sessionService.getSearchOptions(sessionId)).toEqual({
+      isStored: false,
+      isRestore: false,
+      sessionId,
+    });
+
+    sessionService.setSearchSessionInfoProvider({
+      getName: async () => 'Name',
+      getUrlGeneratorData: async () => ({
+        urlGeneratorId: 'id',
+        initialState: {},
+        restoreState: {},
+      }),
+    });
+    await sessionService.save();
+
+    expect(sessionService.getSearchOptions(someOtherId)).toEqual({
+      isStored: false,
+      isRestore: false,
+      sessionId: someOtherId,
+    });
+    expect(sessionService.getSearchOptions(sessionId)).toEqual({
+      isStored: true,
+      isRestore: false,
+      sessionId,
+    });
+
+    await sessionService.restore(sessionId);
+
+    expect(sessionService.getSearchOptions(someOtherId)).toEqual({
+      isStored: false,
+      isRestore: false,
+      sessionId: someOtherId,
+    });
+    expect(sessionService.getSearchOptions(sessionId)).toEqual({
+      isStored: true,
+      isRestore: true,
+      sessionId,
+    });
+  });
+  test('isCurrentSession', () => {
+    expect(sessionService.isCurrentSession()).toBeFalsy();
+
+    const sessionId = sessionService.start();
+
+    expect(sessionService.isCurrentSession()).toBeFalsy();
+    expect(sessionService.isCurrentSession('some-other')).toBeFalsy();
+    expect(sessionService.isCurrentSession(sessionId)).toBeTruthy();
+  });
 });

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

@@ -29,6 +29,7 @@ import {
   SessionStateContainer,
 } from './search_session_state';
 import { ISessionsClient } from './sessions_client';
+import { ISearchOptions } from '../../../common';
 import { NowProviderInternalContract } from '../../now_provider';
 
 export type ISessionService = PublicContract<SessionService>;
@@ -256,4 +257,27 @@ export class SessionService {
       this.state.transitions.store();
     }
   }
+
+  /**
+   * Checks if passed sessionId is a current sessionId
+   * @param sessionId
+   */
+  public isCurrentSession(sessionId?: string): boolean {
+    return !!sessionId && this.getSessionId() === sessionId;
+  }
+
+  /**
+   * Infers search session options for sessionId using current session state
+   * @param sessionId
+   */
+  public getSearchOptions(
+    sessionId: string
+  ): Required<Pick<ISearchOptions, 'sessionId' | 'isRestore' | 'isStored'>> {
+    const isCurrentSession = this.isCurrentSession(sessionId);
+    return {
+      sessionId,
+      isRestore: isCurrentSession ? this.isRestore() : false,
+      isStored: isCurrentSession ? this.isStored() : false,
+    };
+  }
 }

src/plugins/vis_type_timelion/public/helpers/timelion_request_handler.ts

@@ -73,12 +73,15 @@ export function getTimelionRequestHandler({
     filters,
     query,
     visParams,
+    searchSessionId,
   }: {
     timeRange: TimeRange;
     filters: Filter[];
     query: Query;
     visParams: TimelionVisParams;
+    searchSessionId?: string;
   }): Promise<TimelionSuccessResponse> {
+    const dataSearch = getDataSearch();
     const expression = visParams.expression;
 
     if (!expression) {
@@ -93,7 +96,13 @@ export function getTimelionRequestHandler({
 
     // parse the time range client side to make sure it behaves like other charts
     const timeRangeBounds = timefilter.calculateBounds(timeRange);
-    const sessionId = getDataSearch().session.getSessionId();
+    const untrackSearch =
+      dataSearch.session.isCurrentSession(searchSessionId) &&
+      dataSearch.session.trackSearch({
+        abort: () => {
+          // TODO: support search cancellations
+        },
+      });
 
     try {
       return await http.post('/api/timelion/run', {
@@ -110,7 +119,9 @@ export function getTimelionRequestHandler({
             interval: visParams.interval,
             timezone,
           },
-          sessionId,
+          ...(searchSessionId && {
+            searchSession: dataSearch.session.getSearchOptions(searchSessionId),
+          }),
         }),
       });
     } catch (e) {
@@ -125,6 +136,11 @@ export function getTimelionRequestHandler({
       } else {
         throw e;
       }
+    } finally {
+      if (untrackSearch && dataSearch.session.isCurrentSession(searchSessionId)) {
+        // call `untrack` if this search still belongs to current session
+        untrackSearch();
+      }
     }
   };
 }

src/plugins/vis_type_timelion/public/timelion_vis_fn.ts

@@ -70,7 +70,7 @@ export const getTimelionVisualizationConfig = (
       help: '',
     },
   },
-  async fn(input, args) {
+  async fn(input, args, { getSearchSessionId }) {
     const timelionRequestHandler = getTimelionRequestHandler(dependencies);
 
     const visParams = { expression: args.expression, interval: args.interval };
@@ -80,6 +80,7 @@ export const getTimelionVisualizationConfig = (
       query: get(input, 'query') as Query,
       filters: get(input, 'filters') as Filter[],
       visParams,
+      searchSessionId: getSearchSessionId(),
     });
 
     response.visType = TIMELION_VIS_NAME;

src/plugins/vis_type_timelion/server/routes/run.ts

@@ -75,7 +75,13 @@ export function runRoute(
               to: schema.maybe(schema.string()),
             })
           ),
-          sessionId: schema.maybe(schema.string()),
+          searchSession: schema.maybe(
+            schema.object({
+              sessionId: schema.string(),
+              isRestore: schema.boolean({ defaultValue: false }),
+              isStored: schema.boolean({ defaultValue: false }),
+            })
+          ),
         }),
       },
     },

src/plugins/vis_type_timelion/server/series_functions/es/es.test.js

@@ -60,19 +60,26 @@ describe('es', () => {
         });
     });
 
-    test('should call data search with sessionId', async () => {
+    test('should call data search with sessionId, isRestore and isStored', async () => {
       tlConfig = {
         ...stubRequestAndServer({ rawResponse: esResponse }),
         request: {
           body: {
-            sessionId: 1,
+            searchSession: {
+              sessionId: '1',
+              isRestore: true,
+              isStored: false,
+            },
           },
         },
       };
 
       await invoke(es, [5], tlConfig);
 
-      expect(tlConfig.context.search.search.mock.calls[0][1]).toHaveProperty('sessionId', 1);
+      const res = tlConfig.context.search.search.mock.calls[0][1];
+      expect(res).toHaveProperty('sessionId', '1');
+      expect(res).toHaveProperty('isRestore', true);
+      expect(res).toHaveProperty('isStored', false);
     });
 
     test('returns a seriesList', () => {

src/plugins/vis_type_timelion/server/series_functions/es/index.js

@@ -133,7 +133,7 @@ export default new Datasource('es', {
       .search(
         body,
         {
-          sessionId: tlConfig.request?.body.sessionId,
+          ...tlConfig.request?.body.searchSession,
         },
         tlConfig.context
       )

src/plugins/vis_type_timeseries/common/vis_schema.ts

@@ -284,5 +284,12 @@ export const visPayloadSchema = schema.object({
     min: stringRequired,
     max: stringRequired,
   }),
-  sessionId: schema.maybe(schema.string()),
+
+  searchSession: schema.maybe(
+    schema.object({
+      sessionId: schema.string(),
+      isRestore: schema.boolean({ defaultValue: false }),
+      isStored: schema.boolean({ defaultValue: false }),
+    })
+  ),
 });

src/plugins/vis_type_timeseries/public/metrics_fn.ts

@@ -65,7 +65,7 @@ export const createMetricsFn = (): TimeseriesExpressionFunctionDefinition => ({
       help: '',
     },
   },
-  async fn(input, args) {
+  async fn(input, args, { getSearchSessionId }) {
     const visParams: TimeseriesVisParams = JSON.parse(args.params);
     const uiState = JSON.parse(args.uiState);
 
@@ -73,6 +73,7 @@ export const createMetricsFn = (): TimeseriesExpressionFunctionDefinition => ({
       input,
       visParams,
       uiState,
+      searchSessionId: getSearchSessionId(),
     });
 
     return {

src/plugins/vis_type_timeseries/public/request_handler.ts

@@ -29,39 +29,57 @@ interface MetricsRequestHandlerParams {
   input: KibanaContext | null;
   uiState: Record<string, any>;
   visParams: TimeseriesVisParams;
+  searchSessionId?: string;
 }
 
 export const metricsRequestHandler = async ({
   input,
   uiState,
   visParams,
+  searchSessionId,
 }: MetricsRequestHandlerParams): Promise<TimeseriesVisData | {}> => {
   const config = getUISettings();
   const timezone = getTimezone(config);
   const uiStateObj = uiState[visParams.type] ?? {};
-  const dataSearch = getDataStart();
+  const data = getDataStart();
-  const parsedTimeRange = dataSearch.query.timefilter.timefilter.calculateBounds(input?.timeRange!);
+  const dataSearch = getDataStart().search;
+  const parsedTimeRange = data.query.timefilter.timefilter.calculateBounds(input?.timeRange!);
 
   if (visParams && visParams.id && !visParams.isModelInvalid) {
     const maxBuckets = config.get(MAX_BUCKETS_SETTING);
 
     validateInterval(parsedTimeRange, visParams, maxBuckets);
 
-    const resp = await getCoreStart().http.post(ROUTES.VIS_DATA, {
+    const untrackSearch =
-      body: JSON.stringify({
+      dataSearch.session.isCurrentSession(searchSessionId) &&
-        timerange: {
+      dataSearch.session.trackSearch({
-          timezone,
+        abort: () => {
-          ...parsedTimeRange,
+          // TODO: support search cancellations
         },
-        query: input?.query,
+      });
-        filters: input?.filters,
-        panels: [visParams],
-        state: uiStateObj,
-        sessionId: dataSearch.search.session.getSessionId(),
-      }),
-    });
 
-    return resp;
+    try {
+      return await getCoreStart().http.post(ROUTES.VIS_DATA, {
+        body: JSON.stringify({
+          timerange: {
+            timezone,
+            ...parsedTimeRange,
+          },
+          query: input?.query,
+          filters: input?.filters,
+          panels: [visParams],
+          state: uiStateObj,
+          ...(searchSessionId && {
+            searchSession: dataSearch.session.getSearchOptions(searchSessionId),
+          }),
+        }),
+      });
+    } finally {
+      if (untrackSearch && dataSearch.session.isCurrentSession(searchSessionId)) {
+        // untrack if this search still belongs to current session
+        untrackSearch();
+      }
+    }
   }
 
   return {};

src/plugins/vis_type_timeseries/server/lib/search_strategies/strategies/abstract_search_strategy.test.js

@@ -66,7 +66,11 @@ describe('AbstractSearchStrategy', () => {
     const responses = await abstractSearchStrategy.search(
       {
         payload: {
-          sessionId: 1,
+          searchSession: {
+            sessionId: '1',
+            isRestore: false,
+            isStored: true,
+          },
         },
         requestContext: {
           search: { search: searchFn },
@@ -85,7 +89,9 @@ describe('AbstractSearchStrategy', () => {
         indexType: undefined,
       },
       {
-        sessionId: 1,
+        sessionId: '1',
+        isRestore: false,
+        isStored: true,
       }
     );
   });

src/plugins/vis_type_timeseries/server/lib/search_strategies/strategies/abstract_search_strategy.ts

@@ -66,7 +66,6 @@ const toSanitizedFieldType = (fields: IFieldType[]) => {
 export abstract class AbstractSearchStrategy {
   async search(req: ReqFacade<VisPayload>, bodies: any[], indexType?: string) {
     const requests: any[] = [];
-    const { sessionId } = req.payload;
 
     bodies.forEach((body) => {
       requests.push(
@@ -78,9 +77,7 @@ export abstract class AbstractSearchStrategy {
                 ...body,
               },
             },
-            {
+            req.payload.searchSession
-              sessionId,
-            }
           )
           .toPromise()
       );

src/plugins/vis_type_vega/public/data_model/search_api.ts

@@ -40,7 +40,8 @@ export class SearchAPI {
   constructor(
     private readonly dependencies: SearchAPIDependencies,
     private readonly abortSignal?: AbortSignal,
-    public readonly inspectorAdapters?: VegaInspectorAdapters
+    public readonly inspectorAdapters?: VegaInspectorAdapters,
+    private readonly searchSessionId?: string
   ) {}
 
   search(searchRequests: SearchRequest[]) {
@@ -60,10 +61,7 @@ export class SearchAPI {
         }
 
         return search
-          .search(
+          .search({ params }, { abortSignal: this.abortSignal, sessionId: this.searchSessionId })
-            { params },
-            { abortSignal: this.abortSignal, sessionId: search.session.getSessionId() }
-          )
           .pipe(
             tap((data) => this.inspectSearchResult(data, requestResponders[requestId])),
             map((data) => ({

src/plugins/vis_type_vega/public/vega_fn.ts

@@ -74,6 +74,7 @@ export const createVegaFn = (
       query: get(input, 'query') as Query,
       filters: get(input, 'filters') as any,
       visParams: { spec: args.spec },
+      searchSessionId: context.getSearchSessionId(),
     });
 
     return {

src/plugins/vis_type_vega/public/vega_request_handler.ts

@@ -32,6 +32,7 @@ interface VegaRequestHandlerParams {
   filters: Filter;
   timeRange: TimeRange;
   visParams: VisParams;
+  searchSessionId?: string;
 }
 
 interface VegaRequestHandlerContext {
@@ -52,6 +53,7 @@ export function createVegaRequestHandler(
     filters,
     query,
     visParams,
+    searchSessionId,
   }: VegaRequestHandlerParams) {
     if (!searchAPI) {
       searchAPI = new SearchAPI(
@@ -61,7 +63,8 @@ export function createVegaRequestHandler(
           injectedMetadata: getInjectedMetadata(),
         },
         context.abortSignal,
-        context.inspectorAdapters
+        context.inspectorAdapters,
+        searchSessionId
       );
     }
 

x-pack/plugins/data_enhanced/public/search/search_interceptor.test.ts

@@ -373,7 +373,7 @@ describe('EnhancedSearchInterceptor', () => {
 
     test('should NOT DELETE a running SAVED async search on abort', async () => {
       const sessionId = 'sessionId';
-      sessionService.getSessionId.mockImplementation(() => sessionId);
+      sessionService.isCurrentSession.mockImplementation((_sessionId) => _sessionId === sessionId);
       const responses = [
         {
           time: 10,
@@ -479,6 +479,7 @@ describe('EnhancedSearchInterceptor', () => {
 
     test('should track searches', async () => {
       const sessionId = 'sessionId';
+      sessionService.isCurrentSession.mockImplementation((_sessionId) => _sessionId === sessionId);
       sessionService.getSessionId.mockImplementation(() => sessionId);
 
       const untrack = jest.fn();
@@ -496,6 +497,7 @@ describe('EnhancedSearchInterceptor', () => {
 
     test('session service should be able to cancel search', async () => {
       const sessionId = 'sessionId';
+      sessionService.isCurrentSession.mockImplementation((_sessionId) => _sessionId === sessionId);
       sessionService.getSessionId.mockImplementation(() => sessionId);
 
       const untrack = jest.fn();
@@ -519,6 +521,7 @@ describe('EnhancedSearchInterceptor', () => {
 
     test("don't track non current session searches", async () => {
       const sessionId = 'sessionId';
+      sessionService.isCurrentSession.mockImplementation((_sessionId) => _sessionId === sessionId);
       sessionService.getSessionId.mockImplementation(() => sessionId);
 
       const untrack = jest.fn();
@@ -539,6 +542,7 @@ describe('EnhancedSearchInterceptor', () => {
 
     test("don't track if no current session", async () => {
       sessionService.getSessionId.mockImplementation(() => undefined);
+      sessionService.isCurrentSession.mockImplementation((_sessionId) => false);
 
       const untrack = jest.fn();
       sessionService.trackSearch.mockImplementation(() => untrack);

x-pack/plugins/data_enhanced/public/search/search_interceptor.ts

@@ -64,20 +64,24 @@ export class EnhancedSearchInterceptor extends SearchInterceptor {
     const search = () => this.runSearch({ id, ...request }, searchOptions);
 
     this.pendingCount$.next(this.pendingCount$.getValue() + 1);
-    const isCurrentSession = () =>
-      !!options.sessionId && options.sessionId === this.deps.session.getSessionId();
 
-    const untrackSearch = isCurrentSession() && this.deps.session.trackSearch({ abort });
+    const untrackSearch =
+      this.deps.session.isCurrentSession(options.sessionId) &&
+      this.deps.session.trackSearch({ abort });
 
     // track if this search's session will be send to background
     // if yes, then we don't need to cancel this search when it is aborted
     let isSavedToBackground = false;
     const savedToBackgroundSub =
-      isCurrentSession() &&
+      this.deps.session.isCurrentSession(options.sessionId) &&
       this.deps.session.state$
         .pipe(
           skip(1), // ignore any state, we are only interested in transition x -> BackgroundLoading
-          filter((state) => isCurrentSession() && state === SearchSessionState.BackgroundLoading),
+          filter(
+            (state) =>
+              this.deps.session.isCurrentSession(options.sessionId) &&
+              state === SearchSessionState.BackgroundLoading
+          ),
           take(1)
         )
         .subscribe(() => {
@@ -93,7 +97,8 @@ export class EnhancedSearchInterceptor extends SearchInterceptor {
       finalize(() => {
         this.pendingCount$.next(this.pendingCount$.getValue() - 1);
         cleanup();
-        if (untrackSearch && isCurrentSession()) {
+        if (untrackSearch && this.deps.session.isCurrentSession(options.sessionId)) {
+          // untrack if this search still belongs to current session
           untrackSearch();
         }
         if (savedToBackgroundSub) {