[9.0] [data.search] Add APM instrumentation to search route (#214280) (#215214)

# Backport This will backport the following commits from `main` to `9.0`: - [[data.search] Add APM instrumentation to search route (#214280)](https://github.com/elastic/kibana/pull/214280)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sorenlouv/backport)  Co-authored-by: Lukas Olson <lukas@elastic.co>
2025-04-24 09:48:58 -04:00 · 2025-03-19 19:21:54 +01:00 · 2025-03-19 19:21:54 +01:00 · d10a4a0112
commit d10a4a0112
parent 483f5ecffa
5 changed files with 71 additions and 35 deletions
--- a/src/platform/plugins/shared/data/public/search/search_interceptor/search_interceptor.ts
+++ b/src/platform/plugins/shared/data/public/search/search_interceptor/search_interceptor.ts
@ -443,7 +443,7 @@ export class SearchInterceptor {
        {
          version: '1',
          signal: abortSignal,
-          context: executionContext,
+          context: this.deps.executionContext.withGlobalContext(executionContext),
          body: JSON.stringify({
            ...request,
            ...searchOptions,
--- a/src/platform/plugins/shared/data/server/search/routes/search.test.ts
+++ b/src/platform/plugins/shared/data/server/search/routes/search.test.ts
@ -9,7 +9,7 @@

 import type { MockedKeys } from '@kbn/utility-types-jest';
 import { from } from 'rxjs';
-import { CoreSetup, RequestHandlerContext } from '@kbn/core/server';
+import { CoreSetup, type Logger, RequestHandlerContext } from '@kbn/core/server';
 import { coreMock, httpServerMock } from '@kbn/core/server/mocks';
 import { registerSearchRoute } from './search';
 import { DataPluginStart } from '../../plugin';
@ -19,13 +19,18 @@ import { KbnSearchError } from '../report_search_error';

 describe('Search service', () => {
  let mockCoreSetup: MockedKeys<CoreSetup<{}, DataPluginStart>>;
+  let mockLogger: Logger;

  function mockEsError(message: string, statusCode: number, errBody?: Record<string, any>) {
    return new KbnSearchError(message, statusCode, errBody);
  }

  async function runMockSearch(mockContext: any, mockRequest: any, mockResponse: any) {
-    registerSearchRoute(mockCoreSetup.http.createRouter());
+    registerSearchRoute(
+      mockCoreSetup.http.createRouter(),
+      mockLogger,
+      mockCoreSetup.executionContext
+    );

    const mockRouter = mockCoreSetup.http.createRouter.mock.results[0].value;
    const handler = mockRouter.versioned.post.mock.results[0].value.addVersion.mock.calls[0][1];
@ -33,7 +38,11 @@ describe('Search service', () => {
  }

  async function runMockDelete(mockContext: any, mockRequest: any, mockResponse: any) {
-    registerSearchRoute(mockCoreSetup.http.createRouter());
+    registerSearchRoute(
+      mockCoreSetup.http.createRouter(),
+      mockLogger,
+      mockCoreSetup.executionContext
+    );

    const mockRouter = mockCoreSetup.http.createRouter.mock.results[0].value;
    const handler = mockRouter.versioned.delete.mock.results[0].value.addVersion.mock.calls[0][1];
@ -43,6 +52,7 @@ describe('Search service', () => {
  beforeEach(() => {
    jest.clearAllMocks();
    mockCoreSetup = coreMock.createSetup();
+    mockLogger = coreMock.createPluginInitializerContext().logger.get();
  });

  it('handler calls context.search.search with the given request and strategy', async () => {
--- a/src/platform/plugins/shared/data/server/search/routes/search.ts
+++ b/src/platform/plugins/shared/data/server/search/routes/search.ts
@ -11,13 +11,21 @@ import { first } from 'rxjs';
 import { schema } from '@kbn/config-schema';
 import { reportServerError } from '@kbn/kibana-utils-plugin/server';
 import { IncomingMessage } from 'http';
+import type { KibanaExecutionContext } from '@kbn/core-execution-context-common';
+import { Logger } from '@kbn/logging';
+import type { ExecutionContextSetup } from '@kbn/core-execution-context-server';
+import apm from 'elastic-apm-node';
 import { reportSearchError } from '../report_search_error';
 import { getRequestAbortedSignal } from '../../lib';
 import type { DataPluginRouter } from '../types';

 export const SEARCH_API_BASE_URL = '/internal/search';

-export function registerSearchRoute(router: DataPluginRouter): void {
+export function registerSearchRoute(
+  router: DataPluginRouter,
+  logger: Logger,
+  executionContextSetup: ExecutionContextSetup
+): void {
  router.versioned
    .post({
      path: `${SEARCH_API_BASE_URL}/{strategy}/{id?}`,
@ -65,39 +73,54 @@ export function registerSearchRoute(router: DataPluginRouter): void {
        const { strategy, id } = request.params;
        const abortSignal = getRequestAbortedSignal(request.events.aborted$);

+        let executionContext: KibanaExecutionContext | undefined;
+        const contextHeader = request.headers['x-kbn-context'];
        try {
-          const search = await context.search;
-          const response = await search
-            .search(
-              { ...searchRequest, id },
-              {
-                abortSignal,
-                strategy,
-                legacyHitsTotal,
-                sessionId,
-                isStored,
-                isRestore,
-                retrieveResults,
-                stream,
-              }
-            )
-            .pipe(first())
-            .toPromise();
-
-          if (response && (response.rawResponse as unknown as IncomingMessage).pipe) {
-            return res.ok({
-              body: response.rawResponse,
-              headers: {
-                'kbn-search-is-restored': response.isRestored ? '?1' : '?0',
-                'kbn-search-request-params': JSON.stringify(response.requestParams),
-              },
-            });
-          } else {
-            return res.ok({ body: response });
+          if (contextHeader != null) {
+            executionContext = JSON.parse(
+              decodeURIComponent(Array.isArray(contextHeader) ? contextHeader[0] : contextHeader)
+            );
          }
        } catch (err) {
-          return reportSearchError(res, err);
+          logger.error(`Error parsing search execution context: ${contextHeader}`);
        }
+
+        return executionContextSetup.withContext(executionContext, async () => {
+          apm.addLabels(executionContextSetup.getAsLabels());
+          try {
+            const search = await context.search;
+            const response = await search
+              .search(
+                { ...searchRequest, id },
+                {
+                  abortSignal,
+                  strategy,
+                  legacyHitsTotal,
+                  sessionId,
+                  isStored,
+                  isRestore,
+                  retrieveResults,
+                  stream,
+                }
+              )
+              .pipe(first())
+              .toPromise();
+
+            if (response && (response.rawResponse as unknown as IncomingMessage).pipe) {
+              return res.ok({
+                body: response.rawResponse,
+                headers: {
+                  'kbn-search-is-restored': response.isRestored ? '?1' : '?0',
+                  'kbn-search-request-params': JSON.stringify(response.requestParams),
+                },
+              });
+            } else {
+              return res.ok({ body: response });
+            }
+          } catch (err) {
+            return reportSearchError(res, err);
+          }
+        });
      }
    );

--- a/src/platform/plugins/shared/data/server/search/search_service.ts
+++ b/src/platform/plugins/shared/data/server/search/search_service.ts
@ -148,7 +148,7 @@ export class SearchService {
    const usage = usageCollection ? usageProvider(core) : undefined;

    const router = core.http.createRouter<DataRequestHandlerContext>();
-    registerSearchRoute(router);
+    registerSearchRoute(router, this.logger, core.executionContext);
    registerSessionRoutes(router, this.logger);

    this.sessionService.setup(core, {});
--- a/src/platform/plugins/shared/data/tsconfig.json
+++ b/src/platform/plugins/shared/data/tsconfig.json
@ -55,6 +55,9 @@
    "@kbn/esql-utils",
    "@kbn/shared-ux-table-persist",
    "@kbn/core-deprecations-common",
+    "@kbn/core-execution-context-common",
+    "@kbn/logging",
+    "@kbn/core-execution-context-server",
  ],
  "exclude": [
    "target/**/*",