feat: Add mutation functions for SETtings

src/ast/mutate/README.md

@@ -66,5 +66,9 @@ console.log(src); // FROM index METADATA _lang, _id
   - `.join`
     - `.list()` — List all `JOIN` commands.
     - `.byIndex()` — Find a `JOIN` command by index.
-
-
+  - `.set`
+    - `.list()` — List all `SET` header commands.
+    - `.findBySettingName()` — Find a `SET` command by its setting name.
+    - `.update()` — Modify the value of an existing `SET` setting.
+    - `.upsert()` — Modify a `SET` setting value, or add a new `SET` command if it does not exist.
+    - `.remove()` — Deletes a SET command by setting name.

src/ast/mutate/commands/index.ts

@@ -7,10 +7,11 @@
 
 import * as from from './from';
 import * as limit from './limit';
+import * as set from './set';
 import * as sort from './sort';
 import * as stats from './stats';
 import * as where from './where';
 import * as join from './join';
 import * as rerank from './rerank';
 
-export { from, limit, sort, stats, where, join, rerank };
+export { from, limit, set, sort, stats, where, join, rerank };

src/ast/mutate/commands/set/index.test.ts

@@ -0,0 +1,229 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { Parser } from '../../../../parser';
+import { BasicPrettyPrinter } from '../../../../pretty_print';
+import * as commands from '..';
+
+describe('commands.set', () => {
+  describe('.list()', () => {
+    it('lists all SET header commands', () => {
+      const src = 'SET approximation = true; SET unmapped_fields = "DEFAULT"; FROM index';
+      const { root } = Parser.parse(src);
+
+      const nodes = [...commands.set.list(root)];
+
+      expect(nodes).toHaveLength(2);
+    });
+
+    it('returns empty iterator when no SET commands exist', () => {
+      const src = 'FROM index | LIMIT 10';
+      const { root } = Parser.parse(src);
+
+      const nodes = [...commands.set.list(root)];
+
+      expect(nodes).toHaveLength(0);
+    });
+  });
+
+  describe('.findBySettingName()', () => {
+    it('finds a SET command by setting name', () => {
+      const src = 'SET approximation = true; SET unmapped_fields = "DEFAULT"; FROM index';
+      const { root } = Parser.parse(src);
+
+      const node = commands.set.findBySettingName(root, 'unmapped_fields');
+
+      expect(node).toMatchObject({
+        type: 'header-command',
+        name: 'set',
+        args: [
+          {
+            type: 'function',
+            subtype: 'binary-expression',
+            name: '=',
+            args: [
+              { type: 'identifier', name: 'unmapped_fields' },
+              { type: 'literal', literalType: 'keyword', valueUnquoted: 'DEFAULT' },
+            ],
+          },
+        ],
+      });
+    });
+
+    it('returns undefined when the setting does not exist', () => {
+      const src = 'SET approximation = true; FROM index';
+      const { root } = Parser.parse(src);
+
+      const node = commands.set.findBySettingName(root, 'unmapped_fields');
+
+      expect(node).toBeUndefined();
+    });
+  });
+
+  describe('.update()', () => {
+    it('modifies the value of an existing setting', () => {
+      const src = 'SET unmapped_fields = "DEFAULT"; FROM index';
+      const { root } = Parser.parse(src);
+
+      commands.set.update(root, 'unmapped_fields', '"LOAD"');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET unmapped_fields = "LOAD"; FROM index');
+    });
+
+    it('returns undefined when the setting does not exist', () => {
+      const src = 'SET approximation = TRUE; FROM index';
+      const { root } = Parser.parse(src);
+
+      const node = commands.set.update(root, 'unmapped_fields', '"LOAD"');
+      const printed = BasicPrettyPrinter.print(root);
+      expect(printed).toBe('SET approximation = TRUE; FROM index');
+      expect(node).toBeUndefined();
+    });
+
+    it('only modifies the targeted setting when multiple exist', () => {
+      const src = 'SET approximation = TRUE; SET unmapped_fields = "DEFAULT"; FROM index';
+      const { root } = Parser.parse(src);
+
+      commands.set.update(root, 'unmapped_fields', '"LOAD"');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET approximation = TRUE; SET unmapped_fields = "LOAD"; FROM index');
+    });
+
+    it('modifies the value of an existing setting with a boolean value', () => {
+      const src = 'SET approximation = TRUE; FROM index';
+      const { root } = Parser.parse(src);
+
+      commands.set.update(root, 'approximation', 'FALSE');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET approximation = FALSE; FROM index');
+    });
+
+    it('modifies the value of an existing setting with a MAP value', () => {
+      const src = 'SET approximation = TRUE; FROM index';
+      const { root } = Parser.parse(src);
+
+      const node = commands.set.update(root, 'approximation', '{ "confidence_level": 0.99 }');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET approximation = {"confidence_level": 0.99}; FROM index');
+      expect(node).toMatchObject({
+        type: 'header-command',
+        name: 'set',
+        args: [
+          {
+            type: 'function',
+            subtype: 'binary-expression',
+            name: '=',
+            args: [
+              { type: 'identifier', name: 'approximation' },
+              {
+                type: 'map',
+                entries: [
+                  {
+                    key: {
+                      type: 'literal',
+                      literalType: 'keyword',
+                      valueUnquoted: 'confidence_level',
+                    },
+                    value: { type: 'literal', literalType: 'double', value: 0.99 },
+                  },
+                ],
+              },
+            ],
+          },
+        ],
+      });
+    });
+  });
+
+  describe('.upsert()', () => {
+    it('modifies the value when the setting already exists', () => {
+      const src = 'SET unmapped_fields = "DEFAULT"; FROM index';
+      const { root } = Parser.parse(src);
+
+      const node = commands.set.upsert(root, 'unmapped_fields', '"LOAD"');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET unmapped_fields = "LOAD"; FROM index');
+      expect(node).toMatchObject({
+        type: 'header-command',
+        name: 'set',
+      });
+    });
+
+    it('adds a new SET command when the setting does not exist', () => {
+      const src = 'FROM index | LIMIT 10';
+      const { root } = Parser.parse(src);
+
+      const node = commands.set.upsert(root, 'unmapped_fields', '"LOAD"');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET unmapped_fields = "LOAD"; FROM index | LIMIT 10');
+      expect(node).toMatchObject({
+        type: 'header-command',
+        name: 'set',
+        args: [
+          {
+            type: 'function',
+            subtype: 'binary-expression',
+            name: '=',
+            args: [
+              { type: 'identifier', name: 'unmapped_fields' },
+              { type: 'literal', literalType: 'keyword', valueUnquoted: 'LOAD' },
+            ],
+          },
+        ],
+      });
+    });
+
+    it('appends alongside existing SET commands', () => {
+      const src = 'SET approximation = TRUE; FROM index';
+      const { root } = Parser.parse(src);
+
+      commands.set.upsert(root, 'unmapped_fields', '"LOAD"');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET approximation = TRUE; SET unmapped_fields = "LOAD"; FROM index');
+    });
+  });
+
+  describe('.remove()', () => {
+    it('removes a SET command by setting name', () => {
+      const src = 'SET unmapped_fields = "DEFAULT"; FROM index';
+      const { root } = Parser.parse(src);
+
+      commands.set.remove(root, 'unmapped_fields');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('FROM index');
+    });
+
+    it('returns undefined when the setting does not exist', () => {
+      const src = 'SET approximation = TRUE; FROM index';
+      const { root } = Parser.parse(src);
+
+      const removed = commands.set.remove(root, 'unmapped_fields');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(removed).toBeUndefined();
+      expect(printed).toBe('SET approximation = TRUE; FROM index');
+    });
+
+    it('only removes the targeted setting when multiple exist', () => {
+      const src = 'SET approximation = TRUE; SET unmapped_fields = "DEFAULT"; FROM index';
+      const { root } = Parser.parse(src);
+
+      commands.set.remove(root, 'unmapped_fields');
+      const printed = BasicPrettyPrinter.print(root);
+
+      expect(printed).toBe('SET approximation = TRUE; FROM index');
+    });
+  });
+});

