From 0846fa8a9e01aff0f1b0177b868b8047d7fb4be7 Mon Sep 17 00:00:00 2001 From: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com> Date: Wed, 30 Oct 2024 21:34:57 +0100 Subject: [PATCH] [ES|QL] Add`SORT` command mutation APIs (#197185) ## Summary Partially addresses https://github.com/elastic/kibana/issues/191812 - Adds traversal and manipulation APIs for `SORT` command. - `commands.sort.listCommands()` - `commands.sort.getCommand()` - `commands.sort.list()` - `commands.sort.findByPredicate()` - `commands.sort.find()` - `commands.sort.remove()` - `commands.sort.insertIntoCommand()` - `commands.sort.insertExpression()` - `commands.sort.insertCommand()` - Refactors "generic" AST manipulation routines into (1) `commands`, (2) `commands.args`, (3) `commands.options`. - `generic.commands.*` - `generic.commands.args.*` - `generic.commands.options.*` ### Checklist Delete any items that are not applicable to this PR. - [x] [Unit or functional tests](https://www.elastic.co/guide/en/kibana/master/development-tests.html) were updated or added to match the most common scenarios ### For maintainers - [x] This was checked for breaking API changes and was [labeled appropriately](https://www.elastic.co/guide/en/kibana/master/contributing.html#_add_your_labels) --------- Co-authored-by: Stratoula Kalafateli (cherry picked from commit 0cc89457124f9cefbdd08766a7a22b1c4bdcc440) --- .../components/esql_inspector/helpers.tsx | 2 + packages/kbn-esql-ast/src/builder/builder.ts | 27 +- .../src/mutate/commands/from/metadata.ts | 10 +- .../src/mutate/commands/from/sources.ts | 6 +- .../kbn-esql-ast/src/mutate/commands/index.ts | 3 +- .../src/mutate/commands/limit/index.ts | 8 +- .../src/mutate/commands/sort/index.test.ts | 527 ++++++++++++++++++ .../src/mutate/commands/sort/index.ts | 313 +++++++++++ packages/kbn-esql-ast/src/mutate/generic.ts | 287 ---------- .../generic/commands/args/index.test.ts | 132 +++++ .../src/mutate/generic/commands/args/index.ts | 86 +++ .../commands/index.test.ts} | 77 +-- .../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/parser/factories.ts | 17 +- .../kbn-esql-ast/src/parser/formatting.ts | 4 + packages/kbn-esql-ast/src/parser/walkers.ts | 2 +- packages/kbn-esql-ast/src/types.ts | 2 +- 20 files changed, 1449 insertions(+), 374 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 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 rename packages/kbn-esql-ast/src/mutate/{generic.test.ts => generic/commands/index.test.ts} (54%) 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/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, diff --git a/packages/kbn-esql-ast/src/builder/builder.ts b/packages/kbn-esql-ast/src/builder/builder.ts index 26b64a6312ee4..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, @@ -20,6 +22,7 @@ import { ESQLIntegerLiteral, ESQLList, ESQLLocation, + ESQLOrderExpression, ESQLSource, } from '../types'; import { AstNodeParserFields, AstNodeTemplate, PartialFields } from './types'; @@ -63,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 { @@ -130,6 +133,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/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/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..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 new file mode 100644 index 0000000000000..d04f79b96541a --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.test.ts @@ -0,0 +1,527 @@ +/* + * 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'; +import { Builder } from '../../../builder'; + +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)].map(([node]) => node); + + 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)].map(([node]) => node); + + 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)].map(([node]) => node); + + 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)].map(([node]) => node); + + 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'], + }); + }); + + 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); + }); + }); + + 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'); + }); + }); + + 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 new file mode 100644 index 0000000000000..d2b2c7cd5f3d4 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/sort/index.ts @@ -0,0 +1,313 @@ +/* + * 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 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. + * + * @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); +}; + +/** + * 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 +): IterableIterator<[sortExpression: SortExpression, sortCommand: ESQLCommand]> => { + return new Visitor() + .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, ctx.node]; + } + } + } + }) + .on('visitCommand', function* (): IterableIterator<[SortExpression, ESQLCommand]> {}) + .on('visitQuery', function* (ctx): IterableIterator<[SortExpression, ESQLCommand]> { + for (const command of ctx.visitCommands()) { + yield* command; + } + }) + .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]>, + index?: number +): [sortExpression: SortExpression, sortCommand: ESQLCommand] | undefined => { + 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[], + index: number = 0 +): [sortExpression: SortExpression, sortCommand: ESQLCommand] | 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; + }); +}; + +/** + * 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[], + index?: number +): [sortExpression: SortExpression, sortCommand: ESQLCommand] | undefined => { + const tuple = find(ast, parts, index); + + if (!tuple) { + return undefined; + } + + const [node] = tuple; + const cmd = generic.commands.args.remove(ast, node); + + if (cmd) { + if (!cmd.args.length) { + generic.commands.remove(ast, cmd); + } + } + + return cmd ? tuple : undefined; +}; + +/** + * 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.ts b/packages/kbn-esql-ast/src/mutate/generic.ts deleted file mode 100644 index f27b0e2ae399f..0000000000000 --- a/packages/kbn-esql-ast/src/mutate/generic.ts +++ /dev/null @@ -1,287 +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 appendCommandArgument = ( - command: ESQLCommand, - expression: ESQLSingleAstItem -): number => { - if (expression.type === 'option') { - command.args.push(expression); - return command.args.length - 1; - } - - const index = command.args.findIndex((arg) => isOptionNode(arg)); - - if (index > -1) { - command.args.splice(index, 0, expression); - return index; - } - - command.args.push(expression); - return command.args.length - 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 true if the argument was removed, false otherwise. - */ -export const removeCommandArgument = ( - ast: ESQLAstQueryExpression, - node: ESQLProperNode -): boolean => { - return new Visitor() - .on('visitCommand', (ctx): boolean => { - 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 false; - }) - .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/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.test.ts b/packages/kbn-esql-ast/src/mutate/generic/commands/index.test.ts similarity index 54% rename from packages/kbn-esql-ast/src/mutate/generic.test.ts rename to packages/kbn-esql-ast/src/mutate/generic/commands/index.test.ts index 0109ff838ffda..b35d0b6415247 100644 --- a/packages/kbn-esql-ast/src/mutate/generic.test.ts +++ b/packages/kbn-esql-ast/src/mutate/generic/commands/index.test.ts @@ -7,26 +7,26 @@ * 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 './generic'; +import { parse } from '../../../parser'; +import { BasicPrettyPrinter } from '../../../pretty_print'; +import * as generic from '..'; -describe('generic', () => { - describe('.listCommands()', () => { +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.listCommands(root)].map((cmd) => cmd.name); + const commands = [...generic.commands.list(root)].map((cmd) => cmd.name); expect(commands).toEqual(['from', 'where', 'limit']); }); }); - describe('.findCommand()', () => { + describe('.find()', () => { 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'); + const command = generic.commands.find(root, (cmd) => cmd.name === 'from'); expect(command).toMatchObject({ type: 'command', @@ -42,7 +42,7 @@ describe('generic', () => { 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'); + const command = generic.commands.find(root, (cmd) => cmd.name === 'limit'); expect(command).toMatchObject({ type: 'command', @@ -58,7 +58,7 @@ describe('generic', () => { 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( + const command = generic.commands.find( root, (cmd) => cmd.name === 'limit' && (cmd.args?.[0] as any).value === 2 ); @@ -76,34 +76,13 @@ describe('generic', () => { }); }); - 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('.removeCommand()', () => { + describe('.remove()', () => { it('can remove the last command', () => { const src = 'FROM index | LIMIT 10'; const { root } = parse(src); - const command = generic.findCommandByName(root, 'limit', 0); + const command = generic.commands.findByName(root, 'limit', 0); - generic.removeCommand(root, command!); + generic.commands.remove(root, command!); const src2 = BasicPrettyPrinter.print(root); @@ -113,9 +92,9 @@ describe('generic', () => { 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); + const command = generic.commands.findByName(root, 'limit', 1); - generic.removeCommand(root, command!); + generic.commands.remove(root, command!); const src2 = BasicPrettyPrinter.print(root); @@ -125,29 +104,15 @@ describe('generic', () => { 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'); + const cmd1 = generic.commands.findByName(root, 'where'); + const cmd2 = generic.commands.findByName(root, 'limit'); + const cmd3 = generic.commands.findByName(root, 'from'); - generic.removeCommand(root, cmd1!); - generic.removeCommand(root, cmd2!); - generic.removeCommand(root, cmd3!); + generic.commands.remove(root, cmd1!); + generic.commands.remove(root, cmd2!); + generic.commands.remove(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/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/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/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); 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( 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;