[ES|QL] Support subqueries in pretty printer (#241473)

## Summary

support subqueries for pretty printer

Author: bartoval

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/__tests__/basic_pretty_printer.comments.test.ts

@@ -285,3 +285,12 @@ describe('commands', () => {
     });
   });
 });
+
+describe('subqueries (parens)', () => {
+  test('can print comments in complex subqueries', () => {
+    const query =
+      'FROM index1, /* before subquery */ (/* inside start */ FROM index2 /* after source */ | WHERE a > 10 /* after where */ | EVAL b = a * 2 | STATS cnt = COUNT(*) BY c | SORT cnt DESC | LIMIT 10) /* after first subquery */, index3, (FROM index4 | STATS COUNT(*)) /* after second */ | WHERE d > 10 | STATS max = MAX(*) BY e | SORT max DESC';
+
+    assertPrint(query);
+  });
+});

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

@@ -1098,3 +1098,14 @@ describe('unary operator precedence and grouping', () => {
     assertReprint('ROW (a + b) * (c + d)');
   });
 });
+
+describe('subqueries (parens)', () => {
+  test('can print 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 expected =
+      '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';
+
+    assertReprint(src, expected);
+  });
+});

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/__tests__/wrapping_pretty_printer.comments.test.ts

@@ -932,4 +932,42 @@ FROM index`;
       });
     });
   });
+
+  describe('subqueries (parens)', () => {
+    test('can print comments in complex subqueries', () => {
+      // This single test covers comment preservation in wrapped output
+      const query = [
+        'FROM index1,',
+        '/* before subquery */ (/* inside start */ FROM index2 /* after source */ |',
+        'WHERE a > 10 /* after where */ |',
+        'EVAL b = a * 2 |',
+        'STATS cnt = COUNT(*) BY c |',
+        'SORT cnt DESC |',
+        'LIMIT 10) /* after first subquery */,',
+        'index3,',
+        '(FROM index4 | STATS COUNT(*)) /* after second */ |',
+        'WHERE d > 10 |',
+        'STATS max = MAX(*) BY e |',
+        'SORT max DESC',
+      ].join(' ');
+
+      const expected = `FROM
+  index1,
+  /* before subquery */ (/* inside start */ FROM index2 /* after source */
+    | WHERE a > 10 /* after where */
+    | EVAL b = a * 2
+    | STATS cnt = COUNT(*)
+          BY c
+    | SORT cnt DESC
+    | LIMIT 10) /* after first subquery */,
+  index3,
+  (FROM index4 | STATS COUNT(*)) /* after second */
+  | WHERE d > 10
+  | STATS max = MAX(*)
+        BY e
+  | SORT max DESC`;
+
+      assertReprint(query, expected);
+    });
+  });
 });

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

@@ -1310,4 +1310,30 @@ describe('unary operator precedence and grouping', () => {
   });
 });
 
+describe('subqueries (parens)', () => {
+  test('can print 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';
+
+    assertReprint(
+      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`
+    );
+  });
+});
+
 test.todo('Idempotence on multiple times pretty printing');

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/basic_pretty_printer.ts

@@ -348,6 +348,13 @@ export class BasicPrettyPrinter {
       return this.decorateWithComments(ctx.node, formatted);
     })
 
+    .on('visitParensExpression', (ctx) => {
+      const child = ctx.visitChild();
+      const formatted = `(${child})`;
+
+      return this.decorateWithComments(ctx.node, formatted);
+    })
+
     .on('visitFunctionCallExpression', (ctx) => {
       const opts = this.opts;
       const node = ctx.node;

src/platform/packages/shared/kbn-esql-ast/src/pretty_print/wrapping_pretty_printer.ts

@@ -679,6 +679,13 @@ export class WrappingPrettyPrinter {
       return this.decorateWithComments(inp, ctx.node, formatted);
     })
 
+    .on('visitParensExpression', (ctx, inp: Input): Output => {
+      const child = ctx.visitChild(inp);
+      const formatted = `(${child.txt.trimStart()})`;
+
+      return this.decorateWithComments(inp, ctx.node, formatted);
+    })
+
     .on('visitFunctionCallExpression', (ctx, inp: Input): Output => {
       const node = ctx.node;
       let operator = ctx.operator();
@@ -919,17 +926,10 @@ export class WrappingPrettyPrinter {
       }
 
       let i = 0;
-      let prevOut: Output | undefined;
       let hasCommands = false;
 
       for (const out of ctx.visitCommands({ indent, remaining: remaining - indent.length })) {
         const isFirstCommand = i === 0;
-        const isSecondCommand = i === 1;
-
-        if (isSecondCommand) {
-          const firstCommandIsMultiline = prevOut?.lines && prevOut.lines > 1;
-          if (firstCommandIsMultiline) text += '\n' + indent;
-        }
 
         const commandIndent = isFirstCommand ? indent : pipedCommandIndent;
         const topDecorations = this.printTopDecorations(commandIndent, commands[i]);
@@ -954,7 +954,6 @@ export class WrappingPrettyPrinter {
 
         text += out.txt;
         i++;
-        prevOut = out;
         hasCommands = true;
       }
 

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

@@ -545,8 +545,12 @@ export class GlobalVisitorContext<
         return this.visitQuery(parent, expressionNode, input as any);
       }
       case 'parens': {
-        if (!this.methods.visitParensExpression) break;
-        return this.visitParensExpression(parent, expressionNode, input as any);
+        if (this.methods.visitParensExpression) {
+          const result = this.visitParensExpression(parent, expressionNode, input as any);
+          if (result) return result;
+        }
+        // Parens wraps subqueries: can return null to let comments attach to nodes inside the subquery
+        break;
       }
     }
     return this.visitExpressionGeneric(parent, expressionNode, input as any);

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

@@ -86,5 +86,9 @@ export function* children(node: ESQLProperNode): Iterable<ESQLAstExpression> {
       }
       break;
     }
+    case 'parens': {
+      yield node.child;
+      break;
+    }
   }
 }

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

@@ -94,6 +94,21 @@ export class Visitor<
         }
         return null;
       })
+      .on('visitParensExpression', (ctx): ESQLProperNode | null => {
+        const parens = ctx.node;
+        const parensLocation = parens.location;
+
+        if (!parensLocation) {
+          return null;
+        }
+
+        // Handle position before opening "("
+        if (parensLocation.min > pos) {
+          return parens;
+        }
+
+        return null;
+      })
       .on('visitCommand', visitCommand)
       .on('visitHeaderCommand', visitCommand)
       .on('visitQuery', (ctx): ESQLProperNode | null => {
@@ -184,6 +199,31 @@ export class Visitor<
 
         return null;
       })
+      .on('visitParensExpression', (ctx): ESQLProperNode | null => {
+        const parens = ctx.node;
+        const parensLocation = parens.location;
+
+        if (!parensLocation) {
+          return null;
+        }
+
+        const childQuery = ctx.child();
+
+        // Handle comments between end of subquery content and closing ")"
+        if (childQuery?.type === 'query' && childQuery.location) {
+          if (childQuery.location.max < pos && pos <= parensLocation.max) {
+            return parens;
+          }
+        }
+
+        // Handle comments immediately after closing ")"
+        if (pos > parensLocation.max) {
+          return parens;
+        }
+
+        // For comments inside parens but before/during content, continue visiting children
+        return null;
+      })
       .on('visitCommand', visitCommand)
       .on('visitHeaderCommand', visitCommand)
       .on('visitQuery', (ctx): ESQLProperNode | null => {