Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[ES|QL] JOIN command Traversal API and prety-printing support #202750

Merged
merged 9 commits into from
Dec 6, 2024
113 changes: 113 additions & 0 deletions packages/kbn-esql-ast/src/parser/__tests__/comments.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
*/

import { parse } from '..';
import { Walker } from '../../walker';

describe('Comments', () => {
describe('can attach "top" comment(s)', () => {
Expand Down Expand Up @@ -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)', () => {
Expand Down Expand Up @@ -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)', () => {
Expand Down Expand Up @@ -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', () => {
Expand Down
30 changes: 26 additions & 4 deletions packages/kbn-esql-ast/src/parser/__tests__/join.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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);
});
});
});
21 changes: 21 additions & 0 deletions packages/kbn-esql-ast/src/parser/factories.ts
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ import type {
InlineCastingType,
ESQLFunctionCallExpression,
ESQLIdentifier,
ESQLBinaryExpression,
BinaryExpressionOperator,
} from '../types';
import { parseIdentifier, getPosition } from './helpers';
import { Builder, type AstNodeParserFields } from '../builder';
Expand Down Expand Up @@ -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) {
Expand Down
20 changes: 14 additions & 6 deletions packages/kbn-esql-ast/src/parser/factories/join.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 = (
Expand All @@ -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;
Expand All @@ -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) {
Expand All @@ -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;
};
Original file line number Diff line number Diff line change
Expand Up @@ -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 };
};
Expand Down Expand Up @@ -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*/');
});
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
Loading
Loading