[ES|QL] Add subqueries support for walker and visitor (#241451)

## Summary
Part of #237401

Parens support for walker and visitor

Author: bartoval

src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/commands.test.ts

@@ -230,3 +230,39 @@ test('can visit COMPLETION command', () => {
 
   expect(list).toEqual(['COMPLETION']);
 });
+
+test('can visit FROM command with complex subqueries', () => {
+  const { ast } = EsqlQuery.fromSrc(`
+    FROM index1,
+         (FROM index2
+          | WHERE a > 10
+          | EVAL b = a * 2
+          | STATS cnt = COUNT(*) BY c
+          | SORT cnt desc
+          | LIMIT 10),
+         index3,
+         (FROM index4 | STATS count(*))
+    | WHERE d > 10
+    | STATS max = max(*) BY e
+    | SORT max desc
+  `);
+  const visitor = new Visitor()
+    .on('visitParensExpression', () => {
+      return 'SUBQUERY';
+    })
+    .on('visitExpression', () => {
+      return null;
+    })
+    .on('visitFromCommand', (ctx) => {
+      return [...ctx.visitArguments()];
+    })
+    .on('visitCommand', (ctx) => {
+      return null;
+    })
+    .on('visitQuery', (ctx) => {
+      return [...ctx.visitCommands()].flat();
+    });
+  const list = visitor.visitQuery(ast).flat().filter(Boolean);
+
+  expect(list).toEqual(['SUBQUERY', 'SUBQUERY']);
+});

src/platform/packages/shared/kbn-esql-ast/src/visitor/__tests__/expressions.test.ts

@@ -263,3 +263,40 @@ test('"visitExpression" does visit identifier nodes', () => {
 
   expect(expressions.sort()).toEqual(['a', 'index']);
 });
+
+test('"visitParensExpression" can traverse complex subqueries with processing', () => {
+  const { ast } = parse(`
+    FROM index1,
+         (FROM index2
+          | WHERE a > 10
+          | EVAL b = a * 2
+          | STATS cnt = COUNT(*) BY c
+          | SORT cnt desc
+          | LIMIT 10),
+         index3,
+         (FROM index4 | STATS count(*))
+    | WHERE d > 10
+    | STATS max = max(*) BY e
+    | SORT max desc
+  `);
+  const visitor = new Visitor()
+    .on('visitParensExpression', (ctx) => {
+      const child = ctx.visitChild(undefined);
+      return `PARENS(${child})`;
+    })
+    .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, PARENS(FROM E | WHERE E | EVAL E | STATS E | SORT E | LIMIT E), E, PARENS(FROM E | STATS E) | WHERE E | STATS E | SORT E'
+  );
+});

src/platform/packages/shared/kbn-esql-ast/src/visitor/contexts.ts

@@ -35,6 +35,7 @@ import type {
   ESQLMap,
   ESQLMapEntry,
   ESQLOrderExpression,
+  ESQLParens,
   ESQLSource,
   ESQLStringLiteral,
 } from '../types';
@@ -725,3 +726,20 @@ export class MapEntryExpressionVisitorContext<
     return this.visitExpression(this.value(), input as any);
   }
 }
+
+export class ParensExpressionVisitorContext<
+  Methods extends VisitorMethods = VisitorMethods,
+  Data extends SharedData = SharedData
+> extends VisitorContext<Methods, Data, ESQLParens> {
+  public child(): ESQLAstExpression {
+    return this.node.child;
+  }
+
+  public visitChild(
+    input: VisitorInput<Methods, 'visitExpression'>
+  ): VisitorOutput<Methods, 'visitExpression'> {
+    this.ctx.assertMethodExists('visitExpression');
+
+    return this.visitExpression(this.child(), input as any);
+  }
+}

src/platform/packages/shared/kbn-esql-ast/src/visitor/global_visitor_context.ts

@@ -25,6 +25,7 @@ import type {
   ESQLMap,
   ESQLMapEntry,
   ESQLOrderExpression,
+  ESQLParens,
   ESQLSource,
 } from '../types';
 import type * as types from './types';
@@ -543,6 +544,10 @@ export class GlobalVisitorContext<
         if (!this.methods.visitQuery || expressionNode.type !== 'query') break;
         return this.visitQuery(parent, expressionNode, input as any);
       }
+      case 'parens': {
+        if (!this.methods.visitParensExpression) break;
+        return this.visitParensExpression(parent, expressionNode, input as any);
+      }
     }
     return this.visitExpressionGeneric(parent, expressionNode, input as any);
   }
@@ -645,6 +650,15 @@ export class GlobalVisitorContext<
     const context = new contexts.MapEntryExpressionVisitorContext(this, node, parent);
     return this.visitWithSpecificContext('visitMapEntryExpression', context, input);
   }