src/ast/mutate/commands/set/index.ts

@@ -0,0 +1,144 @@
+/*
+ * 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; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { Builder } from '../../../builder';
+import { isAssignment } from '../../../is';
+import { Parser } from '../../../../parser';
+import type {
+  ESQLAstQueryExpression,
+  ESQLAstSetHeaderCommand,
+  ESQLBinaryExpression,
+  BinaryExpressionAssignmentOperator,
+  ESQLIdentifier,
+} from '../../../../types';
+
+/**
+ * Lists all SET header commands in the query AST.
+ *
+ * @param ast The root AST node.
+ * @returns An iterable of SET header commands.
+ */
+export const list = function* (
+  ast: ESQLAstQueryExpression
+): IterableIterator<ESQLAstSetHeaderCommand> {
+  if (!ast.header) return;
+  for (const cmd of ast.header) {
+    if (cmd.name === 'set') {
+      yield cmd as ESQLAstSetHeaderCommand;
+    }
+  }
+};
+
+/**
+ * Finds a SET header command by its setting name.
+ *
+ * @param ast The root AST node.
+ * @param settingName The name of the setting to find (e.g. "unmapped_fields").
+ * @returns The matching SET header command, or undefined.
+ */
+export const findBySettingName = (
+  ast: ESQLAstQueryExpression,
+  settingName: string
+): ESQLAstSetHeaderCommand | undefined => {
+  for (const cmd of list(ast)) {
+    const assignment = cmd.args[0];
+    if (isAssignment(assignment)) {
+      const identifier = assignment.args[0] as ESQLIdentifier;
+      if (identifier.name === settingName) {
+        return cmd;
+      }
+    }
+  }
+  return undefined;
+};
+
+/**
+ * Updates the value of an existing SET setting. Returns the modified header
+ * command, or `undefined` if no matching setting was found.
+ *
+ * @param ast The root AST node.
+ * @param settingName The name of the setting to modify.
+ * @param value The new value, parsed as an ES|QL expression. For example:
+ *   - `'"LOAD"'` → string literal `"LOAD"`
+ *   - `'true'` → boolean literal `TRUE`
+ *   - `'42'` → integer literal `42`
+ *   - `'{ "key": "value" }'` → map expression
+ * @returns The modified SET header command, or undefined if not found.
+ */
+export const update = (
+  ast: ESQLAstQueryExpression,
+  settingName: string,
+  value: string
+): ESQLAstSetHeaderCommand | undefined => {
+  const cmd = findBySettingName(ast, settingName);
+  if (!cmd) return undefined;
+
+  const newNode = Parser.parseExpression(value).root;
+
+  const assignment = cmd.args[0] as ESQLBinaryExpression<BinaryExpressionAssignmentOperator>;
+  assignment.args[1] = newNode;
+
+  return cmd;
+};
+
+/**
+ * Updates the value of an existing SET setting, or inserts a new SET header
+ * command if the setting does not exist. The new command is inserted as the last
+ * command in the header.
+ *
+ * @param ast The root AST node.
+ * @param settingName The name of the setting (e.g. "unmapped_fields").
+ * @param value The new value, parsed as an ES|QL expression. For example:
+ *   - `'"LOAD"'` → string literal `"LOAD"`
+ *   - `'true'` → boolean literal `TRUE`
+ *   - `'42'` → integer literal `42`
+ *   - `'{ "key": "value" }'` → map expression
+ * @returns The modified or newly created SET header command.
+ */
+export const upsert = (
+  ast: ESQLAstQueryExpression,
+  settingName: string,
+  value: string
+): ESQLAstSetHeaderCommand => {
+  const existing = update(ast, settingName, value);
+  if (existing) return existing;
+
+  const identifier = Builder.identifier(settingName);
+  const newNode = Parser.parseExpression(value).root;
+  const assignment = Builder.expression.func.binary('=', [identifier, newNode]);
+  const cmd = Builder.header.command.set([
+    assignment as ESQLBinaryExpression<BinaryExpressionAssignmentOperator>,
+  ]);
+
+  if (!ast.header) {
+    ast.header = [];
+  }
+  ast.header.push(cmd);
+
+  return cmd;
+};
+
+/**
+ * Removes a SET header command by its setting name.
+ *
+ * @param ast The root AST node.
+ * @param settingName The name of the setting to remove.
+ * @returns The removed SET header command, or undefined if not found.
+ */
+export const remove = (
+  ast: ESQLAstQueryExpression,
+  settingName: string
+): ESQLAstSetHeaderCommand | undefined => {
+  const cmd = findBySettingName(ast, settingName);
+  if (!cmd || !ast.header) return;
+
+  const index = ast.header.indexOf(cmd);
+  if (index === -1) return;
+
+  ast.header.splice(index, 1);
+  return cmd;
+};