[8.x] [ES|QL] More AST mutation APIs (#196240) (#196355)

# Backport This will backport the following commits from `main` to `8.x`: - [[ES|QL] More AST mutation APIs (#196240)](https://github.com/elastic/kibana/pull/196240)  ### Questions ? Please refer to the [Backport tool documentation](https://github.com/sqren/backport)  Co-authored-by: Vadim Kibana <82822460+vadimkibana@users.noreply.github.com>
2025-04-24 17:59:23 -04:00 · 2024-10-16 04:53:55 +11:00 · 2024-10-16 04:53:55 +11:00 · 79f9d05a78
commit 79f9d05a78
parent 11f535a422
12 changed files with 1003 additions and 13 deletions
--- a/packages/kbn-esql-ast/src/ast/util.ts
+++ 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';
+};
--- 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<AstNodeTemplate<ESQLSource>, 'name' | 'index' | 'cluster'>,
+      fromParser?: Partial<AstNodeParserFields>
+    ): ESQLSource => {
+      return {
+        ...template,
+        ...Builder.parserFields(fromParser),
+        index,
+        cluster,
+        name: (cluster ? cluster + ':' : '') + index,
+        sourceType: 'index',
+        type: 'source',
+      };
+    };
+
    export const column = (
      template: Omit<AstNodeTemplate<ESQLColumn>, 'name' | 'quoted'>,
      fromParser?: Partial<AstNodeParserFields>
--- 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()` &mdash; List all `METADATA` fields.
- `.commands.from.metadata.find()` &mdash; Find a `METADATA` field by name.
- `.commands.from.metadata.removeByPredicate()` &mdash; Remove a `METADATA`
-  field by matching a predicate.
- `.commands.from.metadata.remove()` &mdash; Remove a `METADATA` field by name.
- `.commands.from.metadata.insert()` &mdash; Insert a `METADATA` field.
- `.commands.from.metadata.upsert()` &mdash; Insert `METADATA` field, if it does
-  not exist.
+- `.generic`
+  - `.listCommands()` &mdash; Lists all commands. Returns an iterator.
+  - `.findCommand()` &mdash; Finds a specific command by a predicate function.
+  - `.findCommandOption()` &mdash; Finds a specific command option by a predicate function.
+  - `.findCommandByName()` &mdash; Finds a specific command by name.
+  - `.findCommandOptionByName()` &mdash; Finds a specific command option by name.
+  - `.appendCommand()` &mdash; Add a new command to the AST.
+  - `.appendCommandOption()` &mdash; Add a new command option to a command.
+  - `.appendCommandArgument()` &mdash; Add a new main command argument to a command.
+  - `.removeCommand()` &mdash; Remove a command from the AST.
+  - `.removeCommandOption()` &mdash; Remove a command option from the AST.
+  - `.removeCommandArgument()` &mdash; Remove a command argument from the AST.
+- `.commands`
+  - `.from`
+    - `.sources`
+      - `.list()` &mdash; List all `FROM` sources.
+      - `.find()` &mdash; Find a source by name.
+      - `.remove()` &mdash; Remove a source by name.
+      - `.insert()` &mdash; Insert a source.
+      - `.upsert()` &mdash; Insert a source, if it does not exist.
+    - `.metadata`
+      - `.list()` &mdash; List all `METADATA` fields.
+      - `.find()` &mdash; Find a `METADATA` field by name.
+      - `.removeByPredicate()` &mdash; Remove a `METADATA` field by matching a predicate function.
+      - `.remove()` &mdash; Remove a `METADATA` field by name.
+      - `.insert()` &mdash; Insert a `METADATA` field.
+      - `.upsert()` &mdash; Insert `METADATA` field, if it does not exist.
+  - `.limit`
+    - `.list()` &mdash; List all `LIMIT` commands.
+    - `.byIndex()` &mdash; Find a `LIMIT` command by index.
+    - `.find()` &mdash; Find a `LIMIT` command by a predicate function.
+    - `.remove()` &mdash; Remove a `LIMIT` command by index.
+    - `.set()` &mdash; Set the limit value of a specific `LIMIT` command.
+    - `.upsert()` &mdash; Insert a `LIMIT` command, or update the limit value if it already exists.
--- 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 };
--- 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;
--- a/packages/kbn-esql-ast/src/mutate/commands/from/sources.test.ts
+++ 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');
+    });
+  });
+});
--- a/packages/kbn-esql-ast/src/mutate/commands/from/sources.ts
+++ 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<ESQLSource> => {
+  return new Visitor()
+    .on('visitFromCommand', function* (ctx): IterableIterator<ESQLSource> {
+      for (const argument of ctx.arguments()) {
+        if (argument.type === 'source') {
+          yield argument;
+        }
+      }
+    })
+    .on('visitCommand', function* (): IterableIterator<ESQLSource> {})
+    .on('visitQuery', function* (ctx): IterableIterator<ESQLSource> {
+      for (const command of ctx.visitCommands()) {
+        yield* command;
+      }
+    })
+    .visitQuery(ast);
+};
+
+export const findByPredicate = (
+  ast: ESQLAstQueryExpression,
+  predicate: Predicate<ESQLSource>
+): 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);
+};
--- 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 };
--- a/packages/kbn-esql-ast/src/mutate/commands/limit/index.test.ts
+++ 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,
+          },
+        ],
+      });
+    });
+  });
+});
--- a/packages/kbn-esql-ast/src/mutate/commands/limit/index.ts
+++ 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<ESQLCommand> => {
+  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>
+): 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<ESQLCommand> = 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<ESQLCommand> = 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;
+};
--- 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';
--- 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);
+};