+
+  public visitParensExpression(
+    parent: contexts.VisitorContext | null,
+    node: ESQLParens,
+    input: types.VisitorInput<Methods, 'visitParensExpression'>
+  ): types.VisitorOutput<Methods, 'visitParensExpression'> {
+    const context = new contexts.ParensExpressionVisitorContext(this, node, parent);
+    return this.visitWithSpecificContext('visitParensExpression', context, input);
+  }
 }
 
 // #endregion

src/platform/packages/shared/kbn-esql-ast/src/visitor/types.ts

@@ -63,7 +63,8 @@ export type ExpressionVisitorInput<Methods extends VisitorMethods> = AnyToVoid<
       VisitorInput<Methods, 'visitOrderExpression'> &
       VisitorInput<Methods, 'visitIdentifierExpression'> &
       VisitorInput<Methods, 'visitMapExpression'> &
-      VisitorInput<Methods, 'visitMapEntryExpression'>
+      VisitorInput<Methods, 'visitMapEntryExpression'> &
+      VisitorInput<Methods, 'visitParensExpression'>
 >;
 
 /**
@@ -81,7 +82,8 @@ export type ExpressionVisitorOutput<Methods extends VisitorMethods> =
   | VisitorOutput<Methods, 'visitOrderExpression'>
   | VisitorOutput<Methods, 'visitIdentifierExpression'>
   | VisitorOutput<Methods, 'visitMapExpression'>
-  | VisitorOutput<Methods, 'visitMapEntryExpression'>;
+  | VisitorOutput<Methods, 'visitMapEntryExpression'>
+  | VisitorOutput<Methods, 'visitParensExpression'>;
 
 /**
  * Input that satisfies any command visitor input constraints.
@@ -237,6 +239,11 @@ export interface VisitorMethods<
     any,
     any
   >;
+  visitParensExpression?: Visitor<
+    contexts.ParensExpressionVisitorContext<Visitors, Data>,
+    any,
+    any
+  >;
 }
 
 /**
@@ -268,6 +275,8 @@ export type AstNodeToVisitorName<Node extends VisitorAstNode> = Node extends ESQ
   ? 'visitMapExpression'
   : Node extends ast.ESQLMapEntry
   ? 'visitMapEntryExpression'
+  : Node extends ast.ESQLParens
+  ? 'visitParensExpression'
   : never;
 
 /**

src/platform/packages/shared/kbn-esql-ast/src/walker/__tests__/walker.test.ts

@@ -1246,4 +1246,38 @@ describe('header commands', () => {
       expect(regularCommands.map((cmd) => cmd.name)).toStrictEqual(['from', 'limit']);
     });
   });
+
+  describe('parens (subquery)', () => {
+    test('can visit complex subqueries with processing', () => {
+      const src = `
+        FROM index1,
+             (FROM index2
+              | WHERE a > 10
+              | EVAL b = a * 2
+              | STATS cnt = COUNT(*) BY c
+              | SORT cnt desc
+              | LIMIT 10),
+             index3,
+             (FROM index4 | STATS count(*))
+        | WHERE d > 10
+        | STATS max = max(*) BY e
+        | SORT max desc
+      `;
+      const { ast } = parse(src);
+      let parensCount = 0;
+      const sources: string[] = [];
+
+      walk(ast, {
+        visitParens: (node) => {
+          parensCount++;
+        },
+        visitSource: (node) => {
+          sources.push(node.name);
+        },
+      });
+
+      expect(parensCount).toBe(2);
+      expect(sources).toEqual(['index1', 'index2', 'index3', 'index4']);
+    });
+  });
 });

src/platform/packages/shared/kbn-esql-ast/src/walker/walker.ts

@@ -90,6 +90,11 @@ export interface WalkerOptions {
     parent: types.ESQLProperNode | undefined,
     walker: WalkerVisitorApi
   ) => void;
+  visitParens?: (
+    node: types.ESQLParens,
+    parent: types.ESQLProperNode | undefined,
+    walker: WalkerVisitorApi
+  ) => void;
 
   /**
    * Called on every expression node.
@@ -703,6 +708,15 @@ export class Walker {
     }
   }
 
+  public walkParens(node: types.ESQLParens, parent: types.ESQLProperNode | undefined): void {
+    const { options } = this;
+    (options.visitParens ?? options.visitAny)?.(node, parent, this);
+
+    if (node.child) {
+      this.walkSingleAstItem(node.child, node);
+    }
+  }
+
   public walkQuery(
     node: types.ESQLAstQueryExpression,
     parent: types.ESQLProperNode | undefined
@@ -770,6 +784,10 @@ export class Walker {
         this.walkInlineCast(node, parent);
         break;
       }
+      case 'parens': {
+        this.walkParens(node as types.ESQLParens, parent);
+        break;
+      }
       case 'identifier': {
         (options.visitIdentifier ?? options.visitAny)?.(node, parent, this);
         break;