From ac8a81d02f57cf96a3d03ef48d3d183d6d2c5f58 Mon Sep 17 00:00:00 2001 From: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com> Date: Fri, 6 Dec 2024 08:44:23 +0100 Subject: [PATCH] [ES|QL] `JOIN` command Traversal API and prety-printing support (#202750) ## Summary Partially addresses https://github.com/elastic/kibana/issues/200858 - Add support for the new `JOIN` command and `AS` expression in Traversal API: `Walker` and `Visitor` - Adds support for the new `JOIN`command and `AS` expression in the pretty-printer. - Fixes some parser bugs related to the `JOIN` command. ### Checklist - [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 --- .../src/parser/__tests__/comments.test.ts | 113 +++++++++++++++++ .../src/parser/__tests__/join.test.ts | 30 ++++- packages/kbn-esql-ast/src/parser/factories.ts | 21 ++++ .../kbn-esql-ast/src/parser/factories/join.ts | 20 ++- .../basic_pretty_printer.comments.test.ts | 14 ++- .../__tests__/basic_pretty_printer.test.ts | 34 +++++ .../wrapping_pretty_printer.comments.test.ts | 43 +++++++ .../__tests__/wrapping_pretty_printer.test.ts | 23 +++- .../src/pretty_print/basic_pretty_printer.ts | 17 ++- .../pretty_print/wrapping_pretty_printer.ts | 20 ++- packages/kbn-esql-ast/src/types.ts | 10 +- .../src/visitor/__tests__/commands.test.ts | 117 ++++++++++++++++++ .../src/visitor/__tests__/expressions.test.ts | 22 ++++ packages/kbn-esql-ast/src/visitor/contexts.ts | 15 ++- .../src/visitor/global_visitor_context.ts | 28 +++++ packages/kbn-esql-ast/src/visitor/types.ts | 20 ++- packages/kbn-esql-ast/src/walker/helpers.ts | 39 +++++- .../kbn-esql-ast/src/walker/walker.test.ts | 49 ++++++++ 18 files changed, 607 insertions(+), 28 deletions(-) create mode 100644 packages/kbn-esql-ast/src/visitor/__tests__/commands.test.ts diff --git a/packages/kbn-esql-ast/src/parser/__tests__/comments.test.ts b/packages/kbn-esql-ast/src/parser/__tests__/comments.test.ts index dedc99ba322ec..73b022457d2f6 100644 --- a/packages/kbn-esql-ast/src/parser/__tests__/comments.test.ts +++ b/packages/kbn-esql-ast/src/parser/__tests__/comments.test.ts @@ -8,6 +8,7 @@ */ import { parse } from '..'; +import { Walker } from '../../walker'; describe('Comments', () => { describe('can attach "top" comment(s)', () => { @@ -442,6 +443,35 @@ FROM index`; ], }); }); + + it('to an identifier', () => { + const text = `FROM index | LEFT JOIN + // comment + abc + ON a = b`; + const { root } = parse(text, { withFormatting: true }); + + expect(root.commands[1]).toMatchObject({ + type: 'command', + name: 'join', + args: [ + { + type: 'identifier', + name: 'abc', + formatting: { + top: [ + { + type: 'comment', + subtype: 'single-line', + text: ' comment', + }, + ], + }, + }, + {}, + ], + }); + }); }); describe('can attach "left" comment(s)', () => { @@ -549,6 +579,34 @@ FROM index`; }, ]); }); + + it('to an identifier', () => { + const text = `FROM index | LEFT JOIN + /* left */ abc + ON a = b`; + const { root } = parse(text, { withFormatting: true }); + + expect(root.commands[1]).toMatchObject({ + type: 'command', + name: 'join', + args: [ + { + type: 'identifier', + name: 'abc', + formatting: { + left: [ + { + type: 'comment', + subtype: 'multi-line', + text: ' left ', + }, + ], + }, + }, + {}, + ], + }); + }); }); describe('can attach "right" comment(s)', () => { @@ -776,6 +834,61 @@ FROM index`; ], }); }); + + it('to an identifier', () => { + const text = `FROM index | LEFT JOIN + abc /* right */ // right 2 + ON a = b`; + const { root } = parse(text, { withFormatting: true }); + + expect(root.commands[1]).toMatchObject({ + type: 'command', + name: 'join', + args: [ + { + type: 'identifier', + name: 'abc', + formatting: { + right: [ + { + type: 'comment', + subtype: 'multi-line', + text: ' right ', + }, + ], + rightSingleLine: { + type: 'comment', + subtype: 'single-line', + text: ' right 2', + }, + }, + }, + {}, + ], + }); + }); + + it('to a column inside ON option', () => { + const text = `FROM index | LEFT JOIN + abc + ON a /* right */ = b`; + const { root } = parse(text, { withFormatting: true }); + const a = Walker.match(root, { type: 'column', name: 'a' }); + + expect(a).toMatchObject({ + type: 'column', + name: 'a', + formatting: { + right: [ + { + type: 'comment', + subtype: 'multi-line', + text: ' right ', + }, + ], + }, + }); + }); }); describe('can attach "right end" comments', () => { diff --git a/packages/kbn-esql-ast/src/parser/__tests__/join.test.ts b/packages/kbn-esql-ast/src/parser/__tests__/join.test.ts index 5784e0c71cb86..baeee1aeac446 100644 --- a/packages/kbn-esql-ast/src/parser/__tests__/join.test.ts +++ b/packages/kbn-esql-ast/src/parser/__tests__/join.test.ts @@ -184,19 +184,41 @@ describe(' JOIN command', () => { const node2 = Walker.match(query.ast, { type: 'identifier', name: 'alias' }); const node3 = Walker.match(query.ast, { type: 'column', name: 'on_1' }); const node4 = Walker.match(query.ast, { type: 'column', name: 'on_2' }); + const node5 = Walker.match(query.ast, { type: 'function', name: 'as' }); expect(query.src.slice(node1?.location.min, node1?.location.max! + 1)).toBe('index'); expect(query.src.slice(node2?.location.min, node2?.location.max! + 1)).toBe('alias'); expect(query.src.slice(node3?.location.min, node3?.location.max! + 1)).toBe('on_1'); expect(query.src.slice(node4?.location.min, node4?.location.max! + 1)).toBe('on_2'); + expect(query.src.slice(node5?.location.min, node5?.location.max! + 1)).toBe('index AS alias'); + }); + + it('correctly extracts JOIN command position', () => { + const text = `FROM employees | LOOKUP JOIN index AS alias ON on_1, on_2 | LIMIT 1`; + const query = EsqlQuery.fromSrc(text); + const join = Walker.match(query.ast, { type: 'command', name: 'join' }); + + expect(query.src.slice(join?.location.min, join?.location.max! + 1)).toBe( + 'LOOKUP JOIN index AS alias ON on_1, on_2' + ); + }); + + it('correctly extracts ON option position', () => { + const text = `FROM employees | LOOKUP JOIN index AS alias ON on_1, on_2 | LIMIT 1`; + const query = EsqlQuery.fromSrc(text); + const on = Walker.match(query.ast, { type: 'option', name: 'on' }); + + expect(query.src.slice(on?.location.min, on?.location.max! + 1)).toBe('ON on_1, on_2'); }); }); describe('incorrectly formatted', () => { - const text = `FROM employees | LOOKUP JOIN index AAS alias ON on_1, on_2 | LIMIT 1`; - const query = EsqlQuery.fromSrc(text); + it('throws error on invalid "AS" keyword', () => { + const text = `FROM employees | LOOKUP JOIN index AAS alias ON on_1, on_2 | LIMIT 1`; + const query = EsqlQuery.fromSrc(text); - expect(query.errors.length > 0).toBe(true); - expect(query.errors[0].message.includes('AAS')).toBe(true); + expect(query.errors.length > 0).toBe(true); + expect(query.errors[0].message.includes('AAS')).toBe(true); + }); }); }); diff --git a/packages/kbn-esql-ast/src/parser/factories.ts b/packages/kbn-esql-ast/src/parser/factories.ts index 311dcced8a617..388469e82ab99 100644 --- a/packages/kbn-esql-ast/src/parser/factories.ts +++ b/packages/kbn-esql-ast/src/parser/factories.ts @@ -55,6 +55,8 @@ import type { InlineCastingType, ESQLFunctionCallExpression, ESQLIdentifier, + ESQLBinaryExpression, + BinaryExpressionOperator, } from '../types'; import { parseIdentifier, getPosition } from './helpers'; import { Builder, type AstNodeParserFields } from '../builder'; @@ -240,6 +242,25 @@ export const createFunctionCall = (ctx: FunctionContext): ESQLFunctionCallExpres return node; }; +export const createBinaryExpression = ( + operator: BinaryExpressionOperator, + ctx: ParserRuleContext, + args: ESQLBinaryExpression['args'] +): ESQLBinaryExpression => { + const node = Builder.expression.func.binary( + operator, + args, + {}, + { + text: ctx.getText(), + location: getPosition(ctx.start, ctx.stop), + incomplete: Boolean(ctx.exception), + } + ) as ESQLBinaryExpression; + + return node; +}; + export const createIdentifierOrParam = (ctx: IdentifierOrParameterContext) => { const identifier = ctx.identifier(); if (identifier) { diff --git a/packages/kbn-esql-ast/src/parser/factories/join.ts b/packages/kbn-esql-ast/src/parser/factories/join.ts index 400313aa39045..9ff5083ac5e28 100644 --- a/packages/kbn-esql-ast/src/parser/factories/join.ts +++ b/packages/kbn-esql-ast/src/parser/factories/join.ts @@ -8,9 +8,13 @@ */ import { JoinCommandContext, JoinTargetContext } from '../../antlr/esql_parser'; -import { Builder } from '../../builder'; import { ESQLAstItem, ESQLBinaryExpression, ESQLCommand, ESQLIdentifier } from '../../types'; -import { createCommand, createIdentifier } from '../factories'; +import { + createBinaryExpression, + createCommand, + createIdentifier, + createOption, +} from '../factories'; import { visitValueExpression } from '../walkers'; const createNodeFromJoinTarget = ( @@ -24,7 +28,7 @@ const createNodeFromJoinTarget = ( } const alias = createIdentifier(aliasCtx); - const renameExpression = Builder.expression.func.binary('as', [ + const renameExpression = createBinaryExpression('as', ctx, [ index, alias, ]) as ESQLBinaryExpression; @@ -39,10 +43,11 @@ export const createJoinCommand = (ctx: JoinCommandContext): ESQLCommand => { command.commandType = (ctx._type_.text ?? '').toLocaleLowerCase(); const joinTarget = createNodeFromJoinTarget(ctx.joinTarget()); - const onOption = Builder.option({ name: 'on' }); + const joinCondition = ctx.joinCondition(); + const onOption = createOption('on', joinCondition); const joinPredicates: ESQLAstItem[] = onOption.args; - for (const joinPredicateCtx of ctx.joinCondition().joinPredicate_list()) { + for (const joinPredicateCtx of joinCondition.joinPredicate_list()) { const expression = visitValueExpression(joinPredicateCtx.valueExpression()); if (expression) { @@ -51,7 +56,10 @@ export const createJoinCommand = (ctx: JoinCommandContext): ESQLCommand => { } command.args.push(joinTarget); - command.args.push(onOption); + + if (onOption.args.length) { + command.args.push(onOption); + } return command; }; diff --git a/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.comments.test.ts b/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.comments.test.ts index 3839c5e52f9bc..d725759172928 100644 --- a/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.comments.test.ts +++ b/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.comments.test.ts @@ -14,7 +14,7 @@ const reprint = (src: string) => { const { root } = parse(src, { withFormatting: true }); const text = BasicPrettyPrinter.print(root); - // console.log(JSON.stringify(ast, null, 2)); + // console.log(JSON.stringify(root.commands, null, 2)); return { text }; }; @@ -184,3 +184,15 @@ describe('rename expressions', () => { assertPrint('FROM a | RENAME /*I*/ a /*II*/ AS /*III*/ b /*IV*/, c AS d'); }); }); + +describe('commands', () => { + describe('JOIN', () => { + test('around JOIN targets', () => { + assertPrint('FROM a | LEFT JOIN /*1*/ a /*2*/ AS /*3*/ b /*4*/ ON c'); + }); + + test('around JOIN conditions', () => { + assertPrint('FROM a | LEFT JOIN a AS b ON /*1*/ c /*2*/, /*3*/ d /*4*/'); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts b/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts index c8880b9bfe678..81cde1ed94e7f 100644 --- a/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts +++ b/packages/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.test.ts @@ -115,6 +115,40 @@ describe('single line query', () => { expect(text).toBe('FROM index | DISSECT input "pattern" APPEND_SEPARATOR = ""'); }); }); + + describe('JOIN', () => { + test('example from docs', () => { + const { text } = reprint(` + FROM employees + | EVAL language_code = languages + | LOOKUP JOIN languages_lookup ON language_code + | WHERE emp_no < 500 + | KEEP emp_no, language_name + | SORT emp_no + | LIMIT 10 + `); + + expect(text).toBe( + 'FROM employees | EVAL language_code = languages | LOOKUP JOIN languages_lookup ON language_code | WHERE emp_no < 500 | KEEP emp_no, language_name | SORT emp_no | LIMIT 10' + ); + }); + + test('supports aliases', () => { + const { text } = reprint(` + FROM employees | LEFT JOIN languages_lookup AS something ON language_code`); + + expect(text).toBe( + 'FROM employees | LEFT JOIN languages_lookup AS something ON language_code' + ); + }); + + test('supports multiple conditions', () => { + const { text } = reprint(` + FROM employees | LEFT JOIN a ON b, c, d.e.f`); + + expect(text).toBe('FROM employees | LEFT JOIN a ON b, c, d.e.f'); + }); + }); }); describe('expressions', () => { diff --git a/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.comments.test.ts b/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.comments.test.ts index 861d274493a42..1b1fa3b01896d 100644 --- a/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.comments.test.ts +++ b/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.comments.test.ts @@ -497,6 +497,49 @@ ROW 1 }); }); + describe('as-expressions', () => { + test('JOIN main arguments surrounded in comments', () => { + const query = ` + FROM index | LEFT JOIN + /* 1 */ + // 2 + /* 3 */ + // 4 + /* 5 */ a /* 6 */ AS /* 7 */ b + ON c`; + const text = reprint(query).text; + expect('\n' + text).toBe(` +FROM index + | LEFT JOIN + /* 1 */ + // 2 + /* 3 */ + // 4 + /* 5 */ a /* 6 */ AS + /* 7 */ b + ON c`); + }); + + test('JOIN "ON" option argument comments', () => { + const query = ` + FROM index | RIGHT JOIN a AS b ON + // c.1 + /* c.2 */ c /* c.3 */, + // d.1 + /* d.2 */ d /* d.3 */`; + const text = reprint(query).text; + expect('\n' + text).toBe(` +FROM index + | RIGHT JOIN + a AS b + ON + // c.1 + /* c.2 */ c, /* c.3 */ + // d.1 + /* d.2 */ d /* d.3 */`); + }); + }); + describe('function call expressions', () => { describe('binary expressions', () => { test('first operand surrounded by inline comments', () => { diff --git a/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.test.ts b/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.test.ts index 6422ae9a451af..b26906e5d2e75 100644 --- a/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.test.ts +++ b/packages/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.test.ts @@ -14,12 +14,33 @@ const reprint = (src: string, opts?: WrappingPrettyPrinterOptions) => { const { root } = parse(src); const text = WrappingPrettyPrinter.print(root, opts); - // console.log(JSON.stringify(ast, null, 2)); + // console.log(JSON.stringify(root.commands, null, 2)); return { text }; }; describe('commands', () => { + describe('JOIN', () => { + test('with short identifiers', () => { + const { text } = reprint('FROM a | RIGHT JOIN b AS c ON d, e'); + + expect(text).toBe('FROM a | RIGHT JOIN b AS c ON d, e'); + }); + + test('with long identifiers', () => { + const { text } = reprint( + 'FROM aaaaaaaaaaaa | RIGHT JOIN bbbbbbbbbbbbbbbbb AS cccccccccccccccccccc ON dddddddddddddddddddddddddddddddddddddddd, eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee' + ); + + expect('\n' + text).toBe(` +FROM aaaaaaaaaaaa + | RIGHT JOIN bbbbbbbbbbbbbbbbb AS cccccccccccccccccccc + ON + dddddddddddddddddddddddddddddddddddddddd, + eeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeeee`); + }); + }); + describe('GROK', () => { test('two basic arguments', () => { const { text } = reprint('FROM search-movies | GROK Awards "text"'); diff --git a/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts b/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts index cf252825c243f..60248774f9157 100644 --- a/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts +++ b/packages/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts @@ -223,6 +223,11 @@ export class BasicPrettyPrinter { return ''; }) + .on('visitIdentifierExpression', (ctx) => { + const formatted = LeafPrinter.identifier(ctx.node); + return this.decorateWithComments(ctx.node, formatted); + }) + .on('visitSourceExpression', (ctx) => { const formatted = LeafPrinter.source(ctx.node); return this.decorateWithComments(ctx.node, formatted); @@ -383,12 +388,16 @@ export class BasicPrettyPrinter { const argsFormatted = args ? `${separator}${args}` : ''; const optionFormatted = `${option}${argsFormatted}`; - return optionFormatted; + return this.decorateWithComments(ctx.node, optionFormatted); }) .on('visitCommand', (ctx) => { const opts = this.opts; - const cmd = opts.lowercaseCommands ? ctx.node.name : ctx.node.name.toUpperCase(); + const node = ctx.node; + const cmd = opts.lowercaseCommands ? node.name : node.name.toUpperCase(); + const cmdType = !node.commandType + ? '' + : (opts.lowercaseCommands ? node.commandType : node.commandType.toUpperCase()) + ' '; let args = ''; let options = ''; @@ -406,9 +415,9 @@ export class BasicPrettyPrinter { const argsFormatted = args ? ` ${args}` : ''; const optionsFormatted = options ? ` ${options}` : ''; - const cmdFormatted = `${cmd}${argsFormatted}${optionsFormatted}`; + const cmdFormatted = `${cmdType}${cmd}${argsFormatted}${optionsFormatted}`; - return cmdFormatted; + return this.decorateWithComments(ctx.node, cmdFormatted); }) .on('visitQuery', (ctx) => { diff --git a/packages/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts b/packages/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts index 2f863524740ee..9ba4ce8b0a5ae 100644 --- a/packages/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts +++ b/packages/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts @@ -427,12 +427,19 @@ export class WrappingPrettyPrinter { return { txt, indented }; } - protected readonly visitor = new Visitor() + protected readonly visitor: Visitor = new Visitor() .on('visitExpression', (ctx, inp: Input): Output => { const txt = ctx.node.text ?? ''; return { txt }; }) + .on('visitIdentifierExpression', (ctx, inp: Input) => { + const formatted = LeafPrinter.identifier(ctx.node); + const { txt, indented } = this.decorateWithComments(inp.indent, ctx.node, formatted); + + return { txt, indented }; + }) + .on('visitSourceExpression', (ctx, inp: Input): Output => { const formatted = LeafPrinter.source(ctx.node) + (inp.suffix ?? ''); const { txt, indented } = this.decorateWithComments(inp.indent, ctx.node, formatted); @@ -570,7 +577,14 @@ export class WrappingPrettyPrinter { .on('visitCommand', (ctx, inp: Input): Output => { const opts = this.opts; - const cmd = opts.lowercaseCommands ? ctx.node.name : ctx.node.name.toUpperCase(); + const node = ctx.node; + let cmd = opts.lowercaseCommands ? node.name : node.name.toUpperCase(); + + if (node.commandType) { + const type = opts.lowercaseCommands ? node.commandType : node.commandType.toUpperCase(); + cmd = `${type} ${cmd}`; + } + const args = this.printArguments(ctx, { indent: inp.indent, remaining: inp.remaining - cmd.length - 1, @@ -678,6 +692,6 @@ export class WrappingPrettyPrinter { }); public print(query: ESQLAstQueryExpression) { - return this.visitor.visitQuery(query); + return this.visitor.visitQuery(query, undefined); } } diff --git a/packages/kbn-esql-ast/src/types.ts b/packages/kbn-esql-ast/src/types.ts index 60da69b04989a..a204cf50901ce 100644 --- a/packages/kbn-esql-ast/src/types.ts +++ b/packages/kbn-esql-ast/src/types.ts @@ -9,7 +9,7 @@ export type ESQLAst = ESQLAstCommand[]; -export type ESQLAstCommand = ESQLCommand | ESQLAstMetricsCommand; +export type ESQLAstCommand = ESQLCommand | ESQLAstMetricsCommand | ESQLAstJoinCommand; export type ESQLAstNode = ESQLAstCommand | ESQLAstExpression | ESQLAstItem; @@ -92,6 +92,10 @@ export interface ESQLAstMetricsCommand extends ESQLCommand<'metrics'> { grouping?: ESQLAstField[]; } +export interface ESQLAstJoinCommand extends ESQLCommand<'join'> { + commandType: 'lookup' | 'left' | 'right'; +} + export interface ESQLCommandOption extends ESQLAstBaseItem { type: 'option'; args: ESQLAstItem[]; @@ -199,12 +203,14 @@ export type BinaryExpressionOperator = | BinaryExpressionArithmeticOperator | BinaryExpressionAssignmentOperator | BinaryExpressionComparisonOperator - | BinaryExpressionRegexOperator; + | BinaryExpressionRegexOperator + | BinaryExpressionRenameOperator; export type BinaryExpressionArithmeticOperator = '+' | '-' | '*' | '/' | '%'; export type BinaryExpressionAssignmentOperator = '='; export type BinaryExpressionComparisonOperator = '==' | '=~' | '!=' | '<' | '<=' | '>' | '>='; export type BinaryExpressionRegexOperator = 'like' | 'not_like' | 'rlike' | 'not_rlike'; +export type BinaryExpressionRenameOperator = 'as'; // from https://github.com/elastic/elasticsearch/blob/122e7288200ee03e9087c98dff6cebbc94e774aa/docs/reference/esql/functions/kibana/inline_cast.json export type InlineCastingType = diff --git a/packages/kbn-esql-ast/src/visitor/__tests__/commands.test.ts b/packages/kbn-esql-ast/src/visitor/__tests__/commands.test.ts new file mode 100644 index 0000000000000..35bfda59070c8 --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/__tests__/commands.test.ts @@ -0,0 +1,117 @@ +/* + * 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'; +import { Visitor } from '../visitor'; + +test('"visitCommand" captures all non-captured commands', () => { + const { ast } = EsqlQuery.fromSrc(` + FROM index + | STATS 1, "str", [true], a = b BY field + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitStatsCommand', (ctx) => { + return ''; + }) + .on('visitCommand', (ctx) => { + return `${ctx.name()}`; + }) + .on('visitQuery', (ctx) => { + return [...ctx.visitCommands()].join(' | '); + }); + const text = visitor.visitQuery(ast); + + expect(text).toBe('FROM | | LIMIT'); +}); + +test('can visit JOIN command', () => { + const { ast } = EsqlQuery.fromSrc(` + FROM index + | STATS 1, "str", [true], a = b BY field + | RIGHT JOIN abc ON xyz + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitJoinCommand', (ctx) => { + return `JOIN[type = ${ctx.node.commandType}]`; + }) + .on('visitCommand', (ctx) => { + return `${ctx.name()}`; + }) + .on('visitQuery', (ctx) => { + return [...ctx.visitCommands()].join(' | '); + }); + const text = visitor.visitQuery(ast); + + expect(text).toBe('FROM | STATS | JOIN[type = right] | LIMIT'); +}); + +test('can visit JOIN command arguments', () => { + const { ast } = EsqlQuery.fromSrc(` + FROM index + | STATS 1, "str", [true], a = b BY field + | RIGHT JOIN abc AS xxx ON xyz + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitFunctionCallExpression', (ctx) => { + if (ctx.node.subtype === 'binary-expression') { + return ctx.node.name; + } else { + return null; + } + }) + .on('visitExpression', (ctx) => { + return null; + }) + .on('visitJoinCommand', (ctx) => { + return [...ctx.visitArgs()]; + }) + .on('visitCommand', (ctx) => { + return null; + }) + .on('visitQuery', (ctx) => { + return [...ctx.visitCommands()]; + }); + const list = visitor.visitQuery(ast).flat().filter(Boolean); + + expect(list).toMatchObject(['as']); +}); + +test('can visit JOIN ON option', () => { + const { ast } = EsqlQuery.fromSrc(` + FROM index + | STATS 1, "str", [true], a = b BY field + | RIGHT JOIN abc AS xxx ON xyz + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitColumnExpression', (ctx) => { + return ctx.node.name; + }) + .on('visitExpression', (ctx) => { + return null; + }) + .on('visitCommandOption', (ctx) => { + return [...ctx.visitArguments()].flat(); + }) + .on('visitJoinCommand', (ctx) => { + return [...ctx.visitOptions()].flat(); + }) + .on('visitCommand', (ctx) => { + return null; + }) + .on('visitQuery', (ctx) => { + return [...ctx.visitCommands()].flat(); + }); + const list = visitor.visitQuery(ast).flat().filter(Boolean); + + expect(list).toMatchObject(['xyz']); +}); diff --git a/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts b/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts index 32f8fff9816a7..2d1d364b204ea 100644 --- a/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts +++ b/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts @@ -158,3 +158,25 @@ test('"visitLiteral" takes over all literal visits', () => { expect(text).toBe('FROM E | STATS , , E, E | LIMIT '); }); + +test('"visitExpression" does visit identifier nodes', () => { + const { ast } = parse(` + FROM index + | RIGHT JOIN a AS b ON c + `); + const expressions: string[] = []; + new Visitor() + .on('visitExpression', (ctx) => { + expressions.push(ctx.node.name); + for (const _ of ctx.visitArguments(undefined)); + }) + .on('visitCommand', (ctx) => { + for (const _ of ctx.visitArguments()); + }) + .on('visitQuery', (ctx) => { + for (const _ of ctx.visitCommands()); + }) + .visitQuery(ast); + + expect(expressions.sort()).toEqual(['a', 'as', 'b', 'index']); +}); diff --git a/packages/kbn-esql-ast/src/visitor/contexts.ts b/packages/kbn-esql-ast/src/visitor/contexts.ts index e1570f7143c07..913abfacd9702 100644 --- a/packages/kbn-esql-ast/src/visitor/contexts.ts +++ b/packages/kbn-esql-ast/src/visitor/contexts.ts @@ -17,6 +17,7 @@ import type { ESQLAstCommand, ESQLAstExpression, ESQLAstItem, + ESQLAstJoinCommand, ESQLAstNodeWithArgs, ESQLAstNodeWithChildren, ESQLAstRenameExpression, @@ -24,6 +25,7 @@ import type { ESQLCommandOption, ESQLDecimalLiteral, ESQLFunction, + ESQLIdentifier, ESQLInlineCast, ESQLIntegerLiteral, ESQLList, @@ -86,7 +88,7 @@ export class VisitorContext< const node = this.node; if (!isNodeWithArgs(node)) { - throw new Error('Node does not have arguments'); + return; } for (const arg of singleItems(node.args)) { @@ -467,6 +469,12 @@ export class MvExpandCommandVisitorContext< Data extends SharedData = SharedData > extends CommandVisitorContext {} +// JOIN ON +export class JoinCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + // Expressions ----------------------------------------------------------------- export class ExpressionVisitorContext< @@ -567,3 +575,8 @@ export class OrderExpressionVisitorContext< Methods extends VisitorMethods = VisitorMethods, Data extends SharedData = SharedData > extends VisitorContext {} + +export class IdentifierExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends VisitorContext {} diff --git a/packages/kbn-esql-ast/src/visitor/global_visitor_context.ts b/packages/kbn-esql-ast/src/visitor/global_visitor_context.ts index 793803bc48f54..5240b4fe2e224 100644 --- a/packages/kbn-esql-ast/src/visitor/global_visitor_context.ts +++ b/packages/kbn-esql-ast/src/visitor/global_visitor_context.ts @@ -10,9 +10,11 @@ import * as contexts from './contexts'; import type { ESQLAstCommand, + ESQLAstJoinCommand, ESQLAstRenameExpression, ESQLColumn, ESQLFunction, + ESQLIdentifier, ESQLInlineCast, ESQLList, ESQLLiteral, @@ -165,6 +167,10 @@ export class GlobalVisitorContext< if (!this.methods.visitMvExpandCommand) break; return this.visitMvExpandCommand(parent, commandNode, input as any); } + case 'join': { + if (!this.methods.visitJoinCommand) break; + return this.visitJoinCommand(parent, commandNode as ESQLAstJoinCommand, input as any); + } } return this.visitCommandGeneric(parent, commandNode, input as any); } @@ -349,6 +355,15 @@ export class GlobalVisitorContext< return this.visitWithSpecificContext('visitMvExpandCommand', context, input); } + public visitJoinCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstJoinCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.JoinCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitJoinCommand', context, input); + } + // Expression visiting ------------------------------------------------------- public visitExpressionGeneric( @@ -405,6 +420,10 @@ export class GlobalVisitorContext< if (!this.methods.visitOrderExpression) break; return this.visitOrderExpression(parent, expressionNode, input as any); } + case 'identifier': { + if (!this.methods.visitIdentifierExpression) break; + return this.visitIdentifierExpression(parent, expressionNode, input as any); + } case 'option': { switch (expressionNode.name) { case 'as': { @@ -501,4 +520,13 @@ export class GlobalVisitorContext< const context = new contexts.OrderExpressionVisitorContext(this, node, parent); return this.visitWithSpecificContext('visitOrderExpression', context, input); } + + public visitIdentifierExpression( + parent: contexts.VisitorContext | null, + node: ESQLIdentifier, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.IdentifierExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitIdentifierExpression', context, input); + } } diff --git a/packages/kbn-esql-ast/src/visitor/types.ts b/packages/kbn-esql-ast/src/visitor/types.ts index 6dd49eabb013e..b471eb67258fe 100644 --- a/packages/kbn-esql-ast/src/visitor/types.ts +++ b/packages/kbn-esql-ast/src/visitor/types.ts @@ -61,7 +61,8 @@ export type ExpressionVisitorInput = AnyToVoid< VisitorInput & VisitorInput & VisitorInput & - VisitorInput + VisitorInput & + VisitorInput >; /** @@ -77,7 +78,8 @@ export type ExpressionVisitorOutput = | VisitorOutput | VisitorOutput | VisitorOutput - | VisitorOutput; + | VisitorOutput + | VisitorOutput; /** * Input that satisfies any command visitor input constraints. @@ -103,7 +105,8 @@ export type CommandVisitorInput = AnyToVoid< VisitorInput & VisitorInput & VisitorInput & - VisitorInput + VisitorInput & + VisitorInput >; /** @@ -130,7 +133,8 @@ export type CommandVisitorOutput = | VisitorOutput | VisitorOutput | VisitorOutput - | VisitorOutput; + | VisitorOutput + | VisitorOutput; export interface VisitorMethods< Visitors extends VisitorMethods = any, @@ -162,6 +166,7 @@ export interface VisitorMethods< visitGrokCommand?: Visitor, any, any>; visitEnrichCommand?: Visitor, any, any>; visitMvExpandCommand?: Visitor, any, any>; + visitJoinCommand?: Visitor, any, any>; visitCommandOption?: Visitor, any, any>; visitExpression?: Visitor, any, any>; visitSourceExpression?: Visitor< @@ -205,6 +210,11 @@ export interface VisitorMethods< any >; visitOrderExpression?: Visitor, any, any>; + visitIdentifierExpression?: Visitor< + contexts.IdentifierExpressionVisitorContext, + any, + any + >; } /** @@ -230,6 +240,8 @@ export type AstNodeToVisitorName = Node extends ESQ ? 'visitTimeIntervalLiteralExpression' : Node extends ast.ESQLInlineCast ? 'visitInlineCastExpression' + : Node extends ast.ESQLIdentifier + ? 'visitIdentifierExpression' : never; /** diff --git a/packages/kbn-esql-ast/src/walker/helpers.ts b/packages/kbn-esql-ast/src/walker/helpers.ts index eb455ec7991dc..8159c9abbde49 100644 --- a/packages/kbn-esql-ast/src/walker/helpers.ts +++ b/packages/kbn-esql-ast/src/walker/helpers.ts @@ -7,11 +7,46 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import { ESQLProperNode } from '../types'; +import { + ESQLAstCommand, + ESQLAstQueryExpression, + ESQLColumn, + ESQLCommandMode, + ESQLCommandOption, + ESQLFunction, + ESQLIdentifier, + ESQLInlineCast, + ESQLList, + ESQLLiteral, + ESQLOrderExpression, + ESQLProperNode, + ESQLSource, + ESQLTimeInterval, + ESQLUnknownItem, +} from '../types'; + +export type NodeMatchKeys = + | keyof ESQLAstCommand + | keyof ESQLAstQueryExpression + | keyof ESQLFunction + | keyof ESQLCommandOption + | keyof ESQLSource + | keyof ESQLColumn + | keyof ESQLTimeInterval + | keyof ESQLList + | keyof ESQLLiteral + | keyof ESQLIdentifier + | keyof ESQLCommandMode + | keyof ESQLInlineCast + | keyof ESQLOrderExpression + | keyof ESQLUnknownItem; export type NodeMatchTemplateKey = V | V[] | RegExp; + export type NodeMatchTemplate = { - [K in keyof ESQLProperNode]?: NodeMatchTemplateKey; + [K in NodeMatchKeys]?: K extends keyof ESQLProperNode + ? NodeMatchTemplateKey + : NodeMatchTemplateKey; }; /** diff --git a/packages/kbn-esql-ast/src/walker/walker.test.ts b/packages/kbn-esql-ast/src/walker/walker.test.ts index 048e7e259eec2..c2db01f719d4b 100644 --- a/packages/kbn-esql-ast/src/walker/walker.test.ts +++ b/packages/kbn-esql-ast/src/walker/walker.test.ts @@ -20,6 +20,7 @@ import { ESQLTimeInterval, ESQLInlineCast, ESQLUnknownItem, + ESQLIdentifier, } from '../types'; import { walk, Walker } from './walker'; @@ -82,6 +83,23 @@ describe('structurally can walk all nodes', () => { ]); }); + test('can traverse JOIN command', () => { + const { ast } = parse('FROM index | LEFT JOIN a AS b ON c, d'); + const commands: ESQLCommand[] = []; + const identifiers: ESQLIdentifier[] = []; + const columns: ESQLColumn[] = []; + + walk(ast, { + visitCommand: (cmd) => commands.push(cmd), + visitIdentifier: (id) => identifiers.push(id), + visitColumn: (col) => columns.push(col), + }); + + expect(commands.map(({ name }) => name).sort()).toStrictEqual(['from', 'join']); + expect(identifiers.map(({ name }) => name).sort()).toStrictEqual(['a', 'as', 'b', 'c', 'd']); + expect(columns.map(({ name }) => name).sort()).toStrictEqual(['c', 'd']); + }); + test('"visitAny" can capture command nodes', () => { const { ast } = parse('FROM index | STATS a = 123 | WHERE 123 | LIMIT 10'); const commands: ESQLCommand[] = []; @@ -1050,6 +1068,37 @@ describe('Walker.match()', () => { name: 'a.b.c', }); }); + + test('can find WHERE command by its type', () => { + const query = 'FROM index | LEFT JOIN a | RIGHT JOIN b'; + const { root } = parse(query); + + const join1 = Walker.match(root, { + type: 'command', + name: 'join', + commandType: 'left', + })!; + const identifier1 = Walker.match(join1, { + type: 'identifier', + name: 'a', + })!; + const join2 = Walker.match(root, { + type: 'command', + name: 'join', + commandType: 'right', + })!; + const identifier2 = Walker.match(join2, { + type: 'identifier', + name: 'b', + })!; + + expect(identifier1).toMatchObject({ + name: 'a', + }); + expect(identifier2).toMatchObject({ + name: 'b', + }); + }); }); describe('Walker.matchAll()', () => {