From e49e89e5753893d9f3bdff0d23b84e45b384bba2 Mon Sep 17 00:00:00 2001 From: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com> Date: Wed, 31 Jul 2024 21:03:42 +0200 Subject: [PATCH] [ES|QL] Implements `Visitor` pattern for ES|QL AST (#189516) ## Summary Partially addresses https://github.com/elastic/kibana/issues/182255 - Implements the `Visitor` pattern for ES|QL AST trees. Unlike the `Walker` (which automatically traverses the whole tree exactly once), the `Visitor` pattern allows to control the traversal. The developer has to manually call children "visitor" routines. This manual handling enables: - The AST tree can be traversed any number of times. - Only a specific subset of the tree can be travered. - Each visitor receives a *context* object, which can provide the global context as well as a linked list to all parent nodes. - The context object also provides node-specific read/write functionality. - Each visitor can receive *input* from its parent node. - Each visitor can return *output* to its parent node. - The visitor nodes are strictly typed: the context object as well as inputs and outputs have specific types. Also the inputs and outputs TypeScript types are inferred automatically from the callback signature the developer specifies and then the correct input/output usage is enforced in other callbacks. - The "scenarios" test file contains real-world usage scenarios, like: - [Changing the `LIMIT`](https://github.com/elastic/kibana/pull/189516/files#diff-571e21fd50dbdb664e71297e2edd72c1a1b2b96f346248f0360558ef8ceb75f7R20) - [Removing a "filter", a `WHERE` command](https://github.com/elastic/kibana/pull/189516/files#diff-571e21fd50dbdb664e71297e2edd72c1a1b2b96f346248f0360558ef8ceb75f7R57) ### Checklist Delete any items that are not applicable to this PR. - [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 ### For maintainers - [x] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#kibana-release-notes-process) --------- Co-authored-by: Elastic Machine --- .../src/__tests__/ast_parser.commands.test.ts | 369 ++++++++++++++ .../src/__tests__/ast_parser.literal.test.ts | 25 + .../src/__tests__/ast_parser.sort.test.ts | 85 ++++ .../src/__tests__/ast_parser.where.test.ts | 38 ++ .../kbn-esql-ast/src/builder/index.test.ts | 20 + packages/kbn-esql-ast/src/builder/index.ts | 43 ++ packages/kbn-esql-ast/src/builder/types.ts | 30 ++ packages/kbn-esql-ast/src/types.ts | 9 + packages/kbn-esql-ast/src/visitor/README.md | 69 +++ .../src/visitor/__tests__/expressions.test.ts | 159 ++++++ .../src/visitor/__tests__/scenarios.test.ts | 193 ++++++++ .../src/visitor/__tests__/visitor.test.ts | 118 +++++ packages/kbn-esql-ast/src/visitor/contexts.ts | 438 ++++++++++++++++ .../src/visitor/global_visitor_context.ts | 468 ++++++++++++++++++ packages/kbn-esql-ast/src/visitor/index.ts | 12 + packages/kbn-esql-ast/src/visitor/types.ts | 256 ++++++++++ packages/kbn-esql-ast/src/visitor/utils.ts | 36 ++ packages/kbn-esql-ast/src/visitor/visitor.ts | 98 ++++ 18 files changed, 2466 insertions(+) create mode 100644 packages/kbn-esql-ast/src/__tests__/ast_parser.commands.test.ts create mode 100644 packages/kbn-esql-ast/src/__tests__/ast_parser.literal.test.ts create mode 100644 packages/kbn-esql-ast/src/__tests__/ast_parser.sort.test.ts create mode 100644 packages/kbn-esql-ast/src/__tests__/ast_parser.where.test.ts create mode 100644 packages/kbn-esql-ast/src/builder/index.test.ts create mode 100644 packages/kbn-esql-ast/src/builder/index.ts create mode 100644 packages/kbn-esql-ast/src/builder/types.ts create mode 100644 packages/kbn-esql-ast/src/visitor/README.md create mode 100644 packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts create mode 100644 packages/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts create mode 100644 packages/kbn-esql-ast/src/visitor/__tests__/visitor.test.ts create mode 100644 packages/kbn-esql-ast/src/visitor/contexts.ts create mode 100644 packages/kbn-esql-ast/src/visitor/global_visitor_context.ts create mode 100644 packages/kbn-esql-ast/src/visitor/index.ts create mode 100644 packages/kbn-esql-ast/src/visitor/types.ts create mode 100644 packages/kbn-esql-ast/src/visitor/utils.ts create mode 100644 packages/kbn-esql-ast/src/visitor/visitor.ts diff --git a/packages/kbn-esql-ast/src/__tests__/ast_parser.commands.test.ts b/packages/kbn-esql-ast/src/__tests__/ast_parser.commands.test.ts new file mode 100644 index 0000000000000..a636f4a448595 --- /dev/null +++ b/packages/kbn-esql-ast/src/__tests__/ast_parser.commands.test.ts @@ -0,0 +1,369 @@ +/* + * 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 { getAstAndSyntaxErrors as parse } from '../ast_parser'; + +describe('commands', () => { + describe('correctly formatted, basic usage', () => { + it('SHOW', () => { + const query = 'SHOW info'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + { + type: 'command', + name: 'show', + args: [ + { + type: 'function', + name: 'info', + }, + ], + }, + ]); + }); + + it('META', () => { + const query = 'META functions'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + { + type: 'command', + name: 'meta', + args: [ + { + type: 'function', + name: 'functions', + }, + ], + }, + ]); + }); + + it('FROM', () => { + const query = 'FROM index'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + { + type: 'command', + name: 'from', + args: [ + { + type: 'source', + name: 'index', + }, + ], + }, + ]); + }); + + it('ROW', () => { + const query = 'ROW 1'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + { + type: 'command', + name: 'row', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }, + ]); + }); + + it('EVAL', () => { + const query = 'FROM index | EVAL 1'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'eval', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }, + ]); + }); + + it('STATS', () => { + const query = 'FROM index | STATS 1'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'stats', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }, + ]); + }); + + it('LIMIT', () => { + const query = 'FROM index | LIMIT 1'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }, + ]); + }); + + it('KEEP', () => { + const query = 'FROM index | KEEP abc'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'keep', + args: [ + { + type: 'column', + name: 'abc', + }, + ], + }, + ]); + }); + + it('SORT', () => { + const query = 'FROM index | SORT 1'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'sort', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }, + ]); + }); + + it('WHERE', () => { + const query = 'FROM index | WHERE 1'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'where', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }, + ]); + }); + + it('DROP', () => { + const query = 'FROM index | DROP abc'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'drop', + args: [ + { + type: 'column', + name: 'abc', + }, + ], + }, + ]); + }); + + it('RENAME', () => { + const query = 'FROM index | RENAME a AS b, c AS d'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'rename', + args: [ + { + type: 'option', + name: 'as', + args: [ + { + type: 'column', + name: 'a', + }, + { + type: 'column', + name: 'b', + }, + ], + }, + { + type: 'option', + name: 'as', + args: [ + { + type: 'column', + name: 'c', + }, + { + type: 'column', + name: 'd', + }, + ], + }, + ], + }, + ]); + }); + + it('DISSECT', () => { + const query = 'FROM index | DISSECT a "b" APPEND_SEPARATOR="c"'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'dissect', + args: [ + { + type: 'column', + name: 'a', + }, + { + type: 'literal', + value: '"b"', + }, + { + type: 'option', + name: 'append_separator', + args: [ + { + type: 'literal', + value: '"c"', + }, + ], + }, + ], + }, + ]); + }); + + it('GROK', () => { + const query = 'FROM index | GROK a "b"'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'grok', + args: [ + { + type: 'column', + name: 'a', + }, + { + type: 'literal', + value: '"b"', + }, + ], + }, + ]); + }); + + it('ENRICH', () => { + const query = 'FROM index | ENRICH a ON b WITH c, d'; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'enrich', + args: [ + { + type: 'source', + name: 'a', + }, + { + type: 'option', + name: 'on', + args: [ + { + type: 'column', + name: 'b', + }, + ], + }, + { + type: 'option', + name: 'with', + }, + ], + }, + ]); + }); + + it('MV_EXPAND', () => { + const query = 'FROM index | MV_EXPAND a '; + const { ast } = parse(query); + + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'mv_expand', + args: [ + { + type: 'column', + name: 'a', + }, + ], + }, + ]); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/__tests__/ast_parser.literal.test.ts b/packages/kbn-esql-ast/src/__tests__/ast_parser.literal.test.ts new file mode 100644 index 0000000000000..9b966905308d7 --- /dev/null +++ b/packages/kbn-esql-ast/src/__tests__/ast_parser.literal.test.ts @@ -0,0 +1,25 @@ +/* + * 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 { getAstAndSyntaxErrors as parse } from '../ast_parser'; +import { ESQLLiteral } from '../types'; + +describe('literal expression', () => { + it('numeric expression captures "value", and "name" fields', () => { + const text = 'ROW 1'; + const { ast } = parse(text); + const literal = ast[0].args[0] as ESQLLiteral; + + expect(literal).toMatchObject({ + type: 'literal', + literalType: 'number', + name: '1', + value: 1, + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/__tests__/ast_parser.sort.test.ts b/packages/kbn-esql-ast/src/__tests__/ast_parser.sort.test.ts new file mode 100644 index 0000000000000..ccfbceb890893 --- /dev/null +++ b/packages/kbn-esql-ast/src/__tests__/ast_parser.sort.test.ts @@ -0,0 +1,85 @@ +/* + * 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 { getAstAndSyntaxErrors as parse } from '../ast_parser'; + +describe('SORT', () => { + describe('correctly formatted', () => { + // Un-skip one https://github.com/elastic/kibana/issues/189491 fixed. + it.skip('example from documentation', () => { + const text = ` + FROM employees + | KEEP first_name, last_name, height + | SORT height DESC + `; + const { ast, errors } = parse(text); + + expect(errors.length).toBe(0); + expect(ast).toMatchObject([ + {}, + {}, + { + type: 'command', + name: 'sort', + args: [ + { + type: 'column', + name: 'height', + }, + ], + }, + ]); + }); + + // Un-skip once https://github.com/elastic/kibana/issues/189491 fixed. + it.skip('can parse various sorting columns with options', () => { + const text = + 'FROM a | SORT a, b ASC, c DESC, d NULLS FIRST, e NULLS LAST, f ASC NULLS FIRST, g DESC NULLS LAST'; + const { ast, errors } = parse(text); + + expect(errors.length).toBe(0); + expect(ast).toMatchObject([ + {}, + { + type: 'command', + name: 'sort', + args: [ + { + type: 'column', + name: 'a', + }, + { + type: 'column', + name: 'b', + }, + { + type: 'column', + name: 'c', + }, + { + type: 'column', + name: 'd', + }, + { + type: 'column', + name: 'e', + }, + { + type: 'column', + name: 'f', + }, + { + type: 'column', + name: 'g', + }, + ], + }, + ]); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/__tests__/ast_parser.where.test.ts b/packages/kbn-esql-ast/src/__tests__/ast_parser.where.test.ts new file mode 100644 index 0000000000000..34148ec1aecd2 --- /dev/null +++ b/packages/kbn-esql-ast/src/__tests__/ast_parser.where.test.ts @@ -0,0 +1,38 @@ +/* + * 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 { getAstAndSyntaxErrors as parse } from '../ast_parser'; + +describe('WHERE', () => { + describe('correctly formatted', () => { + it('example from documentation', () => { + const text = ` + FROM employees + | KEEP first_name, last_name, still_hired + | WHERE still_hired == true + `; + const { ast, errors } = parse(text); + + expect(errors.length).toBe(0); + expect(ast).toMatchObject([ + {}, + {}, + { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + }, + ], + }, + ]); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/builder/index.test.ts b/packages/kbn-esql-ast/src/builder/index.test.ts new file mode 100644 index 0000000000000..f54ab2f90a9ca --- /dev/null +++ b/packages/kbn-esql-ast/src/builder/index.test.ts @@ -0,0 +1,20 @@ +/* + * 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 { Builder } from '.'; + +test('can mint a numeric literal', () => { + const node = Builder.numericLiteral({ value: 42 }); + + expect(node).toMatchObject({ + type: 'literal', + literalType: 'number', + name: '42', + value: 42, + }); +}); diff --git a/packages/kbn-esql-ast/src/builder/index.ts b/packages/kbn-esql-ast/src/builder/index.ts new file mode 100644 index 0000000000000..d389caf40fab2 --- /dev/null +++ b/packages/kbn-esql-ast/src/builder/index.ts @@ -0,0 +1,43 @@ +/* + * 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 { ESQLNumberLiteral } from '../types'; +import { AstNodeParserFields, AstNodeTemplate } from './types'; + +export class Builder { + /** + * Constructs fields which are only available when the node is minted by + * the parser. + */ + public static readonly parserFields = ({ + location = { min: 0, max: 0 }, + text = '', + incomplete = false, + }: Partial): AstNodeParserFields => ({ + location, + text, + incomplete, + }); + + /** + * Constructs a number literal node. + */ + public static readonly numericLiteral = ( + template: Omit, 'literalType' | 'name'> + ): ESQLNumberLiteral => { + const node: ESQLNumberLiteral = { + ...template, + ...Builder.parserFields(template), + type: 'literal', + literalType: 'number', + name: template.value.toString(), + }; + + return node; + }; +} diff --git a/packages/kbn-esql-ast/src/builder/types.ts b/packages/kbn-esql-ast/src/builder/types.ts new file mode 100644 index 0000000000000..60575c0d00994 --- /dev/null +++ b/packages/kbn-esql-ast/src/builder/types.ts @@ -0,0 +1,30 @@ +/* + * 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 type { ESQLProperNode, ESQLAstBaseItem } from '../types'; + +/** + * Node fields which are available only when the node is minted by the parser. + * When creating nodes manually, these fields are not available. + */ +export type AstNodeParserFields = Pick; + +/** + * The node *template* transforms ES|QL AST nodes into a permissive shape, with + * the aim to: + * + * - Remove the `type` property, as the builder will set it. + * - Make properties like `text`, `location`, and `incomplete` optional, as they + * are a available only when the AST node is minted by the parser. + * - Make all other properties optional, for easy node creation. + */ +export type AstNodeTemplate = Omit< + Node, + 'type' | 'text' | 'location' | 'incomplete' +> & + Partial>; diff --git a/packages/kbn-esql-ast/src/types.ts b/packages/kbn-esql-ast/src/types.ts index 257a004e78f10..5bc1a02ffd2ae 100644 --- a/packages/kbn-esql-ast/src/types.ts +++ b/packages/kbn-esql-ast/src/types.ts @@ -35,6 +35,15 @@ export type ESQLAstField = ESQLFunction | ESQLColumn; */ export type ESQLAstItem = ESQLSingleAstItem | ESQLAstItem[]; +export type ESQLAstNodeWithArgs = ESQLCommand | ESQLCommandOption | ESQLFunction; + +/** + * *Proper* are nodes which are objects with `type` property, once we get rid + * of the nodes which are plain arrays, all nodes will be *proper* and we can + * remove this type. + */ +export type ESQLProperNode = ESQLSingleAstItem | ESQLAstCommand; + export interface ESQLLocation { min: number; max: number; diff --git a/packages/kbn-esql-ast/src/visitor/README.md b/packages/kbn-esql-ast/src/visitor/README.md new file mode 100644 index 0000000000000..71729bf56a0ab --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/README.md @@ -0,0 +1,69 @@ +## High-level AST structure + +Broadly, there are two AST node types: (1) commands (say `FROM ...`, like +*statements* in other languages), and (2) expressions (say `a + b`, or `fn()`). + + +### Commands + +Commands in ES|QL are like *statements* in other languages. They are the top +level nodes in the AST. + +The root node of the AST is considered to bye the "query" node. It contains a +list of commands. + +``` +Quey = Command[] +``` + +Each command receives a list of positional arguments. For example: + +``` +COMMAND arg1, arg2, arg3 +``` + +A command may also receive additional lists of *named* arguments, we refer to +them as `option`s. For example: + +``` +COMMAND arg1, arg2, arg3 OPTION1 arg4, arg5 OPTION2 arg6, arg7 +``` + +Essentially, one can of command arguments as a list of expressions, with the +ability to add named arguments to the command. For example, the above command +arguments can be represented as: + +```js +{ + '': [arg1, arg2, arg3], + 'option1': [arg4, arg5], + 'option2': [arg6, arg7] +} +``` + +Each command has a command specific `visitCommandX` callback, where `X` is the +name of the command. If a command-specific callback is not found, the generic +`visitCommand` callback is called. + + +### Expressions + +Expressions just like expressions in other languages. Expressions can be deeply +nested, as one expression can contain other expressions. For example, math +expressions `1 + 2`, function call expressions `fn()`, identifier expressions +`my.index` and so on. + +As of this writing, the following expressions are defined: + +- Column identifier expression, `{type: "column"}`, like `@timestamp` +- Source identifier expression, `{type: "source"}`, like `tsdb_index` +- Function call expression, `{type: "function"}`, like `fn(123)` +- Literal expression, `{type: "literal"}`, like `123`, `"hello"` +- List literal expression, `{type: "list"}`, like `[1, 2, 3]`, `["a", "b", "c"]`, `[true, false]` +- Time interval expression, `{type: "interval"}`, like `1h`, `1d`, `1w` +- Inline cast expression, `{type: "cast"}`, like `abc::int`, `def::string` +- Unknown node, `{type: "unknown"}` + +Each expression has a `visitExpressionX` callback, where `X` is the type of the +expression. If a expression-specific callback is not found, the generic +`visitExpression` callback is called. diff --git a/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts b/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts new file mode 100644 index 0000000000000..efd30f035e7ca --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts @@ -0,0 +1,159 @@ +/* + * 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 { getAstAndSyntaxErrors } from '../../ast_parser'; +import { Visitor } from '../visitor'; + +test('"visitExpression" captures all non-captured expressions', () => { + const { ast } = getAstAndSyntaxErrors(` + FROM index + | STATS 1, "str", [true], a = b BY field + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitExpression', (ctx) => { + return ''; + }) + .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 | STATS , , , | LIMIT ' + ); +}); + +test('can terminate walk early, does not visit all literals', () => { + const numbers: number[] = []; + const { ast } = getAstAndSyntaxErrors(` + FROM index + | STATS 0, 1, 2, 3 + | LIMIT 123 + `); + const result = new Visitor() + .on('visitExpression', (ctx) => { + return 0; + }) + .on('visitLiteralExpression', (ctx) => { + numbers.push(ctx.node.value as number); + return ctx.node.value; + }) + .on('visitCommand', (ctx) => { + for (const res of ctx.visitArguments()) if (res) return res; + }) + .on('visitQuery', (ctx) => { + for (const res of ctx.visitCommands()) if (res) return res; + }) + .visitQuery(ast); + + expect(result).toBe(1); + expect(numbers).toEqual([0, 1]); +}); + +test('"visitColumnExpression" takes over all column visits', () => { + const { ast } = getAstAndSyntaxErrors(` + FROM index | STATS a + `); + const visitor = new Visitor() + .on('visitColumnExpression', (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 '); +}); + +test('"visitSourceExpression" takes over all source visits', () => { + const { ast } = getAstAndSyntaxErrors(` + FROM index + | STATS 1, "str", [true], a = b BY field + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitSourceExpression', (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 | STATS E, E, E, E | LIMIT E'); +}); + +test('"visitFunctionCallExpression" takes over all literal visits', () => { + const { ast } = getAstAndSyntaxErrors(` + FROM index + | STATS 1, "str", [true], a = b BY field + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitFunctionCallExpression', (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 E, E, E, | LIMIT E'); +}); + +test('"visitLiteral" takes over all literal visits', () => { + const { ast } = getAstAndSyntaxErrors(` + FROM index + | STATS 1, "str", [true], a = b BY field + | LIMIT 123 + `); + const visitor = new Visitor() + .on('visitLiteralExpression', (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 , , E, E | LIMIT '); +}); diff --git a/packages/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts b/packages/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts new file mode 100644 index 0000000000000..ce338e8bd72ba --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/__tests__/scenarios.test.ts @@ -0,0 +1,193 @@ +/* + * 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. + */ + +/** + * @category Visitor Real-world Scenarios + * + * This test suite contains real-world scenarios that demonstrate how to use the + * visitor to traverse the AST and make changes to it, or how to extract useful + */ + +import { getAstAndSyntaxErrors } from '../../ast_parser'; +import { ESQLAstQueryNode } from '../types'; +import { Visitor } from '../visitor'; + +test('change LIMIT from 24 to 42', () => { + const { ast } = getAstAndSyntaxErrors(` + FROM index + | STATS 1, "str", [true], a = b BY field + | LIMIT 24 + `); + + // Find the LIMIT node + const limit = () => + new Visitor() + .on('visitLimitCommand', (ctx) => ctx.numeric()) + .on('visitCommand', () => null) + .on('visitQuery', (ctx) => [...ctx.visitCommands()]) + .visitQuery(ast) + .filter(Boolean)[0]; + + expect(limit()).toBe(24); + + // Change LIMIT to 42 + new Visitor() + .on('visitLimitCommand', (ctx) => { + ctx.setLimit(42); + }) + .on('visitCommand', () => {}) + .on('visitQuery', (ctx) => [...ctx.visitCommands()]) + .visitQuery(ast); + + expect(limit()).toBe(42); +}); + +/** + * Implement this once sorting order expressions are available: + * + * - https://github.com/elastic/kibana/issues/189491 + */ +test.todo('can modify sorting orders'); + +test('can remove a specific WHERE command', () => { + const query = getAstAndSyntaxErrors(` + FROM employees + | KEEP first_name, last_name, still_hired + | WHERE still_hired == true + | WHERE last_name == "Jeo" + | WHERE 123 == salary + `); + + const print = () => + new Visitor() + .on('visitColumnExpression', (ctx) => ctx.node.name) + .on( + 'visitFunctionCallExpression', + (ctx) => `${ctx.node.name}(${[...ctx.visitArguments()].join(', ')})` + ) + .on('visitExpression', (ctx) => '') + .on('visitCommand', (ctx) => { + if (ctx.node.name === 'where') { + const args = [...ctx.visitArguments()].join(', '); + return `${ctx.name()}${args ? ` ${args}` : ''}`; + } else { + return ''; + } + }) + .on('visitQuery', (ctx) => [...ctx.visitCommands()].filter(Boolean).join(' | ')) + .visitQuery(query.ast); + + const removeFilter = (field: string) => { + query.ast = new Visitor() + .on('visitColumnExpression', (ctx) => (ctx.node.name === field ? null : ctx.node)) + .on('visitFunctionCallExpression', (ctx) => { + const args = [...ctx.visitArguments()]; + return args.some((arg) => arg === null) ? null : ctx.node; + }) + .on('visitExpression', (ctx) => ctx.node) + .on('visitCommand', (ctx) => { + if (ctx.node.name === 'where') { + ctx.node.args = [...ctx.visitArguments()].filter(Boolean); + } + return ctx.node; + }) + .on('visitQuery', (ctx) => [...ctx.visitCommands()].filter((cmd) => cmd.args.length)) + .visitQuery(query.ast); + }; + + expect(print()).toBe( + 'WHERE ==(still_hired, ) | WHERE ==(last_name, ) | WHERE ==(, salary)' + ); + removeFilter('last_name'); + expect(print()).toBe('WHERE ==(still_hired, ) | WHERE ==(, salary)'); + removeFilter('still_hired'); + removeFilter('still_hired'); + expect(print()).toBe('WHERE ==(, salary)'); + removeFilter('still_hired'); + removeFilter('salary'); + removeFilter('salary'); + expect(print()).toBe(''); +}); + +export const prettyPrint = (ast: ESQLAstQueryNode) => + new Visitor() + .on('visitSourceExpression', (ctx) => { + return ctx.node.name; + }) + .on('visitColumnExpression', (ctx) => { + return ctx.node.name; + }) + .on('visitFunctionCallExpression', (ctx) => { + let args = ''; + for (const arg of ctx.visitArguments()) { + args += (args ? ', ' : '') + arg; + } + return `${ctx.node.name.toUpperCase()}${args ? `(${args})` : ''}`; + }) + .on('visitLiteralExpression', (ctx) => { + return ctx.node.value; + }) + .on('visitListLiteralExpression', (ctx) => { + return ''; + }) + .on('visitTimeIntervalLiteralExpression', (ctx) => { + return ''; + }) + .on('visitInlineCastExpression', (ctx) => { + return ''; + }) + .on('visitExpression', (ctx) => { + return ''; + }) + .on('visitCommandOption', (ctx) => { + let args = ''; + for (const arg of ctx.visitArguments()) { + args += (args ? ', ' : '') + arg; + } + return ctx.node.name.toUpperCase() + (args ? ` ${args}` : ''); + }) + .on('visitCommand', (ctx) => { + let args = ''; + for (const source of ctx.visitArguments()) { + args += (args ? ', ' : '') + source; + } + return `${ctx.node.name.toUpperCase()}${args ? ` ${args}` : ''}`; + }) + .on('visitFromCommand', (ctx) => { + let sources = ''; + for (const source of ctx.visitSources()) { + sources += (sources ? ', ' : '') + source; + } + let options = ''; + for (const option of ctx.visitOptions()) { + options += ' ' + option; + } + return `FROM ${sources}${options}`; + }) + .on('visitLimitCommand', (ctx) => { + return `LIMIT ${ctx.numeric() ?? 0}`; + }) + .on('visitQuery', (ctx) => { + let text = ''; + for (const cmd of ctx.visitCommands()) { + text += (text ? ' | ' : '') + cmd; + } + return text; + }) + .visitQuery(ast); + +test('can print a query to text', () => { + const { ast } = getAstAndSyntaxErrors( + 'FROM index METADATA _id, asdf, 123 | STATS fn([1,2], 1d, 1::string, x in (1, 2)), a = b | LIMIT 1000' + ); + const text = prettyPrint(ast); + + expect(text).toBe( + 'FROM index METADATA _id, asdf, 123 | STATS FN(, , , IN(x, 1, 2)), =(a, b) | LIMIT 1000' + ); +}); diff --git a/packages/kbn-esql-ast/src/visitor/__tests__/visitor.test.ts b/packages/kbn-esql-ast/src/visitor/__tests__/visitor.test.ts new file mode 100644 index 0000000000000..24944f635ee44 --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/__tests__/visitor.test.ts @@ -0,0 +1,118 @@ +/* + * 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 { getAstAndSyntaxErrors } from '../../ast_parser'; +import { CommandVisitorContext, WhereCommandVisitorContext } from '../contexts'; +import { Visitor } from '../visitor'; + +test('can collect all command names in type safe way', () => { + const visitor = new Visitor() + .on('visitCommand', (ctx) => { + return ctx.node.name; + }) + .on('visitQuery', (ctx) => { + const cmds = []; + for (const cmd of ctx.visitCommands()) { + cmds.push(cmd); + } + return cmds; + }); + + const { ast } = getAstAndSyntaxErrors('FROM index | LIMIT 123'); + const res = visitor.visitQuery(ast); + + expect(res).toEqual(['from', 'limit']); +}); + +test('can pass inputs to visitors', () => { + const visitor = new Visitor() + .on('visitCommand', (ctx, prefix: string) => { + return prefix + ctx.node.name; + }) + .on('visitQuery', (ctx) => { + const cmds = []; + for (const cmd of ctx.visitCommands('pfx:')) { + cmds.push(cmd); + } + return cmds; + }); + + const { ast } = getAstAndSyntaxErrors('FROM index | LIMIT 123'); + const res = visitor.visitQuery(ast); + + expect(res).toEqual(['pfx:from', 'pfx:limit']); +}); + +test('can specify specific visitors for commands', () => { + const { ast } = getAstAndSyntaxErrors( + 'FROM index | SORT asfd | WHERE 1 | ENRICH adsf | LIMIT 123' + ); + const res = new Visitor() + .on('visitWhereCommand', () => 'where') + .on('visitSortCommand', () => 'sort') + .on('visitEnrichCommand', () => 'very rich') + .on('visitCommand', () => 'DEFAULT') + .on('visitQuery', (ctx) => [...ctx.visitCommands()]) + .visitQuery(ast); + + expect(res).toEqual(['DEFAULT', 'sort', 'where', 'very rich', 'DEFAULT']); +}); + +test('a command can access parent query node', () => { + const { ast } = getAstAndSyntaxErrors( + 'FROM index | SORT asfd | WHERE 1 | ENRICH adsf | LIMIT 123' + ); + new Visitor() + .on('visitWhereCommand', (ctx) => { + if (ctx.parent!.node !== ast) { + throw new Error('Expected parent to be query node'); + } + }) + .on('visitCommand', (ctx) => { + if (ctx.parent!.node !== ast) { + throw new Error('Expected parent to be query node'); + } + }) + .on('visitQuery', (ctx) => [...ctx.visitCommands()]) + .visitQuery(ast); +}); + +test('specific commands receive specific visitor contexts', () => { + const { ast } = getAstAndSyntaxErrors( + 'FROM index | SORT asfd | WHERE 1 | ENRICH adsf | LIMIT 123' + ); + + new Visitor() + .on('visitWhereCommand', (ctx) => { + if (!(ctx instanceof WhereCommandVisitorContext)) { + throw new Error('Expected WhereCommandVisitorContext'); + } + if (!(ctx instanceof CommandVisitorContext)) { + throw new Error('Expected WhereCommandVisitorContext'); + } + }) + .on('visitCommand', (ctx) => { + if (!(ctx instanceof CommandVisitorContext)) { + throw new Error('Expected CommandVisitorContext'); + } + }) + .on('visitQuery', (ctx) => [...ctx.visitCommands()]) + .visitQuery(ast); + + new Visitor() + .on('visitCommand', (ctx) => { + if (!(ctx instanceof CommandVisitorContext)) { + throw new Error('Expected CommandVisitorContext'); + } + if (ctx instanceof WhereCommandVisitorContext) { + throw new Error('Did not expect WhereCommandVisitorContext'); + } + }) + .on('visitQuery', (ctx) => [...ctx.visitCommands()]) + .visitQuery(ast); +}); diff --git a/packages/kbn-esql-ast/src/visitor/contexts.ts b/packages/kbn-esql-ast/src/visitor/contexts.ts new file mode 100644 index 0000000000000..ca6044c017aa6 --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/contexts.ts @@ -0,0 +1,438 @@ +/* + * 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. + */ + +/* eslint-disable max-classes-per-file */ +// Splitting classes across files runs into issues with circular dependencies +// and makes it harder to understand the code structure. + +import { type GlobalVisitorContext, SharedData } from './global_visitor_context'; +import { firstItem, singleItems } from './utils'; +import type { + ESQLAstCommand, + ESQLAstItem, + ESQLAstNodeWithArgs, + ESQLColumn, + ESQLCommandOption, + ESQLFunction, + ESQLInlineCast, + ESQLList, + ESQLLiteral, + ESQLNumberLiteral, + ESQLSource, + ESQLTimeInterval, +} from '../types'; +import type { + CommandVisitorInput, + ESQLAstExpressionNode, + ESQLAstQueryNode, + ExpressionVisitorInput, + ExpressionVisitorOutput, + UndefinedToVoid, + VisitorAstNode, + VisitorMethods, +} from './types'; +import { Builder } from '../builder'; + +const isNodeWithArgs = (x: unknown): x is ESQLAstNodeWithArgs => + !!x && typeof x === 'object' && Array.isArray((x as any).args); + +export class VisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData, + Node extends VisitorAstNode = VisitorAstNode +> { + constructor( + /** + * Global visitor context. + */ + public readonly ctx: GlobalVisitorContext, + + /** + * ES|QL AST node which is currently being visited. + */ + public readonly node: Node, + + /** + * Context of the parent node, from which the current node was reached + * during the AST traversal. + */ + public readonly parent: VisitorContext | null = null + ) {} + + public *visitArguments( + input: ExpressionVisitorInput + ): Iterable> { + this.ctx.assertMethodExists('visitExpression'); + + const node = this.node; + + if (!isNodeWithArgs(node)) { + throw new Error('Node does not have arguments'); + } + + for (const arg of singleItems(node.args)) { + yield this.visitExpression(arg, input as any); + } + } + + public visitExpression( + expressionNode: ESQLAstExpressionNode, + input: ExpressionVisitorInput + ): ExpressionVisitorOutput { + return this.ctx.visitExpression(this, expressionNode, input); + } + + public visitCommand( + commandNode: ESQLAstCommand, + input: CommandVisitorInput + ): ExpressionVisitorOutput { + return this.ctx.visitCommand(this, commandNode, input); + } +} + +export class QueryVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends VisitorContext { + public *visitCommands( + input: UndefinedToVoid>[1]> + ): Iterable< + | ReturnType> + | ReturnType> + > { + this.ctx.assertMethodExists('visitCommand'); + + for (const cmd of this.node) { + yield this.visitCommand(cmd, input as any); + } + } +} + +// Commands -------------------------------------------------------------------- + +export class CommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData, + Node extends ESQLAstCommand = ESQLAstCommand +> extends VisitorContext { + public name(): string { + return this.node.name.toUpperCase(); + } + + public *options(): Iterable { + for (const arg of this.node.args) { + if (Array.isArray(arg)) { + continue; + } + if (arg.type === 'option') { + yield arg; + } + } + } + + public *visitOptions( + input: UndefinedToVoid>[1]> + ): Iterable>> { + this.ctx.assertMethodExists('visitCommandOption'); + + for (const option of this.options()) { + const sourceContext = new CommandOptionVisitorContext(this.ctx, option, this); + const result = this.ctx.methods.visitCommandOption!(sourceContext, input); + + yield result; + } + } + + public *arguments(option: '' | string = ''): Iterable { + option = option.toLowerCase(); + + if (!option) { + for (const arg of this.node.args) { + if (Array.isArray(arg)) { + yield arg; + continue; + } + if (arg.type !== 'option') { + yield arg; + } + } + } + + const optionNode = this.node.args.find( + (arg) => !Array.isArray(arg) && arg.type === 'option' && arg.name === option + ); + + if (optionNode) { + yield* (optionNode as ESQLCommandOption).args; + } + } + + public *visitArguments( + input: ExpressionVisitorInput, + option: '' | string = '' + ): Iterable> { + this.ctx.assertMethodExists('visitExpression'); + + const node = this.node; + + if (!isNodeWithArgs(node)) { + throw new Error('Node does not have arguments'); + } + + for (const arg of singleItems(this.arguments(option))) { + yield this.visitExpression(arg, input as any); + } + } + + public *visitSources( + input: UndefinedToVoid>[1]> + ): Iterable>> { + this.ctx.assertMethodExists('visitSourceExpression'); + + for (const arg of singleItems(this.node.args)) { + if (arg.type === 'source') { + const sourceContext = new SourceExpressionVisitorContext(this.ctx, arg, this); + const result = this.ctx.methods.visitSourceExpression!(sourceContext, input); + + yield result; + } + } + } +} + +export class CommandOptionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends VisitorContext {} + +// FROM [ METADATA ] +export class FromCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext { + /** + * Visit the METADATA part of the FROM command. + * + * FROM [ METADATA ] + * + * @param input Input object to pass to all "visitColumn" children methods. + * @returns An iterable of results of all the "visitColumn" visitor methods. + */ + public *visitMetadataColumns( + input: UndefinedToVoid>[1]> + ): Iterable>> { + this.ctx.assertMethodExists('visitColumnExpression'); + + let metadataOption: ESQLCommandOption | undefined; + + for (const arg of singleItems(this.node.args)) { + if (arg.type === 'option' && arg.name === 'metadata') { + metadataOption = arg; + break; + } + } + + if (!metadataOption) { + return; + } + + for (const arg of singleItems(metadataOption.args)) { + if (arg.type === 'column') { + const columnContext = new ColumnExpressionVisitorContext(this.ctx, arg, this); + const result = this.ctx.methods.visitColumnExpression!(columnContext, input); + + yield result; + } + } + } +} + +// LIMIT +export class LimitCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext { + /** + * @returns The first numeric literal argument of the command. + */ + public numericLiteral(): ESQLNumberLiteral | undefined { + const arg = firstItem(this.node.args); + + if (arg && arg.type === 'literal' && arg.literalType === 'number') { + return arg; + } + } + + /** + * @returns The value of the first numeric literal argument of the command. + */ + public numeric(): number | undefined { + const literal = this.numericLiteral(); + + return literal?.value; + } + + public setLimit(value: number): void { + const literalNode = Builder.numericLiteral({ value }); + + this.node.args = [literalNode]; + } +} + +// EXPLAIN +export class ExplainCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// ROW +export class RowCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// METRICS +export class MetricsCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// SHOW +export class ShowCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// META +export class MetaCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// EVAL +export class EvalCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// STATS [ BY ] +export class StatsCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// INLINESTATS [ BY ] +export class InlineStatsCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// LOOKUP ON +export class LookupCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// KEEP +export class KeepCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// SORT +export class SortCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// WHERE +export class WhereCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// DROP +export class DropCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// RENAME AS +export class RenameCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// DISSECT [ APPEND_SEPARATOR = ] +export class DissectCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// GROK +export class GrokCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// ENRICH [ ON ] [ WITH ] +export class EnrichCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// MV_EXPAND +export class MvExpandCommandVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends CommandVisitorContext {} + +// Expressions ----------------------------------------------------------------- + +export class ExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData, + Node extends ESQLAstExpressionNode = ESQLAstExpressionNode +> extends VisitorContext {} + +export class ColumnExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends VisitorContext {} + +export class SourceExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends VisitorContext {} + +export class FunctionCallExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends VisitorContext {} + +export class LiteralExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData, + Node extends ESQLLiteral = ESQLLiteral +> extends ExpressionVisitorContext {} + +export class ListLiteralExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData, + Node extends ESQLList = ESQLList +> extends ExpressionVisitorContext {} + +export class TimeIntervalLiteralExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends ExpressionVisitorContext {} + +export class InlineCastExpressionVisitorContext< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> extends ExpressionVisitorContext {} diff --git a/packages/kbn-esql-ast/src/visitor/global_visitor_context.ts b/packages/kbn-esql-ast/src/visitor/global_visitor_context.ts new file mode 100644 index 0000000000000..9cae41f36dde5 --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/global_visitor_context.ts @@ -0,0 +1,468 @@ +/* + * 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 * as contexts from './contexts'; +import type { + ESQLAstCommand, + ESQLColumn, + ESQLFunction, + ESQLInlineCast, + ESQLList, + ESQLLiteral, + ESQLSource, + ESQLTimeInterval, +} from '../types'; +import type * as types from './types'; + +export type SharedData = Record; + +/** + * Global shared visitor context available to all visitors when visiting the AST. + * It contains the shared data, which can be accessed and modified by all visitors. + */ +export class GlobalVisitorContext< + Methods extends types.VisitorMethods = types.VisitorMethods, + Data extends SharedData = SharedData +> { + constructor( + /** + * Visitor methods, used internally by the visitor to traverse the AST. + * @protected + */ + public readonly methods: Methods, + + /** + * Shared data, which can be accessed and modified by all visitors. + */ + public data: Data + ) {} + + public assertMethodExists(name: K) { + if (!this.methods[name]) { + throw new Error(`${name}() method is not defined`); + } + } + + private visitWithSpecificContext< + Method extends keyof types.VisitorMethods, + Context extends contexts.VisitorContext + >( + method: Method, + context: Context, + input: types.VisitorInput + ): types.VisitorOutput { + this.assertMethodExists(method); + return this.methods[method]!(context as any, input); + } + + // Command visiting ---------------------------------------------------------- + + public visitCommandGeneric( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + this.assertMethodExists('visitCommand'); + + const context = new contexts.CommandVisitorContext(this, node, parent); + const output = this.methods.visitCommand!(context, input); + + return output; + } + + public visitCommand( + parent: contexts.VisitorContext | null, + commandNode: ESQLAstCommand, + input: types.CommandVisitorInput + ): types.CommandVisitorOutput { + switch (commandNode.name) { + case 'from': { + if (!this.methods.visitFromCommand) break; + return this.visitFromCommand(parent, commandNode, input as any); + } + case 'limit': { + if (!this.methods.visitLimitCommand) break; + return this.visitLimitCommand(parent, commandNode, input as any); + } + case 'explain': { + if (!this.methods.visitExplainCommand) break; + return this.visitExplainCommand(parent, commandNode, input as any); + } + case 'row': { + if (!this.methods.visitRowCommand) break; + return this.visitRowCommand(parent, commandNode, input as any); + } + // TODO: uncomment this when the command is implemented + // case 'metrics': { + // if (!this.methods.visitMetricsCommand) break; + // return this.visitMetricsCommand(parent, commandNode, input as any); + // } + case 'show': { + if (!this.methods.visitShowCommand) break; + return this.visitShowCommand(parent, commandNode, input as any); + } + case 'meta': { + if (!this.methods.visitMetaCommand) break; + return this.visitMetaCommand(parent, commandNode, input as any); + } + case 'eval': { + if (!this.methods.visitEvalCommand) break; + return this.visitEvalCommand(parent, commandNode, input as any); + } + case 'stats': { + if (!this.methods.visitStatsCommand) break; + return this.visitStatsCommand(parent, commandNode, input as any); + } + // TODO: uncomment this when the command is implemented + // case 'inline_stats': { + // if (!this.methods.visitInlineStatsCommand) break; + // return this.visitInlineStatsCommand(parent, commandNode, input as any); + // } + case 'lookup': { + if (!this.methods.visitLookupCommand) break; + return this.visitLookupCommand(parent, commandNode, input as any); + } + case 'keep': { + if (!this.methods.visitKeepCommand) break; + return this.visitKeepCommand(parent, commandNode, input as any); + } + case 'sort': { + if (!this.methods.visitSortCommand) break; + return this.visitSortCommand(parent, commandNode, input as any); + } + case 'where': { + if (!this.methods.visitWhereCommand) break; + return this.visitWhereCommand(parent, commandNode, input as any); + } + case 'drop': { + if (!this.methods.visitDropCommand) break; + return this.visitDropCommand(parent, commandNode, input as any); + } + case 'rename': { + if (!this.methods.visitRenameCommand) break; + return this.visitRenameCommand(parent, commandNode, input as any); + } + case 'dissect': { + if (!this.methods.visitDissectCommand) break; + return this.visitDissectCommand(parent, commandNode, input as any); + } + case 'grok': { + if (!this.methods.visitGrokCommand) break; + return this.visitGrokCommand(parent, commandNode, input as any); + } + case 'enrich': { + if (!this.methods.visitEnrichCommand) break; + return this.visitEnrichCommand(parent, commandNode, input as any); + } + case 'mv_expand': { + if (!this.methods.visitMvExpandCommand) break; + return this.visitMvExpandCommand(parent, commandNode, input as any); + } + } + return this.visitCommandGeneric(parent, commandNode, input as any); + } + + public visitFromCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.FromCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitFromCommand', context, input); + } + + public visitLimitCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.LimitCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitLimitCommand', context, input); + } + + public visitExplainCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.ExplainCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitExplainCommand', context, input); + } + + public visitRowCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.RowCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitRowCommand', context, input); + } + + public visitMetricsCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.MetricsCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitMetricsCommand', context, input); + } + + public visitShowCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.ShowCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitShowCommand', context, input); + } + + public visitMetaCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.MetaCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitMetaCommand', context, input); + } + + public visitEvalCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.EvalCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitEvalCommand', context, input); + } + + public visitStatsCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.StatsCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitStatsCommand', context, input); + } + + public visitInlineStatsCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.InlineStatsCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitInlineStatsCommand', context, input); + } + + public visitLookupCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.LookupCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitLookupCommand', context, input); + } + + public visitKeepCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.KeepCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitKeepCommand', context, input); + } + + public visitSortCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.SortCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitSortCommand', context, input); + } + + public visitWhereCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.WhereCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitWhereCommand', context, input); + } + + public visitDropCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.DropCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitDropCommand', context, input); + } + + public visitRenameCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.RenameCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitRenameCommand', context, input); + } + + public visitDissectCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.DissectCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitDissectCommand', context, input); + } + + public visitGrokCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.GrokCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitGrokCommand', context, input); + } + + public visitEnrichCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.EnrichCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitEnrichCommand', context, input); + } + + public visitMvExpandCommand( + parent: contexts.VisitorContext | null, + node: ESQLAstCommand, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.MvExpandCommandVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitMvExpandCommand', context, input); + } + + // Expression visiting ------------------------------------------------------- + + public visitExpressionGeneric( + parent: contexts.VisitorContext | null, + node: types.ESQLAstExpressionNode, + input: types.VisitorInput + ): types.VisitorOutput { + this.assertMethodExists('visitExpression'); + + const context = new contexts.ExpressionVisitorContext(this, node, parent); + const output = this.methods.visitExpression!(context, input); + + return output; + } + + public visitExpression( + parent: contexts.VisitorContext | null, + expressionNode: types.ESQLAstExpressionNode, + input: types.ExpressionVisitorInput + ): types.ExpressionVisitorOutput { + if (Array.isArray(expressionNode)) { + throw new Error('should not happen'); + } + switch (expressionNode.type) { + case 'column': { + if (!this.methods.visitColumnExpression) break; + return this.visitColumnExpression(parent, expressionNode, input as any); + } + case 'source': { + if (!this.methods.visitSourceExpression) break; + return this.visitSourceExpression(parent, expressionNode, input as any); + } + case 'function': { + if (!this.methods.visitFunctionCallExpression) break; + return this.visitFunctionCallExpression(parent, expressionNode, input as any); + } + case 'literal': { + if (!this.methods.visitLiteralExpression) break; + return this.visitLiteralExpression(parent, expressionNode, input as any); + } + case 'list': { + if (!this.methods.visitListLiteralExpression) break; + return this.visitListLiteralExpression(parent, expressionNode, input as any); + } + case 'timeInterval': { + if (!this.methods.visitTimeIntervalLiteralExpression) break; + return this.visitTimeIntervalLiteralExpression(parent, expressionNode, input as any); + } + case 'inlineCast': { + if (!this.methods.visitInlineCastExpression) break; + return this.visitInlineCastExpression(parent, expressionNode, input as any); + } + } + return this.visitExpressionGeneric(parent, expressionNode, input as any); + } + + public visitColumnExpression( + parent: contexts.VisitorContext | null, + node: ESQLColumn, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.ColumnExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitColumnExpression', context, input); + } + + public visitSourceExpression( + parent: contexts.VisitorContext | null, + node: ESQLSource, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.SourceExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitSourceExpression', context, input); + } + + public visitFunctionCallExpression( + parent: contexts.VisitorContext | null, + node: ESQLFunction, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.FunctionCallExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitFunctionCallExpression', context, input); + } + + public visitLiteralExpression( + parent: contexts.VisitorContext | null, + node: ESQLLiteral, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.LiteralExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitLiteralExpression', context, input); + } + + public visitListLiteralExpression( + parent: contexts.VisitorContext | null, + node: ESQLList, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.ListLiteralExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitListLiteralExpression', context, input); + } + + public visitTimeIntervalLiteralExpression( + parent: contexts.VisitorContext | null, + node: ESQLTimeInterval, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.TimeIntervalLiteralExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitTimeIntervalLiteralExpression', context, input); + } + + public visitInlineCastExpression( + parent: contexts.VisitorContext | null, + node: ESQLInlineCast, + input: types.VisitorInput + ): types.VisitorOutput { + const context = new contexts.InlineCastExpressionVisitorContext(this, node, parent); + return this.visitWithSpecificContext('visitInlineCastExpression', context, input); + } +} diff --git a/packages/kbn-esql-ast/src/visitor/index.ts b/packages/kbn-esql-ast/src/visitor/index.ts new file mode 100644 index 0000000000000..9e46616a50002 --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/index.ts @@ -0,0 +1,12 @@ +/* + * 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. + */ + +export * from './types'; +export { Visitor, type VisitorOptions } from './visitor'; +export { GlobalVisitorContext, type SharedData } from './global_visitor_context'; +export * from './contexts'; diff --git a/packages/kbn-esql-ast/src/visitor/types.ts b/packages/kbn-esql-ast/src/visitor/types.ts new file mode 100644 index 0000000000000..a8ec5e9bd1785 --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/types.ts @@ -0,0 +1,256 @@ +/* + * 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 type { SharedData } from './global_visitor_context'; +import type * as ast from '../types'; +import type * as contexts from './contexts'; + +/** + * We don't have a dedicated "query" AST node, so - for now - we use the root + * array of commands as the "query" node. + */ +export type ESQLAstQueryNode = ast.ESQLAst; + +/** + * Represents an "expression" node in the AST. + */ +// export type ESQLAstExpressionNode = ESQLAstItem; +export type ESQLAstExpressionNode = ast.ESQLSingleAstItem; + +/** + * All possible AST nodes supported by the visitor. + */ +export type VisitorAstNode = ESQLAstQueryNode | ast.ESQLAstNode; + +export type Visitor = ( + ctx: Ctx, + input: Input +) => Output; + +/** + * Retrieves the `Input` of a {@link Visitor} function. + */ +export type VisitorInput< + Methods extends VisitorMethods, + Method extends keyof Methods +> = UndefinedToVoid>>[1]>; + +/** + * Retrieves the `Output` of a {@link Visitor} function. + */ +export type VisitorOutput< + Methods extends VisitorMethods, + Method extends keyof Methods +> = ReturnType>>; + +/** + * Input that satisfies any expression visitor input constraints. + */ +export type ExpressionVisitorInput = AnyToVoid< + | VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput +>; + +/** + * Input that satisfies any expression visitor output constraints. + */ +export type ExpressionVisitorOutput = + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput; + +/** + * Input that satisfies any command visitor input constraints. + */ +export type CommandVisitorInput = AnyToVoid< + | VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput & + VisitorInput +>; + +/** + * Input that satisfies any command visitor output constraints. + */ +export type CommandVisitorOutput = + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput + | VisitorOutput; + +export interface VisitorMethods< + Visitors extends VisitorMethods = any, + Data extends SharedData = SharedData +> { + visitQuery?: Visitor, any, any>; + visitCommand?: Visitor, any, any>; + visitFromCommand?: Visitor, any, any>; + visitLimitCommand?: Visitor, any, any>; + visitExplainCommand?: Visitor, any, any>; + visitRowCommand?: Visitor, any, any>; + visitMetricsCommand?: Visitor, any, any>; + visitShowCommand?: Visitor, any, any>; + visitMetaCommand?: Visitor, any, any>; + visitEvalCommand?: Visitor, any, any>; + visitStatsCommand?: Visitor, any, any>; + visitInlineStatsCommand?: Visitor< + contexts.InlineStatsCommandVisitorContext, + any, + any + >; + visitLookupCommand?: Visitor, any, any>; + visitKeepCommand?: Visitor, any, any>; + visitSortCommand?: Visitor, any, any>; + visitWhereCommand?: Visitor, any, any>; + visitDropCommand?: Visitor, any, any>; + visitRenameCommand?: Visitor, any, any>; + visitDissectCommand?: Visitor, any, any>; + visitGrokCommand?: Visitor, any, any>; + visitEnrichCommand?: Visitor, any, any>; + visitMvExpandCommand?: Visitor, any, any>; + visitCommandOption?: Visitor, any, any>; + visitExpression?: Visitor, any, any>; + visitSourceExpression?: Visitor< + contexts.SourceExpressionVisitorContext, + any, + any + >; + visitColumnExpression?: Visitor< + contexts.ColumnExpressionVisitorContext, + any, + any + >; + visitFunctionCallExpression?: Visitor< + contexts.FunctionCallExpressionVisitorContext, + any, + any + >; + visitLiteralExpression?: Visitor< + contexts.LiteralExpressionVisitorContext, + any, + any + >; + visitListLiteralExpression?: Visitor< + contexts.ListLiteralExpressionVisitorContext, + any, + any + >; + visitTimeIntervalLiteralExpression?: Visitor< + contexts.TimeIntervalLiteralExpressionVisitorContext, + any, + any + >; + visitInlineCastExpression?: Visitor< + contexts.InlineCastExpressionVisitorContext, + any, + any + >; +} + +/** + * Maps any AST node to the corresponding visitor context. + */ +export type AstNodeToVisitorName = Node extends ESQLAstQueryNode + ? 'visitQuery' + : Node extends ast.ESQLCommand + ? 'visitCommand' + : Node extends ast.ESQLCommandOption + ? 'visitCommandOption' + : Node extends ast.ESQLSource + ? 'visitSourceExpression' + : Node extends ast.ESQLColumn + ? 'visitColumnExpression' + : Node extends ast.ESQLFunction + ? 'visitFunctionCallExpression' + : Node extends ast.ESQLLiteral + ? 'visitLiteralExpression' + : Node extends ast.ESQLList + ? 'visitListLiteralExpression' + : Node extends ast.ESQLTimeInterval + ? 'visitTimeIntervalLiteralExpression' + : Node extends ast.ESQLInlineCast + ? 'visitInlineCastExpression' + : never; + +/** + * Maps any AST node to the corresponding visitor context. + */ +export type AstNodeToVisitor< + Node extends VisitorAstNode, + Methods extends VisitorMethods = VisitorMethods +> = Methods[AstNodeToVisitorName]; + +/** + * Maps any AST node to its corresponding visitor context. + */ +export type AstNodeToContext< + Node extends VisitorAstNode, + Methods extends VisitorMethods = VisitorMethods +> = Parameters>>[0]; + +/** + * Asserts that a type is a function. + */ +export type EnsureFunction = T extends (...args: any[]) => any ? T : never; + +/** + * Converts `undefined` to `void`. This allows to make optional a function + * parameter or the return value. + */ +export type UndefinedToVoid = T extends undefined ? void : T; + +/** Returns `Y` if `T` is `any`, or `N` otherwise. */ +export type IfAny = 0 extends 1 & T ? Y : N; + +/** Converts `any` type to `void`. */ +export type AnyToVoid = IfAny; diff --git a/packages/kbn-esql-ast/src/visitor/utils.ts b/packages/kbn-esql-ast/src/visitor/utils.ts new file mode 100644 index 0000000000000..d79cc6fd5ed1a --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/utils.ts @@ -0,0 +1,36 @@ +/* + * 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 { ESQLAstItem, ESQLSingleAstItem } from '../types'; + +/** + * Normalizes AST "item" list to only contain *single* items. + * + * @param items A list of single or nested items. + */ +export function* singleItems(items: Iterable): Iterable { + for (const item of items) { + if (Array.isArray(item)) { + yield* singleItems(item); + } else { + yield item; + } + } +} + +/** + * Returns the first normalized "single item" from the "item" list. + * + * @param items Returns the first "single item" from the "item" list. + * @returns A "single item", if any. + */ +export const firstItem = (items: ESQLAstItem[]): ESQLSingleAstItem | undefined => { + for (const item of singleItems(items)) { + return item; + } +}; diff --git a/packages/kbn-esql-ast/src/visitor/visitor.ts b/packages/kbn-esql-ast/src/visitor/visitor.ts new file mode 100644 index 0000000000000..3956fe126723e --- /dev/null +++ b/packages/kbn-esql-ast/src/visitor/visitor.ts @@ -0,0 +1,98 @@ +/* + * 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 { GlobalVisitorContext, SharedData } from './global_visitor_context'; +import { QueryVisitorContext } from './contexts'; +import { VisitorContext } from './contexts'; +import type { + AstNodeToVisitorName, + EnsureFunction, + ESQLAstQueryNode, + UndefinedToVoid, + VisitorMethods, +} from './types'; + +export interface VisitorOptions< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> { + visitors?: Methods; + data?: Data; +} + +export class Visitor< + Methods extends VisitorMethods = VisitorMethods, + Data extends SharedData = SharedData +> { + public readonly ctx: GlobalVisitorContext; + + constructor(protected readonly options: VisitorOptions = {}) { + this.ctx = new GlobalVisitorContext( + options.visitors ?? ({} as Methods), + options.data ?? ({} as Data) + ); + } + + public visitors>( + visitors: NewMethods + ): Visitor { + Object.assign(this.ctx.methods, visitors); + return this as any; + } + + public on< + K extends keyof VisitorMethods, + F extends VisitorMethods[K] + >(visitor: K, fn: F): Visitor { + (this.ctx.methods as any)[visitor] = fn; + return this as any; + } + + /** + * Traverse any AST node given any visitor context. + * + * @param node AST node to traverse. + * @param ctx Traversal context. + * @returns Result of the visitor callback. + */ + public visit>( + ctx: Ctx, + input: UndefinedToVoid]>>[1]> + ): ReturnType]>> { + const node = ctx.node; + if (node instanceof Array) { + this.ctx.assertMethodExists('visitQuery'); + return this.ctx.methods.visitQuery!(ctx as any, input) as ReturnType< + NonNullable + >; + } else if (node && typeof node === 'object') { + switch (node.type) { + case 'command': + this.ctx.assertMethodExists('visitCommand'); + return this.ctx.methods.visitCommand!(ctx as any, input) as ReturnType< + NonNullable + >; + } + } + throw new Error(`Unsupported node type: ${typeof node}`); + } + + /** + * Traverse the root node of ES|QL query with default context. + * + * @param node Query node to traverse. + * @returns The result of the query visitor. + */ + public visitQuery( + node: ESQLAstQueryNode, + input: UndefinedToVoid>[1]> + ) { + const queryContext = new QueryVisitorContext(this.ctx, node, null); + return this.visit(queryContext, input); + } +}