[skip-ci] [data.search.session] Add extend functionality and expose route

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

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

src/plugins/data/server/index.ts

@@ -246,6 +246,7 @@ export {
   SearchUsage,
   SessionService,
   ISessionService,
+  ISessionServiceDependencies,
 } from './search';
 
 // Search namespace

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

@@ -77,7 +77,7 @@ import { KbnServerError } from '../../../kibana_utils/server';
 
 declare module 'src/core/server' {
   interface RequestHandlerContext {
-    search?: ISearchClient & { session: IScopedSessionService };
+    search?: ISearchClient & { session: ISearchClient };
   }
 }
 
@@ -131,6 +131,8 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     registerSearchRoute(router);
     registerMsearchRoute(router, routeDependencies);
 
+    this.sessionService.setup(core, {});
+
     core.getStartServices().then(([coreStart]) => {
       this.coreStart = coreStart;
     });
@@ -293,15 +295,14 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     SearchStrategyRequest extends IKibanaSearchRequest = IEsSearchRequest,
     SearchStrategyResponse extends IKibanaSearchResponse = IEsSearchResponse
   >(
-    session: IScopedSessionService,
     request: SearchStrategyRequest,
     options: ISearchOptions,
     deps: SearchStrategyDependencies
   ) => {
     const strategy = this.getSearchStrategy<SearchStrategyRequest, SearchStrategyResponse>(
       options.strategy
     );
-    return session.search(strategy, request, options, deps);
+    return strategy.search(request, options, deps);
   };
 
   private cancel = (id: string, options: ISearchOptions, deps: SearchStrategyDependencies) => {
@@ -346,17 +347,23 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
     const { elasticsearch, savedObjects, uiSettings } = core;
     const getSessionAsScoped = this.sessionService.asScopedProvider(core);
     return (request: KibanaRequest): ISearchClient => {
-      const scopedSession = getSessionAsScoped(request);
       const savedObjectsClient = savedObjects.getScopedClient(request);
       const deps = {
         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),
+        search: (searchRequest, options = {}) => {
+          return options.sessionId
+            ? getSessionAsScoped(request).search(searchRequest, options)
+            : this.search(searchRequest, options, deps);
+        },
+        cancel: (id, options = {}) => {
+          return options.sessionId
+            ? getSessionAsScoped(request).cancel(id, options)
+            : this.cancel(id, options, deps);
+        },
         extend: (id, keepAlive, options = {}) => this.extend(id, keepAlive, options, deps),
       };
     };

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

@@ -17,27 +17,69 @@
  * under the License.
  */
 
