[data.search] Add search session methods to search service contract

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

@@ -79,6 +79,12 @@ export class SessionsClient {
     });
   }
 
+  public extend(sessionId: string, keepAlive: string): Promise<SavedObjectsFindResponse> {
+    return this.http!.post(`/internal/session/${encodeURIComponent(sessionId)}/_extend`, {
+      body: JSON.stringify({ keepAlive }),
+    });
+  }
+
   public delete(sessionId: string): Promise<void> {
     return this.http!.delete(`/internal/session/${encodeURIComponent(sessionId)}`);
   }

src/plugins/data/server/index.ts

@@ -244,8 +244,8 @@ export {
   searchUsageObserver,
   shimAbortSignal,
   SearchUsage,
-  SessionService,
-  ISessionService,
+  SearchSessionService,
+  ISearchSessionService,
 } from './search';
 
 // Search namespace

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

@@ -17,7 +17,7 @@
  * under the License.
  */
 
-import { BehaviorSubject, Observable } from 'rxjs';
+import { BehaviorSubject, from, Observable } from 'rxjs';
 import { pick } from 'lodash';
 import {
   CoreSetup,
@@ -29,10 +29,11 @@ import {
   SharedGlobalConfig,
   StartServicesAccessor,
 } from 'src/core/server';
-import { catchError, first, map } from 'rxjs/operators';
+import { catchError, first, map, switchMap } from 'rxjs/operators';
 import { BfetchServerSetup } from 'src/plugins/bfetch/server';
 import { ExpressionsServerSetup } from 'src/plugins/expressions/server';
 import {
+  IScopedSearchClient,
   ISearchSetup,
   ISearchStart,
   ISearchStrategy,
@@ -56,7 +57,6 @@ import {
   IEsSearchResponse,
   IKibanaSearchRequest,
   IKibanaSearchResponse,
-  ISearchClient,
   ISearchOptions,
   kibana,
   kibanaContext,
@@ -72,12 +72,13 @@ import {
 } from '../../common/search/aggs/buckets/shard_delay';
 import { aggShardDelay } from '../../common/search/aggs/buckets/shard_delay_fn';
 import { ConfigSchema } from '../../config';
-import { SessionService, IScopedSessionService, ISessionService } from './session';
+import { ISearchSessionService, SearchSessionService } from './session';
 import { KbnServerError } from '../../../kibana_utils/server';
+import { tapFirst } from '../../common';
 
 declare module 'src/core/server' {
   interface RequestHandlerContext {
-    search?: ISearchClient & { session: IScopedSessionService };
+    search?: IScopedSearchClient;
   }
 }
 
@@ -107,14 +108,14 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
   private readonly searchSourceService = new SearchSourceService();
   private defaultSearchStrategyName: string = ES_SEARCH_STRATEGY;
   private searchStrategies: StrategyMap = {};
-  private sessionService: ISessionService;
-  private coreStart?: CoreStart;
+  private sessionService: ISearchSessionService;
+  private asScoped!: (request: KibanaRequest) => IScopedSearchClient;
 
   constructor(
     private initializerContext: PluginInitializerContext<ConfigSchema>,
     private readonly logger: Logger
   ) {
-    this.sessionService = new SessionService();
+    this.sessionService = new SearchSessionService();
   }
 
   public setup(
@@ -131,14 +132,8 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     registerSearchRoute(router);
     registerMsearchRoute(router, routeDependencies);
 
-    core.getStartServices().then(([coreStart]) => {
-      this.coreStart = coreStart;
-    });
-
     core.http.registerRouteHandlerContext('search', async (context, request) => {
-      const search = this.asScopedProvider(this.coreStart!)(request);
-      const session = this.sessionService.asScopedProvider(this.coreStart!)(request);
-      return { ...search, session };
+      return this.asScoped(request);
     });
 
     this.registerSearchStrategy(
@@ -154,7 +149,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
       { request: IKibanaSearchResponse; options?: ISearchOptions },
       any
     >('/internal/bsearch', (request) => {
-      const search = this.asScopedProvider(this.coreStart!)(request);
+      const search = this.asScoped(request);
 
       return {
         onBatchItem: async ({ request: requestData, options }) => {
@@ -229,15 +224,15 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     { fieldFormats, indexPatterns }: SearchServiceStartDependencies
   ): ISearchStart {
     const { elasticsearch, savedObjects, uiSettings } = core;
-    const asScoped = this.asScopedProvider(core);
+    this.asScoped = this.asScopedProvider(core);
     return {
       aggs: this.aggsService.start({
         fieldFormats,
         uiSettings,
         indexPatterns,
       }),
       getSearchStrategy: this.getSearchStrategy,
-      asScoped,
+      asScoped: this.asScoped,
       searchSource: {
         asScoped: async (request: KibanaRequest) => {
           const esClient = elasticsearch.client.asScoped(request);
@@ -256,7 +251,7 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
 
           const searchSourceDependencies: SearchSourceDependencies = {
             getConfig: <T = any>(key: string): T => uiSettingsCache[key],
-            search: asScoped(request).search,
+            search: this.asScoped(request).search,
             onResponse: (req, res) => shimHitsTotal(res),
             legacy: {
               callMsearch: getCallMsearch({
@@ -289,22 +284,50 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     this.searchStrategies[name] = strategy;
   };
 
-  private search = <
+  private getSearchStrategy = <
     SearchStrategyRequest extends IKibanaSearchRequest = IEsSearchRequest,
     SearchStrategyResponse extends IKibanaSearchResponse = IEsSearchResponse
   >(
-    session: IScopedSessionService,
+    name: string = this.defaultSearchStrategyName
+  ): ISearchStrategy<SearchStrategyRequest, SearchStrategyResponse> => {
+    this.logger.debug(`Get strategy ${name}`);
+    const strategy = this.searchStrategies[name];
+    if (!strategy) {
+      throw new KbnServerError(`Search strategy ${name} not found`, 404);
+    }
+    return strategy;
+  };
+
+  private search = <
+    SearchStrategyRequest extends IKibanaSearchRequest,
+    SearchStrategyResponse extends IKibanaSearchResponse
+  >(
+    deps: SearchStrategyDependencies,
     request: SearchStrategyRequest,
-    options: ISearchOptions,
-    deps: SearchStrategyDependencies
+    options: ISearchOptions
   ) => {
     const strategy = this.getSearchStrategy<SearchStrategyRequest, SearchStrategyResponse>(
       options.strategy
     );
-    return session.search(strategy, request, options, deps);
+
+    const getSearchRequest = async () =>
+      !options.sessionId || !options.isRestore || request.id
+        ? request
+        : {
+            ...request,
+            id: await deps.searchSessionsClient.getId(request, options),
+          };
+
+    return from(getSearchRequest()).pipe(
+      switchMap((searchRequest) => strategy.search(searchRequest, options, deps)),
+      tapFirst((response) => {
+        if (request.id || !options.sessionId || !response.id || options.isRestore) return;
+        deps.searchSessionsClient.trackId(request, response.id, options);
+      })
+    );
   };
 
-  private cancel = (id: string, options: ISearchOptions, deps: SearchStrategyDependencies) => {
+  private cancel = (deps: SearchStrategyDependencies, id: string, options: ISearchOptions = {}) => {
     const strategy = this.getSearchStrategy(options.strategy);
     if (!strategy.cancel) {
       throw new KbnServerError(
@@ -316,10 +339,10 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
   };
 
   private extend = (
+    deps: SearchStrategyDependencies,
     id: string,
     keepAlive: string,
-    options: ISearchOptions,
-    deps: SearchStrategyDependencies
+    options: ISearchOptions = {}
   ) => {
     const strategy = this.getSearchStrategy(options.strategy);
     if (!strategy.extend) {
@@ -328,36 +351,69 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     return strategy.extend(id, keepAlive, options, deps);
   };
 
-  private getSearchStrategy = <
-    SearchStrategyRequest extends IKibanaSearchRequest = IEsSearchRequest,
-    SearchStrategyResponse extends IKibanaSearchResponse = IEsSearchResponse
-  >(
-    name: string = this.defaultSearchStrategyName
-  ): ISearchStrategy<SearchStrategyRequest, SearchStrategyResponse> => {
-    this.logger.debug(`Get strategy ${name}`);
-    const strategy = this.searchStrategies[name];
-    if (!strategy) {
-      throw new KbnServerError(`Search strategy ${name} not found`, 404);
+  private deleteSession = async (deps: SearchStrategyDependencies, sessionId: string) => {
+    const searchIdMapping = await deps.searchSessionsClient.getSearchIdMapping(sessionId);
+
+    for (const [searchId, strategyName] of searchIdMapping.entries()) {
+      const searchOptions = {
+        sessionId,
+        strategy: strategyName,
+        isStored: true,
+      };
+      await this.cancel(deps, searchId, searchOptions);
     }
-    return strategy;
+
+    return deps.searchSessionsClient.delete(sessionId);
+  };
+
+  private extendSession = async (
+    deps: SearchStrategyDependencies,
+    sessionId: string,
+    keepAlive: string
+  ) => {
+    const searchIdMapping = await deps.searchSessionsClient.getSearchIdMapping(sessionId);
+
+    for (const [searchId, strategyName] of searchIdMapping.entries()) {
+      const searchOptions = {
+        sessionId,
+        strategy: strategyName,
+        isStored: true,
+      };
+      await this.extend(deps, searchId, keepAlive, searchOptions);
+    }
+
+    return deps.searchSessionsClient.extend(sessionId, keepAlive);
   };
 
   private asScopedProvider = (core: CoreStart) => {
     const { elasticsearch, savedObjects, uiSettings } = core;
     const getSessionAsScoped = this.sessionService.asScopedProvider(core);
-    return (request: KibanaRequest): ISearchClient => {
-      const scopedSession = getSessionAsScoped(request);
+    return (request: KibanaRequest): IScopedSearchClient => {
       const savedObjectsClient = savedObjects.getScopedClient(request);
+      const searchSessionsClient = getSessionAsScoped(request);
       const deps = {
+        searchSessionsClient,
         savedObjectsClient,
         esClient: elasticsearch.client.asScoped(request),
         uiSettingsClient: uiSettings.asScopedToClient(savedObjectsClient),
       };
       return {
-        search: (searchRequest, options = {}) =>
-          this.search(scopedSession, searchRequest, options, deps),
-        cancel: (id, options = {}) => this.cancel(id, options, deps),
-        extend: (id, keepAlive, options = {}) => this.extend(id, keepAlive, options, deps),
+        search: <
+          SearchStrategyRequest extends IKibanaSearchRequest = IEsSearchRequest,
+          SearchStrategyResponse extends IKibanaSearchResponse = IEsSearchResponse
+        >(
+          searchRequest: SearchStrategyRequest,
+          options: ISearchOptions = {}
+        ) =>
+          this.search<SearchStrategyRequest, SearchStrategyResponse>(deps, searchRequest, options),
+        cancel: this.cancel.bind(this, deps),
+        extend: this.extend.bind(this, deps),
+        saveSession: searchSessionsClient.save,
+        getSession: searchSessionsClient.get,
+        findSessions: searchSessionsClient.find,
+        updateSession: searchSessionsClient.update,
+        extendSession: this.extendSession.bind(this, deps),
+        deleteSession: this.deleteSession.bind(this, deps),
       };
     };
   };

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

@@ -17,22 +17,6 @@
  * under the License.
  */
 
-import { of } from 'rxjs';
-import { SearchStrategyDependencies } from '../types';
-import { SessionService } from './session_service';
+// import { SearchSessionService } from './session_service';
 
-describe('SessionService', () => {
-  it('search invokes `strategy.search`', async () => {
-    const service = new SessionService();
-    const mockSearch = jest.fn().mockReturnValue(of({}));
-    const mockStrategy = { search: mockSearch };
-    const mockRequest = { id: 'bar' };
-    const mockOptions = { sessionId: '1234' };
-    const mockDeps = { savedObjectsClient: {} } as SearchStrategyDependencies;
-
-    await service.search(mockStrategy, mockRequest, mockOptions, mockDeps);
-
-    expect(mockSearch).toHaveBeenCalled();
-    expect(mockSearch).toHaveBeenCalledWith(mockRequest, mockOptions, mockDeps);
-  });
-});
+describe('SearchSessionService', () => {});

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

@@ -17,27 +17,40 @@
  * under the License.
  */
 
-import { CoreStart, KibanaRequest } from 'kibana/server';
-import { IKibanaSearchRequest, IKibanaSearchResponse } from '../../../common';
-import { ISearchStrategy } from '../types';
-import { ISessionService } from './types';
+import { ISearchSessionService } from './types';
 
 /**
- * The OSS session service. See data_enhanced in X-Pack for the search session service.
+ * The OSS session service, which leaves most search session-related logic unimplemented.
+ * @see x-pack/plugins/data_enhanced/server/search/session/session_service.ts
  */
-export class SessionService implements ISessionService {
+export class SearchSessionService implements ISearchSessionService {
   constructor() {}
 
-  public search<Request extends IKibanaSearchRequest, Response extends IKibanaSearchResponse>(
-    strategy: ISearchStrategy<Request, Response>,
-    ...args: Parameters<ISearchStrategy<Request, Response>['search']>
-  ) {
-    return strategy.search(...args);
-  }
-
-  public asScopedProvider(core: CoreStart) {
-    return (request: KibanaRequest) => ({
-      search: this.search,
+  public asScopedProvider() {
+    return () => ({
+      getId: () => {
+        throw new Error('getId not implemented in OSS search session service');
+      },
+      trackId: async () => {},
+      getSearchIdMapping: async () => new Map<string, string>(),
+      save: async () => {
+        throw new Error('save not implemented in OSS search session service');
+      },
+      get: async () => {
+        throw new Error('get not implemented in OSS search session service');
+      },
+      find: async () => {
+        throw new Error('find not implemented in OSS search session service');
+      },
+      update: async () => {
+        throw new Error('update not implemented in OSS search session service');
+      },
+      extend: async () => {
+        throw new Error('extend not implemented in OSS search session service');
+      },
+      delete: async () => {
+        throw new Error('delete not implemented in OSS search session service');
+      },
     });
   }
 }