[Search Session] Fix integration in vega, timelion and TSVB

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.isRestore.mockImplementation(() => isRestore);
-        sessionServiceMock.isStored.mockImplementation(() => isStored);
+        sessionServiceMock.getSearchOptions.mockImplementation(() => ({
+          sessionId,
+          isRestore,
+          isStored,
+        }));
         fetchMock.mockResolvedValue({ result: 200 });
       };
 
@@ -130,72 +132,28 @@ describe('SearchInterceptor', () => {
 
       afterEach(() => {
         const sessionServiceMock = searchMock.session as jest.Mocked<ISessionService>;
-        sessionServiceMock.getSessionId.mockReset();
-        sessionServiceMock.isRestore.mockReset();
-        sessionServiceMock.isStored.mockReset();
+        sessionServiceMock.getSearchOptions.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,
         });
 
         await searchInterceptor.search(mockRequest, { sessionId }).toPromise();
         expect(fetchMock.mock.calls[0][0]).toEqual(
           expect.objectContaining({
-            options: { sessionId: 'sid', isStored: true, isRestore: false },
-          })
-        );
-      });
-
-      test('skips isRestore & isStore in case not a current session Id', async () => {
-        setup({
-          isStored: true,
-          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 },
+            options: { sessionId, isStored: true, isRestore: true },
           })
         );
-      });
 
-      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 },
-          })
-        );
+        expect(
+          (searchMock.session as jest.Mocked<ISessionService>).getSearchOptions
+        ).toHaveBeenCalledWith(sessionId);
       });
     });
 

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,
-          isRestore: isCurrentSession ? this.deps.session.isRestore() : false,
+          ...(options?.sessionId && this.deps.session.getSearchOptions(options.sessionId)),
         },
       },
       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 }),
+    })
+  ),
 });