support "field" AST node

src/platform/packages/shared/kbn-esql-ast/src/builder/builder.ts

@@ -38,6 +38,7 @@ import {
   ESQLTimeInterval,
   ESQLBooleanLiteral,
   ESQLNullLiteral,
+  ESQLField,
 } from '../types';
 import { AstNodeParserFields, AstNodeTemplate, PartialFields } from './types';
 
@@ -56,10 +57,10 @@ export namespace Builder {
     incomplete,
   });
 
-  export const command = (
-    template: PartialFields<AstNodeTemplate<ESQLCommand>, 'args'>,
+  export const command = <Name extends string>(
+    template: PartialFields<AstNodeTemplate<ESQLCommand<Name>>, 'args'>,
     fromParser?: Partial<AstNodeParserFields>
-  ): ESQLCommand => {
+  ): ESQLCommand<Name> => {
     return {
       ...template,
       ...Builder.parserFields(fromParser),
@@ -175,6 +176,20 @@ export namespace Builder {
       return node;
     };
 
+    export const field = (
+      template: Omit<AstNodeTemplate<ESQLField>, 'name' | 'args'>,
+      fromParser?: Partial<AstNodeParserFields>
+    ): ESQLField => {
+      const node: ESQLField = {
+        ...template,
+        ...Builder.parserFields(fromParser),
+        name: '',
+        type: 'field',
+      };
+
+      return node;
+    };
+
     export const order = (
       operand: ESQLColumn,
       template: Omit<AstNodeTemplate<ESQLOrderExpression>, 'name' | 'args'>,

src/platform/packages/shared/kbn-esql-ast/src/parser/__tests__/stats.test.ts

@@ -0,0 +1,167 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { EsqlQuery } from '../../query';
+
+describe('STATS', () => {
+  describe('correctly formatted', () => {
+    it('a simple single aggregate expression', () => {
+      const src = `
+        FROM employees
+          | STATS 123
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {
+              type: 'field',
+              column: {
+                type: 'column',
+                args: [
+                  {
+                    type: 'identifier',
+                    name: '123',
+                  },
+                ],
+              },
+              value: {
+                type: 'literal',
+                value: 123,
+              },
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+
+    it('aggregation function with escaped values', () => {
+      const src = `
+        FROM employees
+          | STATS 123, agg("salary")
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {
+              type: 'field',
+              column: {
+                type: 'column',
+                args: [
+                  {
+                    type: 'identifier',
+                    name: '123',
+                  },
+                ],
+              },
+              value: {
+                type: 'literal',
+                value: 123,
+              },
+            },
+            {
+              type: 'field',
+              column: {
+                type: 'column',
+                args: [
+                  {
+                    type: 'identifier',
+                    name: 'agg("salary")',
+                  },
+                ],
+              },
+              value: {
+                type: 'function',
+                name: 'agg',
+              },
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+
+    it('field column name defined', () => {
+      const src = `
+        FROM employees
+          | STATS my_field = agg("salary")
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {
+              type: 'field',
+              column: {
+                type: 'column',
+                args: [
+                  {
+                    type: 'identifier',
+                    name: 'my_field',
+                  },
+                ],
+              },
+              value: {
+                type: 'function',
+                name: 'agg',
+              },
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+
+    it('parses BY clause', () => {
+      const src = `
+        FROM employees
+          | STATS my_field = agg("salary") BY department
+          | WHERE still_hired == true
+          `;
+      const query = EsqlQuery.fromSrc(src);
+
+      expect(query.errors.length).toBe(0);
+      expect(query.ast.commands).toMatchObject([
+        {},
+        {
+          type: 'command',
+          name: 'stats',
+          args: [
+            {},
+            {
+              type: 'option',
+              name: 'by',
+            },
+          ],
+        },
+        {},
+      ]);
+    });
+  });
+});

src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts

@@ -60,10 +60,13 @@ import type { ESQLAst, ESQLAstMetricsCommand } from '../types';
 import { createJoinCommand } from './factories/join';
 import { createDissectCommand } from './factories/dissect';
 import { createGrokCommand } from './factories/grok';
+import { createStatsCommand } from './factories/stats';
 
 export class ESQLAstBuilderListener implements ESQLParserListener {
   private ast: ESQLAst = [];
 
+  constructor(public src: string) {}
+
   public getAst() {
     return { ast: this.ast };
   }
@@ -173,16 +176,9 @@ export class ESQLAstBuilderListener implements ESQLParserListener {
    * @param ctx the parse tree
    */
   exitStatsCommand(ctx: StatsCommandContext) {
-    const command = createCommand('stats', ctx);
-    this.ast.push(command);
+    const command = createStatsCommand(ctx, this.src);
 
-    // STATS expression is optional
-    if (ctx._stats) {
-      command.args.push(...collectAllAggFields(ctx.aggFields()));
-    }
-    if (ctx._grouping) {
-      command.args.push(...visitByOption(ctx, ctx.fields()));
-    }
+    this.ast.push(command);
   }
 
   /**

src/platform/packages/shared/kbn-esql-ast/src/parser/factories.ts

@@ -84,7 +84,7 @@ const createParserFields = (ctx: ParserRuleContext): AstNodeParserFields => ({
   incomplete: Boolean(ctx.exception),
 });
 
-export const createCommand = (name: string, ctx: ParserRuleContext) =>
+export const createCommand = <Name extends string>(name: Name, ctx: ParserRuleContext) =>
   Builder.command({ name, args: [] }, createParserFields(ctx));
 
 export const createInlineCast = (ctx: InlineCastContext, value: ESQLInlineCast['value']) =>

src/platform/packages/shared/kbn-esql-ast/src/parser/factories/stats.ts

@@ -0,0 +1,42 @@
+/*
+ * 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", the "GNU Affero General Public License v3.0 only", 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", the "GNU Affero General Public
+ * License v3.0 only", or the "Server Side Public License, v 1".
+ */
+
+import { AggFieldContext, StatsCommandContext } from '../../antlr/esql_parser';
+import { ESQLCommand, ESQLField } from '../../types';
+import { createCommand } from '../factories';
+import { createField, visitByOption } from '../walkers';
+
+const createAggField = (ctx: AggFieldContext, src: string): ESQLField => {
+  const fieldCtx = ctx.field();
+  const field = createField(fieldCtx, src);
+
+  return field;
+};
+
+export const createStatsCommand = (ctx: StatsCommandContext, src: string): ESQLCommand<'stats'> => {
+  const command = createCommand('stats', ctx);
+
+  if (ctx._stats) {
+    const fields = ctx.aggFields();
+
+    for (const fieldCtx of fields.aggField_list()) {
+      const node = createAggField(fieldCtx, src);
+
+      command.args.push(node);
+    }
+  }
+
+  if (ctx._grouping) {
+    const options = visitByOption(ctx, ctx.fields());
+
+    command.args.push(...options);
+  }
+
+  return command;
+};

src/platform/packages/shared/kbn-esql-ast/src/parser/parser.ts

@@ -55,11 +55,11 @@ export const getParser = (
   };
 };
 
-export const createParser = (text: string) => {
+export const createParser = (src: string) => {
   const errorListener = new ESQLErrorListener();
-  const parseListener = new ESQLAstBuilderListener();
+  const parseListener = new ESQLAstBuilderListener(src);
 
-  return getParser(CharStreams.fromString(text), errorListener, parseListener);
+  return getParser(CharStreams.fromString(src), errorListener, parseListener);
 };
 
 // These will need to be manually updated whenever the relevant grammar changes.
@@ -106,7 +106,7 @@ export const parse = (text: string | undefined, options: ParseOptions = {}): Par
       return { ast: commands, root: Builder.expression.query(commands), errors: [], tokens: [] };
     }
     const errorListener = new ESQLErrorListener();
-    const parseListener = new ESQLAstBuilderListener();
+    const parseListener = new ESQLAstBuilderListener(text);
     const { tokens, parser } = getParser(
       CharStreams.fromString(text),
       errorListener,

src/platform/packages/shared/kbn-esql-ast/src/parser/walkers.ts

@@ -94,8 +94,10 @@ import {
   ESQLAstField,
   ESQLInlineCast,
   ESQLOrderExpression,
+  ESQLField,
 } from '../types';
 import { firstItem, lastItem } from '../visitor/utils';
+import { Builder } from '../builder';
 
 export function collectAllSourceIdentifiers(ctx: FromCommandContext): ESQLAstItem[] {
   const fromContexts = ctx.getTypedRuleContexts(IndexPatternContext);
@@ -625,3 +627,28 @@ export function visitOrderExpressions(
 
   return ast;
 }
+
+export const createField = (ctx: FieldContext, src: string): ESQLField => {
+  const qualifiedName = ctx.qualifiedName();
+  const hasColumn = !!qualifiedName && !!ctx.ASSIGN();
+  const value = firstItem(collectBooleanExpression(ctx.booleanExpression()))!;
+
+  let column: ESQLColumn | undefined;
+
+  if (hasColumn && qualifiedName) {
+    column = createColumn(qualifiedName);
+  } else {
+    const fieldName = src.slice(value.location.min, value.location.max + 1);
+
+    column = Builder.expression.column({
+      args: [Builder.identifier(fieldName)],
+    });
+  }
+
+  const field = Builder.expression.field({
+    column,
+    value,
+  });
+
+  return field;
+};

src/platform/packages/shared/kbn-esql-ast/src/types.ts

@@ -30,6 +30,7 @@ export type ESQLSingleAstItem =
   | ESQLCommandMode
   | ESQLInlineCast
   | ESQLOrderExpression
+  | ESQLField
   | ESQLUnknownItem;
 
 export type ESQLAstField = ESQLFunction | ESQLColumn;
@@ -317,6 +318,29 @@ export interface ESQLColumn extends ESQLAstBaseItem {
   quoted: boolean;
 }
 
+/**
+ * An ES|QL field definition.
+ */
+export interface ESQLField extends ESQLAstBaseItem {
+  type: 'field';
+
+  /**
+   * The column name of the field. If not specified, the raw text of the
+   * expression is used. Hence, this property is always present.
+   */
+  column: ESQLColumn;
+
+  /**
+   * The expression which defines the value of the field.
+   */
+  value: ESQLAstExpression;
+
+  /**
+   * The WHERE clause of the field, used in STATS command aggregation fields.
+   */
+  where?: ESQLAstExpression;
+}
+
 export interface ESQLList extends ESQLAstBaseItem {
   type: 'list';
   values: ESQLLiteral[];