-import { CoreStart, KibanaRequest } from 'kibana/server';
-import { IKibanaSearchRequest, IKibanaSearchResponse } from '../../../common';
-import { ISearchStrategy } from '../types';
-import { ISessionService } from './types';
+import { CoreSetup, KibanaRequest } from 'kibana/server';
+import { ISessionService, ISessionServiceDependencies } from './types';
+import { DataPluginStart } from '../../plugin';
+import {
+  IKibanaSearchRequest,
+  IKibanaSearchResponse,
+  ISearchClient,
+  ISearchOptions,
+} from '../../../common';
 
 /**
  * The OSS session service. See data_enhanced in X-Pack for the background session service.
  */
 export class SessionService implements ISessionService {
+  protected searchAsScoped!: (request: KibanaRequest) => ISearchClient;
+
   constructor() {}
 
-  public search<Request extends IKibanaSearchRequest, Response extends IKibanaSearchResponse>(
-    strategy: ISearchStrategy<Request, Response>,
-    ...args: Parameters<ISearchStrategy<Request, Response>['search']>
+  public setup(core: CoreSetup<{}, DataPluginStart>) {
+    core.getStartServices().then(([, , dataStart]) => {
+      this.searchAsScoped = dataStart.search.asScoped;
+    });
+  }
+
+  public start() {}
+
+  public stop() {}
+
+  /**
+   * Forward this search request to the search service, removing the sessionId (so the search client
+   * doesn't forward back to this service).
+   */
+  private search<Request extends IKibanaSearchRequest, Response extends IKibanaSearchResponse>(
+    { searchClient }: ISessionServiceDependencies,
+    request: Request,
+    { sessionId, ...options }: ISearchOptions
   ) {
-    return strategy.search(...args);
+    return searchClient.search<Request, Response>(request, options);
   }
 
-  public asScopedProvider(core: CoreStart) {
-    return (request: KibanaRequest) => ({
-      search: this.search,
-    });
+  /**
+   * Forward this search request to the search service, removing the sessionId (so the search client
+   * doesn't forward back to this service).
+   */
+  private cancel(
+    { searchClient }: ISessionServiceDependencies,
+    id: string,
+    { sessionId, ...options }: ISearchOptions = {}
+  ) {
+    return searchClient.cancel(id, options);
+  }
+
+  public asScopedProvider() {
+    return (request: KibanaRequest) => {
+      const searchClient = this.searchAsScoped(request);
+      const deps = { searchClient };
+      return {
+        search: <Request extends IKibanaSearchRequest, Response extends IKibanaSearchResponse>(
+          searchRequest: Request,
+          options: ISearchOptions = {}
+        ) => this.search<Request, Response>(deps, searchRequest, options),
+        cancel: this.cancel.bind(this, deps),
+      };
+    };
   }
 }

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

@@ -17,19 +17,13 @@
  * under the License.
  */
 
-import { Observable } from 'rxjs';
-import { CoreStart, KibanaRequest } from 'kibana/server';
-import { ISearchStrategy } from '../types';
-import { IKibanaSearchRequest, IKibanaSearchResponse } from '../../../common/search';
+import { CoreStart, KibanaRequest, Plugin } from 'kibana/server';
+import { ISearchClient } from '../../../common';
 
-export interface IScopedSessionService {
-  search: <Request extends IKibanaSearchRequest, Response extends IKibanaSearchResponse>(
-    strategy: ISearchStrategy<Request, Response>,
-    ...args: Parameters<ISearchStrategy<Request, Response>['search']>
-  ) => Observable<Response>;
-  [prop: string]: any;
+export interface ISessionService<T extends ISearchClient = ISearchClient> extends Plugin {
+  asScopedProvider: (core: CoreStart) => (request: KibanaRequest) => T;
 }
 
-export interface ISessionService {
-  asScopedProvider: (core: CoreStart) => (request: KibanaRequest) => IScopedSessionService;
+export interface ISessionServiceDependencies {
+  searchClient: ISearchClient;
 }

x-pack/plugins/data_enhanced/server/plugin.ts

@@ -15,17 +15,29 @@ import { ENHANCED_ES_SEARCH_STRATEGY, EQL_SEARCH_STRATEGY } from '../common';
 import { registerSessionRoutes } from './routes';
 import { backgroundSessionMapping } from './saved_objects';
 import {
+  BackgroundSessionClient,
   BackgroundSessionService,
   enhancedEsSearchStrategyProvider,
   eqlSearchStrategyProvider,
 } from './search';
 import { getUiSettings } from './ui_settings';
+import { ISearchClient } from '../../../../src/plugins/data/common';
 
 interface SetupDependencies {
   data: DataPluginSetup;
   usageCollection?: UsageCollectionSetup;
 }
 
+export interface DataEnhancedPluginsStart {
+  data: DataPluginStart;
+}
+
+declare module 'src/core/server' {
+  interface RequestHandlerContext {
+    search?: ISearchClient & { session: BackgroundSessionClient };
+  }
+}
+
 export class EnhancedDataServerPlugin implements Plugin<void, void, SetupDependencies> {
   private readonly logger: Logger;
   private sessionService!: BackgroundSessionService;
@@ -34,7 +46,7 @@ export class EnhancedDataServerPlugin implements Plugin<void, void, SetupDepende
     this.logger = initializerContext.logger.get('data_enhanced');
   }
 
-  public setup(core: CoreSetup<DataPluginStart>, deps: SetupDependencies) {
+  public setup(core: CoreSetup<DataEnhancedPluginsStart>, deps: SetupDependencies) {
     const usage = deps.usageCollection ? usageProvider(core) : undefined;
 
     core.uiSettings.register(getUiSettings());
@@ -55,6 +67,7 @@ export class EnhancedDataServerPlugin implements Plugin<void, void, SetupDepende
     );
 
     this.sessionService = new BackgroundSessionService(this.logger);
+    this.sessionService.setup(core);
 
     deps.data.__enhance({
       search: {

x-pack/plugins/data_enhanced/server/routes/session.ts

@@ -160,4 +160,39 @@ export function registerSessionRoutes(router: IRouter): void {
       }
     }
   );
+
+  router.post(
+    {
+      path: '/internal/session/{id}/_extend',
+      validate: {
+        params: schema.object({
+          id: schema.string(),
+        }),
+        body: schema.object({
+          keepAlive: schema.string(),
+        }),
+      },
+    },
+    async (context, request, res) => {
+      const { id } = request.params;
+      const { keepAlive } = request.body;
+      try {
+        const response = await context.search!.session.extend(id, keepAlive);
+
+        return res.ok({
+          body: response,
+        });
+      } catch (err) {
+        return res.customError({
+          statusCode: err.statusCode || 500,
+          body: {
+            message: err.message,
+            attributes: {
+              error: err.body?.error || err.message,
+            },
+          },
+        });
+      }
+    }
+  );
 }