[KQL] Add util for getting field names from KQL expression (#183573)

## Summary

Resolves https://github.com/elastic/kibana/issues/180555.

Adds a utility to kbn-es-query for getting the field names associated
with a KQL expression.

This utility already (mostly) existed in
x-pack/plugins/observability_solution/apm/common/utils/get_kuery_fields.ts
but didn't have test coverage for things like wildcards and nested
fields. This also updates the utility to be a little more robust in
checking the KQL node types.

### Checklist

- [ ]
[Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html)
was added for features that require explanation or tutorials
- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

---------

Co-authored-by: Matthew Kime <matt@mattki.me>
Co-authored-by: Kibana Machine <42973632+kibanamachine@users.noreply.github.com>

packages/kbn-es-query/index.ts

@@ -118,6 +118,8 @@ export {
   toElasticsearchQuery,
   escapeKuery,
   escapeQuotes,
+  getKqlFieldNames,
+  getKqlFieldNamesFromExpression,
 } from './src/kuery';
 
 export {

packages/kbn-es-query/src/kuery/index.ts

@@ -23,5 +23,10 @@ export const toElasticsearchQuery = (...params: Parameters<typeof astToElasticse
 export { KQLSyntaxError } from './kuery_syntax_error';
 export { nodeTypes, nodeBuilder } from './node_types';
 export { fromKueryExpression, toKqlExpression } from './ast';
-export { escapeKuery, escapeQuotes } from './utils';
+export {
+  escapeKuery,
+  escapeQuotes,
+  getKqlFieldNames,
+  getKqlFieldNamesFromExpression,
+} from './utils';
 export type { DslQuery, KueryNode, KueryQueryOptions, KueryParseOptions } from './types';

packages/kbn-es-query/src/kuery/utils/get_kql_fields.test.ts

@@ -0,0 +1,84 @@
+/*
+ * 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 { getKqlFieldNamesFromExpression } from './get_kql_fields';
+
+describe('getKqlFieldNames', () => {
+  it('returns single kuery field', () => {
+    const expression = 'service.name: my-service';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['service.name']);
+  });
+
+  it('returns kuery fields with wildcard', () => {
+    const expression = 'service.name: *';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['service.name']);
+  });
+
+  it('returns multiple fields used AND operator', () => {
+    const expression = 'service.name: my-service AND service.environment: production';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual([
+      'service.name',
+      'service.environment',
+    ]);
+  });
+
+  it('returns multiple kuery fields with OR operator', () => {
+    const expression = 'network.carrier.mcc: test or child.id: 33';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['network.carrier.mcc', 'child.id']);
+  });
+
+  it('returns multiple kuery fields with wildcard', () => {
+    const expression = 'network.carrier.mcc:* or child.id: *';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['network.carrier.mcc', 'child.id']);
+  });
+
+  it('returns single kuery fields with gt operator', () => {
+    const expression = 'transaction.duration.aggregate > 10';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['transaction.duration.aggregate']);
+  });
+
+  it('returns duplicate fields', () => {
+    const expression = 'service.name: my-service and service.name: my-service and trace.id: trace';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual([
+      'service.name',
+      'service.name',
+      'trace.id',
+    ]);
+  });
+
+  it('returns multiple fields with multiple logical operators', () => {
+    const expression =
+      '(service.name:opbeans-* OR service.name:kibana) and (service.environment:production)';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual([
+      'service.name',
+      'service.name',
+      'service.environment',
+    ]);
+  });
+
+  it('returns nested fields', () => {
+    const expression = 'user.names:{ first: "Alice" and last: "White" }';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['user.names']);
+  });
+
+  it('returns wildcard fields', () => {
+    const expression = 'server.*: kibana';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['server.*']);
+  });
+
+  // _field_caps doesn't allow escaped wildcards, so for now this behavior is what we want
+  it('returns escaped fields', () => {
+    const expression = 'server.\\*: kibana';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual(['server.*']);
+  });
+
+  it('do not return if kuery field is null', () => {
+    const expression = 'opbean';
+    expect(getKqlFieldNamesFromExpression(expression)).toEqual([]);
+  });
+});

packages/kbn-es-query/src/kuery/utils/get_kql_fields.ts

@@ -0,0 +1,45 @@
+/*
+ * 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 { fromKueryExpression, KueryNode } from '../../..';
+import { nodeTypes } from '../node_types';
+import { functions } from '../functions';
+
+export function getKqlFieldNamesFromExpression(expression: string): string[] {
+  const node = fromKueryExpression(expression);
+  return getKqlFieldNames(node);
+}
+
+export function getKqlFieldNames(node: KueryNode): string[] {
+  if (nodeTypes.function.isNode(node)) {
+    if (functions.and.isNode(node) || functions.or.isNode(node)) {
+      return node.arguments.reduce<string[]>((result, child) => {
+        return result.concat(getKqlFieldNames(child));
+      }, []);
+    } else if (
+      functions.not.isNode(node) ||
+      functions.exists.isNode(node) ||
+      functions.is.isNode(node) ||
+      functions.nested.isNode(node) ||
+      functions.range.isNode(node)
+    ) {
+      // For each of these field types, we only need to look at the first argument to determine the fields
+      const [fieldNode] = node.arguments;
+      return getKqlFieldNames(fieldNode);
+    } else {
+      throw new Error(`KQL function ${node.function} not supported in getKqlFieldNames`);
+    }
+  } else if (nodeTypes.literal.isNode(node)) {
+    if (node.value === null) return [];
+    return [`${nodeTypes.literal.toElasticsearchQuery(node)}`];
+  } else if (nodeTypes.wildcard.isNode(node)) {
+    return [nodeTypes.wildcard.toElasticsearchQuery(node)];
+  } else {
+    throw new Error(`KQL node type ${node.type} not supported in getKqlFieldNames`);
+  }
+}

packages/kbn-es-query/src/kuery/utils/index.ts

@@ -7,3 +7,4 @@
  */
 
 export { escapeKuery, escapeQuotes } from './escape_kuery';
+export { getKqlFieldNames, getKqlFieldNamesFromExpression } from './get_kql_fields';

x-pack/plugins/observability_solution/apm/common/service_groups.ts

@@ -5,9 +5,8 @@
  * 2.0.
  */
 
-import { fromKueryExpression } from '@kbn/es-query';
+import { getKqlFieldNamesFromExpression } from '@kbn/es-query';
 import { i18n } from '@kbn/i18n';
-import { getKueryFields } from './utils/get_kuery_fields';
 import {
   AGENT_NAME,
   SERVICE_NAME,
@@ -51,7 +50,7 @@ export function validateServiceGroupKuery(kuery: string): {
   message?: string;
 } {
   try {
-    const kueryFields = getKueryFields([fromKueryExpression(kuery)]);
+    const kueryFields = getKqlFieldNamesFromExpression(kuery);
     const unsupportedKueryFields = kueryFields.filter((fieldName) => !isSupportedField(fieldName));
     if (unsupportedKueryFields.length === 0) {
       return { isValidFields: true, isValidSyntax: true };

x-pack/plugins/observability_solution/apm/common/utils/get_kuery_fields.test.ts

@@ -1,70 +0,0 @@
-/*
- * 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; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-import { getKueryFields } from './get_kuery_fields';
-import { fromKueryExpression } from '@kbn/es-query';
-
-describe('get kuery fields', () => {
-  it('returns single kuery field', () => {
-    const kuery = 'service.name: my-service';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual(['service.name']);
-  });
-
-  it('returns kuery fields with wildcard', () => {
-    const kuery = 'service.name: *';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual(['service.name']);
-  });
-
-  it('returns multiple fields used AND operator', () => {
-    const kuery = 'service.name: my-service AND service.environment: production';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual(['service.name', 'service.environment']);
-  });
-
-  it('returns multiple kuery fields with OR operator', () => {
-    const kuery = 'network.carrier.mcc: test or child.id: 33';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual(['network.carrier.mcc', 'child.id']);
-  });
-
-  it('returns multiple kuery fields with wildcard', () => {
-    const kuery = 'network.carrier.mcc:* or child.id: *';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual(['network.carrier.mcc', 'child.id']);
-  });
-
-  it('returns single kuery fields with gt operator', () => {
-    const kuery = 'transaction.duration.aggregate > 10';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual(['transaction.duration.aggregate']);
-  });
-
-  it('returns dublicate fields', () => {
-    const kueries = ['service.name: my-service', 'service.name: my-service and trace.id: trace'];
-
-    const kueryNodes = kueries.map((kuery) => fromKueryExpression(kuery));
-    expect(getKueryFields(kueryNodes)).toEqual(['service.name', 'service.name', 'trace.id']);
-  });
-
-  it('returns multiple fields with multiple logical operators', () => {
-    const kuery =
-      '(service.name:opbeans-* OR service.name:kibana) and (service.environment:production)';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual([
-      'service.name',
-      'service.name',
-      'service.environment',
-    ]);
-  });
-
-  it('do not return if kuery field is null', () => {
-    const kuery = 'opbean';
-    const kueryNode = fromKueryExpression(kuery);
-    expect(getKueryFields([kueryNode])).toEqual([]);
-  });
-});

x-pack/plugins/observability_solution/apm/common/utils/get_kuery_fields.ts

@@ -1,27 +0,0 @@
-/*
- * 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; you may not use this file except in compliance with the Elastic License
- * 2.0.
- */
-
-import { KueryNode } from '@kbn/es-query';
-import { compact } from 'lodash';
-
-export function getKueryFields(nodes: KueryNode[]): string[] {
-  const allFields = nodes
-    .map((node) => {
-      const {
-        arguments: [fieldNameArg],
-      } = node;
-
-      if (fieldNameArg.type === 'function') {
-        return getKueryFields(node.arguments);
-      }
-
-      return fieldNameArg.value;
-    })
-    .flat();
-
-  return compact(allFields);
-}

x-pack/plugins/observability_solution/apm/public/components/shared/unified_search_bar/index.tsx

@@ -6,7 +6,14 @@
  */
 import React, { useCallback, useEffect } from 'react';
 import { i18n } from '@kbn/i18n';
-import { Filter, fromKueryExpression, Query, TimeRange, toElasticsearchQuery } from '@kbn/es-query';
+import {
+  Filter,
+  fromKueryExpression,
+  getKqlFieldNamesFromExpression,
+  Query,
+  TimeRange,
+  toElasticsearchQuery,
+} from '@kbn/es-query';
 import { useHistory, useLocation } from 'react-router-dom';
 import deepEqual from 'fast-deep-equal';
 import { useKibana } from '@kbn/kibana-react-plugin/public';
@@ -27,7 +34,6 @@ import { useLegacyUrlParams } from '../../../context/url_params_context/use_url_
 import { clearCache } from '../../../services/rest/call_api';
 import { useTimeRangeId } from '../../../context/time_range_id/use_time_range_id';
 import { toBoolean, toNumber } from '../../../context/url_params_context/helpers';
-import { getKueryFields } from '../../../../common/utils/get_kuery_fields';
 import { SearchQueryActions } from '../../../services/telemetry';
 
 export const DEFAULT_REFRESH_INTERVAL = 60000;
@@ -228,7 +234,7 @@ export function UnifiedSearchBar({
       if (!res) {
         return;
       }
-      const kueryFields = getKueryFields([fromKueryExpression(query?.query as string)]);
+      const kueryFields = getKqlFieldNamesFromExpression(query?.query as string);
 
       const existingQueryParams = toQuery(location.search);
       const updatedQueryWithTime = {

x-pack/plugins/observability_solution/apm/server/lib/apm_telemetry/collect_data_telemetry/tasks.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 import type * as estypes from '@elastic/elasticsearch/lib/api/typesWithBodyKey';
-import { fromKueryExpression } from '@kbn/es-query';
+import { getKqlFieldNamesFromExpression } from '@kbn/es-query';
 import { ProcessorEvent } from '@kbn/observability-plugin/common';
 import { createHash } from 'crypto';
 import { flatten, merge, pickBy, sortBy, sum, uniq } from 'lodash';
@@ -54,7 +54,6 @@ import {
   SavedServiceGroup,
 } from '../../../../common/service_groups';
 import { asMutableArray } from '../../../../common/utils/as_mutable_array';
-import { getKueryFields } from '../../../../common/utils/get_kuery_fields';
 import { APMError } from '../../../../typings/es_schemas/ui/apm_error';
 import { AgentName } from '../../../../typings/es_schemas/ui/fields/agent';
 import { Span } from '../../../../typings/es_schemas/ui/span';
@@ -1409,11 +1408,10 @@ export const tasks: TelemetryTask[] = [
         namespaces: ['*'],
       });
 
-      const kueryNodes = response.saved_objects.map(({ attributes: { kuery } }) =>
-        fromKueryExpression(kuery)
-      );
-
-      const kueryFields = getKueryFields(kueryNodes);
+      const kueryExpressions = response.saved_objects.map(({ attributes: { kuery } }) => kuery);
+      const kueryFields = kueryExpressions
+        .map(getKqlFieldNamesFromExpression)
+        .reduce((a, b) => a.concat(b), []);
 
       return {
         service_groups: {
@@ -1435,11 +1433,12 @@ export const tasks: TelemetryTask[] = [
         namespaces: ['*'],
       });
 
-      const kueryNodes = response.saved_objects.map(({ attributes: { kuery } }) =>
-        fromKueryExpression(kuery ?? '')
+      const kueryExpressions = response.saved_objects.map(
+        ({ attributes: { kuery } }) => kuery ?? ''
       );
-
-      const kueryFields = getKueryFields(kueryNodes);
+      const kueryFields = kueryExpressions
+        .map(getKqlFieldNamesFromExpression)
+        .reduce((a, b) => a.concat(b), []);
 
       return {
         custom_dashboards: {

x-pack/plugins/observability_solution/observability/public/components/custom_threshold/components/validation.test.ts

@@ -15,9 +15,7 @@ import { EQUATION_REGEX, validateCustomThreshold } from './validation';
 const errorReason = 'this should appear as error reason';
 
 jest.mock('@kbn/es-query', () => {
-  const actual = jest.requireActual('@kbn/es-query');
   return {
-    ...actual,
     buildEsQuery: jest.fn(() => {
       // eslint-disable-next-line no-throw-literal
       throw { shortMessage: errorReason };