From 888a7176c6b225249a8aea2e0ba9d74225e598f9 Mon Sep 17 00:00:00 2001 From: vadimkibana Date: Tue, 22 Oct 2024 10:57:00 +0200 Subject: [PATCH 1/6] add sort command mutation api --- packages/kbn-esql-ast/src/builder/builder.ts | 15 + .../kbn-esql-ast/src/mutate/commands/index.ts | 3 +- .../src/mutate/commands/limit/index.ts | 2 +- .../src/mutate/commands/sort/index.test.ts | 357 ++++++++++++++++++ .../src/mutate/commands/sort/index.ts | 159 ++++++++ packages/kbn-esql-ast/src/mutate/generic.ts | 20 +- packages/kbn-esql-ast/src/parser/factories.ts | 17 +- packages/kbn-esql-ast/src/parser/walkers.ts | 2 +- 8 files changed, 551 insertions(+), 24 deletions(-) create mode 100644 packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/sort/index.ts diff --git a/packages/kbn-esql-ast/src/builder/builder.ts b/packages/kbn-esql-ast/src/builder/builder.ts index 26b64a6312ee4..accfd82e2e33c 100644 --- a/packages/kbn-esql-ast/src/builder/builder.ts +++ b/packages/kbn-esql-ast/src/builder/builder.ts @@ -20,6 +20,7 @@ import { ESQLIntegerLiteral, ESQLList, ESQLLocation, + ESQLOrderExpression, ESQLSource, } from '../types'; import { AstNodeParserFields, AstNodeTemplate, PartialFields } from './types'; @@ -130,6 +131,20 @@ export namespace Builder { }; }; + export const order = ( + operand: ESQLColumn, + template: Omit, 'name' | 'args'>, + fromParser?: Partial + ): ESQLOrderExpression => { + return { + ...template, + ...Builder.parserFields(fromParser), + name: '', + args: [operand], + type: 'order', + }; + }; + export const inlineCast = ( template: Omit, 'name'>, fromParser?: Partial diff --git a/packages/kbn-esql-ast/src/mutate/commands/index.ts b/packages/kbn-esql-ast/src/mutate/commands/index.ts index 0a779292e6eca..9e2599c493459 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/index.ts @@ -9,5 +9,6 @@ import * as from from './from'; import * as limit from './limit'; +import * as sort from './sort'; -export { from, limit }; +export { from, limit, sort }; diff --git a/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts b/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts index 937538e848328..534b8c4dd28b9 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts @@ -61,7 +61,7 @@ export const remove = (ast: ESQLAstQueryExpression, index: number = 0): ESQLComm return; } - const success = generic.removeCommand(ast, command); + const success = !!generic.removeCommand(ast, command); if (!success) { return; diff --git a/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts b/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts new file mode 100644 index 0000000000000..54b4dafe133bb --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts @@ -0,0 +1,357 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { parse } from '../../../parser'; +import * as commands from '..'; +import { BasicPrettyPrinter } from '../../../pretty_print'; + +describe('commands.sort', () => { + describe('.listCommands()', () => { + it('returns empty array, if there are no sort commands', () => { + const src = 'FROM index METADATA a'; + const { root } = parse(src); + const list = [...commands.sort.listCommands(root)]; + + expect(list.length).toBe(0); + }); + + it('returns all sort commands', () => { + const src = + 'FROM index | SORT a ASC, b DESC, c | LIMIT 123 | SORT d | EVAL 1 | SORT e NULLS FIRST, f NULLS LAST'; + const { root } = parse(src); + const list = [...commands.sort.listCommands(root)]; + + expect(list.length).toBe(3); + }); + + it('can skip given number of sort commands', () => { + const src = + 'FROM index | SORT a ASC, b DESC, c | LIMIT 123 | SORT d | EVAL 1 | SORT e NULLS FIRST, f NULLS LAST'; + const { root } = parse(src); + const list1 = [...commands.sort.listCommands(root, 1)]; + const list2 = [...commands.sort.listCommands(root, 2)]; + const list3 = [...commands.sort.listCommands(root, 3)]; + const list4 = [...commands.sort.listCommands(root, 111)]; + + expect(list1.length).toBe(2); + expect(list2.length).toBe(1); + expect(list3.length).toBe(0); + expect(list4.length).toBe(0); + }); + }); + + describe('.list()', () => { + it('returns empty array, if there are no sort commands', () => { + const src = 'FROM index METADATA a'; + const { root } = parse(src); + const list = [...commands.sort.list(root)]; + + expect(list.length).toBe(0); + }); + + it('returns a single column expression', () => { + const src = 'FROM index | SORT a'; + const { root } = parse(src); + const list = [...commands.sort.list(root)]; + + expect(list.length).toBe(1); + expect(list[0]).toMatchObject({ + type: 'column', + name: 'a', + }); + }); + + it('returns a single order expression', () => { + const src = 'FROM index | SORT a ASC'; + const { root } = parse(src); + const list = [...commands.sort.list(root)]; + + expect(list.length).toBe(1); + expect(list[0]).toMatchObject({ + type: 'order', + args: [ + { + type: 'column', + name: 'a', + }, + ], + }); + }); + + it('returns all sort command expressions', () => { + const src = + 'FROM index | SORT a ASC, b DESC, c | LIMIT 123 | SORT d | EVAL 1 | SORT e NULLS FIRST, f NULLS LAST'; + const { root } = parse(src); + const list = [...commands.sort.list(root)]; + + expect(list).toMatchObject([ + { + type: 'order', + args: [ + { + type: 'column', + name: 'a', + }, + ], + }, + { + type: 'order', + args: [ + { + type: 'column', + name: 'b', + }, + ], + }, + { + type: 'column', + name: 'c', + }, + { + type: 'column', + name: 'd', + }, + { + type: 'order', + args: [ + { + type: 'column', + name: 'e', + }, + ], + }, + { + type: 'order', + args: [ + { + type: 'column', + name: 'f', + }, + ], + }, + ]); + }); + + it('can skip one order expression', () => { + const src = 'FROM index | SORT b DESC, a ASC'; + const { root } = parse(src); + const list = [...commands.sort.list(root, 1)]; + + expect(list.length).toBe(1); + expect(list[0]).toMatchObject({ + type: 'order', + args: [ + { + type: 'column', + name: 'a', + }, + ], + }); + }); + }); + + describe('.find()', () => { + it('returns undefined if sort expression is not found', () => { + const src = 'FROM index | WHERE a = b | LIMIT 123'; + const { root } = parse(src); + const node = commands.sort.find(root, 'abc'); + + expect(node).toBe(undefined); + }); + + it('can find a single sort expression', () => { + const src = 'FROM index | SORT a'; + const { root } = parse(src); + const node = commands.sort.find(root, 'a')!; + + expect(node).toMatchObject({ + type: 'column', + name: 'a', + }); + }); + + it('can find a single sort (order) expression', () => { + const src = 'FROM index | SORT b ASC'; + const { root } = parse(src); + const node = commands.sort.find(root, 'b')!; + + expect(node).toMatchObject({ + type: 'order', + args: [ + { + type: 'column', + name: 'b', + }, + ], + }); + }); + + it('can find a column and specific order expressions among other such expressions', () => { + const src = + 'FROM index | SORT a, b ASC | STATS agg() | SORT c DESC, d, e NULLS FIRST | LIMIT 10'; + const { root } = parse(src); + const node1 = commands.sort.find(root, 'b')!; + const node2 = commands.sort.find(root, 'd')!; + + expect(node1).toMatchObject({ + type: 'order', + args: [ + { + type: 'column', + name: 'b', + }, + ], + }); + expect(node2).toMatchObject({ + type: 'column', + name: 'd', + }); + }); + + it('can select second order expression with the same name', () => { + const src = 'FROM index | SORT b ASC | STATS agg() | SORT b DESC'; + const { root } = parse(src); + const node = commands.sort.find(root, 'b', 1)!; + + expect(node).toMatchObject({ + type: 'order', + order: 'DESC', + args: [ + { + type: 'column', + name: 'b', + }, + ], + }); + }); + }); + + it('can find multipart columns', () => { + const src = 'FROM index | SORT hello, b.a ASC, a.b, c, c.d | STATS agg() | SORT b DESC'; + const { root } = parse(src); + const node1 = commands.sort.find(root, ['b', 'a'])!; + const node2 = commands.sort.find(root, ['a', 'b'])!; + + expect(node1).toMatchObject({ + type: 'order', + order: 'ASC', + args: [ + { + type: 'column', + parts: ['b', 'a'], + }, + ], + }); + expect(node2).toMatchObject({ + type: 'column', + parts: ['a', 'b'], + }); + }); + + describe('.remove()', () => { + it('can remove a column from a list', () => { + const src1 = 'FROM a, b, c | SORT a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a, b, c'); + + commands.sort.remove(root, 'b'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a, c'); + }); + + it('can remove an order expression from a list', () => { + const src1 = 'FROM a, b, c | SORT a, b ASC, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a, b ASC, c'); + + commands.sort.remove(root, 'b'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a, c'); + }); + + it('does nothing if column does not exist', () => { + const src1 = 'FROM a, b, c | SORT a, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a, c'); + + commands.sort.remove(root, 'b'); + commands.sort.remove(root, 'd'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a, c'); + }); + + it('can remove the sort expression at specific index', () => { + const src1 = 'FROM index | SORT a, b, c | LIMIT 1 | SORT a, b, c | LIMIT 2 | SORT a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe( + 'FROM index | SORT a, b, c | LIMIT 1 | SORT a, b, c | LIMIT 2 | SORT a, b, c' + ); + + commands.sort.remove(root, 'a', 1); + commands.sort.remove(root, 'c', 1); + commands.sort.remove(root, 'b', 2); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index | SORT a, b, c | LIMIT 1 | SORT b | LIMIT 2 | SORT a, c'); + }); + + it('removes SORT command, if it is left empty', () => { + const src1 = 'FROM index | SORT a, b, c | LIMIT 1 | SORT a, b, c | LIMIT 2 | SORT a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe( + 'FROM index | SORT a, b, c | LIMIT 1 | SORT a, b, c | LIMIT 2 | SORT a, b, c' + ); + + commands.sort.remove(root, 'c', 1); + commands.sort.remove(root, 'b', 1); + commands.sort.remove(root, 'a', 1); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM index | SORT a, b, c | LIMIT 1 | LIMIT 2 | SORT a, b, c'); + }); + + it('can remove by matching parts', () => { + const src1 = 'FROM a, b, c | SORT a, b.c, d.e NULLS FIRST, e'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a, b.c, d.e NULLS FIRST, e'); + + commands.sort.remove(root, ['b', 'c']); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a, d.e NULLS FIRST, e'); + + commands.sort.remove(root, ['d', 'e']); + + const src4 = BasicPrettyPrinter.print(root); + + expect(src4).toBe('FROM a, b, c | SORT a, e'); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts b/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts new file mode 100644 index 0000000000000..04b7408a61010 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts @@ -0,0 +1,159 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { Builder } from '../../../builder'; +import { + ESQLAstQueryExpression, + ESQLColumn, + ESQLCommand, + ESQLOrderExpression, +} from '../../../types'; +import { Visitor } from '../../../visitor'; +import { Predicate } from '../../types'; +import * as util from '../../util'; +import * as generic from '../../generic'; + +export type SortExpressions = ESQLOrderExpression | ESQLColumn; + +/** + * Iterates through all sort commands starting from the beginning of the query. + * You can specify the `skip` parameter to skip a given number of sort commands. + * + * @param ast The root of the AST. + * @param skip Number of sort commands to skip. + * @returns Iterator through all sort commands. + */ +export const listCommands = ( + ast: ESQLAstQueryExpression, + skip: number = 0 +): IterableIterator => { + return new Visitor() + .on('visitSortCommand', function* (ctx): IterableIterator { + if (skip) { + skip--; + } else { + yield ctx.node; + } + }) + .on('visitCommand', function* (): IterableIterator {}) + .on('visitQuery', function* (ctx): IterableIterator { + for (const command of ctx.visitCommands()) { + yield* command; + } + }) + .visitQuery(ast); +}; + +export const list = ( + ast: ESQLAstQueryExpression, + skip: number = 0 +): IterableIterator => { + return new Visitor() + .on('visitSortCommand', function* (ctx): IterableIterator { + for (const argument of ctx.arguments()) { + if (argument.type === 'order' || argument.type === 'column') { + if (skip) { + skip--; + } else { + yield argument; + } + } + } + }) + .on('visitCommand', function* (): IterableIterator {}) + .on('visitQuery', function* (ctx): IterableIterator { + for (const command of ctx.visitCommands()) { + yield* command; + } + }) + .visitQuery(ast); +}; + +export const findByPredicate = ( + ast: ESQLAstQueryExpression, + predicate: Predicate, + index?: number +): SortExpressions | undefined => { + return util.findByPredicate(list(ast, index), predicate); +}; + +export const find = ( + ast: ESQLAstQueryExpression, + parts: string | string[], + index: number = 0 +): SortExpressions | undefined => { + const arrParts = typeof parts === 'string' ? [parts] : parts; + + return findByPredicate(ast, (node) => { + let isMatch = false; + if (node.type === 'column') { + isMatch = util.cmpArr(node.parts, arrParts); + } else if (node.type === 'order') { + const columnParts = (node.args[0] as ESQLColumn)?.parts; + + if (Array.isArray(columnParts)) { + isMatch = util.cmpArr(columnParts, arrParts); + } + } + + if (isMatch) { + index--; + if (index < 0) { + return true; + } + } + + return false; + }); +}; + +export const remove = ( + ast: ESQLAstQueryExpression, + parts: string | string[], + index?: number +): SortExpressions | undefined => { + const node = find(ast, parts, index); + + if (!node) { + return undefined; + } + + const cmd = generic.removeCommandArgument(ast, node); + + if (cmd) { + if (!cmd.args.length) { + generic.removeCommand(ast, cmd); + } + } + + return cmd ? node : undefined; +}; + +// export const insert = ( +// ast: ESQLAstQueryExpression, +// indexName: string, +// clusterName?: string, +// index: number = -1 +// ): ESQLSource | undefined => { +// const command = generic.findCommandByName(ast, 'from'); + +// if (!command) { +// return; +// } + +// const source = Builder.expression.indexSource(indexName, clusterName); + +// if (index === -1) { +// generic.appendCommandArgument(command, source); +// } else { +// command.args.splice(index, 0, source); +// } + +// return source; +// }; diff --git a/packages/kbn-esql-ast/src/mutate/generic.ts b/packages/kbn-esql-ast/src/mutate/generic.ts index f27b0e2ae399f..9702e5d2fd9e6 100644 --- a/packages/kbn-esql-ast/src/mutate/generic.ts +++ b/packages/kbn-esql-ast/src/mutate/generic.ts @@ -254,34 +254,34 @@ export const removeCommandOption = ( * * @param ast The root AST node to search for command arguments. * @param node The argument AST node to remove. - * @returns Returns true if the argument was removed, false otherwise. + * @returns Returns the command that the argument was removed from, if any. */ export const removeCommandArgument = ( ast: ESQLAstQueryExpression, node: ESQLProperNode -): boolean => { +): ESQLCommand | undefined => { return new Visitor() - .on('visitCommand', (ctx): boolean => { + .on('visitCommand', (ctx): ESQLCommand | undefined => { const args = ctx.node.args; const length = args.length; for (let i = 0; i < length; i++) { if (args[i] === node) { args.splice(i, 1); - return true; + return ctx.node; } } - return false; + return undefined; }) - .on('visitQuery', (ctx): boolean => { - for (const success of ctx.visitCommands()) { - if (success) { - return true; + .on('visitQuery', (ctx): ESQLCommand | undefined => { + for (const cmd of ctx.visitCommands()) { + if (cmd) { + return cmd; } } - return false; + return undefined; }) .visitQuery(ast); }; diff --git a/packages/kbn-esql-ast/src/parser/factories.ts b/packages/kbn-esql-ast/src/parser/factories.ts index 0fffb3a970e4c..246a62747ee6e 100644 --- a/packages/kbn-esql-ast/src/parser/factories.ts +++ b/packages/kbn-esql-ast/src/parser/factories.ts @@ -203,20 +203,15 @@ export function createFunction( export const createOrderExpression = ( ctx: ParserRuleContext, - arg: ESQLAstItem, + arg: ESQLColumn, order: ESQLOrderExpression['order'], nulls: ESQLOrderExpression['nulls'] ) => { - const node: ESQLOrderExpression = { - type: 'order', - name: '', - order, - nulls, - args: [arg], - text: ctx.getText(), - location: getPosition(ctx.start, ctx.stop), - incomplete: Boolean(ctx.exception), - }; + const node = Builder.expression.order( + arg as ESQLColumn, + { order, nulls }, + createParserFields(ctx) + ); return node; }; diff --git a/packages/kbn-esql-ast/src/parser/walkers.ts b/packages/kbn-esql-ast/src/parser/walkers.ts index df10161f68bf8..268c90417078b 100644 --- a/packages/kbn-esql-ast/src/parser/walkers.ts +++ b/packages/kbn-esql-ast/src/parser/walkers.ts @@ -671,7 +671,7 @@ const visitOrderExpression = (ctx: OrderExpressionContext): ESQLOrderExpression return arg; } - return createOrderExpression(ctx, arg, order, nulls); + return createOrderExpression(ctx, arg as ESQLColumn, order, nulls); }; export function visitOrderExpressions( From d88b3bc187b621119f0760822fde933de6d43630 Mon Sep 17 00:00:00 2001 From: vadimkibana Date: Tue, 22 Oct 2024 23:12:54 +0200 Subject: [PATCH 2/6] return the source command of the found nodes --- .../src/mutate/commands/sort/index.test.ts | 65 +++++++++++-------- .../src/mutate/commands/sort/index.ts | 47 +++++++++----- 2 files changed, 67 insertions(+), 45 deletions(-) diff --git a/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts b/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts index 54b4dafe133bb..e12e4e330e432 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts @@ -58,7 +58,7 @@ describe('commands.sort', () => { it('returns a single column expression', () => { const src = 'FROM index | SORT a'; const { root } = parse(src); - const list = [...commands.sort.list(root)]; + const list = [...commands.sort.list(root)].map(([node]) => node); expect(list.length).toBe(1); expect(list[0]).toMatchObject({ @@ -70,7 +70,7 @@ describe('commands.sort', () => { it('returns a single order expression', () => { const src = 'FROM index | SORT a ASC'; const { root } = parse(src); - const list = [...commands.sort.list(root)]; + const list = [...commands.sort.list(root)].map(([node]) => node); expect(list.length).toBe(1); expect(list[0]).toMatchObject({ @@ -88,7 +88,7 @@ describe('commands.sort', () => { const src = 'FROM index | SORT a ASC, b DESC, c | LIMIT 123 | SORT d | EVAL 1 | SORT e NULLS FIRST, f NULLS LAST'; const { root } = parse(src); - const list = [...commands.sort.list(root)]; + const list = [...commands.sort.list(root)].map(([node]) => node); expect(list).toMatchObject([ { @@ -141,7 +141,7 @@ describe('commands.sort', () => { it('can skip one order expression', () => { const src = 'FROM index | SORT b DESC, a ASC'; const { root } = parse(src); - const list = [...commands.sort.list(root, 1)]; + const list = [...commands.sort.list(root, 1)].map(([node]) => node); expect(list.length).toBe(1); expect(list[0]).toMatchObject({ @@ -168,7 +168,7 @@ describe('commands.sort', () => { it('can find a single sort expression', () => { const src = 'FROM index | SORT a'; const { root } = parse(src); - const node = commands.sort.find(root, 'a')!; + const [node] = commands.sort.find(root, 'a')!; expect(node).toMatchObject({ type: 'column', @@ -179,7 +179,7 @@ describe('commands.sort', () => { it('can find a single sort (order) expression', () => { const src = 'FROM index | SORT b ASC'; const { root } = parse(src); - const node = commands.sort.find(root, 'b')!; + const [node] = commands.sort.find(root, 'b')!; expect(node).toMatchObject({ type: 'order', @@ -196,8 +196,8 @@ describe('commands.sort', () => { const src = 'FROM index | SORT a, b ASC | STATS agg() | SORT c DESC, d, e NULLS FIRST | LIMIT 10'; const { root } = parse(src); - const node1 = commands.sort.find(root, 'b')!; - const node2 = commands.sort.find(root, 'd')!; + const [node1] = commands.sort.find(root, 'b')!; + const [node2] = commands.sort.find(root, 'd')!; expect(node1).toMatchObject({ type: 'order', @@ -217,7 +217,7 @@ describe('commands.sort', () => { it('can select second order expression with the same name', () => { const src = 'FROM index | SORT b ASC | STATS agg() | SORT b DESC'; const { root } = parse(src); - const node = commands.sort.find(root, 'b', 1)!; + const [node] = commands.sort.find(root, 'b', 1)!; expect(node).toMatchObject({ type: 'order', @@ -230,27 +230,38 @@ describe('commands.sort', () => { ], }); }); - }); - it('can find multipart columns', () => { - const src = 'FROM index | SORT hello, b.a ASC, a.b, c, c.d | STATS agg() | SORT b DESC'; - const { root } = parse(src); - const node1 = commands.sort.find(root, ['b', 'a'])!; - const node2 = commands.sort.find(root, ['a', 'b'])!; + it('can find multipart columns', () => { + const src = 'FROM index | SORT hello, b.a ASC, a.b, c, c.d | STATS agg() | SORT b DESC'; + const { root } = parse(src); + const [node1] = commands.sort.find(root, ['b', 'a'])!; + const [node2] = commands.sort.find(root, ['a', 'b'])!; - expect(node1).toMatchObject({ - type: 'order', - order: 'ASC', - args: [ - { - type: 'column', - parts: ['b', 'a'], - }, - ], + expect(node1).toMatchObject({ + type: 'order', + order: 'ASC', + args: [ + { + type: 'column', + parts: ['b', 'a'], + }, + ], + }); + expect(node2).toMatchObject({ + type: 'column', + parts: ['a', 'b'], + }); }); - expect(node2).toMatchObject({ - type: 'column', - parts: ['a', 'b'], + + it('returns the parent sort command of the found order expression', () => { + const src = 'FROM index | SORT hello, b.a ASC, a.b, c, c.d | STATS agg() | SORT b DESC'; + const { root } = parse(src); + const [node1, command1] = commands.sort.find(root, ['b', 'a'])!; + const [node2, command2] = commands.sort.find(root, ['a', 'b'])!; + + expect(command1).toBe(command2); + expect(!!command1.args.find((arg) => arg === node1)).toBe(true); + expect(!!command2.args.find((arg) => arg === node2)).toBe(true); }); }); diff --git a/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts b/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts index 04b7408a61010..7f4f3e4edc153 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts @@ -7,7 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -import { Builder } from '../../../builder'; +// import { Builder } from '../../../builder'; import { ESQLAstQueryExpression, ESQLColumn, @@ -19,7 +19,7 @@ import { Predicate } from '../../types'; import * as util from '../../util'; import * as generic from '../../generic'; -export type SortExpressions = ESQLOrderExpression | ESQLColumn; +export type SortExpression = ESQLOrderExpression | ESQLColumn; /** * Iterates through all sort commands starting from the beginning of the query. @@ -53,21 +53,21 @@ export const listCommands = ( export const list = ( ast: ESQLAstQueryExpression, skip: number = 0 -): IterableIterator => { +): IterableIterator<[sortExpression: SortExpression, sortCommand: ESQLCommand]> => { return new Visitor() - .on('visitSortCommand', function* (ctx): IterableIterator { + .on('visitSortCommand', function* (ctx): IterableIterator<[SortExpression, ESQLCommand]> { for (const argument of ctx.arguments()) { if (argument.type === 'order' || argument.type === 'column') { if (skip) { skip--; } else { - yield argument; + yield [argument, ctx.node]; } } } }) - .on('visitCommand', function* (): IterableIterator {}) - .on('visitQuery', function* (ctx): IterableIterator { + .on('visitCommand', function* (): IterableIterator<[SortExpression, ESQLCommand]> {}) + .on('visitQuery', function* (ctx): IterableIterator<[SortExpression, ESQLCommand]> { for (const command of ctx.visitCommands()) { yield* command; } @@ -77,9 +77,9 @@ export const list = ( export const findByPredicate = ( ast: ESQLAstQueryExpression, - predicate: Predicate, + predicate: Predicate<[sortExpression: SortExpression, sortCommand: ESQLCommand]>, index?: number -): SortExpressions | undefined => { +): [sortExpression: SortExpression, sortCommand: ESQLCommand] | undefined => { return util.findByPredicate(list(ast, index), predicate); }; @@ -87,10 +87,10 @@ export const find = ( ast: ESQLAstQueryExpression, parts: string | string[], index: number = 0 -): SortExpressions | undefined => { +): [sortExpression: SortExpression, sortCommand: ESQLCommand] | undefined => { const arrParts = typeof parts === 'string' ? [parts] : parts; - return findByPredicate(ast, (node) => { + return findByPredicate(ast, ([node]) => { let isMatch = false; if (node.type === 'column') { isMatch = util.cmpArr(node.parts, arrParts); @@ -117,13 +117,14 @@ export const remove = ( ast: ESQLAstQueryExpression, parts: string | string[], index?: number -): SortExpressions | undefined => { - const node = find(ast, parts, index); +): [sortExpression: SortExpression, sortCommand: ESQLCommand] | undefined => { + const tuple = find(ast, parts, index); - if (!node) { + if (!tuple) { return undefined; } + const [node] = tuple; const cmd = generic.removeCommandArgument(ast, node); if (cmd) { @@ -132,14 +133,24 @@ export const remove = ( } } - return cmd ? node : undefined; + return cmd ? tuple : undefined; }; +// /** +// * +// * ``` +// * FROM index | SORT a, b | LIMIT 10 | SORT c, d +// * ``` +// * +// * @param ast +// * @param parts +// * @param index +// * @returns +// */ // export const insert = ( // ast: ESQLAstQueryExpression, -// indexName: string, -// clusterName?: string, -// index: number = -1 +// parts: string | string[], +// index?: number // ): ESQLSource | undefined => { // const command = generic.findCommandByName(ast, 'from'); From 19cbba9faf03c5695600fe9edf556b4a1269d4b0 Mon Sep 17 00:00:00 2001 From: vadimkibana Date: Wed, 23 Oct 2024 09:54:23 +0200 Subject: [PATCH 3/6] add ability to insert command arguments --- .../kbn-esql-ast/src/mutate/generic.test.ts | 118 ++++++++++++++++++ packages/kbn-esql-ast/src/mutate/generic.ts | 30 +++-- 2 files changed, 140 insertions(+), 8 deletions(-) diff --git a/packages/kbn-esql-ast/src/mutate/generic.test.ts b/packages/kbn-esql-ast/src/mutate/generic.test.ts index 0109ff838ffda..6108ff8d81146 100644 --- a/packages/kbn-esql-ast/src/mutate/generic.test.ts +++ b/packages/kbn-esql-ast/src/mutate/generic.test.ts @@ -7,6 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ +import { Builder } from '../builder'; import { parse } from '../parser'; import { BasicPrettyPrinter } from '../pretty_print'; import * as generic from './generic'; @@ -97,6 +98,123 @@ describe('generic', () => { }); }); + describe('.insertCommandArgument()', () => { + it('can insert at the end of the list', () => { + const src = 'FROM index | LIMIT 10'; + const { root } = parse(src); + const command = generic.findCommandByName(root, 'from', 0); + + generic.insertCommandArgument( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 123 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, test | LIMIT 10'); + }); + + it('can insert at the beginning of the list', () => { + const src = 'FROM index | LIMIT 10'; + const { root } = parse(src); + const command = generic.findCommandByName(root, 'from', 0); + + generic.insertCommandArgument( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 0 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM test, index | LIMIT 10'); + }); + + it('can insert in the middle of the list', () => { + const src = 'FROM index1, index2 | LIMIT 10'; + const { root } = parse(src); + const command = generic.findCommandByName(root, 'from', 0); + + generic.insertCommandArgument( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 1 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index1, test, index2 | LIMIT 10'); + }); + + describe('with option present', () => { + it('can insert at the end of the list', () => { + const src = 'FROM index METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.findCommandByName(root, 'from', 0); + + generic.insertCommandArgument( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 123 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, test METADATA _id | LIMIT 10'); + }); + + it('can insert at the beginning of the list', () => { + const src = 'FROM index METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.findCommandByName(root, 'from', 0); + + generic.insertCommandArgument( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 0 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM test, index METADATA _id | LIMIT 10'); + }); + + it('can insert in the middle of the list', () => { + const src = 'FROM index1, index2 METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.findCommandByName(root, 'from', 0); + + generic.insertCommandArgument( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 1 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index1, test, index2 METADATA _id | LIMIT 10'); + }); + }); + }); + + describe('.appendCommandArgument()', () => { + it('can append and argument', () => { + const src = 'FROM index METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.findCommandByName(root, 'from', 0); + + generic.appendCommandArgument( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }) + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, test METADATA _id | LIMIT 10'); + }); + }); + describe('.removeCommand()', () => { it('can remove the last command', () => { const src = 'FROM index | LIMIT 10'; diff --git a/packages/kbn-esql-ast/src/mutate/generic.ts b/packages/kbn-esql-ast/src/mutate/generic.ts index 9702e5d2fd9e6..2809e69d9eb1b 100644 --- a/packages/kbn-esql-ast/src/mutate/generic.ts +++ b/packages/kbn-esql-ast/src/mutate/generic.ts @@ -162,24 +162,38 @@ export const appendCommandOption = ( return option; }; -export const appendCommandArgument = ( +export const insertCommandArgument = ( command: ESQLCommand, - expression: ESQLSingleAstItem + expression: ESQLSingleAstItem, + index: number = -1 ): number => { if (expression.type === 'option') { command.args.push(expression); return command.args.length - 1; } - const index = command.args.findIndex((arg) => isOptionNode(arg)); + let mainArgumentCount = command.args.findIndex((arg) => isOptionNode(arg)); - if (index > -1) { - command.args.splice(index, 0, expression); - return index; + if (mainArgumentCount < 0) { + mainArgumentCount = command.args.length; + } + if (index === -1) { + index = mainArgumentCount; } + if (index > mainArgumentCount) { + index = mainArgumentCount; + } + + command.args.splice(index, 0, expression); - command.args.push(expression); - return command.args.length - 1; + return mainArgumentCount + 1; +}; + +export const appendCommandArgument = ( + command: ESQLCommand, + expression: ESQLSingleAstItem +): number => { + return insertCommandArgument(command, expression, -1); }; export const removeCommand = (ast: ESQLAstQueryExpression, command: ESQLCommand): boolean => { From e92a24e21b0d70bea005e54fdf8b2ab6cc4b5806 Mon Sep 17 00:00:00 2001 From: vadimkibana Date: Thu, 24 Oct 2024 16:06:09 +0200 Subject: [PATCH 4/6] finalize sort command api --- packages/kbn-esql-ast/src/builder/builder.ts | 12 +- .../src/mutate/commands/from/metadata.ts | 10 +- .../src/mutate/commands/from/sources.ts | 6 +- .../src/mutate/commands/limit/index.ts | 8 +- .../src/mutate/commands/sort/index.test.ts | 159 +++++++++ .../src/mutate/commands/sort/index.ts | 213 +++++++++++-- .../kbn-esql-ast/src/mutate/generic.test.ts | 271 ---------------- packages/kbn-esql-ast/src/mutate/generic.ts | 301 ------------------ .../generic/commands/args/index.test.ts | 132 ++++++++ .../src/mutate/generic/commands/args/index.ts | 86 +++++ .../src/mutate/generic/commands/index.test.ts | 118 +++++++ .../src/mutate/generic/commands/index.ts | 131 ++++++++ .../generic/commands/options/index.test.ts | 49 +++ .../mutate/generic/commands/options/index.ts | 130 ++++++++ .../kbn-esql-ast/src/mutate/generic/index.ts | 10 + packages/kbn-esql-ast/src/types.ts | 2 +- 16 files changed, 1013 insertions(+), 625 deletions(-) delete mode 100644 packages/kbn-esql-ast/src/mutate/generic.test.ts delete mode 100644 packages/kbn-esql-ast/src/mutate/generic.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic/commands/args/index.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic/commands/args/index.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic/commands/index.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic/commands/index.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic/commands/options/index.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic/commands/options/index.ts create mode 100644 packages/kbn-esql-ast/src/mutate/generic/index.ts diff --git a/packages/kbn-esql-ast/src/builder/builder.ts b/packages/kbn-esql-ast/src/builder/builder.ts index accfd82e2e33c..d033e177bd4b5 100644 --- a/packages/kbn-esql-ast/src/builder/builder.ts +++ b/packages/kbn-esql-ast/src/builder/builder.ts @@ -11,6 +11,8 @@ import { ESQLAstComment, + ESQLAstCommentMultiLine, + ESQLAstCommentSingleLine, ESQLAstQueryExpression, ESQLColumn, ESQLCommand, @@ -64,17 +66,17 @@ export namespace Builder { }; }; - export const comment = ( - subtype: ESQLAstComment['subtype'], + export const comment = ( + subtype: S, text: string, - location: ESQLLocation - ): ESQLAstComment => { + location?: ESQLLocation + ): S extends 'multi-line' ? ESQLAstCommentMultiLine : ESQLAstCommentSingleLine => { return { type: 'comment', subtype, text, location, - }; + } as S extends 'multi-line' ? ESQLAstCommentMultiLine : ESQLAstCommentSingleLine; }; export namespace expression { diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts index 7f08fa2a5e946..5160ab65954cb 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts @@ -106,7 +106,7 @@ export const removeByPredicate = ( option.args.splice(index, 1); if (option.args.length === 0) { - generic.removeCommandOption(ast, option); + generic.commands.options.remove(ast, option); } return tuple; @@ -148,16 +148,16 @@ export const insert = ( fieldName: string | string[], index: number = -1 ): [column: ESQLColumn, option: ESQLCommandOption] | undefined => { - let option = generic.findCommandOptionByName(ast, 'from', 'metadata'); + let option = generic.commands.options.findByName(ast, 'from', 'metadata'); if (!option) { - const command = generic.findCommandByName(ast, 'from'); + const command = generic.commands.findByName(ast, 'from'); if (!command) { return; } - option = generic.appendCommandOption(command, 'metadata'); + option = generic.commands.options.append(command, 'metadata'); } const parts: string[] = typeof fieldName === 'string' ? [fieldName] : fieldName; @@ -189,7 +189,7 @@ export const upsert = ( fieldName: string | string[], index: number = -1 ): [column: ESQLColumn, option: ESQLCommandOption] | undefined => { - const option = generic.findCommandOptionByName(ast, 'from', 'metadata'); + const option = generic.commands.options.findByName(ast, 'from', 'metadata'); if (option) { const parts = Array.isArray(fieldName) ? fieldName : [fieldName]; diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts b/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts index da67500b5b0bd..c10096cec38d9 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts @@ -67,7 +67,7 @@ export const remove = ( return undefined; } - const success = generic.removeCommandArgument(ast, node); + const success = generic.commands.args.remove(ast, node); return success ? node : undefined; }; @@ -78,7 +78,7 @@ export const insert = ( clusterName?: string, index: number = -1 ): ESQLSource | undefined => { - const command = generic.findCommandByName(ast, 'from'); + const command = generic.commands.findByName(ast, 'from'); if (!command) { return; @@ -87,7 +87,7 @@ export const insert = ( const source = Builder.expression.indexSource(indexName, clusterName); if (index === -1) { - generic.appendCommandArgument(command, source); + generic.commands.args.append(command, source); } else { command.args.splice(index, 0, source); } diff --git a/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts b/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts index 534b8c4dd28b9..f181a1d5f0cd4 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts @@ -19,7 +19,7 @@ import { Predicate } from '../../types'; * @returns A collection of "LIMIT" commands. */ export const list = (ast: ESQLAstQueryExpression): IterableIterator => { - return generic.listCommands(ast, (cmd) => cmd.name === 'limit'); + return generic.commands.list(ast, (cmd) => cmd.name === 'limit'); }; /** @@ -55,13 +55,13 @@ export const find = ( * @returns The removed "LIMIT" command, if any. */ export const remove = (ast: ESQLAstQueryExpression, index: number = 0): ESQLCommand | undefined => { - const command = generic.findCommandByName(ast, 'limit', index); + const command = generic.commands.findByName(ast, 'limit', index); if (!command) { return; } - const success = !!generic.removeCommand(ast, command); + const success = !!generic.commands.remove(ast, command); if (!success) { return; @@ -128,7 +128,7 @@ export const upsert = ( args: [literal], }); - generic.appendCommand(ast, command); + generic.commands.append(ast, command); return command; }; diff --git a/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts b/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts index e12e4e330e432..d04f79b96541a 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts @@ -10,6 +10,7 @@ import { parse } from '../../../parser'; import * as commands from '..'; import { BasicPrettyPrinter } from '../../../pretty_print'; +import { Builder } from '../../../builder'; describe('commands.sort', () => { describe('.listCommands()', () => { @@ -365,4 +366,162 @@ describe('commands.sort', () => { expect(src4).toBe('FROM a, b, c | SORT a, e'); }); }); + + describe('.insertIntoCommand()', () => { + it('can insert a sorting condition into the first existing SORT command', () => { + const src1 = 'FROM a, b, c | SORT s1, s2'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT s1, s2'); + + const command = commands.sort.getCommand(root)!; + commands.sort.insertIntoCommand(command, 's3'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT s1, s2, s3'); + }); + + it('can prepend a sorting condition with options into the first existing SORT command', () => { + const src1 = 'FROM a, b, c | SORT s1, s2'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT s1, s2'); + + const command = commands.sort.getCommand(root)!; + commands.sort.insertIntoCommand( + command, + { parts: ['address', 'street🙃'], order: 'ASC', nulls: 'NULLS FIRST' }, + 0 + ); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT address.`street🙃` ASC NULLS FIRST, s1, s2'); + }); + + it('can insert a sorting condition into specific sorting command into specific position', () => { + const src1 = 'FROM a, b, c | SORT a1, a2 | SORT b1, /* HERE */ b3 | SORT c1, c2'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a1, a2 | SORT b1, b3 | SORT c1, c2'); + + const command = commands.sort.getCommand(root, 1)!; + commands.sort.insertIntoCommand(command, 'b2', 1); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a1, a2 | SORT b1, b2, b3 | SORT c1, c2'); + }); + }); + + describe('.insertExpression()', () => { + it('can insert a sorting condition into the first existing SORT command', () => { + const src1 = 'FROM a, b, c | SORT s1, s2'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT s1, s2'); + + commands.sort.insertExpression(root, 's3'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT s1, s2, s3'); + }); + + it('can insert a sorting condition into specific sorting command into specific position', () => { + const src1 = 'FROM a, b, c | SORT a1, a2 | SORT b1, /* HERE */ b3 | SORT c1, c2'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a1, a2 | SORT b1, b3 | SORT c1, c2'); + + commands.sort.insertExpression(root, 'b2', 1, 1); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a1, a2 | SORT b1, b2, b3 | SORT c1, c2'); + }); + + it('when no positional arguments are provided append the column to the first SORT command', () => { + const src1 = 'FROM a, b, c | SORT a1, a2 | SORT b1, b2 | SORT c1, c2'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a1, a2 | SORT b1, b2 | SORT c1, c2'); + + commands.sort.insertExpression(root, 'a3'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a1, a2, a3 | SORT b1, b2 | SORT c1, c2'); + }); + + it('when no SORT command found, inserts a new SORT command', () => { + const src1 = 'FROM a, b, c | LIMIT 10'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | LIMIT 10'); + + commands.sort.insertExpression(root, ['i18n', 'language', 'locale']); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | LIMIT 10 | SORT i18n.language.locale'); + }); + + it('can change the sorting order', () => { + const src1 = 'FROM a, b, c | SORT a ASC'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT a ASC'); + + commands.sort.insertExpression(root, { parts: 'a', order: 'DESC' }); + commands.sort.remove(root, 'a', 0); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT a DESC'); + }); + }); + + describe('.insertCommand()', () => { + it('can append a new SORT command', () => { + const src1 = 'FROM a, b, c | SORT s1, s2'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | SORT s1, s2'); + + commands.sort.insertCommand(root, 's3'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT s1, s2 | SORT s3'); + }); + + it('can insert a SORT command before a LIMIT command (and add a comment)', () => { + const src1 = 'FROM a, b, c | LIMIT 10'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c | LIMIT 10'); + + const [_, column] = commands.sort.insertCommand(root, 'b', 1); + + column.formatting = { + right: [Builder.comment('multi-line', ' we sort by "b" ')], + }; + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c | SORT b /* we sort by "b" */ | LIMIT 10'); + }); + }); }); diff --git a/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts b/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts index 7f4f3e4edc153..d2b2c7cd5f3d4 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts @@ -7,7 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ -// import { Builder } from '../../../builder'; +import { Builder } from '../../../builder'; import { ESQLAstQueryExpression, ESQLColumn, @@ -21,6 +21,57 @@ import * as generic from '../../generic'; export type SortExpression = ESQLOrderExpression | ESQLColumn; +/** + * This "template" allows the developer to easily specify a new sort expression + * AST node, for example: + * + * ```ts + * // as a simple string + * 'column_name' + * + * // column with nested fields + * ['column_name', 'nested_field'] + * + * // as an object with additional options + * { parts: 'column_name', order: 'ASC', nulls: 'NULLS FIRST' } + * { parts: ['column_name', 'nested_field'], order: 'DESC', nulls: 'NULLS LAST' } + * ``` + */ +export type NewSortExpressionTemplate = + | string + | string[] + | { + parts: string | string[]; + order?: ESQLOrderExpression['order']; + nulls?: ESQLOrderExpression['nulls']; + }; + +const createSortExpression = ( + template: string | string[] | NewSortExpressionTemplate +): SortExpression => { + const column = Builder.expression.column({ + parts: + typeof template === 'string' + ? [template] + : Array.isArray(template) + ? template + : typeof template.parts === 'string' + ? [template.parts] + : template.parts, + }); + + if (typeof template === 'string' || Array.isArray(template)) { + return column; + } + + const order = Builder.expression.order(column, { + order: template.order ?? '', + nulls: template.nulls ?? '', + }); + + return order; +}; + /** * Iterates through all sort commands starting from the beginning of the query. * You can specify the `skip` parameter to skip a given number of sort commands. @@ -50,6 +101,31 @@ export const listCommands = ( .visitQuery(ast); }; +/** + * Returns the Nth SORT command found in the query. + * + * @param ast The root of the AST. + * @param index The index (N) of the sort command to return. + * @returns The sort command found in the AST, if any. + */ +export const getCommand = ( + ast: ESQLAstQueryExpression, + index: number = 0 +): ESQLCommand | undefined => { + for (const command of listCommands(ast, index)) { + return command; + } +}; + +/** + * Returns an iterator for all sort expressions (columns and order expressions) + * in the query. You can specify the `skip` parameter to skip a given number of + * expressions. + * + * @param ast The root of the AST. + * @param skip Number of sort expressions to skip. + * @returns Iterator through sort expressions (columns and order expressions). + */ export const list = ( ast: ESQLAstQueryExpression, skip: number = 0 @@ -75,6 +151,17 @@ export const list = ( .visitQuery(ast); }; +/** + * Finds the Nts sort expression that matches the predicate. + * + * @param ast The root of the AST. + * @param predicate A function that returns true if the sort expression matches + * the predicate. + * @param index The index of the sort expression to return. If not specified, + * the first sort expression that matches the predicate will be returned. + * @returns The sort expressions and sort command 2-tuple that matches the + * predicate, if any. + */ export const findByPredicate = ( ast: ESQLAstQueryExpression, predicate: Predicate<[sortExpression: SortExpression, sortCommand: ESQLCommand]>, @@ -83,6 +170,15 @@ export const findByPredicate = ( return util.findByPredicate(list(ast, index), predicate); }; +/** + * Finds the Nth sort expression that matches the sort expression by column + * name. The `parts` argument allows to specify an array of nested field names. + * + * @param ast The root of the AST. + * @param parts A string or an array of strings representing the column name. + * @returns The sort expressions and sort command 2-tuple that matches the + * predicate, if any. + */ export const find = ( ast: ESQLAstQueryExpression, parts: string | string[], @@ -113,6 +209,15 @@ export const find = ( }); }; +/** + * Removes the Nth sort expression that matches the sort expression by column + * name. The `parts` argument allows to specify an array of nested field names. + * + * @param ast The root of the AST. + * @param parts A string or an array of strings representing the column name. + * @param index The index of the sort expression to remove. + * @returns The sort expressions and sort command 2-tuple that was removed, if any. + */ export const remove = ( ast: ESQLAstQueryExpression, parts: string | string[], @@ -125,46 +230,84 @@ export const remove = ( } const [node] = tuple; - const cmd = generic.removeCommandArgument(ast, node); + const cmd = generic.commands.args.remove(ast, node); if (cmd) { if (!cmd.args.length) { - generic.removeCommand(ast, cmd); + generic.commands.remove(ast, cmd); } } return cmd ? tuple : undefined; }; -// /** -// * -// * ``` -// * FROM index | SORT a, b | LIMIT 10 | SORT c, d -// * ``` -// * -// * @param ast -// * @param parts -// * @param index -// * @returns -// */ -// export const insert = ( -// ast: ESQLAstQueryExpression, -// parts: string | string[], -// index?: number -// ): ESQLSource | undefined => { -// const command = generic.findCommandByName(ast, 'from'); - -// if (!command) { -// return; -// } - -// const source = Builder.expression.indexSource(indexName, clusterName); - -// if (index === -1) { -// generic.appendCommandArgument(command, source); -// } else { -// command.args.splice(index, 0, source); -// } - -// return source; -// }; +/** + * Inserts a new sort expression into the specified SORT command at the + * specified argument position. + * + * @param sortCommand The SORT command to insert the new sort expression into. + * @param template The sort expression template. + * @param index Argument position in the command argument list. + * @returns The inserted sort expression. + */ +export const insertIntoCommand = ( + sortCommand: ESQLCommand, + template: NewSortExpressionTemplate, + index?: number +): SortExpression => { + const expression = createSortExpression(template); + + generic.commands.args.insert(sortCommand, expression, index); + + return expression; +}; + +/** + * Creates a new sort expression node and inserts it into the specified SORT + * command at the specified argument position. If not sort command is found, a + * new one is created and appended to the end of the query. + * + * @param ast The root AST node. + * @param parts ES|QL column name parts. + * @param index The new column name position in command argument list. + * @param sortCommandIndex The index of the SORT command in the AST. E.g. 0 is the + * first SORT command in the AST. + * @returns The inserted column AST node. + */ +export const insertExpression = ( + ast: ESQLAstQueryExpression, + template: NewSortExpressionTemplate, + index: number = -1, + sortCommandIndex: number = 0 +): SortExpression => { + let command: ESQLCommand | undefined = getCommand(ast, sortCommandIndex); + + if (!command) { + command = Builder.command({ name: 'sort' }); + generic.commands.append(ast, command); + } + + return insertIntoCommand(command, template, index); +}; + +/** + * Inserts a new SORT command with a single sort expression as its sole argument. + * You can specify the position to insert the command at. + * + * @param ast The root of the AST. + * @param template The sort expression template. + * @param index The position to insert the sort expression at. + * @returns The inserted sort expression and the command it was inserted into. + */ +export const insertCommand = ( + ast: ESQLAstQueryExpression, + template: NewSortExpressionTemplate, + index: number = -1 +): [ESQLCommand, SortExpression] => { + const expression = createSortExpression(template); + const command = Builder.command({ name: 'sort', args: [expression] }); + + generic.commands.insert(ast, command, index); + + return [command, expression]; +}; diff --git a/packages/kbn-esql-ast/src/mutate/generic.test.ts b/packages/kbn-esql-ast/src/mutate/generic.test.ts deleted file mode 100644 index 6108ff8d81146..0000000000000 --- a/packages/kbn-esql-ast/src/mutate/generic.test.ts +++ /dev/null @@ -1,271 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the "Elastic License - * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side - * Public License v 1"; you may not use this file except in compliance with, at - * your election, the "Elastic License 2.0", the "GNU Affero General Public - * License v3.0 only", or the "Server Side Public License, v 1". - */ - -import { Builder } from '../builder'; -import { parse } from '../parser'; -import { BasicPrettyPrinter } from '../pretty_print'; -import * as generic from './generic'; - -describe('generic', () => { - describe('.listCommands()', () => { - it('lists all commands', () => { - const src = 'FROM index | WHERE a == b | LIMIT 123'; - const { root } = parse(src); - const commands = [...generic.listCommands(root)].map((cmd) => cmd.name); - - expect(commands).toEqual(['from', 'where', 'limit']); - }); - }); - - describe('.findCommand()', () => { - it('can the first command', () => { - const src = 'FROM index | WHERE a == b | LIMIT 123'; - const { root } = parse(src); - const command = generic.findCommand(root, (cmd) => cmd.name === 'from'); - - expect(command).toMatchObject({ - type: 'command', - name: 'from', - args: [ - { - type: 'source', - }, - ], - }); - }); - - it('can the last command', () => { - const src = 'FROM index | WHERE a == b | LIMIT 123'; - const { root } = parse(src); - const command = generic.findCommand(root, (cmd) => cmd.name === 'limit'); - - expect(command).toMatchObject({ - type: 'command', - name: 'limit', - args: [ - { - type: 'literal', - }, - ], - }); - }); - - it('find the specific of multiple commands', () => { - const src = 'FROM index | WHERE a == b | LIMIT 1 | LIMIT 2 | LIMIT 3'; - const { root } = parse(src); - const command = generic.findCommand( - root, - (cmd) => cmd.name === 'limit' && (cmd.args?.[0] as any).value === 2 - ); - - expect(command).toMatchObject({ - type: 'command', - name: 'limit', - args: [ - { - type: 'literal', - value: 2, - }, - ], - }); - }); - }); - - describe('.findCommandOptionByName()', () => { - it('can the find a command option', () => { - const src = 'FROM index METADATA _score'; - const { root } = parse(src); - const option = generic.findCommandOptionByName(root, 'from', 'metadata'); - - expect(option).toMatchObject({ - type: 'option', - name: 'metadata', - }); - }); - - it('returns undefined if there is no option', () => { - const src = 'FROM index'; - const { root } = parse(src); - const option = generic.findCommandOptionByName(root, 'from', 'metadata'); - - expect(option).toBe(undefined); - }); - }); - - describe('.insertCommandArgument()', () => { - it('can insert at the end of the list', () => { - const src = 'FROM index | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'from', 0); - - generic.insertCommandArgument( - command!, - Builder.expression.source({ name: 'test', sourceType: 'index' }), - 123 - ); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index, test | LIMIT 10'); - }); - - it('can insert at the beginning of the list', () => { - const src = 'FROM index | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'from', 0); - - generic.insertCommandArgument( - command!, - Builder.expression.source({ name: 'test', sourceType: 'index' }), - 0 - ); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM test, index | LIMIT 10'); - }); - - it('can insert in the middle of the list', () => { - const src = 'FROM index1, index2 | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'from', 0); - - generic.insertCommandArgument( - command!, - Builder.expression.source({ name: 'test', sourceType: 'index' }), - 1 - ); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index1, test, index2 | LIMIT 10'); - }); - - describe('with option present', () => { - it('can insert at the end of the list', () => { - const src = 'FROM index METADATA _id | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'from', 0); - - generic.insertCommandArgument( - command!, - Builder.expression.source({ name: 'test', sourceType: 'index' }), - 123 - ); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index, test METADATA _id | LIMIT 10'); - }); - - it('can insert at the beginning of the list', () => { - const src = 'FROM index METADATA _id | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'from', 0); - - generic.insertCommandArgument( - command!, - Builder.expression.source({ name: 'test', sourceType: 'index' }), - 0 - ); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM test, index METADATA _id | LIMIT 10'); - }); - - it('can insert in the middle of the list', () => { - const src = 'FROM index1, index2 METADATA _id | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'from', 0); - - generic.insertCommandArgument( - command!, - Builder.expression.source({ name: 'test', sourceType: 'index' }), - 1 - ); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index1, test, index2 METADATA _id | LIMIT 10'); - }); - }); - }); - - describe('.appendCommandArgument()', () => { - it('can append and argument', () => { - const src = 'FROM index METADATA _id | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'from', 0); - - generic.appendCommandArgument( - command!, - Builder.expression.source({ name: 'test', sourceType: 'index' }) - ); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index, test METADATA _id | LIMIT 10'); - }); - }); - - describe('.removeCommand()', () => { - it('can remove the last command', () => { - const src = 'FROM index | LIMIT 10'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'limit', 0); - - generic.removeCommand(root, command!); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index'); - }); - - it('can remove the second command out of 3 with the same name', () => { - const src = 'FROM index | LIMIT 1 | LIMIT 2 | LIMIT 3'; - const { root } = parse(src); - const command = generic.findCommandByName(root, 'limit', 1); - - generic.removeCommand(root, command!); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index | LIMIT 1 | LIMIT 3'); - }); - - it('can remove all commands', () => { - const src = 'FROM index | WHERE a == b | LIMIT 123'; - const { root } = parse(src); - const cmd1 = generic.findCommandByName(root, 'where'); - const cmd2 = generic.findCommandByName(root, 'limit'); - const cmd3 = generic.findCommandByName(root, 'from'); - - generic.removeCommand(root, cmd1!); - generic.removeCommand(root, cmd2!); - generic.removeCommand(root, cmd3!); - - expect(root.commands.length).toBe(0); - }); - }); - - describe('.removeCommandOption()', () => { - it('can remove existing command option', () => { - const src = 'FROM index METADATA _score'; - const { root } = parse(src); - const option = generic.findCommandOptionByName(root, 'from', 'metadata'); - - generic.removeCommandOption(root, option!); - - const src2 = BasicPrettyPrinter.print(root); - - expect(src2).toBe('FROM index'); - }); - }); -}); diff --git a/packages/kbn-esql-ast/src/mutate/generic.ts b/packages/kbn-esql-ast/src/mutate/generic.ts deleted file mode 100644 index 2809e69d9eb1b..0000000000000 --- a/packages/kbn-esql-ast/src/mutate/generic.ts +++ /dev/null @@ -1,301 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the "Elastic License - * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side - * Public License v 1"; you may not use this file except in compliance with, at - * your election, the "Elastic License 2.0", the "GNU Affero General Public - * License v3.0 only", or the "Server Side Public License, v 1". - */ - -import { isOptionNode } from '../ast/util'; -import { Builder } from '../builder'; -import { - ESQLAstQueryExpression, - ESQLCommand, - ESQLCommandOption, - ESQLProperNode, - ESQLSingleAstItem, -} from '../types'; -import { Visitor } from '../visitor'; -import { Predicate } from './types'; - -/** - * Returns an iterator for all command AST nodes in the query. If a predicate is - * provided, only commands that satisfy the predicate will be returned. - * - * @param ast Root AST node to search for commands. - * @param predicate Optional predicate to filter commands. - * @returns A list of commands found in the AST. - */ -export const listCommands = ( - ast: ESQLAstQueryExpression, - predicate?: Predicate -): IterableIterator => { - return new Visitor() - .on('visitQuery', function* (ctx): IterableIterator { - for (const cmd of ctx.commands()) { - if (!predicate || predicate(cmd)) { - yield cmd; - } - } - }) - .visitQuery(ast); -}; - -/** - * Returns the first command AST node at a given index in the query that - * satisfies the predicate. If no index is provided, the first command found - * will be returned. - * - * @param ast Root AST node to search for commands. - * @param predicate Optional predicate to filter commands. - * @param index The index of the command to return. - * @returns The command found in the AST, if any. - */ -export const findCommand = ( - ast: ESQLAstQueryExpression, - predicate?: Predicate, - index: number = 0 -): ESQLCommand | undefined => { - for (const cmd of listCommands(ast, predicate)) { - if (!index) { - return cmd; - } - - index--; - } - - return undefined; -}; - -/** - * Returns the first command option AST node that satisfies the predicate. - * - * @param command The command AST node to search for options. - * @param predicate The predicate to filter options. - * @returns The option found in the command, if any. - */ -export const findCommandOption = ( - command: ESQLCommand, - predicate: Predicate -): ESQLCommandOption | undefined => { - return new Visitor() - .on('visitCommand', (ctx): ESQLCommandOption | undefined => { - for (const opt of ctx.options()) { - if (predicate(opt)) { - return opt; - } - } - - return undefined; - }) - .visitCommand(command); -}; - -/** - * Returns the first command AST node at a given index with a given name in the - * query. If no index is provided, the first command found will be returned. - * - * @param ast Root AST node to search for commands. - * @param commandName The name of the command to find. - * @param index The index of the command to return. - * @returns The command found in the AST, if any. - */ -export const findCommandByName = ( - ast: ESQLAstQueryExpression, - commandName: string, - index: number = 0 -): ESQLCommand | undefined => { - return findCommand(ast, (cmd) => cmd.name === commandName, index); -}; - -/** - * Returns the first command option AST node with a given name in the query. - * - * @param ast The root AST node to search for command options. - * @param commandName Command name to search for. - * @param optionName Option name to search for. - * @returns The option found in the command, if any. - */ -export const findCommandOptionByName = ( - ast: ESQLAstQueryExpression, - commandName: string, - optionName: string -): ESQLCommandOption | undefined => { - const command = findCommand(ast, (cmd) => cmd.name === commandName); - - if (!command) { - return undefined; - } - - return findCommandOption(command, (opt) => opt.name === optionName); -}; - -/** - * Adds a new command to the query AST node. - * - * @param ast The root AST node to append the command to. - * @param command The command AST node to append. - */ -export const appendCommand = (ast: ESQLAstQueryExpression, command: ESQLCommand): void => { - ast.commands.push(command); -}; - -/** - * Inserts a command option into the command's arguments list. The option can - * be specified as a string or an AST node. - * - * @param command The command AST node to insert the option into. - * @param option The option to insert. - * @returns The inserted option. - */ -export const appendCommandOption = ( - command: ESQLCommand, - option: string | ESQLCommandOption -): ESQLCommandOption => { - if (typeof option === 'string') { - option = Builder.option({ name: option }); - } - - command.args.push(option); - - return option; -}; - -export const insertCommandArgument = ( - command: ESQLCommand, - expression: ESQLSingleAstItem, - index: number = -1 -): number => { - if (expression.type === 'option') { - command.args.push(expression); - return command.args.length - 1; - } - - let mainArgumentCount = command.args.findIndex((arg) => isOptionNode(arg)); - - if (mainArgumentCount < 0) { - mainArgumentCount = command.args.length; - } - if (index === -1) { - index = mainArgumentCount; - } - if (index > mainArgumentCount) { - index = mainArgumentCount; - } - - command.args.splice(index, 0, expression); - - return mainArgumentCount + 1; -}; - -export const appendCommandArgument = ( - command: ESQLCommand, - expression: ESQLSingleAstItem -): number => { - return insertCommandArgument(command, expression, -1); -}; - -export const removeCommand = (ast: ESQLAstQueryExpression, command: ESQLCommand): boolean => { - const cmds = ast.commands; - const length = cmds.length; - - for (let i = 0; i < length; i++) { - if (cmds[i] === command) { - cmds.splice(i, 1); - return true; - } - } - - return false; -}; - -/** - * Removes the first command option from the command's arguments list that - * satisfies the predicate. - * - * @param command The command AST node to remove the option from. - * @param predicate The predicate to filter options. - * @returns The removed option, if any. - */ -export const removeCommandOption = ( - ast: ESQLAstQueryExpression, - option: ESQLCommandOption -): boolean => { - return new Visitor() - .on('visitCommandOption', (ctx): boolean => { - return ctx.node === option; - }) - .on('visitCommand', (ctx): boolean => { - let target: undefined | ESQLCommandOption; - - for (const opt of ctx.options()) { - if (opt === option) { - target = opt; - break; - } - } - - if (!target) { - return false; - } - - const index = ctx.node.args.indexOf(target); - - if (index === -1) { - return false; - } - - ctx.node.args.splice(index, 1); - - return true; - }) - .on('visitQuery', (ctx): boolean => { - for (const success of ctx.visitCommands()) { - if (success) { - return true; - } - } - - return false; - }) - .visitQuery(ast); -}; - -/** - * Searches all command arguments in the query AST node and removes the node - * from the command's arguments list. - * - * @param ast The root AST node to search for command arguments. - * @param node The argument AST node to remove. - * @returns Returns the command that the argument was removed from, if any. - */ -export const removeCommandArgument = ( - ast: ESQLAstQueryExpression, - node: ESQLProperNode -): ESQLCommand | undefined => { - return new Visitor() - .on('visitCommand', (ctx): ESQLCommand | undefined => { - const args = ctx.node.args; - const length = args.length; - - for (let i = 0; i < length; i++) { - if (args[i] === node) { - args.splice(i, 1); - return ctx.node; - } - } - - return undefined; - }) - .on('visitQuery', (ctx): ESQLCommand | undefined => { - for (const cmd of ctx.visitCommands()) { - if (cmd) { - return cmd; - } - } - - return undefined; - }) - .visitQuery(ast); -}; diff --git a/packages/kbn-esql-ast/src/mutate/generic/commands/args/index.test.ts b/packages/kbn-esql-ast/src/mutate/generic/commands/args/index.test.ts new file mode 100644 index 0000000000000..e687c4528dd7d --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic/commands/args/index.test.ts @@ -0,0 +1,132 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { Builder } from '../../../../builder'; +import { parse } from '../../../../parser'; +import { BasicPrettyPrinter } from '../../../../pretty_print'; +import * as generic from '../..'; + +describe('generic.commands.args', () => { + describe('.insert()', () => { + it('can insert at the end of the list', () => { + const src = 'FROM index | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'from', 0); + + generic.commands.args.insert( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 123 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, test | LIMIT 10'); + }); + + it('can insert at the beginning of the list', () => { + const src = 'FROM index | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'from', 0); + + generic.commands.args.insert( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 0 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM test, index | LIMIT 10'); + }); + + it('can insert in the middle of the list', () => { + const src = 'FROM index1, index2 | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'from', 0); + + generic.commands.args.insert( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 1 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index1, test, index2 | LIMIT 10'); + }); + + describe('with option present', () => { + it('can insert at the end of the list', () => { + const src = 'FROM index METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'from', 0); + + generic.commands.args.insert( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 123 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, test METADATA _id | LIMIT 10'); + }); + + it('can insert at the beginning of the list', () => { + const src = 'FROM index METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'from', 0); + + generic.commands.args.insert( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 0 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM test, index METADATA _id | LIMIT 10'); + }); + + it('can insert in the middle of the list', () => { + const src = 'FROM index1, index2 METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'from', 0); + + generic.commands.args.insert( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }), + 1 + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index1, test, index2 METADATA _id | LIMIT 10'); + }); + }); + }); + + describe('.append()', () => { + it('can append and argument', () => { + const src = 'FROM index METADATA _id | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'from', 0); + + generic.commands.args.append( + command!, + Builder.expression.source({ name: 'test', sourceType: 'index' }) + ); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, test METADATA _id | LIMIT 10'); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/generic/commands/args/index.ts b/packages/kbn-esql-ast/src/mutate/generic/commands/args/index.ts new file mode 100644 index 0000000000000..7072c38a5f1a8 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic/commands/args/index.ts @@ -0,0 +1,86 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { isOptionNode } from '../../../../ast/util'; +import { + ESQLAstQueryExpression, + ESQLCommand, + ESQLProperNode, + ESQLSingleAstItem, +} from '../../../../types'; +import { Visitor } from '../../../../visitor'; + +export const insert = ( + command: ESQLCommand, + expression: ESQLSingleAstItem, + index: number = -1 +): number => { + if (expression.type === 'option') { + command.args.push(expression); + return command.args.length - 1; + } + + let mainArgumentCount = command.args.findIndex((arg) => isOptionNode(arg)); + + if (mainArgumentCount < 0) { + mainArgumentCount = command.args.length; + } + if (index === -1) { + index = mainArgumentCount; + } + if (index > mainArgumentCount) { + index = mainArgumentCount; + } + + command.args.splice(index, 0, expression); + + return mainArgumentCount + 1; +}; + +export const append = (command: ESQLCommand, expression: ESQLSingleAstItem): number => { + return insert(command, expression, -1); +}; + +/** + * Searches all command arguments in the query AST node and removes the node + * from the command's arguments list. + * + * @param ast The root AST node to search for command arguments. + * @param node The argument AST node to remove. + * @returns Returns the command that the argument was removed from, if any. + */ +export const remove = ( + ast: ESQLAstQueryExpression, + node: ESQLProperNode +): ESQLCommand | undefined => { + return new Visitor() + .on('visitCommand', (ctx): ESQLCommand | undefined => { + const args = ctx.node.args; + const length = args.length; + + for (let i = 0; i < length; i++) { + if (args[i] === node) { + args.splice(i, 1); + return ctx.node; + } + } + + return undefined; + }) + .on('visitQuery', (ctx): ESQLCommand | undefined => { + for (const cmd of ctx.visitCommands()) { + if (cmd) { + return cmd; + } + } + + return undefined; + }) + .visitQuery(ast); +}; diff --git a/packages/kbn-esql-ast/src/mutate/generic/commands/index.test.ts b/packages/kbn-esql-ast/src/mutate/generic/commands/index.test.ts new file mode 100644 index 0000000000000..b35d0b6415247 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic/commands/index.test.ts @@ -0,0 +1,118 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { parse } from '../../../parser'; +import { BasicPrettyPrinter } from '../../../pretty_print'; +import * as generic from '..'; + +describe('generic.commands', () => { + describe('.list()', () => { + it('lists all commands', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + const commands = [...generic.commands.list(root)].map((cmd) => cmd.name); + + expect(commands).toEqual(['from', 'where', 'limit']); + }); + }); + + describe('.find()', () => { + it('can the first command', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + const command = generic.commands.find(root, (cmd) => cmd.name === 'from'); + + expect(command).toMatchObject({ + type: 'command', + name: 'from', + args: [ + { + type: 'source', + }, + ], + }); + }); + + it('can the last command', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + const command = generic.commands.find(root, (cmd) => cmd.name === 'limit'); + + expect(command).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + }, + ], + }); + }); + + it('find the specific of multiple commands', () => { + const src = 'FROM index | WHERE a == b | LIMIT 1 | LIMIT 2 | LIMIT 3'; + const { root } = parse(src); + const command = generic.commands.find( + root, + (cmd) => cmd.name === 'limit' && (cmd.args?.[0] as any).value === 2 + ); + + expect(command).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 2, + }, + ], + }); + }); + }); + + describe('.remove()', () => { + it('can remove the last command', () => { + const src = 'FROM index | LIMIT 10'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'limit', 0); + + generic.commands.remove(root, command!); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index'); + }); + + it('can remove the second command out of 3 with the same name', () => { + const src = 'FROM index | LIMIT 1 | LIMIT 2 | LIMIT 3'; + const { root } = parse(src); + const command = generic.commands.findByName(root, 'limit', 1); + + generic.commands.remove(root, command!); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index | LIMIT 1 | LIMIT 3'); + }); + + it('can remove all commands', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + const cmd1 = generic.commands.findByName(root, 'where'); + const cmd2 = generic.commands.findByName(root, 'limit'); + const cmd3 = generic.commands.findByName(root, 'from'); + + generic.commands.remove(root, cmd1!); + generic.commands.remove(root, cmd2!); + generic.commands.remove(root, cmd3!); + + expect(root.commands.length).toBe(0); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/generic/commands/index.ts b/packages/kbn-esql-ast/src/mutate/generic/commands/index.ts new file mode 100644 index 0000000000000..0582bb592edb8 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic/commands/index.ts @@ -0,0 +1,131 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { ESQLAstQueryExpression, ESQLCommand } from '../../../types'; +import { Visitor } from '../../../visitor'; +import { Predicate } from '../../types'; + +export * as args from './args'; +export * as options from './options'; + +/** + * Returns an iterator for all command AST nodes in the query. If a predicate is + * provided, only commands that satisfy the predicate will be returned. + * + * @param ast Root AST node to search for commands. + * @param predicate Optional predicate to filter commands. + * @returns A list of commands found in the AST. + */ +export const list = ( + ast: ESQLAstQueryExpression, + predicate?: Predicate +): IterableIterator => { + return new Visitor() + .on('visitQuery', function* (ctx): IterableIterator { + for (const cmd of ctx.commands()) { + if (!predicate || predicate(cmd)) { + yield cmd; + } + } + }) + .visitQuery(ast); +}; + +/** + * Returns the first command AST node at a given index in the query that + * satisfies the predicate. If no index is provided, the first command found + * will be returned. + * + * @param ast Root AST node to search for commands. + * @param predicate Optional predicate to filter commands. + * @param index The index of the command to return. + * @returns The command found in the AST, if any. + */ +export const find = ( + ast: ESQLAstQueryExpression, + predicate?: Predicate, + index: number = 0 +): ESQLCommand | undefined => { + for (const cmd of list(ast, predicate)) { + if (!index) { + return cmd; + } + + index--; + } + + return undefined; +}; + +/** + * Returns the first command AST node at a given index with a given name in the + * query. If no index is provided, the first command found will be returned. + * + * @param ast Root AST node to search for commands. + * @param commandName The name of the command to find. + * @param index The index of the command to return. + * @returns The command found in the AST, if any. + */ +export const findByName = ( + ast: ESQLAstQueryExpression, + commandName: string, + index: number = 0 +): ESQLCommand | undefined => { + return find(ast, (cmd) => cmd.name === commandName, index); +}; + +/** + * Inserts a new command into the query AST node at the specified index. If the + * `index` is out of bounds, the command will be appended to the end of the + * command list. + * + * @param ast The root AST node. + * @param command The command AST node to insert. + * @param index The index to insert the command at. + * @returns The index the command was inserted at. + */ +export const insert = ( + ast: ESQLAstQueryExpression, + command: ESQLCommand, + index: number = Infinity +): number => { + const commands = ast.commands; + + if (index > commands.length || index < 0) { + index = commands.length; + } + + commands.splice(index, 0, command); + + return index; +}; + +/** + * Adds a new command to the query AST node. + * + * @param ast The root AST node to append the command to. + * @param command The command AST node to append. + */ +export const append = (ast: ESQLAstQueryExpression, command: ESQLCommand): void => { + ast.commands.push(command); +}; + +export const remove = (ast: ESQLAstQueryExpression, command: ESQLCommand): boolean => { + const cmds = ast.commands; + const length = cmds.length; + + for (let i = 0; i < length; i++) { + if (cmds[i] === command) { + cmds.splice(i, 1); + return true; + } + } + + return false; +}; diff --git a/packages/kbn-esql-ast/src/mutate/generic/commands/options/index.test.ts b/packages/kbn-esql-ast/src/mutate/generic/commands/options/index.test.ts new file mode 100644 index 0000000000000..00c3ee90eccdd --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic/commands/options/index.test.ts @@ -0,0 +1,49 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { parse } from '../../../../parser'; +import { BasicPrettyPrinter } from '../../../../pretty_print'; +import * as generic from '../..'; + +describe('generic.commands.options', () => { + describe('.findByName()', () => { + it('can the find a command option', () => { + const src = 'FROM index METADATA _score'; + const { root } = parse(src); + const option = generic.commands.options.findByName(root, 'from', 'metadata'); + + expect(option).toMatchObject({ + type: 'option', + name: 'metadata', + }); + }); + + it('returns undefined if there is no option', () => { + const src = 'FROM index'; + const { root } = parse(src); + const option = generic.commands.options.findByName(root, 'from', 'metadata'); + + expect(option).toBe(undefined); + }); + }); + + describe('.remove()', () => { + it('can remove existing command option', () => { + const src = 'FROM index METADATA _score'; + const { root } = parse(src); + const option = generic.commands.options.findByName(root, 'from', 'metadata'); + + generic.commands.options.remove(root, option!); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index'); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/generic/commands/options/index.ts b/packages/kbn-esql-ast/src/mutate/generic/commands/options/index.ts new file mode 100644 index 0000000000000..b9b2bac452e31 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic/commands/options/index.ts @@ -0,0 +1,130 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +import { Builder } from '../../../../builder'; +import { ESQLAstQueryExpression, ESQLCommand, ESQLCommandOption } from '../../../../types'; +import { Visitor } from '../../../../visitor'; +import { Predicate } from '../../../types'; +import * as commands from '..'; + +/** + * Returns the first command option AST node that satisfies the predicate. + * + * @param command The command AST node to search for options. + * @param predicate The predicate to filter options. + * @returns The option found in the command, if any. + */ +export const find = ( + command: ESQLCommand, + predicate: Predicate +): ESQLCommandOption | undefined => { + return new Visitor() + .on('visitCommand', (ctx): ESQLCommandOption | undefined => { + for (const opt of ctx.options()) { + if (predicate(opt)) { + return opt; + } + } + + return undefined; + }) + .visitCommand(command); +}; + +/** + * Returns the first command option AST node with a given name in the query. + * + * @param ast The root AST node to search for command options. + * @param commandName Command name to search for. + * @param optionName Option name to search for. + * @returns The option found in the command, if any. + */ +export const findByName = ( + ast: ESQLAstQueryExpression, + commandName: string, + optionName: string +): ESQLCommandOption | undefined => { + const command = commands.find(ast, (cmd) => cmd.name === commandName); + + if (!command) { + return undefined; + } + + return find(command, (opt) => opt.name === optionName); +}; + +/** + * Inserts a command option into the command's arguments list. The option can + * be specified as a string or an AST node. + * + * @param command The command AST node to insert the option into. + * @param option The option to insert. + * @returns The inserted option. + */ +export const append = ( + command: ESQLCommand, + option: string | ESQLCommandOption +): ESQLCommandOption => { + if (typeof option === 'string') { + option = Builder.option({ name: option }); + } + + command.args.push(option); + + return option; +}; + +/** + * Removes the first command option from the command's arguments list that + * satisfies the predicate. + * + * @param command The command AST node to remove the option from. + * @param predicate The predicate to filter options. + * @returns The removed option, if any. + */ +export const remove = (ast: ESQLAstQueryExpression, option: ESQLCommandOption): boolean => { + return new Visitor() + .on('visitCommandOption', (ctx): boolean => { + return ctx.node === option; + }) + .on('visitCommand', (ctx): boolean => { + let target: undefined | ESQLCommandOption; + + for (const opt of ctx.options()) { + if (opt === option) { + target = opt; + break; + } + } + + if (!target) { + return false; + } + + const index = ctx.node.args.indexOf(target); + + if (index === -1) { + return false; + } + + ctx.node.args.splice(index, 1); + + return true; + }) + .on('visitQuery', (ctx): boolean => { + for (const success of ctx.visitCommands()) { + if (success) { + return true; + } + } + + return false; + }) + .visitQuery(ast); +}; diff --git a/packages/kbn-esql-ast/src/mutate/generic/index.ts b/packages/kbn-esql-ast/src/mutate/generic/index.ts new file mode 100644 index 0000000000000..e7f26b9340af7 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/generic/index.ts @@ -0,0 +1,10 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the "Elastic License + * 2.0", the "GNU Affero General Public License v3.0 only", and the "Server Side + * Public License v 1"; you may not use this file except in compliance with, at + * your election, the "Elastic License 2.0", the "GNU Affero General Public + * License v3.0 only", or the "Server Side Public License, v 1". + */ + +export * as commands from './commands'; diff --git a/packages/kbn-esql-ast/src/types.ts b/packages/kbn-esql-ast/src/types.ts index 0df75ee2e8f24..eabdefa5a401a 100644 --- a/packages/kbn-esql-ast/src/types.ts +++ b/packages/kbn-esql-ast/src/types.ts @@ -404,7 +404,7 @@ export interface ESQLAstGenericComment; From a1a29b4ebbf31b74cb7f9b9bcb331facc3a7ae36 Mon Sep 17 00:00:00 2001 From: vadimkibana Date: Thu, 24 Oct 2024 21:16:41 +0200 Subject: [PATCH 5/6] do not attempt to attach comment nodes if location information not available --- packages/kbn-esql-ast/src/parser/formatting.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/kbn-esql-ast/src/parser/formatting.ts b/packages/kbn-esql-ast/src/parser/formatting.ts index 492e8a76ddeac..f7c556da63008 100644 --- a/packages/kbn-esql-ast/src/parser/formatting.ts +++ b/packages/kbn-esql-ast/src/parser/formatting.ts @@ -173,6 +173,10 @@ const attachCommentDecoration = ( ) => { const commentConsumesWholeLine = !comment.hasContentToLeft && !comment.hasContentToRight; + if (!comment.node.location) { + return; + } + if (commentConsumesWholeLine) { const node = Visitor.findNodeAtOrAfter(ast, comment.node.location.max - 1); From 79e64bba4d2b4e044bdd8f90f99d7cb623ed4418 Mon Sep 17 00:00:00 2001 From: vadimkibana Date: Fri, 25 Oct 2024 13:41:11 +0200 Subject: [PATCH 6/6] fix typescript type error --- .../public/components/esql_inspector/helpers.tsx | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/esql_ast_inspector/public/components/esql_inspector/helpers.tsx b/examples/esql_ast_inspector/public/components/esql_inspector/helpers.tsx index a117062f7efa9..19a0c54a722c6 100644 --- a/examples/esql_ast_inspector/public/components/esql_inspector/helpers.tsx +++ b/examples/esql_ast_inspector/public/components/esql_inspector/helpers.tsx @@ -82,6 +82,8 @@ export const highlight = (query: EsqlQuery): Annotation[] => { }); Walker.visitComments(query.ast, (comment) => { + if (!comment.location) return; + annotations.push([ comment.location.min, comment.location.max,