diff --git a/packages/kbn-esql-ast/src/builder/builder.ts b/packages/kbn-esql-ast/src/builder/builder.ts index fae1981b454c2..07b9f14875abb 100644 --- a/packages/kbn-esql-ast/src/builder/builder.ts +++ b/packages/kbn-esql-ast/src/builder/builder.ts @@ -29,6 +29,9 @@ import { ESQLPositionalParamLiteral, ESQLOrderExpression, ESQLSource, + ESQLParamLiteral, + ESQLFunction, + ESQLAstItem, } from '../types'; import { AstNodeParserFields, AstNodeTemplate, PartialFields } from './types'; @@ -171,6 +174,53 @@ export namespace Builder { }; }; + export namespace func { + export const node = ( + template: AstNodeTemplate, + fromParser?: Partial + ): ESQLFunction => { + return { + ...template, + ...Builder.parserFields(fromParser), + type: 'function', + }; + }; + + export const call = ( + nameOrOperator: string | ESQLIdentifier | ESQLParamLiteral, + args: ESQLAstItem[], + template?: Omit, 'subtype' | 'name' | 'operator' | 'args'>, + fromParser?: Partial + ): ESQLFunction => { + let name: string; + let operator: ESQLIdentifier | ESQLParamLiteral; + if (typeof nameOrOperator === 'string') { + name = nameOrOperator; + operator = Builder.identifier({ name }); + } else { + operator = nameOrOperator; + name = LeafPrinter.print(operator); + } + return Builder.expression.func.node( + { ...template, name, operator, args, subtype: 'variadic-call' }, + fromParser + ); + }; + + export const binary = ( + name: string, + args: [left: ESQLAstItem, right: ESQLAstItem], + template?: Omit, 'subtype' | 'name' | 'operator' | 'args'>, + fromParser?: Partial + ): ESQLFunction => { + const operator = Builder.identifier({ name }); + return Builder.expression.func.node( + { ...template, name, operator, args, subtype: 'binary-expression' }, + fromParser + ); + }; + } + export namespace literal { /** * Constructs an integer literal node. @@ -189,6 +239,21 @@ export namespace Builder { return node; }; + export const integer = ( + value: number, + template?: Omit, 'name'>, + fromParser?: Partial + ): ESQLIntegerLiteral | ESQLDecimalLiteral => { + return Builder.expression.literal.numeric( + { + ...template, + value, + literalType: 'integer', + }, + fromParser + ); + }; + export const list = ( template: Omit, 'name'>, fromParser?: Partial @@ -262,5 +327,25 @@ export namespace Builder { return node; }; + + export const build = ( + name: string, + options: Partial = {}, + fromParser?: Partial + ): ESQLParam => { + const value: string = name.startsWith('?') ? name.slice(1) : name; + + if (!value) { + return Builder.param.unnamed(options); + } + + const isNumeric = !isNaN(Number(value)) && String(Number(value)) === value; + + if (isNumeric) { + return Builder.param.positional({ ...options, value: Number(value) }, fromParser); + } else { + return Builder.param.named({ ...options, value }, fromParser); + } + }; } } diff --git a/packages/kbn-esql-ast/src/mutate/commands/index.ts b/packages/kbn-esql-ast/src/mutate/commands/index.ts index 8068aa5d3ef94..3de5669a43eaf 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/index.ts @@ -11,5 +11,6 @@ import * as from from './from'; import * as limit from './limit'; import * as sort from './sort'; import * as stats from './stats'; +import * as where from './where'; -export { from, limit, sort, stats }; +export { from, limit, sort, stats, where }; diff --git a/packages/kbn-esql-ast/src/mutate/commands/where/index.test.ts b/packages/kbn-esql-ast/src/mutate/commands/where/index.test.ts new file mode 100644 index 0000000000000..c2ac4407b2baa --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/where/index.test.ts @@ -0,0 +1,333 @@ +/* + * 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 * as commands from '..'; +import { EsqlQuery } from '../../../query'; +import { Builder } from '../../../builder'; + +describe('commands.where', () => { + describe('.list()', () => { + it('lists all "WHERE" commands', () => { + const src = 'FROM index | LIMIT 1 | WHERE a == 1 | LIMIT 2 | WHERE b == 2'; + const query = EsqlQuery.fromSrc(src); + + const nodes = [...commands.where.list(query.ast)]; + + expect(nodes).toMatchObject([ + { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + }, + ], + }, + { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + }, + ], + }, + ]); + }); + }); + + describe('.byIndex()', () => { + it('retrieves the specific "WHERE" command by index', () => { + const src = 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 2 == b'; + const query = EsqlQuery.fromSrc(src); + + const node1 = commands.where.byIndex(query.ast, 1); + const node2 = commands.where.byIndex(query.ast, 0); + + expect(node1).toMatchObject({ + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 2, + }, + {}, + ], + }, + ], + }); + expect(node2).toMatchObject({ + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 1, + }, + {}, + ], + }, + ], + }); + }); + }); + + describe('.byField()', () => { + it('retrieves the specific "WHERE" command by field', () => { + const src = 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 2 == b'; + const query = EsqlQuery.fromSrc(src); + + const [command1] = commands.where.byField(query.ast, 'b')!; + const [command2] = commands.where.byField(query.ast, 'a')!; + + expect(command1).toMatchObject({ + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 2, + }, + {}, + ], + }, + ], + }); + expect(command2).toMatchObject({ + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 1, + }, + {}, + ], + }, + ], + }); + }); + + it('can find command by nested field', () => { + const src = 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 2 == a.b.c'; + const query = EsqlQuery.fromSrc(src); + + const [command] = commands.where.byField(query.ast, ['a', 'b', 'c'])!; + + expect(command).toMatchObject({ + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 2, + }, + {}, + ], + }, + ], + }); + }); + + it('can find command by param', () => { + const src = 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE ?param == 123'; + const query = EsqlQuery.fromSrc(src); + + const [command1] = commands.where.byField(query.ast, ['?param'])!; + const [command2] = commands.where.byField(query.ast, '?param')!; + + const expected = { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + {}, + { + type: 'literal', + value: 123, + }, + ], + }, + ], + }; + + expect(command1).toMatchObject(expected); + expect(command2).toMatchObject(expected); + }); + + it('can find command by nested param', () => { + const src = 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE a.b.?param == 123'; + const query = EsqlQuery.fromSrc(src); + + const [command] = commands.where.byField(query.ast, ['a', 'b', '?param'])!; + + const expected = { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + {}, + { + type: 'literal', + value: 123, + }, + ], + }, + ], + }; + + expect(command).toMatchObject(expected); + }); + + it('can find command when field is used in function', () => { + const src = 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 123 == fn(a.b.c)'; + const query = EsqlQuery.fromSrc(src); + + const [command] = commands.where.byField(query.ast, ['a', 'b', 'c'])!; + + const expected = { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 123, + }, + {}, + ], + }, + ], + }; + + expect(command).toMatchObject(expected); + }); + + it('can find command when various decorations are applied to the field', () => { + const src = + 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 123 == add(1 + fn(NOT -(a.b.c::ip)::INTEGER /* comment */))'; + const query = EsqlQuery.fromSrc(src); + + const [command1] = commands.where.byField(query.ast, ['a', 'b', 'c'])!; + const command2 = commands.where.byField(query.ast, 'a.b.c'); + + const expected = { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 123, + }, + {}, + ], + }, + ], + }; + + expect(command1).toMatchObject(expected); + expect(command2).toBe(undefined); + }); + + it('can construct field template using Builder', () => { + const src = + 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 123 == add(1 + fn(NOT -(a.b.c::ip)::INTEGER /* comment */))'; + const query = EsqlQuery.fromSrc(src); + + const [command] = commands.where.byField( + query.ast, + Builder.expression.column({ + args: [ + Builder.identifier({ name: 'a' }), + Builder.identifier({ name: 'b' }), + Builder.identifier({ name: 'c' }), + ], + }) + )!; + + const expected = { + type: 'command', + name: 'where', + args: [ + { + type: 'function', + name: '==', + args: [ + { + type: 'literal', + value: 123, + }, + {}, + ], + }, + ], + }; + + expect(command).toMatchObject(expected); + }); + + it('returns the found column node', () => { + const src = + 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 123 == add(1 + fn(NOT -(a.b.c::ip)::INTEGER /* comment */))'; + const query = EsqlQuery.fromSrc(src); + + const [_, column] = commands.where.byField( + query.ast, + Builder.expression.column({ + args: [ + Builder.identifier({ name: 'a' }), + Builder.identifier({ name: 'b' }), + Builder.identifier({ name: 'c' }), + ], + }) + )!; + + expect(column).toMatchObject({ + type: 'column', + name: 'a.b.c', + }); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/commands/where/index.ts b/packages/kbn-esql-ast/src/mutate/commands/where/index.ts new file mode 100644 index 0000000000000..32b61a733dd4f --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/where/index.ts @@ -0,0 +1,131 @@ +/* + * 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 { Walker } from '../../../walker'; +import { LeafPrinter } from '../../../pretty_print'; +import { Builder } from '../../../builder'; +import type { + ESQLAstQueryExpression, + ESQLColumn, + ESQLCommand, + ESQLIdentifier, + ESQLParamLiteral, + ESQLProperNode, +} from '../../../types'; +import * as generic from '../../generic'; + +/** + * Lists all "WHERE" commands in the query AST. + * + * @param ast The root AST node to search for "WHERE" commands. + * @returns A collection of "WHERE" commands. + */ +export const list = (ast: ESQLAstQueryExpression): IterableIterator => { + return generic.commands.list(ast, (cmd) => cmd.name === 'where'); +}; + +/** + * Retrieves the "WHERE" command at the specified index in order of appearance. + * + * @param ast The root AST node to search for "WHERE" commands. + * @param index The index of the "WHERE" command to retrieve. + * @returns The "WHERE" command at the specified index, if any. + */ +export const byIndex = (ast: ESQLAstQueryExpression, index: number): ESQLCommand | undefined => { + return [...list(ast)][index]; +}; + +export type ESQLAstField = ESQLColumn | ESQLIdentifier | ESQLParamLiteral; +export type ESQLAstFieldTemplate = string | string[] | ESQLAstField; + +const fieldTemplateToField = (template: ESQLAstFieldTemplate): ESQLAstField => { + if (typeof template === 'string') { + const part = template.startsWith('?') + ? Builder.param.build(template) + : Builder.identifier({ name: template }); + const column = Builder.expression.column({ args: [part] }); + return column; + } else if (Array.isArray(template)) { + const identifiers = template.map((name) => { + if (name.startsWith('?')) { + return Builder.param.build(name); + } else { + return Builder.identifier({ name }); + } + }); + const column = Builder.expression.column({ args: identifiers }); + return column; + } + + return template; +}; + +const matchNodeAgainstField = (node: ESQLProperNode, field: ESQLAstField): boolean => { + return LeafPrinter.print(node) === LeafPrinter.print(field); +}; + +/** + * Finds the first "WHERE" command which contains the specified text as one of + * its comparison operands. The text can represent a field (including nested + * fields or a single identifier), or a param. If the text starts with "?", it + * is assumed to be a param. + * + * Examples: + * + * ```ts + * byField(ast, 'field'); + * byField(ast, ['nested', 'field']); + * byField(ast, '?param'); + * byField(ast, ['nested', '?param']); + * byField(ast, ['nested', 'positional', 'param', '?123']); + * byField(ast, '?'); + * ``` + * + * Alternatively you can build your own field template using the builder: + * + * ```ts + * byField(ast, Builder.expression.column({ + * args: [Builder.identifier({ name: 'field' })] + * })); + * ``` + * + * @param ast The root AST node search for "WHERE" commands. + * @param text The text or nested column name texts to search for. + */ +export const byField = ( + ast: ESQLAstQueryExpression, + template: ESQLAstFieldTemplate +): [command: ESQLCommand, field: ESQLProperNode] | undefined => { + const field = fieldTemplateToField(template); + + for (const command of list(ast)) { + let found: ESQLProperNode | undefined; + + const matchNode = (node: ESQLProperNode) => { + if (found) { + return; + } + if (matchNodeAgainstField(node, field)) { + found = node; + } + }; + + Walker.walk(command, { + visitColumn: matchNode, + visitIdentifier: matchNode, + visitLiteral: matchNode, + }); + + if (found) { + return [command, found]; + } + } + + return undefined; +}; diff --git a/packages/kbn-esql-ast/src/mutate/commands/where/index_scenarios.test.ts b/packages/kbn-esql-ast/src/mutate/commands/where/index_scenarios.test.ts new file mode 100644 index 0000000000000..eb0c54a933807 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/where/index_scenarios.test.ts @@ -0,0 +1,104 @@ +/* + * 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 { BasicPrettyPrinter } from '../../../pretty_print'; +import * as mutate from '../..'; +import { EsqlQuery } from '../../../query'; +import { Builder } from '../../../builder'; +import { ESQLFunction } from '../../../types'; + +describe('scenarios', () => { + it('can remove the found WHERE command', () => { + const src = + 'FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2 | WHERE 123 == add(1 + fn(NOT -(a.b.c::ip)::INTEGER /* comment */))'; + const query = EsqlQuery.fromSrc(src); + + const [command1] = mutate.commands.where.byField(query.ast, ['a', 'b', 'c'])!; + mutate.generic.commands.remove(query.ast, command1); + + const text1 = BasicPrettyPrinter.print(query.ast); + + expect(text1).toBe('FROM index | LIMIT 1 | WHERE 1 == a | LIMIT 2'); + + const [command2] = mutate.commands.where.byField(query.ast, 'a')!; + mutate.generic.commands.remove(query.ast, command2); + + const text2 = BasicPrettyPrinter.print(query.ast); + + expect(text2).toBe('FROM index | LIMIT 1 | LIMIT 2'); + }); + + it('can insert a new WHERE command', () => { + const src = 'FROM index | LIMIT 1'; + const query = EsqlQuery.fromSrc(src); + const command = Builder.command({ + name: 'where', + args: [ + Builder.expression.func.binary('==', [ + Builder.expression.column({ + args: [Builder.identifier({ name: 'a' })], + }), + Builder.expression.literal.numeric({ + value: 1, + literalType: 'integer', + }), + ]), + ], + }); + + mutate.generic.commands.insert(query.ast, command, 1); + + const text = BasicPrettyPrinter.print(query.ast); + + expect(text).toBe('FROM index | WHERE a == 1 | LIMIT 1'); + }); + + it('can insert a new WHERE command with function call condition and param in column name', () => { + const src = 'FROM index | LIMIT 1'; + const query = EsqlQuery.fromSrc(src); + const command = Builder.command({ + name: 'where', + args: [ + Builder.expression.func.binary('==', [ + Builder.expression.func.call('add', [ + Builder.expression.literal.integer(1), + Builder.expression.literal.integer(2), + Builder.expression.literal.integer(3), + ]), + Builder.expression.column({ + args: [ + Builder.identifier({ name: 'a' }), + Builder.identifier({ name: 'b' }), + Builder.param.build('?param'), + ], + }), + ]), + ], + }); + + mutate.generic.commands.insert(query.ast, command, 1); + + const text = BasicPrettyPrinter.print(query.ast); + + expect(text).toBe('FROM index | WHERE ADD(1, 2, 3) == a.b.?param | LIMIT 1'); + }); + + it('can update WHERE command condition', () => { + const src = 'FROM index | WHERE a /* important field */ == 1 | LIMIT 1'; + const query = EsqlQuery.fromSrc(src, { withFormatting: true }); + const [command] = mutate.commands.where.byField(query.ast, ['a'])!; + const fn = command.args[0] as ESQLFunction; + + fn.args[1] = Builder.expression.literal.integer(2); + + const text = BasicPrettyPrinter.print(query.ast); + + expect(text).toBe('FROM index | WHERE a /* important field */ == 2 | LIMIT 1'); + }); +}); diff --git a/packages/kbn-esql-ast/src/pretty_print/leaf_printer.ts b/packages/kbn-esql-ast/src/pretty_print/leaf_printer.ts index b413234cbe263..bb768c1e6a738 100644 --- a/packages/kbn-esql-ast/src/pretty_print/leaf_printer.ts +++ b/packages/kbn-esql-ast/src/pretty_print/leaf_printer.ts @@ -11,8 +11,10 @@ import { ESQLAstComment, ESQLAstCommentMultiLine, ESQLColumn, + ESQLIdentifier, ESQLLiteral, ESQLParamLiteral, + ESQLProperNode, ESQLSource, ESQLTimeInterval, } from '../types'; @@ -27,6 +29,18 @@ const regexUnquotedIdPattern = /^([a-z\*_\@]{1})[a-z0-9_\*]*$/i; export const LeafPrinter = { source: (node: ESQLSource) => node.name, + identifier: (node: ESQLIdentifier) => { + const name = node.name; + + if (regexUnquotedIdPattern.test(name)) { + return name; + } else { + // Escape backticks "`" with double backticks "``". + const escaped = name.replace(/`/g, '``'); + return '`' + escaped + '`'; + } + }, + column: (node: ESQLColumn) => { const args = node.args; @@ -35,18 +49,11 @@ export const LeafPrinter = { for (const arg of args) { switch (arg.type) { case 'identifier': { - const name = arg.name; - if (formatted.length > 0) { formatted += '.'; } - if (regexUnquotedIdPattern.test(name)) { - formatted += name; - } else { - // Escape backticks "`" with double backticks "``". - const escaped = name.replace(/`/g, '``'); - formatted += '`' + escaped + '`'; - } + + formatted += LeafPrinter.identifier(arg); break; } @@ -136,4 +143,22 @@ export const LeafPrinter = { } return text; }, + + print: (node: ESQLProperNode): string => { + switch (node.type) { + case 'identifier': { + return LeafPrinter.identifier(node); + } + case 'column': { + return LeafPrinter.column(node); + } + case 'literal': { + return LeafPrinter.literal(node); + } + case 'timeInterval': { + return LeafPrinter.timeInterval(node); + } + } + return ''; + }, }; diff --git a/packages/kbn-esql-ast/src/walker/walker.test.ts b/packages/kbn-esql-ast/src/walker/walker.test.ts index 49c50a0f7fa5d..048e7e259eec2 100644 --- a/packages/kbn-esql-ast/src/walker/walker.test.ts +++ b/packages/kbn-esql-ast/src/walker/walker.test.ts @@ -709,6 +709,25 @@ describe('structurally can walk all nodes', () => { ]); }); + test('can visit a column inside a deeply nested inline cast', () => { + const query = + 'FROM index | WHERE 123 == add(1 + fn(NOT -(a.b.c)::INTEGER /* comment */))'; + const { root } = parse(query); + + const columns: ESQLColumn[] = []; + + walk(root, { + visitColumn: (node) => columns.push(node), + }); + + expect(columns).toMatchObject([ + { + type: 'column', + name: 'a.b.c', + }, + ]); + }); + test('"visitAny" can capture cast expression', () => { const query = 'FROM index | STATS a = 123::integer'; const { ast } = parse(query); @@ -1016,6 +1035,21 @@ describe('Walker.match()', () => { ], }); }); + + test('can find a deeply nested column', () => { + const query = + 'FROM index | WHERE 123 == add(1 + fn(NOT 10 + -(a.b.c::ip)::INTEGER /* comment */))'; + const { root } = parse(query); + const res = Walker.match(root, { + type: 'column', + name: 'a.b.c', + }); + + expect(res).toMatchObject({ + type: 'column', + name: 'a.b.c', + }); + }); }); describe('Walker.matchAll()', () => { diff --git a/packages/kbn-esql-ast/src/walker/walker.ts b/packages/kbn-esql-ast/src/walker/walker.ts index f3b6de91649b7..0e6811c02efef 100644 --- a/packages/kbn-esql-ast/src/walker/walker.ts +++ b/packages/kbn-esql-ast/src/walker/walker.ts @@ -361,6 +361,12 @@ export class Walker { } } + public walkInlineCast(node: ESQLInlineCast): void { + const { options } = this; + (options.visitInlineCast ?? options.visitAny)?.(node); + this.walkAstItem(node.value); + } + public walkFunction(node: ESQLFunction): void { const { options } = this; (options.visitFunction ?? options.visitAny)?.(node); @@ -427,7 +433,7 @@ export class Walker { break; } case 'inlineCast': { - (options.visitInlineCast ?? options.visitAny)?.(node); + this.walkInlineCast(node); break; } case 'identifier': {