[Data] Add support of expression AST generation to the search source (#…

…132660) * Update the `kibana_context` function to support multiple queries in arguments * Add the possibility to build an expression AST from the search source * Update visualizations embeddable to reuse search source to expression AST converter * Add the possibility to build an expression AST from the agg configs
elastic · Jun 3, 2022 · 14ba9b3 · 14ba9b3
1 parent 64f005b
commit 14ba9b3
Show file tree

Hide file tree

Showing 59 changed files with 465 additions and 568 deletions.
diff --git a/src/plugins/data/common/search/aggs/agg_configs.test.ts b/src/plugins/data/common/search/aggs/agg_configs.test.ts
@@ -7,6 +7,7 @@
  */
 
 import { keyBy } from 'lodash';
+import { ExpressionAstExpression, buildExpression } from '@kbn/expressions-plugin/common';
 import { AggConfig } from './agg_config';
 import { AggConfigs } from './agg_configs';
 import { AggTypesRegistryStart } from './agg_types_registry';
@@ -787,4 +788,62 @@ describe('AggConfigs', () => {
       });
     });
   });
+
+  describe('#toExpressionAst', () => {
+    function toString(ast: ExpressionAstExpression) {
+      return buildExpression(ast).toString();
+    }
+
+    it('should generate the `index` argument', () => {
+      const ac = new AggConfigs(indexPattern, [], { typesRegistry });
+
+      expect(toString(ac.toExpressionAst())).toMatchInlineSnapshot(
+        `"esaggs index={indexPatternLoad id=\\"logstash-*\\"}"`
+      );
+    });
+
+    it('should generate the `metricsAtAllLevels` if hierarchical', () => {
+      const ac = new AggConfigs(indexPattern, [], { typesRegistry });
+      ac.hierarchical = true;
+
+      expect(toString(ac.toExpressionAst())).toMatchInlineSnapshot(
+        `"esaggs index={indexPatternLoad id=\\"logstash-*\\"} metricsAtAllLevels=true"`
+      );
+    });
+
+    it('should generate the `partialRows` argument', () => {
+      const ac = new AggConfigs(indexPattern, [], { typesRegistry });
+      ac.partialRows = true;
+
+      expect(toString(ac.toExpressionAst())).toMatchInlineSnapshot(
+        `"esaggs index={indexPatternLoad id=\\"logstash-*\\"} partialRows=true"`
+      );
+    });
+
+    it('should generate the `aggs` argument', () => {
+      const configStates = [
+        {
+          enabled: true,
+          type: 'date_histogram',
+          schema: 'segment',
+          params: { field: '@timestamp', interval: '10s' },
+        },
+        { enabled: true, type: 'avg', schema: 'metric', params: { field: 'bytes' } },
+        { enabled: true, type: 'sum', schema: 'metric', params: { field: 'bytes' } },
+        { enabled: true, type: 'min', schema: 'metric', params: { field: 'bytes' } },
+        { enabled: true, type: 'max', schema: 'metric', params: { field: 'bytes' } },
+      ];
+
+      const ac = new AggConfigs(indexPattern, configStates, { typesRegistry });
+
+      expect(toString(ac.toExpressionAst())).toMatchInlineSnapshot(`
+        "esaggs index={indexPatternLoad id=\\"logstash-*\\"} 
+          aggs={aggDateHistogram field=\\"@timestamp\\" useNormalizedEsInterval=true extendToTimeRange=false scaleMetricValues=false interval=\\"10s\\" drop_partials=false min_doc_count=1 extended_bounds={extendedBounds} id=\\"1\\" enabled=true schema=\\"segment\\"}
+          aggs={aggAvg field=\\"bytes\\" id=\\"2\\" enabled=true schema=\\"metric\\"}
+          aggs={aggSum field=\\"bytes\\" emptyAsNull=false id=\\"3\\" enabled=true schema=\\"metric\\"}
+          aggs={aggMin field=\\"bytes\\" id=\\"4\\" enabled=true schema=\\"metric\\"}
+          aggs={aggMax field=\\"bytes\\" id=\\"5\\" enabled=true schema=\\"metric\\"}"
+      `);
+    });
+  });
 });
