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

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,15 @@ 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 isCurrentSession = () =>
+      !!searchSessionId && searchSessionId === dataSearch.session.getSessionId();
+    const untrackSearch =
+      isCurrentSession() &&
+      dataSearch.session.trackSearch({
+        abort: () => {
+          // TODO: support search cancellations
+        },
+      });
 
     try {
       return await http.post('/api/timelion/run', {
@@ -110,7 +121,9 @@ export function getTimelionRequestHandler({
             interval: visParams.interval,
             timezone,
           },
-          sessionId,
+          sessionId: searchSessionId,
+          isRestore: isCurrentSession() ? dataSearch.session.isRestore() : false,
+          isStored: isCurrentSession() ? dataSearch.session.isStored() : false,
         }),
       });
     } catch (e) {
@@ -125,6 +138,10 @@ export function getTimelionRequestHandler({
       } else {
         throw e;
       }
+    } finally {
+      if (untrackSearch && isCurrentSession()) {
+        untrackSearch();
+      }
     }
   };
 }

src/plugins/vis_type_timelion/public/timelion_vis_fn.ts

@@ -72,7 +72,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 };
@@ -82,6 +82,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

@@ -76,6 +76,8 @@ export function runRoute(
             })
           ),
           sessionId: schema.maybe(schema.string()),
+          isRestore: schema.maybe(schema.boolean()),
+          isStored: schema.maybe(schema.boolean()),
         }),
       },
     },

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

@@ -60,19 +60,24 @@ 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,
+            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

@@ -134,6 +134,8 @@ export default new Datasource('es', {
         body,
         {
           sessionId: tlConfig.request?.body.sessionId,
+          isRestore: tlConfig.request?.body.isRestore,
+          isStored: tlConfig.request?.body.isStored,
         },
         tlConfig.context
       )

src/plugins/vis_type_timeseries/common/vis_schema.ts

@@ -281,5 +281,10 @@ export const visPayloadSchema = schema.object({
     min: stringRequired,
     max: stringRequired,
   }),
+
+  // search sessions integration start
   sessionId: schema.maybe(schema.string()),
+  isRestore: schema.maybe(schema.boolean()),
+  isStored: schema.maybe(schema.boolean()),
+  // search sessions integration end
 });

src/plugins/vis_type_timeseries/public/metrics_fn.ts

@@ -65,14 +65,15 @@ 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);
 
     const response = await metricsRequestHandler({
       input,
       visParams,
       uiState,
+      searchSessionId: getSearchSessionId(),
     });
 
     return {

src/plugins/vis_type_timeseries/public/request_handler.ts

@@ -29,12 +29,14 @@ 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);
@@ -47,21 +49,37 @@ export const metricsRequestHandler = async ({
 
     validateInterval(parsedTimeRange, visParams, maxBuckets);
 
-    const resp = await getCoreStart().http.post(ROUTES.VIS_DATA, {
-      body: JSON.stringify({
-        timerange: {
-          timezone,
-          ...parsedTimeRange,
+    const isCurrentSession = () =>
+      !!searchSessionId && searchSessionId === dataSearch.search.session.getSessionId();
+    const untrackSearch =
+      isCurrentSession() &&
+      dataSearch.search.session.trackSearch({
+        abort: () => {
+          // 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,
+          sessionId: searchSessionId,
+          isRestore: isCurrentSession() ? dataSearch.search.session.isRestore() : false,
+          isStored: isCurrentSession() ? dataSearch.search.session.isStored() : false,
+        }),
+      });
+    } finally {
+      if (untrackSearch && isCurrentSession()) {
+        untrackSearch();
+      }
+    }
   }
 
   return {};

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

@@ -49,7 +49,7 @@ export interface ReqFacade<T = unknown> extends FakeRequest {
 export abstract class AbstractSearchStrategy {
   async search(req: ReqFacade<VisPayload>, bodies: any[], indexType?: string) {
     const requests: any[] = [];
-    const { sessionId } = req.payload;
+    const { sessionId, isRestore, isStored } = req.payload;
 
     bodies.forEach((body) => {
       requests.push(
@@ -63,6 +63,8 @@ export abstract class AbstractSearchStrategy {
             },
             {
               sessionId,
+              isRestore,
+              isStored,
             }
           )
           .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(
-            { params },
-            { abortSignal: this.abortSignal, sessionId: search.session.getSessionId() }
-          )
+          .search({ params }, { abortSignal: this.abortSignal, sessionId: this.searchSessionId })
           .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
       );
     }