From 4f6e4ba40fdbcc044158dfaf6a608b67e981ec98 Mon Sep 17 00:00:00 2001
From: vadimkibana <vadimkibana@gmail.com>
Date: Fri, 12 Jul 2024 13:52:09 +0200
Subject: [PATCH] add AST helper for find functions

---
 .../kbn-esql-ast/src/walker/walker.test.ts    | 33 ++++++++++++++++-
 packages/kbn-esql-ast/src/walker/walker.ts    | 36 +++++++++++++++++++
 2 files changed, 68 insertions(+), 1 deletion(-)

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 {