diff --git a/src/plugins/data/common/search/aggs/agg_configs.ts b/src/plugins/data/common/search/aggs/agg_configs.ts
@@ -13,6 +13,8 @@ import { Assign } from '@kbn/utility-types';
 import { isRangeFilter } from '@kbn/es-query';
 import type { DataView } from '@kbn/data-views-plugin/common';
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
+import type { IndexPatternLoadExpressionFunctionDefinition } from '@kbn/data-views-plugin/common';
+import { buildExpression, buildExpressionFunction } from '@kbn/expressions-plugin/common';
 
 import {
   IEsSearchResponse,
@@ -21,6 +23,7 @@ import {
   RangeFilter,
   // eslint-disable-next-line @kbn/eslint/no-restricted-paths
 } from '../../../public';
+import type { EsaggsExpressionFunctionDefinition } from '../expressions';
 import { AggConfig, AggConfigSerialized, IAggConfig } from './agg_config';
 import { IAggType } from './agg_type';
 import { AggTypesRegistryStart } from './agg_types_registry';
@@ -55,6 +58,7 @@ function parseParentAggs(dslLvlCursor: any, dsl: any) {
 export interface AggConfigsOptions {
   typesRegistry: AggTypesRegistryStart;
   hierarchical?: boolean;
+  partialRows?: boolean;
 }
 
 export type CreateAggConfigParams = Assign<AggConfigSerialized, { type: string | IAggType }>;
@@ -83,6 +87,7 @@ export class AggConfigs {
   public timeFields?: string[];
   public forceNow?: Date;
   public hierarchical?: boolean = false;
+  public partialRows?: boolean = false;
 
   private readonly typesRegistry: AggTypesRegistryStart;
 
@@ -100,6 +105,7 @@ export class AggConfigs {
     this.aggs = [];
     this.indexPattern = indexPattern;
     this.hierarchical = opts.hierarchical;
+    this.partialRows = opts.partialRows;
 
     configStates.forEach((params: any) => this.createAggConfig(params));
   }
@@ -493,4 +499,26 @@ export class AggConfigs {
       this.getRequestAggs().map((agg: AggConfig) => agg.onSearchRequestStart(searchSource, options))
     );
   }
+
+  /**
+   * Generates an expression abstract syntax tree using the `esaggs` expression function.
+   * @returns The expression AST.
+   */
+  toExpressionAst() {
+    return buildExpression([
+      buildExpressionFunction<EsaggsExpressionFunctionDefinition>('esaggs', {
+        index: buildExpression([
+          buildExpressionFunction<IndexPatternLoadExpressionFunctionDefinition>(
+            'indexPatternLoad',
+            {
+              id: this.indexPattern.id!,
+            }
+          ),
+        ]),
+        metricsAtAllLevels: this.hierarchical,
+        partialRows: this.partialRows,
+        aggs: this.aggs.map((agg) => buildExpression(agg.toExpressionAst())),
+      }),
+    ]).toAst();
+  }
 }
diff --git a/src/plugins/data/common/search/expressions/esaggs/request_handler.test.ts b/src/plugins/data/common/search/expressions/esaggs/request_handler.test.ts
@@ -40,6 +40,7 @@ describe('esaggs expression function - public', () => {
       abortSignal: jest.fn() as unknown as jest.Mocked<AbortSignal>,
       aggs: {
         aggs: [{ type: { name: 'terms', postFlightRequest: jest.fn().mockResolvedValue({}) } }],
+        partialRows: false,
         setTimeRange: jest.fn(),
         toDsl: jest.fn().mockReturnValue({ aggs: {} }),
         onSearchRequestStart: jest.fn(),
@@ -49,7 +50,6 @@ describe('esaggs expression function - public', () => {
       filters: undefined,
       indexPattern: { id: 'logstash-*' } as unknown as jest.Mocked<DataView>,
       inspectorAdapters: {},
-      partialRows: false,
       query: undefined,
       searchSessionId: 'abc123',
       searchSourceService: searchSourceCommonMock,
@@ -147,7 +147,7 @@ describe('esaggs expression function - public', () => {
       mockParams.aggs,
       {},
       {
-        partialRows: mockParams.partialRows,
+        partialRows: mockParams.aggs.partialRows,
         timeRange: mockParams.timeRange,
       }
     );

diff --git a/src/plugins/data/common/search/expressions/esaggs/request_handler.ts b/src/plugins/data/common/search/expressions/esaggs/request_handler.ts
@@ -24,8 +24,6 @@ interface RequestHandlerParams {
   filters?: Filter[];
   indexPattern?: DataView;
   inspectorAdapters: Adapters;
-  metricsAtAllLevels?: boolean;
-  partialRows?: boolean;
   query?: Query;
   searchSessionId?: string;
   searchSourceService: ISearchStartSearchSource;
@@ -41,7 +39,6 @@ export const handleRequest = ({
   filters,
   indexPattern,
   inspectorAdapters,
-  partialRows,
   query,
   searchSessionId,
   searchSourceService,
@@ -131,7 +128,7 @@ export const handleRequest = ({
             const parsedTimeRange = timeRange ? calculateBounds(timeRange, { forceNow }) : null;
             const tabifyParams = {
               metricsAtAllLevels: aggs.hierarchical,
-              partialRows,
+              partialRows: aggs.partialRows,
               timeRange: parsedTimeRange
                 ? { from: parsedTimeRange.min, to: parsedTimeRange.max, timeFields: allTimeFields }
                 : undefined,

diff --git a/src/plugins/data/common/search/expressions/kibana_context.test.ts b/src/plugins/data/common/search/expressions/kibana_context.test.ts
@@ -89,16 +89,28 @@ describe('kibanaContextFn', () => {
     } as any);
     const args = {
       ...emptyArgs,
-      q: {
-        type: 'kibana_query' as 'kibana_query',
-        language: 'test',
-        query: {
-          type: 'test',
-          match_phrase: {
-            test: 'something2',
+      q: [
+        {
+          type: 'kibana_query' as 'kibana_query',
+          language: 'test',
+          query: {
+            type: 'test',
+            match_phrase: {
+              test: 'something2',
+            },
           },
         },
-      },
+        {
+          type: 'kibana_query' as 'kibana_query',
+          language: 'test',
+          query: {
+            type: 'test',
+            match_phrase: {
+              test: 'something3',
+            },
+          },
+        },
+      ],
       savedSearchId: 'test',
     };
     const input: KibanaContext = {
@@ -183,6 +195,16 @@ describe('kibanaContextFn', () => {
           },
         },
       },
+      {
+        type: 'kibana_query',
+        language: 'test',
+        query: {
+          type: 'test',
+          match_phrase: {
+            test: 'something3',
+          },
+        },
+      },
       {
         language: 'kuery',
         query: {

diff --git a/src/plugins/data/common/search/expressions/kibana_context.ts b/src/plugins/data/common/search/expressions/kibana_context.ts
@@ -24,7 +24,7 @@ export interface KibanaContextStartDependencies {
 }
 
 interface Arguments {
-  q?: KibanaQueryOutput | null;
+  q?: KibanaQueryOutput[] | null;
   filters?: KibanaFilter[] | null;
   timeRange?: KibanaTimerangeOutput | null;
   savedSearchId?: string | null;
@@ -62,8 +62,8 @@ export const getKibanaContextFn = (
     args: {
       q: {
         types: ['kibana_query', 'null'],
+        multi: true,
         aliases: ['query', '_'],
-        default: null,
         help: i18n.translate('data.search.functions.kibana_context.q.help', {
           defaultMessage: 'Specify Kibana free form text query',
         }),
@@ -123,7 +123,7 @@ export const getKibanaContextFn = (
       const { savedObjectsClient } = await getStartDependencies(getKibanaRequest);
 
       const timeRange = args.timeRange || input?.timeRange;
-      let queries = mergeQueries(input?.query, args?.q || []);
+      let queries = mergeQueries(input?.query, args?.q?.filter(Boolean) || []);
       let filters = [
         ...(input?.filters || []),
         ...((args?.filters?.map(unboxExpressionValue) || []) as Filter[]),

diff --git a/src/plugins/data/common/search/search_source/create_search_source.test.ts b/src/plugins/data/common/search/search_source/create_search_source.test.ts
@@ -19,6 +19,7 @@ describe('createSearchSource', () => {
 
   beforeEach(() => {
     dependencies = {
+      aggs: {} as SearchSourceDependencies['aggs'],
       getConfig: jest.fn(),
       search: jest.fn(),
       onResponse: (req, res) => res,

diff --git a/src/plugins/data/common/search/search_source/mocks.ts b/src/plugins/data/common/search/search_source/mocks.ts
@@ -10,7 +10,7 @@ import { of } from 'rxjs';
 import type { MockedKeys } from '@kbn/utility-types/jest';
 import { uiSettingsServiceMock } from '@kbn/core/public/mocks';
 
-import { SearchSource } from './search_source';
+import { SearchSource, SearchSourceDependencies } from './search_source';
 import { ISearchStartSearchSource, ISearchSource, SearchSourceFields } from './types';
 
 export const searchSourceInstanceMock: MockedKeys<ISearchSource> = {
@@ -35,6 +35,7 @@ export const searchSourceInstanceMock: MockedKeys<ISearchSource> = {
   history: [],
   getSerializedFields: jest.fn(),
   serialize: jest.fn(),
+  toExpressionAst: jest.fn(),
 };
 
 export const searchSourceCommonMock: jest.Mocked<ISearchStartSearchSource> = {
@@ -48,6 +49,9 @@ export const searchSourceCommonMock: jest.Mocked<ISearchStartSearchSource> = {
 
 export const createSearchSourceMock = (fields?: SearchSourceFields, response?: any) =>
   new SearchSource(fields, {
+    aggs: {
+      createAggConfigs: jest.fn(),
+    } as unknown as SearchSourceDependencies['aggs'],
     getConfig: uiSettingsServiceMock.createStartContract().get,
     search: jest.fn().mockReturnValue(
       of(