[ES|QL] JOIN command Traversal API and prety-printing support (elas…

…tic#202750)

## Summary

Partially addresses elastic#200858

- Add support for the new `JOIN` command and `AS` expression in
Traversal API: `Walker` and `Visitor`
- Adds support for the new `JOIN`command and `AS` expression in the
pretty-printer.
- Fixes some parser bugs related to the `JOIN` command.


### Checklist

- [x] [Unit or functional
tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html)
were updated or added to match the most common scenarios

packages/kbn-esql-ast/src/parser/__tests__/comments.test.ts

@@ -8,6 +8,7 @@
  */
 
 import { parse } from '..';
+import { Walker } from '../../walker';
 
 describe('Comments', () => {
   describe('can attach "top" comment(s)', () => {
@@ -442,6 +443,35 @@ FROM index`;
         ],
       });
     });
+
+    it('to an identifier', () => {
+      const text = `FROM index | LEFT JOIN
+        // comment
+        abc
+        ON a = b`;
+      const { root } = parse(text, { withFormatting: true });
+
+      expect(root.commands[1]).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'abc',
+            formatting: {
+              top: [
+                {
+                  type: 'comment',
+                  subtype: 'single-line',
+                  text: ' comment',
+                },
+              ],
+            },
+          },
+          {},
+        ],
+      });
+    });
   });
 
   describe('can attach "left" comment(s)', () => {
@@ -549,6 +579,34 @@ FROM index`;
         },
       ]);
     });
