diff --git a/src/platform/packages/shared/kbn-esql-ast/src/builder/builder.ts b/src/platform/packages/shared/kbn-esql-ast/src/builder/builder.ts index 7988b0e953559..369321a272d64 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/builder/builder.ts +++ b/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, 'args'>, + export const command = ( + template: PartialFields>, 'args'>, fromParser?: Partial - ): ESQLCommand => { + ): ESQLCommand => { return { ...template, ...Builder.parserFields(fromParser), @@ -175,6 +176,20 @@ export namespace Builder { return node; }; + export const field = ( + template: Omit, 'name' | 'args'>, + fromParser?: Partial + ): ESQLField => { + const node: ESQLField = { + ...template, + ...Builder.parserFields(fromParser), + name: '', + type: 'field', + }; + + return node; + }; + export const order = ( operand: ESQLColumn, template: Omit, 'name' | 'args'>, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/__tests__/stats.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/__tests__/stats.test.ts new file mode 100644 index 0000000000000..0fdadc0fb6a7b --- /dev/null +++ b/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', + }, + ], + }, + {}, + ]); + }); + }); +}); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts index ea9201b800721..9c8c524d8c340 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/parser/esql_ast_builder_listener.ts +++ b/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); } /** diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/factories.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/factories.ts index 0db09c0f9dfa7..a63ce09301c26 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/parser/factories.ts +++ b/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: Name, ctx: ParserRuleContext) => Builder.command({ name, args: [] }, createParserFields(ctx)); export const createInlineCast = (ctx: InlineCastContext, value: ESQLInlineCast['value']) => diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/factories/stats.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/factories/stats.ts new file mode 100644 index 0000000000000..59a24b6a57b7e --- /dev/null +++ b/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; +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/parser.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/parser.ts index f99e00e92d1e0..dc17607b9dc8e 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/parser/parser.ts +++ b/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, diff --git a/src/platform/packages/shared/kbn-esql-ast/src/parser/walkers.ts b/src/platform/packages/shared/kbn-esql-ast/src/parser/walkers.ts index a9f3ac7157a1f..288731fa4b02b 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/parser/walkers.ts +++ b/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; +}; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/types.ts index 9adb7cbe36c5c..6ea62dd77ef39 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/types.ts +++ b/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[]; diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts index 2d1d364b204ea..684c3c7401fef 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts @@ -43,7 +43,7 @@ test('can terminate walk early, does not visit all literals', () => { `); const result = new Visitor() .on('visitExpression', (ctx) => { - return 0; + for (const res of ctx.visitArguments(undefined)) if (res) return res; }) .on('visitLiteralExpression', (ctx) => { numbers.push(ctx.node.value as number); @@ -69,8 +69,9 @@ test('"visitColumnExpression" takes over all column visits', () => { .on('visitColumnExpression', (ctx) => { return ''; }) - .on('visitExpression', (ctx) => { - return 'E'; + .on('visitExpression', (ctx) => 'E') + .on('visitFieldExpression', (ctx) => { + return [...ctx.visitArguments()].join(':'); }) .on('visitCommand', (ctx) => { const args = [...ctx.visitArguments()].join(', '); @@ -81,7 +82,7 @@ test('"visitColumnExpression" takes over all column visits', () => { }); const text = visitor.visitQuery(ast); - expect(text).toBe('FROM E | STATS '); + expect(text).toBe('FROM E | STATS :'); }); test('"visitSourceExpression" takes over all source visits', () => { @@ -109,18 +110,46 @@ test('"visitSourceExpression" takes over all source visits', () => { expect(text).toBe('FROM | STATS E, E, E, E | LIMIT E'); }); +test('"visitFieldExpression" takes over all "field" visits', () => { + const { ast } = parse(` + FROM index + | STATS 1 + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitFieldExpression', (ctx) => { + return ''; + }) + .on('visitExpression', (ctx) => { + return 'E'; + }) + .on('visitCommand', (ctx) => { + const args = [...ctx.visitArguments()].join(', '); + return `${ctx.name()}${args ? ` ${args}` : ''}`; + }) + .on('visitQuery', (ctx) => { + return [...ctx.visitCommands()].join(' | '); + }); + const text = visitor.visitQuery(ast); + + expect(text).toBe('FROM E | STATS | LIMIT E'); +}); + test('"visitFunctionCallExpression" takes over all literal visits', () => { const { ast } = parse(` FROM index - | STATS 1, "str", [true], a = b BY field + | STATS 1, "str", [true], a = b(c) BY field | LIMIT 123 `); const visitor = new Visitor() + .on('visitExpression', (ctx) => { + return 'E'; + }) .on('visitFunctionCallExpression', (ctx) => { return ''; }) - .on('visitExpression', (ctx) => { - return 'E'; + .on('visitFieldExpression', (ctx) => { + return [...ctx.visitArguments()].join(':'); }) .on('visitCommand', (ctx) => { const args = [...ctx.visitArguments()].join(', '); @@ -131,7 +160,7 @@ test('"visitFunctionCallExpression" takes over all literal visits', () => { }); const text = visitor.visitQuery(ast); - expect(text).toBe('FROM E | STATS E, E, E, | LIMIT E'); + expect(text).toBe('FROM E | STATS E:E, E:E, E:E, E: | LIMIT E'); }); test('"visitLiteral" takes over all literal visits', () => { @@ -141,11 +170,14 @@ test('"visitLiteral" takes over all literal visits', () => { | LIMIT 123 `); const visitor = new Visitor() + .on('visitExpression', (ctx) => { + return 'E'; + }) .on('visitLiteralExpression', (ctx) => { return ''; }) - .on('visitExpression', (ctx) => { - return 'E'; + .on('visitFieldExpression', (ctx) => { + return [...ctx.visitArguments()].join(':'); }) .on('visitCommand', (ctx) => { const args = [...ctx.visitArguments()].join(', '); @@ -156,7 +188,7 @@ test('"visitLiteral" takes over all literal visits', () => { }); const text = visitor.visitQuery(ast); - expect(text).toBe('FROM E | STATS , , E, E | LIMIT '); + expect(text).toBe('FROM E | STATS E:, E:, E:E, E:E | LIMIT '); }); test('"visitExpression" does visit identifier nodes', () => { diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts index 0b2b5e4a4cb5c..f7fa2ce238f52 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts @@ -126,6 +126,9 @@ export const prettyPrint = (ast: ESQLAstQueryExpression | ESQLAstQueryExpression .on('visitColumnExpression', (ctx) => { return ctx.node.name; }) + .on('visitFieldExpression', (ctx) => { + return `${ctx.visitArgument(1)}`; + }) .on('visitFunctionCallExpression', (ctx) => { let args = ''; for (const arg of ctx.visitArguments()) { @@ -189,6 +192,6 @@ test('can print a query to text', () => { const text = prettyPrint(ast); expect(text).toBe( - 'FROM index METADATA _id, asdf, 123 | STATS FN(, , , IN(x, 1, 2)), =(a, b) | LIMIT 1000' + 'FROM index METADATA _id, asdf, 123 | STATS FN(, , , IN(x, 1, 2)), b | LIMIT 1000' ); }); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts index 913abfacd9702..69739fcf332d1 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts @@ -24,6 +24,7 @@ import type { ESQLColumn, ESQLCommandOption, ESQLDecimalLiteral, + ESQLField, ESQLFunction, ESQLIdentifier, ESQLInlineCast, @@ -46,6 +47,7 @@ import type { VisitorOutput, } from './types'; import { Builder } from '../builder'; +import { isProperNode } from '../ast/helpers'; const isNodeWithArgs = (x: unknown): x is ESQLAstNodeWithArgs => !!x && typeof x === 'object' && Array.isArray((x as any).args); @@ -85,13 +87,7 @@ export class VisitorContext< ): Iterable> { this.ctx.assertMethodExists('visitExpression'); - const node = this.node; - - if (!isNodeWithArgs(node)) { - return; - } - - for (const arg of singleItems(node.args)) { + for (const arg of this.arguments()) { if (arg.type === 'option' && arg.name !== 'as') { continue; } @@ -107,7 +103,7 @@ export class VisitorContext< public arguments(): ESQLAstExpressionNode[] { const node = this.node; - if (!isNodeWithChildren(node)) { + if (!isProperNode(node)) { return []; } @@ -128,12 +124,12 @@ export class VisitorContext< const node = this.node; - if (!isNodeWithArgs(node)) { + if (!isProperNode(node)) { throw new Error('Node does not have arguments'); } let i = 0; - for (const arg of singleItems(node.args)) { + for (const arg of this.arguments()) { if (i === index) { return this.visitExpression(arg, input as any); } @@ -580,3 +576,8 @@ export class IdentifierExpressionVisitorContext< Methods extends VisitorMethods = VisitorMethods, Data extends SharedData = SharedData > extends VisitorContext {} + +export class FieldExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends VisitorContext {} diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts index 5240b4fe2e224..892be88989408 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts @@ -13,6 +13,7 @@ import type { ESQLAstJoinCommand, ESQLAstRenameExpression, ESQLColumn, + ESQLField, ESQLFunction, ESQLIdentifier, ESQLInlineCast, @@ -424,6 +425,10 @@ export class GlobalVisitorContext< if (!this.methods.visitIdentifierExpression) break; return this.visitIdentifierExpression(parent, expressionNode, input as any); } + case 'field': { + if (!this.methods.visitFieldExpression) break; + return this.visitFieldExpression(parent, expressionNode, input as any); + } case 'option': { switch (expressionNode.name) { case 'as': { @@ -529,4 +534,13 @@ export class GlobalVisitorContext< const context = new contexts.IdentifierExpressionVisitorContext(this, node, parent); return this.visitWithSpecificContext('visitIdentifierExpression', context, input); } + + public visitFieldExpression( + parent: contexts.VisitorContext | null, + node: ESQLField, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.FieldExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitFieldExpression', context, input); + } } diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts index b471eb67258fe..fa3c8eb510137 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts @@ -62,7 +62,8 @@ export type ExpressionVisitorInput = AnyToVoid< VisitorInput & VisitorInput & VisitorInput & - VisitorInput + VisitorInput & + VisitorInput >; /** @@ -79,7 +80,8 @@ export type ExpressionVisitorOutput = | VisitorOutput | VisitorOutput | VisitorOutput - | VisitorOutput; + | VisitorOutput + | VisitorOutput; /** * Input that satisfies any command visitor input constraints. @@ -215,6 +217,7 @@ export interface VisitorMethods< any, any >; + visitFieldExpression?: Visitor, any, any>; } /** @@ -242,6 +245,8 @@ export type AstNodeToVisitorName = Node extends ESQ ? 'visitInlineCastExpression' : Node extends ast.ESQLIdentifier ? 'visitIdentifierExpression' + : Node extends ast.ESQLField + ? 'visitFieldExpression' : never; /** diff --git a/src/platform/packages/shared/kbn-esql-ast/src/visitor/utils.ts b/src/platform/packages/shared/kbn-esql-ast/src/visitor/utils.ts index 00b7a4541ccdd..c50aac677bcd7 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/visitor/utils.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/visitor/utils.ts @@ -81,5 +81,14 @@ export function* children(node: ESQLProperNode): Iterable { } break; } + case 'field': { + yield node.column; + yield node.value; + + if (node.where) { + yield node.where; + } + break; + } } } diff --git a/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts b/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts index c2db01f719d4b..35e779ab3f841 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.test.ts @@ -21,6 +21,7 @@ import { ESQLInlineCast, ESQLUnknownItem, ESQLIdentifier, + ESQLField, } from '../types'; import { walk, Walker } from './walker'; @@ -279,6 +280,33 @@ describe('structurally can walk all nodes', () => { }); }); + describe('fields', () => { + test('can walk through "field" node', () => { + const query = 'FROM index | STATS a = 123'; + const { ast } = parse(query); + const fields: ESQLField[] = []; + + walk(ast, { + visitField: (node) => fields.push(node), + }); + + expect(fields).toMatchObject([ + { + type: 'field', + column: { + type: 'column', + name: 'a', + }, + value: { + type: 'literal', + literalType: 'integer', + value: 123, + }, + }, + ]); + }); + }); + describe('functions', () => { test('can walk through functions', () => { const query = 'FROM a | STATS fn(1), agg(true)'; @@ -1189,11 +1217,11 @@ describe('Walker.matchAll()', () => { }); describe('Walker.hasFunction()', () => { - test('can find assignment expression', () => { + test('can find binary expression expression', () => { const query1 = 'FROM a | STATS bucket(bytes, 1 hour)'; - const query2 = 'FROM b | STATS var0 = bucket(bytes, 1 hour)'; - const has1 = Walker.hasFunction(parse(query1).ast!, '='); - const has2 = Walker.hasFunction(parse(query2).ast!, '='); + const query2 = 'FROM b | STATS var0 == bucket(bytes, 1 hour)'; + const has1 = Walker.hasFunction(parse(query1).ast!, '=='); + const has2 = Walker.hasFunction(parse(query2).ast!, '=='); expect(has1).toBe(false); expect(has2).toBe(true); diff --git a/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.ts b/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.ts index 0e6811c02efef..d3bf99fce6757 100644 --- a/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.ts +++ b/src/platform/packages/shared/kbn-esql-ast/src/walker/walker.ts @@ -19,6 +19,7 @@ import type { ESQLCommand, ESQLCommandMode, ESQLCommandOption, + ESQLField, ESQLFunction, ESQLIdentifier, ESQLInlineCast, @@ -51,6 +52,7 @@ export interface WalkerOptions { visitInlineCast?: (node: ESQLInlineCast) => void; visitUnknown?: (node: ESQLUnknownItem) => void; visitIdentifier?: (node: ESQLIdentifier) => void; + visitField?: (node: ESQLField) => void; /** * Called for any node type that does not have a specific visitor. @@ -65,20 +67,6 @@ export type WalkerAstNode = ESQLAstNode | ESQLAstNode[]; /** * Iterates over all nodes in the AST and calls the appropriate visitor * functions. - * - * AST nodes supported: - * - * - [x] command - * - [x] option - * - [x] mode - * - [x] function - * - [x] source - * - [x] column - * - [x] literal - * - [x] list literal - * - [x] timeInterval - * - [x] inlineCast - * - [x] unknown */ export class Walker { /** @@ -325,7 +313,7 @@ export class Walker { } } - public walkAstItem(node: ESQLAstItem): void { + public walkAstItem(node: ESQLAstItem | ESQLAstExpression): void { if (node instanceof Array) { const list = node as ESQLAstItem[]; for (const item of list) this.walkAstItem(item); @@ -367,13 +355,25 @@ export class Walker { this.walkAstItem(node.value); } + public walkField(node: ESQLField): void { + const { options } = this; + + (options.visitField ?? options.visitAny)?.(node); + this.walkColumn(node.column); + this.walkSingleAstItem(node.value); + + if (node.where) { + this.walkSingleAstItem(node.where); + } + } + public walkFunction(node: ESQLFunction): void { const { options } = this; (options.visitFunction ?? options.visitAny)?.(node); const args = node.args; const length = args.length; - if (node.operator) this.walkAstItem(node.operator); + if (node.operator) this.walkSingleAstItem(node.operator); for (let i = 0; i < length; i++) { const arg = args[i]; @@ -440,6 +440,10 @@ export class Walker { (options.visitIdentifier ?? options.visitAny)?.(node); break; } + case 'field': { + this.walkField(node); + break; + } case 'unknown': { (options.visitUnknown ?? options.visitAny)?.(node); break;