From f0845d55d2257a353643015efbac09f74e1426d4 Mon Sep 17 00:00:00 2001 From: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com> Date: Tue, 15 Oct 2024 17:45:03 +0200 Subject: [PATCH] [ES|QL] More AST mutation APIs (#196240) ## Summary Partially addresses https://github.com/elastic/kibana/issues/191812 Implements the following high-level ES|QL AST manipulation methods: - `.generic` - `.appendCommandArgument()` — Add a new main command argument to a command. - `.removeCommandArgument()` — Remove a command argument from the AST. - `.commands` - `.from` - `.sources` - `.list()` — List all `FROM` sources. - `.find()` — Find a source by name. - `.remove()` — Remove a source by name. - `.insert()` — Insert a source. - `.upsert()` — Insert a source, if it does not exist. - `.limit` - `.list()` — List all `LIMIT` commands. - `.byIndex()` — Find a `LIMIT` command by index. - `.find()` — Find a `LIMIT` command by a predicate function. - `.remove()` — Remove a `LIMIT` command by index. - `.set()` — Set the limit value of a specific `LIMIT` command. - `.upsert()` — Insert a `LIMIT` command, or update the limit value if it already exists. ### Checklist Delete any items that are not applicable to this PR. - [x] [Documentation](https://www.elastic.co/guide/en/kibana/master/development-documentation.html) was added for features that require explanation or tutorials - [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) (cherry picked from commit 10364fba2db8bb2080a97173c76a9d1aef1e80ed) --- packages/kbn-esql-ast/src/ast/util.ts | 14 + packages/kbn-esql-ast/src/builder/builder.ts | 17 + packages/kbn-esql-ast/src/mutate/README.md | 42 ++- .../src/mutate/commands/from/index.ts | 3 +- .../src/mutate/commands/from/metadata.ts | 2 +- .../src/mutate/commands/from/sources.test.ts | 246 ++++++++++++++ .../src/mutate/commands/from/sources.ts | 111 +++++++ .../kbn-esql-ast/src/mutate/commands/index.ts | 3 +- .../src/mutate/commands/limit/index.test.ts | 311 ++++++++++++++++++ .../src/mutate/commands/limit/index.ts | 134 ++++++++ .../kbn-esql-ast/src/mutate/generic.test.ts | 40 +++ packages/kbn-esql-ast/src/mutate/generic.ts | 93 +++++- 12 files changed, 1003 insertions(+), 13 deletions(-) create mode 100644 packages/kbn-esql-ast/src/ast/util.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/from/sources.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/from/sources.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/limit/index.test.ts create mode 100644 packages/kbn-esql-ast/src/mutate/commands/limit/index.ts diff --git a/packages/kbn-esql-ast/src/ast/util.ts b/packages/kbn-esql-ast/src/ast/util.ts new file mode 100644 index 0000000000000..0cd94aba85cf1 --- /dev/null +++ b/packages/kbn-esql-ast/src/ast/util.ts @@ -0,0 +1,14 @@ +/* + * 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 { ESQLAstNode, ESQLCommandOption } from '../types'; + +export const isOptionNode = (node: ESQLAstNode): node is ESQLCommandOption => { + return !!node && typeof node === 'object' && !Array.isArray(node) && node.type === 'option'; +}; diff --git a/packages/kbn-esql-ast/src/builder/builder.ts b/packages/kbn-esql-ast/src/builder/builder.ts index ece92fbcd7d5e..26b64a6312ee4 100644 --- a/packages/kbn-esql-ast/src/builder/builder.ts +++ b/packages/kbn-esql-ast/src/builder/builder.ts @@ -100,6 +100,23 @@ export namespace Builder { }; }; + export const indexSource = ( + index: string, + cluster?: string, + template?: Omit, 'name' | 'index' | 'cluster'>, + fromParser?: Partial + ): ESQLSource => { + return { + ...template, + ...Builder.parserFields(fromParser), + index, + cluster, + name: (cluster ? cluster + ':' : '') + index, + sourceType: 'index', + type: 'source', + }; + }; + export const column = ( template: Omit, 'name' | 'quoted'>, fromParser?: Partial diff --git a/packages/kbn-esql-ast/src/mutate/README.md b/packages/kbn-esql-ast/src/mutate/README.md index 8c38bb72ca226..7dfd3d77a1395 100644 --- a/packages/kbn-esql-ast/src/mutate/README.md +++ b/packages/kbn-esql-ast/src/mutate/README.md @@ -26,11 +26,37 @@ console.log(src); // FROM index METADATA _lang, _id ## API -- `.commands.from.metadata.list()` — List all `METADATA` fields. -- `.commands.from.metadata.find()` — Find a `METADATA` field by name. -- `.commands.from.metadata.removeByPredicate()` — Remove a `METADATA` - field by matching a predicate. -- `.commands.from.metadata.remove()` — Remove a `METADATA` field by name. -- `.commands.from.metadata.insert()` — Insert a `METADATA` field. -- `.commands.from.metadata.upsert()` — Insert `METADATA` field, if it does - not exist. +- `.generic` + - `.listCommands()` — Lists all commands. Returns an iterator. + - `.findCommand()` — Finds a specific command by a predicate function. + - `.findCommandOption()` — Finds a specific command option by a predicate function. + - `.findCommandByName()` — Finds a specific command by name. + - `.findCommandOptionByName()` — Finds a specific command option by name. + - `.appendCommand()` — Add a new command to the AST. + - `.appendCommandOption()` — Add a new command option to a command. + - `.appendCommandArgument()` — Add a new main command argument to a command. + - `.removeCommand()` — Remove a command from the AST. + - `.removeCommandOption()` — Remove a command option from the AST. + - `.removeCommandArgument()` — Remove a command argument from the AST. +- `.commands` + - `.from` + - `.sources` + - `.list()` — List all `FROM` sources. + - `.find()` — Find a source by name. + - `.remove()` — Remove a source by name. + - `.insert()` — Insert a source. + - `.upsert()` — Insert a source, if it does not exist. + - `.metadata` + - `.list()` — List all `METADATA` fields. + - `.find()` — Find a `METADATA` field by name. + - `.removeByPredicate()` — Remove a `METADATA` field by matching a predicate function. + - `.remove()` — Remove a `METADATA` field by name. + - `.insert()` — Insert a `METADATA` field. + - `.upsert()` — Insert `METADATA` field, if it does not exist. + - `.limit` + - `.list()` — List all `LIMIT` commands. + - `.byIndex()` — Find a `LIMIT` command by index. + - `.find()` — Find a `LIMIT` command by a predicate function. + - `.remove()` — Remove a `LIMIT` command by index. + - `.set()` — Set the limit value of a specific `LIMIT` command. + - `.upsert()` — Insert a `LIMIT` command, or update the limit value if it already exists. diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/index.ts b/packages/kbn-esql-ast/src/mutate/commands/from/index.ts index df76e072b346e..2a86a43dbe8d1 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/from/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/from/index.ts @@ -7,6 +7,7 @@ * License v3.0 only", or the "Server Side Public License, v 1". */ +import * as sources from './sources'; import * as metadata from './metadata'; -export { metadata }; +export { sources, metadata }; 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 5892b028823aa..7f08fa2a5e946 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/from/metadata.ts @@ -157,7 +157,7 @@ export const insert = ( return; } - option = generic.insertCommandOption(command, 'metadata'); + option = generic.appendCommandOption(command, 'metadata'); } const parts: string[] = typeof fieldName === 'string' ? [fieldName] : fieldName; diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/sources.test.ts b/packages/kbn-esql-ast/src/mutate/commands/from/sources.test.ts new file mode 100644 index 0000000000000..866a6dd8bdb20 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/from/sources.test.ts @@ -0,0 +1,246 @@ +/* + * 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 commands from '..'; + +describe('commands.from.sources', () => { + describe('.list()', () => { + it('returns empty array, if there are no sources', () => { + const src = 'ROW 123'; + const { root } = parse(src); + const list = [...commands.from.sources.list(root)]; + + expect(list.length).toBe(0); + }); + + it('returns a single source', () => { + const src = 'FROM index METADATA a'; + const { root } = parse(src); + const list = [...commands.from.sources.list(root)]; + + expect(list.length).toBe(1); + expect(list[0]).toMatchObject({ + type: 'source', + }); + }); + + it('returns all source fields', () => { + const src = 'FROM index, index2, cl:index3 METADATA a | LIMIT 88'; + const { root } = parse(src); + const list = [...commands.from.sources.list(root)]; + + expect(list).toMatchObject([ + { + type: 'source', + index: 'index', + }, + { + type: 'source', + index: 'index2', + }, + { + type: 'source', + index: 'index3', + cluster: 'cl', + }, + ]); + }); + }); + + describe('.find()', () => { + it('returns undefined if source is not found', () => { + const src = 'FROM index | WHERE a = b | LIMIT 123'; + const { root } = parse(src); + const source = commands.from.sources.find(root, 'abc'); + + expect(source).toBe(undefined); + }); + + it('can find a single source', () => { + const src = 'FROM index METADATA a'; + const { root } = parse(src); + const source = commands.from.sources.find(root, 'index')!; + + expect(source).toMatchObject({ + type: 'source', + name: 'index', + index: 'index', + }); + }); + + it('can find a source withing other sources', () => { + const src = 'FROM index, a, b, c:s1, s1, s2 METADATA a, b, c, _lang, _id'; + const { root } = parse(src); + const source1 = commands.from.sources.find(root, 's2')!; + const source2 = commands.from.sources.find(root, 's1', 'c')!; + + expect(source1).toMatchObject({ + type: 'source', + name: 's2', + index: 's2', + }); + expect(source2).toMatchObject({ + type: 'source', + name: 'c:s1', + index: 's1', + cluster: 'c', + }); + }); + }); + + describe('.remove()', () => { + it('can remove a source from a list', () => { + const src1 = 'FROM a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c'); + + commands.from.sources.remove(root, 'b'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, c'); + }); + + it('does nothing if source-to-delete does not exist', () => { + const src1 = 'FROM a, b, c'; + const { root } = parse(src1); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a, b, c'); + + commands.from.sources.remove(root, 'd'); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM a, b, c'); + }); + }); + + describe('.insert()', () => { + it('can append a source', () => { + const src1 = 'FROM index METADATA a'; + const { root } = parse(src1); + + commands.from.sources.insert(root, 'index2'); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, index2 METADATA a'); + }); + + it('can insert at specified position', () => { + const src1 = 'FROM a1, a2, a3'; + const { root } = parse(src1); + + commands.from.sources.insert(root, 'x', '', 0); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM x, a1, a2, a3'); + + commands.from.sources.insert(root, 'y', '', 2); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM x, a1, y, a2, a3'); + + commands.from.sources.insert(root, 'z', '', 4); + + const src4 = BasicPrettyPrinter.print(root); + + expect(src4).toBe('FROM x, a1, y, a2, z, a3'); + }); + + it('appends element, when insert position too high', () => { + const src1 = 'FROM a1, a2, a3'; + const { root } = parse(src1); + + commands.from.sources.insert(root, 'x', '', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a1, a2, a3, x'); + }); + + it('can inset the same source twice', () => { + const src1 = 'FROM index'; + const { root } = parse(src1); + + commands.from.sources.insert(root, 'x', '', 999); + commands.from.sources.insert(root, 'x', '', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, x, x'); + }); + }); + + describe('.upsert()', () => { + it('can append a source', () => { + const src1 = 'FROM index METADATA a'; + const { root } = parse(src1); + + commands.from.sources.upsert(root, 'index2'); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, index2 METADATA a'); + }); + + it('can upsert at specified position', () => { + const src1 = 'FROM a1, a2, a3'; + const { root } = parse(src1); + + commands.from.sources.upsert(root, 'x', '', 0); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM x, a1, a2, a3'); + + commands.from.sources.upsert(root, 'y', '', 2); + + const src3 = BasicPrettyPrinter.print(root); + + expect(src3).toBe('FROM x, a1, y, a2, a3'); + + commands.from.sources.upsert(root, 'z', '', 4); + + const src4 = BasicPrettyPrinter.print(root); + + expect(src4).toBe('FROM x, a1, y, a2, z, a3'); + }); + + it('appends element, when upsert position too high', () => { + const src1 = 'FROM a1, a2, a3'; + const { root } = parse(src1); + + commands.from.sources.upsert(root, 'x', '', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM a1, a2, a3, x'); + }); + + it('inserting already existing source is a no-op', () => { + const src1 = 'FROM index'; + const { root } = parse(src1); + + commands.from.sources.upsert(root, 'x', '', 999); + commands.from.sources.upsert(root, 'x', '', 999); + + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index, x'); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts b/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts new file mode 100644 index 0000000000000..da67500b5b0bd --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts @@ -0,0 +1,111 @@ +/* + * 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, ESQLSource } from '../../../types'; +import { Visitor } from '../../../visitor'; +import * as generic from '../../generic'; +import * as util from '../../util'; +import type { Predicate } from '../../types'; + +export const list = (ast: ESQLAstQueryExpression): IterableIterator => { + return new Visitor() + .on('visitFromCommand', function* (ctx): IterableIterator { + for (const argument of ctx.arguments()) { + if (argument.type === 'source') { + 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 +): ESQLSource | undefined => { + return util.findByPredicate(list(ast), predicate); +}; + +export const find = ( + ast: ESQLAstQueryExpression, + index: string, + cluster?: string +): ESQLSource | undefined => { + return findByPredicate(ast, (source) => { + if (index !== source.index) { + return false; + } + if (typeof cluster === 'string' && cluster !== source.cluster) { + return false; + } + + return true; + }); +}; + +export const remove = ( + ast: ESQLAstQueryExpression, + index: string, + cluster?: string +): ESQLSource | undefined => { + const node = find(ast, index, cluster); + + if (!node) { + return undefined; + } + + const success = generic.removeCommandArgument(ast, node); + + return success ? 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; +}; + +export const upsert = ( + ast: ESQLAstQueryExpression, + indexName: string, + clusterName?: string, + index: number = -1 +): ESQLSource | undefined => { + const source = find(ast, indexName, clusterName); + + if (source) { + return source; + } + + return insert(ast, indexName, clusterName, index); +}; diff --git a/packages/kbn-esql-ast/src/mutate/commands/index.ts b/packages/kbn-esql-ast/src/mutate/commands/index.ts index cc3b7f446fa88..0a779292e6eca 100644 --- a/packages/kbn-esql-ast/src/mutate/commands/index.ts +++ b/packages/kbn-esql-ast/src/mutate/commands/index.ts @@ -8,5 +8,6 @@ */ import * as from from './from'; +import * as limit from './limit'; -export { from }; +export { from, limit }; diff --git a/packages/kbn-esql-ast/src/mutate/commands/limit/index.test.ts b/packages/kbn-esql-ast/src/mutate/commands/limit/index.test.ts new file mode 100644 index 0000000000000..9d734055cfeff --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/limit/index.test.ts @@ -0,0 +1,311 @@ +/* + * 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 commands from '..'; + +describe('commands.limit', () => { + describe('.list()', () => { + it('lists all "LIMIT" commands', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const nodes = [...commands.limit.list(root)]; + + expect(nodes).toMatchObject([ + { + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }, + { + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 2, + }, + ], + }, + { + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 3, + }, + ], + }, + ]); + }); + }); + + describe('.byIndex()', () => { + it('retrieves the specific "LIMIT" command by index', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const node = commands.limit.byIndex(root, 1); + + expect(node).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 2, + }, + ], + }); + }); + }); + + describe('.find()', () => { + it('can find a limit command by predicate', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const node = commands.limit.find(root, (cmd) => (cmd.args?.[0] as any).value === 3); + + expect(node).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 3, + }, + ], + }); + }); + }); + + describe('.remove()', () => { + it('can remove the only limit command', () => { + const src = 'FROM index | WHERE a == b | LIMIT 123'; + const { root } = parse(src); + + const node = commands.limit.remove(root); + const src2 = BasicPrettyPrinter.print(root); + + expect(node).toMatchObject({ + type: 'command', + name: 'limit', + }); + expect(src2).toBe('FROM index | WHERE a == b'); + }); + + it('can remove the specific limit node', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const node1 = commands.limit.remove(root, 1); + const src1 = BasicPrettyPrinter.print(root); + + expect(node1).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 2, + }, + ], + }); + expect(src1).toBe('FROM index | LIMIT 1 | STATS AGG() | WHERE a == b | LIMIT 3'); + + const node2 = commands.limit.remove(root); + const src2 = BasicPrettyPrinter.print(root); + + expect(node2).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 1, + }, + ], + }); + expect(src2).toBe('FROM index | STATS AGG() | WHERE a == b | LIMIT 3'); + + const node3 = commands.limit.remove(root); + const src3 = BasicPrettyPrinter.print(root); + + expect(node3).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 3, + }, + ], + }); + expect(src3).toBe('FROM index | STATS AGG() | WHERE a == b'); + + const node4 = commands.limit.remove(root); + + expect(node4).toBe(undefined); + }); + }); + + describe('.set()', () => { + it('can update a specific LIMIT command', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const node1 = commands.limit.set(root, 2222, 1); + const node2 = commands.limit.set(root, 3333, 2); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe( + 'FROM index | LIMIT 1 | STATS AGG() | LIMIT 2222 | WHERE a == b | LIMIT 3333' + ); + expect(node1).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 2222, + }, + ], + }); + expect(node2).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 3333, + }, + ], + }); + }); + + it('by default, updates the first LIMIT command', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const node = commands.limit.set(root, 99999999); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe( + 'FROM index | LIMIT 99999999 | STATS AGG() | LIMIT 2 | WHERE a == b | LIMIT 3' + ); + expect(node).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 99999999, + }, + ], + }); + }); + + it('does nothing if there is no existing limit command', () => { + const src = 'FROM index | STATS agg() | WHERE a == b'; + const { root } = parse(src); + + const node = commands.limit.set(root, 99999999); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index | STATS AGG() | WHERE a == b'); + expect(node).toBe(undefined); + }); + }); + + describe('.upsert()', () => { + it('can update a specific LIMIT command', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const node1 = commands.limit.upsert(root, 2222, 1); + const node2 = commands.limit.upsert(root, 3333, 2); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe( + 'FROM index | LIMIT 1 | STATS AGG() | LIMIT 2222 | WHERE a == b | LIMIT 3333' + ); + expect(node1).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 2222, + }, + ], + }); + expect(node2).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 3333, + }, + ], + }); + }); + + it('by default, updates the first LIMIT command', () => { + const src = 'FROM index | LIMIT 1 | STATS agg() | LIMIT 2 | WHERE a == b | LIMIT 3'; + const { root } = parse(src); + + const node = commands.limit.upsert(root, 99999999); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe( + 'FROM index | LIMIT 99999999 | STATS AGG() | LIMIT 2 | WHERE a == b | LIMIT 3' + ); + expect(node).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 99999999, + }, + ], + }); + }); + + it('inserts a new LIMIT command, if there is none existing', () => { + const src = 'FROM index | STATS agg() | WHERE a == b'; + const { root } = parse(src); + + const node = commands.limit.upsert(root, 99999999); + const src2 = BasicPrettyPrinter.print(root); + + expect(src2).toBe('FROM index | STATS AGG() | WHERE a == b | LIMIT 99999999'); + expect(node).toMatchObject({ + type: 'command', + name: 'limit', + args: [ + { + type: 'literal', + value: 99999999, + }, + ], + }); + }); + }); +}); diff --git a/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts b/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts new file mode 100644 index 0000000000000..937538e848328 --- /dev/null +++ b/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts @@ -0,0 +1,134 @@ +/* + * 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 type { ESQLAstQueryExpression, ESQLCommand } from '../../../types'; +import * as generic from '../../generic'; +import { Predicate } from '../../types'; + +/** + * Lists all "LIMIT" commands in the query AST. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @returns A collection of "LIMIT" commands. + */ +export const list = (ast: ESQLAstQueryExpression): IterableIterator => { + return generic.listCommands(ast, (cmd) => cmd.name === 'limit'); +}; + +/** + * Retrieves the "LIMIT" command at the specified index in order of appearance. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @param index The index of the "LIMIT" command to retrieve. + * @returns The "LIMIT" command at the specified index, if any. + */ +export const byIndex = (ast: ESQLAstQueryExpression, index: number): ESQLCommand | undefined => { + return [...list(ast)][index]; +}; + +/** + * Finds the first "LIMIT" command that satisfies the provided predicate. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @param predicate The predicate function to apply to each "LIMIT" command. + * @returns The first "LIMIT" command that satisfies the predicate, if any. + */ +export const find = ( + ast: ESQLAstQueryExpression, + predicate: Predicate +): ESQLCommand | undefined => { + return [...list(ast)].find(predicate); +}; + +/** + * Deletes the specified "LIMIT" command from the query AST. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @param index The index of the "LIMIT" command to remove. + * @returns The removed "LIMIT" command, if any. + */ +export const remove = (ast: ESQLAstQueryExpression, index: number = 0): ESQLCommand | undefined => { + const command = generic.findCommandByName(ast, 'limit', index); + + if (!command) { + return; + } + + const success = generic.removeCommand(ast, command); + + if (!success) { + return; + } + + return command; +}; + +/** + * Sets the value of the specified "LIMIT" command. If `indexOrPredicate` is not + * specified will update the first "LIMIT" command found, if any. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @param value The new value to set. + * @param indexOrPredicate The index of the "LIMIT" command to update, or a + * predicate function. + * @returns The updated "LIMIT" command, if any. + */ +export const set = ( + ast: ESQLAstQueryExpression, + value: number, + indexOrPredicate: number | Predicate = 0 +): ESQLCommand | undefined => { + const node = + typeof indexOrPredicate === 'number' + ? byIndex(ast, indexOrPredicate) + : find(ast, indexOrPredicate); + + if (!node) { + return; + } + + const literal = Builder.expression.literal.numeric({ literalType: 'integer', value }); + + node.args = [literal]; + + return node; +}; + +/** + * Updates the value of the specified "LIMIT" command. If the "LIMIT" command + * is not found, a new one will be created and appended to the query AST. + * + * @param ast The root AST node to search for "LIMIT" commands. + * @param value The new value to set. + * @param indexOrPredicate The index of the "LIMIT" command to update, or a + * predicate function. + * @returns The updated or newly created "LIMIT" command. + */ +export const upsert = ( + ast: ESQLAstQueryExpression, + value: number, + indexOrPredicate: number | Predicate = 0 +): ESQLCommand => { + const node = set(ast, value, indexOrPredicate); + + if (node) { + return node; + } + + const literal = Builder.expression.literal.numeric({ literalType: 'integer', value }); + const command = Builder.command({ + name: 'limit', + args: [literal], + }); + + generic.appendCommand(ast, command); + + return command; +}; diff --git a/packages/kbn-esql-ast/src/mutate/generic.test.ts b/packages/kbn-esql-ast/src/mutate/generic.test.ts index 14d951db1bccb..0109ff838ffda 100644 --- a/packages/kbn-esql-ast/src/mutate/generic.test.ts +++ b/packages/kbn-esql-ast/src/mutate/generic.test.ts @@ -97,6 +97,46 @@ describe('generic', () => { }); }); + 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'; diff --git a/packages/kbn-esql-ast/src/mutate/generic.ts b/packages/kbn-esql-ast/src/mutate/generic.ts index 968eaf84f4a46..f27b0e2ae399f 100644 --- a/packages/kbn-esql-ast/src/mutate/generic.ts +++ b/packages/kbn-esql-ast/src/mutate/generic.ts @@ -7,8 +7,15 @@ * 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 } from '../types'; +import { + ESQLAstQueryExpression, + ESQLCommand, + ESQLCommandOption, + ESQLProperNode, + ESQLSingleAstItem, +} from '../types'; import { Visitor } from '../visitor'; import { Predicate } from './types'; @@ -124,6 +131,16 @@ export const findCommandOptionByName = ( 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. @@ -132,7 +149,7 @@ export const findCommandOptionByName = ( * @param option The option to insert. * @returns The inserted option. */ -export const insertCommandOption = ( +export const appendCommandOption = ( command: ESQLCommand, option: string | ESQLCommandOption ): ESQLCommandOption => { @@ -145,6 +162,40 @@ export const insertCommandOption = ( 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. @@ -196,3 +247,41 @@ export const removeCommandOption = ( }) .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); +};