[Vega] Add the ability to override runtime_mappings (#107875) (#108021)

* Add the ability to override runtime_mappings

Closes: #107855

* runtime_mappings not correctly logged in Inspect panel.

src/plugins/vis_type_vega/public/data_model/search_api.test.ts

@@ -0,0 +1,79 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+import { extendSearchParamsWithRuntimeFields } from './search_api';
+import { dataPluginMock } from '../../../data/public/mocks';
+
+import { getSearchParamsFromRequest, DataPublicPluginStart } from '../../../data/public';
+
+const mockComputedFields = (
+  dataStart: DataPublicPluginStart,
+  index: string,
+  runtimeFields: Record<string, unknown>
+) => {
+  dataStart.indexPatterns.find = jest.fn().mockReturnValue([
+    {
+      title: index,
+      getComputedFields: () => ({
+        runtimeFields,
+      }),
+    },
+  ]);
+};
+
+describe('extendSearchParamsWithRuntimeFields', () => {
+  let dataStart: DataPublicPluginStart;
+
+  beforeEach(() => {
+    dataStart = dataPluginMock.createStartContract();
+  });
+
+  test('should inject default runtime_mappings for known indexes', async () => {
+    const requestParams = {};
+    const runtimeFields = { foo: {} };
+
+    mockComputedFields(dataStart, 'index', runtimeFields);
+
+    expect(
+      await extendSearchParamsWithRuntimeFields(dataStart.indexPatterns, requestParams, 'index')
+    ).toMatchInlineSnapshot(`
+      Object {
+        "body": Object {
+          "runtime_mappings": Object {
+            "foo": Object {},
+          },
+        },
+      }
+    `);
+  });
+
+  test('should use runtime mappings from spec if it is specified', async () => {
+    const requestParams = ({
+      body: {
+        runtime_mappings: {
+          test: {},
+        },
+      },
+    } as unknown) as ReturnType<typeof getSearchParamsFromRequest>;
+    const runtimeFields = { foo: {} };
+
+    mockComputedFields(dataStart, 'index', runtimeFields);
+
+    expect(
+      await extendSearchParamsWithRuntimeFields(dataStart.indexPatterns, requestParams, 'index')
+    ).toMatchInlineSnapshot(`
+      Object {
+        "body": Object {
+          "runtime_mappings": Object {
+            "test": Object {},
+          },
+        },
+      }
+    `);
+  });
+});

src/plugins/vis_type_vega/public/data_model/search_api.ts

@@ -8,31 +8,36 @@
 
 import { combineLatest, from } from 'rxjs';
 import { map, tap, switchMap } from 'rxjs/operators';
-import { CoreStart, IUiSettingsClient } from 'kibana/public';
-import { getData } from '../services';
+import type { CoreStart, IUiSettingsClient } from 'kibana/public';
 import {
   getSearchParamsFromRequest,
   SearchRequest,
   DataPublicPluginStart,
   IEsSearchResponse,
 } from '../../../data/public';
 import { search as dataPluginSearch } from '../../../data/public';
-import { VegaInspectorAdapters } from '../vega_inspector';
-import { RequestResponder } from '../../../inspector/public';
+import type { VegaInspectorAdapters } from '../vega_inspector';
+import type { RequestResponder } from '../../../inspector/public';
 
-const extendSearchParamsWithRuntimeFields = async (
+/** @internal **/
+export const extendSearchParamsWithRuntimeFields = async (
+  indexPatterns: SearchAPIDependencies['indexPatterns'],
   requestParams: ReturnType<typeof getSearchParamsFromRequest>,
   indexPatternString?: string
 ) => {
   if (indexPatternString) {
-    const indexPattern = (await getData().indexPatterns.find(indexPatternString)).find(
-      (index) => index.title === indexPatternString
-    );
-    const runtimeFields = indexPattern?.getComputedFields().runtimeFields;
+    let runtimeMappings = requestParams.body?.runtime_mappings;
+
+    if (!runtimeMappings) {
+      const indexPattern = (await indexPatterns.find(indexPatternString)).find(
+        (index) => index.title === indexPatternString
+      );
+      runtimeMappings = indexPattern?.getComputedFields().runtimeFields;
+    }
 
     return {
       ...requestParams,
-      body: { ...requestParams.body, runtime_mappings: runtimeFields },
+      body: { ...requestParams.body, runtime_mappings: runtimeMappings },
     };
   }
 
@@ -43,6 +48,7 @@ export interface SearchAPIDependencies {
   uiSettings: IUiSettingsClient;
   injectedMetadata: CoreStart['injectedMetadata'];
   search: DataPublicPluginStart['search'];
+  indexPatterns: DataPublicPluginStart['indexPatterns'];
 }
 
 export class SearchAPI {
@@ -54,7 +60,7 @@ export class SearchAPI {
   ) {}
 
   search(searchRequests: SearchRequest[]) {
-    const { search } = this.dependencies;
+    const { search, indexPatterns } = this.dependencies;
     const requestResponders: any = {};
 
     return combineLatest(
@@ -64,15 +70,19 @@ export class SearchAPI {
           getConfig: this.dependencies.uiSettings.get.bind(this.dependencies.uiSettings),
         });
 
-        if (this.inspectorAdapters) {
-          requestResponders[requestId] = this.inspectorAdapters.requests.start(requestId, {
-            ...request,
-            searchSessionId: this.searchSessionId,
-          });
-          requestResponders[requestId].json(requestParams.body);
-        }
-
-        return from(extendSearchParamsWithRuntimeFields(requestParams, request.index)).pipe(
+        return from(
+          extendSearchParamsWithRuntimeFields(indexPatterns, requestParams, request.index)
+        ).pipe(
+          tap((params) => {
+            /** inspect request data **/
+            if (this.inspectorAdapters) {
+              requestResponders[requestId] = this.inspectorAdapters.requests.start(requestId, {
+                ...request,
+                searchSessionId: this.searchSessionId,
+              });
+              requestResponders[requestId].json(params.body);
+            }
+          }),
           switchMap((params) =>
             search
               .search(

src/plugins/vis_type_vega/public/vega_request_handler.ts

@@ -45,10 +45,13 @@ export function createVegaRequestHandler(
     searchSessionId,
   }: VegaRequestHandlerParams) {
     if (!searchAPI) {
+      const { search, indexPatterns } = getData();
+
       searchAPI = new SearchAPI(
         {
           uiSettings,
-          search: getData().search,
+          search,
+          indexPatterns,
           injectedMetadata: getInjectedMetadata(),
         },
         context.abortSignal,

src/plugins/vis_type_vega/public/vega_view/vega_map_view/view.test.ts

@@ -129,6 +129,7 @@ describe('vega_map_view/view', () => {
         JSON.stringify(vegaMap),
         new SearchAPI({
           search: dataPluginStart.search,
+          indexPatterns: dataPluginStart.indexPatterns,
           uiSettings: coreStart.uiSettings,
           injectedMetadata: coreStart.injectedMetadata,
         }),

src/plugins/vis_type_vega/public/vega_visualization.test.js

@@ -92,6 +92,7 @@ describe('VegaVisualizations', () => {
           JSON.stringify(vegaliteGraph),
           new SearchAPI({
             search: dataPluginStart.search,
+            indexPatterns: dataPluginStart.indexPatterns,
             uiSettings: coreStart.uiSettings,
             injectedMetadata: coreStart.injectedMetadata,
           }),
@@ -123,6 +124,7 @@ describe('VegaVisualizations', () => {
           JSON.stringify(vegaGraph),
           new SearchAPI({
             search: dataPluginStart.search,
+            indexPatterns: dataPluginStart.indexPatterns,
             uiSettings: coreStart.uiSettings,
             injectedMetadata: coreStart.injectedMetadata,
           }),