+
+    it('to an identifier', () => {
+      const text = `FROM index | LEFT JOIN
+        /* left */ abc
+        ON a = b`;
+      const { root } = parse(text, { withFormatting: true });
+
+      expect(root.commands[1]).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'abc',
+            formatting: {
+              left: [
+                {
+                  type: 'comment',
+                  subtype: 'multi-line',
+                  text: ' left ',
+                },
+              ],
+            },
+          },
+          {},
+        ],
+      });
+    });
   });
 
   describe('can attach "right" comment(s)', () => {
@@ -776,6 +834,61 @@ FROM index`;
         ],
       });
     });
+
+    it('to an identifier', () => {
+      const text = `FROM index | LEFT JOIN
+        abc /* right */ // right 2
+        ON a = b`;
+      const { root } = parse(text, { withFormatting: true });
+
+      expect(root.commands[1]).toMatchObject({
+        type: 'command',
+        name: 'join',
+        args: [
+          {
+            type: 'identifier',
+            name: 'abc',
+            formatting: {
+              right: [
+                {
+                  type: 'comment',
+                  subtype: 'multi-line',
+                  text: ' right ',
+                },
+              ],
+              rightSingleLine: {
+                type: 'comment',
+                subtype: 'single-line',
+                text: ' right 2',
+              },
+            },
+          },
+          {},
+        ],
+      });
+    });
+
+    it('to a column inside ON option', () => {
+      const text = `FROM index | LEFT JOIN
+        abc
+        ON a /* right */ = b`;
+      const { root } = parse(text, { withFormatting: true });
+      const a = Walker.match(root, { type: 'column', name: 'a' });
+
+      expect(a).toMatchObject({
+        type: 'column',
+        name: 'a',
+        formatting: {
+          right: [
+            {
+              type: 'comment',
+              subtype: 'multi-line',
+              text: ' right ',
+            },
+          ],
+        },
+      });
+    });
   });
 
   describe('can attach "right end" comments', () => {

packages/kbn-esql-ast/src/parser/__tests__/join.test.ts

@@ -184,19 +184,41 @@ describe('<TYPE> JOIN command', () => {
       const node2 = Walker.match(query.ast, { type: 'identifier', name: 'alias' });
       const node3 = Walker.match(query.ast, { type: 'column', name: 'on_1' });
       const node4 = Walker.match(query.ast, { type: 'column', name: 'on_2' });
+      const node5 = Walker.match(query.ast, { type: 'function', name: 'as' });
 
       expect(query.src.slice(node1?.location.min, node1?.location.max! + 1)).toBe('index');
       expect(query.src.slice(node2?.location.min, node2?.location.max! + 1)).toBe('alias');
       expect(query.src.slice(node3?.location.min, node3?.location.max! + 1)).toBe('on_1');
       expect(query.src.slice(node4?.location.min, node4?.location.max! + 1)).toBe('on_2');
+      expect(query.src.slice(node5?.location.min, node5?.location.max! + 1)).toBe('index AS alias');
+    });
+
+    it('correctly extracts JOIN command position', () => {
+      const text = `FROM employees | LOOKUP JOIN index AS alias ON on_1, on_2 | LIMIT 1`;
+      const query = EsqlQuery.fromSrc(text);
+      const join = Walker.match(query.ast, { type: 'command', name: 'join' });
+
+      expect(query.src.slice(join?.location.min, join?.location.max! + 1)).toBe(
+        'LOOKUP JOIN index AS alias ON on_1, on_2'
+      );
+    });
+
+    it('correctly extracts ON option position', () => {
+      const text = `FROM employees | LOOKUP JOIN index AS alias ON on_1, on_2 | LIMIT 1`;
+      const query = EsqlQuery.fromSrc(text);
+      const on = Walker.match(query.ast, { type: 'option', name: 'on' });
+
+      expect(query.src.slice(on?.location.min, on?.location.max! + 1)).toBe('ON on_1, on_2');
     });
   });
 
   describe('incorrectly formatted', () => {
-    const text = `FROM employees | LOOKUP JOIN index AAS alias ON on_1, on_2 | LIMIT 1`;
-    const query = EsqlQuery.fromSrc(text);
+    it('throws error on invalid "AS" keyword', () => {
+      const text = `FROM employees | LOOKUP JOIN index AAS alias ON on_1, on_2 | LIMIT 1`;
+      const query = EsqlQuery.fromSrc(text);
 
-    expect(query.errors.length > 0).toBe(true);
-    expect(query.errors[0].message.includes('AAS')).toBe(true);
+      expect(query.errors.length > 0).toBe(true);
+      expect(query.errors[0].message.includes('AAS')).toBe(true);
+    });
   });
 });

packages/kbn-esql-ast/src/parser/factories.ts

@@ -55,6 +55,8 @@ import type {
   InlineCastingType,
   ESQLFunctionCallExpression,
   ESQLIdentifier,
+  ESQLBinaryExpression,
+  BinaryExpressionOperator,
 } from '../types';
 import { parseIdentifier, getPosition } from './helpers';
 import { Builder, type AstNodeParserFields } from '../builder';
@@ -240,6 +242,25 @@ export const createFunctionCall = (ctx: FunctionContext): ESQLFunctionCallExpres
   return node;
 };
 
+export const createBinaryExpression = (
+  operator: BinaryExpressionOperator,
+  ctx: ParserRuleContext,
+  args: ESQLBinaryExpression['args']
+): ESQLBinaryExpression => {
+  const node = Builder.expression.func.binary(
+    operator,
+    args,
+    {},
+    {
+      text: ctx.getText(),
+      location: getPosition(ctx.start, ctx.stop),
+      incomplete: Boolean(ctx.exception),
+    }
+  ) as ESQLBinaryExpression;
+
+  return node;
+};
+
 export const createIdentifierOrParam = (ctx: IdentifierOrParameterContext) => {
   const identifier = ctx.identifier();
   if (identifier) {

packages/kbn-esql-ast/src/parser/factories/join.ts

@@ -8,9 +8,13 @@
  */
 
 import { JoinCommandContext, JoinTargetContext } from '../../antlr/esql_parser';
-import { Builder } from '../../builder';
 import { ESQLAstItem, ESQLBinaryExpression, ESQLCommand, ESQLIdentifier } from '../../types';
-import { createCommand, createIdentifier } from '../factories';
+import {
+  createBinaryExpression,
+  createCommand,
+  createIdentifier,
+  createOption,
+} from '../factories';
 import { visitValueExpression } from '../walkers';
 
 const createNodeFromJoinTarget = (
@@ -24,7 +28,7 @@ const createNodeFromJoinTarget = (
   }
 
   const alias = createIdentifier(aliasCtx);
-  const renameExpression = Builder.expression.func.binary('as', [
+  const renameExpression = createBinaryExpression('as', ctx, [
     index,
     alias,
   ]) as ESQLBinaryExpression;
@@ -39,10 +43,11 @@ export const createJoinCommand = (ctx: JoinCommandContext): ESQLCommand => {
   command.commandType = (ctx._type_.text ?? '').toLocaleLowerCase();
 
   const joinTarget = createNodeFromJoinTarget(ctx.joinTarget());
-  const onOption = Builder.option({ name: 'on' });
+  const joinCondition = ctx.joinCondition();
+  const onOption = createOption('on', joinCondition);
   const joinPredicates: ESQLAstItem[] = onOption.args;
 
-  for (const joinPredicateCtx of ctx.joinCondition().joinPredicate_list()) {
+  for (const joinPredicateCtx of joinCondition.joinPredicate_list()) {
     const expression = visitValueExpression(joinPredicateCtx.valueExpression());
 
     if (expression) {
@@ -51,7 +56,10 @@ export const createJoinCommand = (ctx: JoinCommandContext): ESQLCommand => {
   }
 
   command.args.push(joinTarget);
-  command.args.push(onOption);
+
+  if (onOption.args.length) {
+    command.args.push(onOption);
+  }
 
   return command;
 };

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

@@ -14,7 +14,7 @@ const reprint = (src: string) => {
   const { root } = parse(src, { withFormatting: true });
   const text = BasicPrettyPrinter.print(root);
 
-  // console.log(JSON.stringify(ast, null, 2));
+  // console.log(JSON.stringify(root.commands, null, 2));
 
   return { text };
 };
@@ -184,3 +184,15 @@ describe('rename expressions', () => {
     assertPrint('FROM a | RENAME /*I*/ a /*II*/ AS /*III*/ b /*IV*/, c AS d');
   });
 });
+
+describe('commands', () => {
+  describe('JOIN', () => {
+    test('around JOIN targets', () => {
+      assertPrint('FROM a | LEFT JOIN /*1*/ a /*2*/ AS /*3*/ b /*4*/ ON c');
+    });
+
+    test('around JOIN conditions', () => {
+      assertPrint('FROM a | LEFT JOIN a AS b ON /*1*/ c /*2*/, /*3*/ d /*4*/');
+    });
+  });
+});

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

@@ -115,6 +115,40 @@ describe('single line query', () => {
         expect(text).toBe('FROM index | DISSECT input "pattern" APPEND_SEPARATOR = "<separator>"');
       });
     });
+
+    describe('JOIN', () => {
+      test('example from docs', () => {
+        const { text } = reprint(`
+          FROM employees
+            | EVAL language_code = languages
+            | LOOKUP JOIN languages_lookup ON language_code
+            | WHERE emp_no < 500
+            | KEEP emp_no, language_name
+            | SORT emp_no
+            | LIMIT 10
+        `);
+
+        expect(text).toBe(
+          'FROM employees | EVAL language_code = languages | LOOKUP JOIN languages_lookup ON language_code | WHERE emp_no < 500 | KEEP emp_no, language_name | SORT emp_no | LIMIT 10'
+        );
+      });
+
+      test('supports aliases', () => {
+        const { text } = reprint(`
+          FROM employees | LEFT JOIN languages_lookup AS something ON language_code`);
+
+        expect(text).toBe(
+          'FROM employees | LEFT JOIN languages_lookup AS something ON language_code'
+        );
+      });
+
+      test('supports multiple conditions', () => {
+        const { text } = reprint(`
+          FROM employees | LEFT JOIN a ON b, c, d.e.f`);
+
+        expect(text).toBe('FROM employees | LEFT JOIN a ON b, c, d.e.f');
+      });
+    });
   });
 
   describe('expressions', () => {

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

@@ -497,6 +497,49 @@ ROW 1
     });
   });
 
+  describe('as-expressions', () => {
+    test('JOIN main arguments surrounded in comments', () => {
+      const query = `
+        FROM index | LEFT JOIN
+        /* 1 */
+        // 2
+        /* 3 */
+        // 4
+        /* 5 */ a /* 6 */ AS /* 7 */ b
+        ON c`;
+      const text = reprint(query).text;
+      expect('\n' + text).toBe(`
+FROM index
+  | LEFT JOIN
+      /* 1 */
+      // 2
+      /* 3 */
+      // 4
+      /* 5 */ a /* 6 */ AS
+        /* 7 */ b
+        ON c`);
+    });
+
+    test('JOIN "ON" option argument comments', () => {
+      const query = `
+        FROM index | RIGHT JOIN a AS b ON
+        // c.1
+        /* c.2 */ c /* c.3 */,
+        // d.1
+        /* d.2 */ d /* d.3 */`;
+      const text = reprint(query).text;
+      expect('\n' + text).toBe(`
+FROM index
+  | RIGHT JOIN
+      a AS b
+        ON
+          // c.1
+          /* c.2 */ c, /* c.3 */
+          // d.1
+          /* d.2 */ d /* d.3 */`);
+    });
+  });
+
   describe('function call expressions', () => {
     describe('binary expressions', () => {
       test('first operand surrounded by inline comments', () => {