diff --git a/packages/kbn-esql-ast/src/walker/walker.test.ts b/packages/kbn-esql-ast/src/walker/walker.test.ts index 11c8a141080d5..ec4e0a8cd701c 100644 --- a/packages/kbn-esql-ast/src/walker/walker.test.ts +++ b/packages/kbn-esql-ast/src/walker/walker.test.ts @@ -6,7 +6,7 @@ * Side Public License, v 1. */ -import { ESQLColumn, ESQLLiteral, getAstAndSyntaxErrors } from '../..'; +import { ESQLColumn, ESQLFunction, ESQLLiteral, getAstAndSyntaxErrors } from '../..'; import { walk, Walker } from './walker'; test('can walk all functions', () => { @@ -67,3 +67,34 @@ test('can collect all params', () => { }, ]); }); + +test('can find assignment expression', () => { + const query = 'METRICS source var0 = bucket(bytes, 1 hour)'; + const { ast } = getAstAndSyntaxErrors(query); + const functions: ESQLFunction[] = []; + + Walker.walk(ast, { + visitFunction: (fn) => { + if (fn.name === '=') { + functions.push(fn); + } + }, + }); + + expect(functions.length).toBe(1); + expect(functions[0].name).toBe('='); + expect(functions[0].args.length).toBe(2); + expect((functions[0].args[0] as any).name).toBe('var0'); +}); + +describe('Walker.hasFunction()', () => { + test('can find assignment expression', () => { + const query1 = 'METRICS source bucket(bytes, 1 hour)'; + const query2 = 'METRICS source var0 = bucket(bytes, 1 hour)'; + const has1 = Walker.hasFunction(getAstAndSyntaxErrors(query1).ast!, '='); + const has2 = Walker.hasFunction(getAstAndSyntaxErrors(query2).ast!, '='); + + expect(has1).toBe(false); + expect(has2).toBe(true); + }); +}); diff --git a/packages/kbn-esql-ast/src/walker/walker.ts b/packages/kbn-esql-ast/src/walker/walker.ts index 8ebf6911599d4..70c78b8a533f5 100644 --- a/packages/kbn-esql-ast/src/walker/walker.ts +++ b/packages/kbn-esql-ast/src/walker/walker.ts @@ -54,6 +54,42 @@ export class Walker { return params; }; + /** + * Returns the first function that matches the predicate. + * + * @param node AST subtree to search in. + * @param predicate Function to test each function with. + * @returns The first function that matches the predicate. + */ + public static readonly findFunction = ( + node: ESQLAstNode | ESQLAstNode[], + predicate: (fn: ESQLFunction) => boolean + ): ESQLFunction | undefined => { + let found: ESQLFunction | undefined; + Walker.walk(node, { + visitFunction: (fn) => { + if (!found && predicate(fn)) { + found = fn; + } + }, + }); + return found; + }; + + /** + * Searches for at least one occurrence of a function or expression in the AST. + * + * @param node AST subtree to search in. + * @param name Function or expression name to search for. + * @returns True if the function or expression is found in the AST. + */ + public static readonly hasFunction = ( + node: ESQLAstNode | ESQLAstNode[], + name: string + ): boolean => { + return !!Walker.findFunction(node, (fn) => fn.name === name); + }; + constructor(protected readonly options: WalkerOptions) {} public walk(node: undefined | ESQLAstNode | ESQLAstNode[]): void {