Skip to content

Commit 568ffc7

Browse files
sddonnesddonne
committed
add mutation functions for SETINGS
1 parent 30cdbef commit 568ffc7

4 files changed

Lines changed: 268 additions & 2 deletions

File tree

‎src/ast/mutate/README.md‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,5 +66,9 @@ console.log(src); // FROM index METADATA _lang, _id
6666
- `.join`
6767
- `.list()` — List all `JOIN` commands.
6868
- `.byIndex()` — Find a `JOIN` command by index.
69-
69+
- `.set`
70+
- `.list()` — List all `SET` header commands.
71+
- `.findBySettingName()` — Find a `SET` command by its setting name.
72+
- `.set()` — Modify the value of an existing `SET` setting.
73+
- `.upsert()` — Modify a `SET` setting value, or add a new `SET` command if it does not exist.
7074

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,11 @@
77

88
import * as from from './from';
99
import * as limit from './limit';
10+
import * as set from './set';
1011
import * as sort from './sort';
1112
import * as stats from './stats';
1213
import * as where from './where';
1314
import * as join from './join';
1415
import * as rerank from './rerank';
1516

16-
export { from, limit, sort, stats, where, join, rerank };
17+
export { from, limit, set, sort, stats, where, join, rerank };

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

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
import { Parser } from '../../../../parser';
9+
import { BasicPrettyPrinter } from '../../../../pretty_print';
10+
import * as commands from '..';
11+
12+
describe('commands.set', () => {
13+
describe('.list()', () => {
14+
it('lists all SET header commands', () => {
15+
const src = 'SET approximation = true; SET unmapped_fields = "DEFAULT"; FROM index';
16+
const { root } = Parser.parse(src);
17+
18+
const nodes = [...commands.set.list(root)];
19+
20+
expect(nodes).toHaveLength(2);
21+
});
22+
23+
it('returns empty iterator when no SET commands exist', () => {
24+
const src = 'FROM index | LIMIT 10';
25+
const { root } = Parser.parse(src);
26+
27+
const nodes = [...commands.set.list(root)];
28+
29+
expect(nodes).toHaveLength(0);
30+
});
31+
});
32+
33+
describe('.findBySettingName()', () => {
34+
it('finds a SET command by setting name', () => {
35+
const src = 'SET approximation = true; SET unmapped_fields = "DEFAULT"; FROM index';
36+
const { root } = Parser.parse(src);
37+
38+
const node = commands.set.findBySettingName(root, 'unmapped_fields');
39+
40+
expect(node).toMatchObject({
41+
type: 'header-command',
42+
name: 'set',
43+
args: [
44+
{
45+
type: 'function',
46+
subtype: 'binary-expression',
47+
name: '=',
48+
args: [
49+
{ type: 'identifier', name: 'unmapped_fields' },
50+
{ type: 'literal', literalType: 'keyword', valueUnquoted: 'DEFAULT' },
51+
],
52+
},
53+
],
54+
});
55+
});
56+
57+
it('returns undefined when the setting does not exist', () => {
58+
const src = 'SET approximation = true; FROM index';
59+
const { root } = Parser.parse(src);
60+
61+
const node = commands.set.findBySettingName(root, 'unmapped_fields');
62+
63+
expect(node).toBeUndefined();
64+
});
65+
});
66+
67+
describe('.set()', () => {
68+
it('modifies the value of an existing setting', () => {
69+
const src = 'SET unmapped_fields = "DEFAULT"; FROM index';
70+
const { root } = Parser.parse(src);
71+
72+
commands.set.set(root, 'unmapped_fields', 'LOAD');
73+
const printed = BasicPrettyPrinter.print(root);
74+
75+
expect(printed).toBe('SET unmapped_fields = "LOAD"; FROM index');
76+
});
77+
78+
it('returns undefined when the setting does not exist', () => {
79+
const src = 'SET approximation = TRUE; FROM index';
80+
const { root } = Parser.parse(src);
81+
82+
const node = commands.set.set(root, 'unmapped_fields', 'LOAD');
83+
const printed = BasicPrettyPrinter.print(root);
84+
expect(printed).toBe('SET approximation = TRUE; FROM index');
85+
expect(node).toBeUndefined();
86+
});
87+
88+
it('only modifies the targeted setting when multiple exist', () => {
89+
const src = 'SET approximation = TRUE; SET unmapped_fields = "DEFAULT"; FROM index';
90+
const { root } = Parser.parse(src);
91+
92+
commands.set.set(root, 'unmapped_fields', 'LOAD');
93+
const printed = BasicPrettyPrinter.print(root);
94+
95+
expect(printed).toBe('SET approximation = TRUE; SET unmapped_fields = "LOAD"; FROM index');
96+
});
97+
});
98+
99+
describe('.upsert()', () => {
100+
it('modifies the value when the setting already exists', () => {
101+
const src = 'SET unmapped_fields = "DEFAULT"; FROM index';
102+
const { root } = Parser.parse(src);
103+
104+
const node = commands.set.upsert(root, 'unmapped_fields', 'LOAD');
105+
const printed = BasicPrettyPrinter.print(root);
106+
107+
expect(printed).toBe('SET unmapped_fields = "LOAD"; FROM index');
108+
expect(node).toMatchObject({
109+
type: 'header-command',
110+
name: 'set',
111+
});
112+
});
113+
114+
it('adds a new SET command when the setting does not exist', () => {
115+
const src = 'FROM index | LIMIT 10';
116+
const { root } = Parser.parse(src);
117+
118+
const node = commands.set.upsert(root, 'unmapped_fields', 'LOAD');
119+
const printed = BasicPrettyPrinter.print(root);
120+
121+
expect(printed).toBe('SET unmapped_fields = "LOAD"; FROM index | LIMIT 10');
122+
expect(node).toMatchObject({
123+
type: 'header-command',
124+
name: 'set',
125+
args: [
126+
{
127+
type: 'function',
128+
subtype: 'binary-expression',
129+
name: '=',
130+
args: [
131+
{ type: 'identifier', name: 'unmapped_fields' },
132+
{ type: 'literal', literalType: 'keyword', valueUnquoted: 'LOAD' },
133+
],
134+
},
135+
],
136+
});
137+
});
138+
139+
it('appends alongside existing SET commands', () => {
140+
const src = 'SET approximation = TRUE; FROM index';
141+
const { root } = Parser.parse(src);
142+
143+
commands.set.upsert(root, 'unmapped_fields', 'LOAD');
144+
const printed = BasicPrettyPrinter.print(root);
145+
146+
expect(printed).toBe('SET approximation = TRUE; SET unmapped_fields = "LOAD"; FROM index');
147+
});
148+
});
149+
});

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

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/*
2+
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
3+
* or more contributor license agreements. Licensed under the Elastic License
4+
* 2.0; you may not use this file except in compliance with the Elastic License
5+
* 2.0.
6+
*/
7+
8+
import { Builder } from '../../../builder';
9+
import { isAssignment } from '../../../is';
10+
import type {
11+
ESQLAstQueryExpression,
12+
ESQLAstSetHeaderCommand,
13+
ESQLBinaryExpression,
14+
BinaryExpressionAssignmentOperator,
15+
ESQLSingleAstItem,
16+
} from '../../../../types';
17+
18+
/**
19+
* Lists all SET header commands in the query AST.
20+
*
21+
* @param ast The root AST node.
22+
* @returns An iterable of SET header commands.
23+
*/
24+
export const list = function* (
25+
ast: ESQLAstQueryExpression
26+
): IterableIterator<ESQLAstSetHeaderCommand> {
27+
if (!ast.header) return;
28+
for (const cmd of ast.header) {
29+
if (cmd.name === 'set') {
30+
yield cmd as ESQLAstSetHeaderCommand;
31+
}
32+
}
33+
};
34+
35+
/**
36+
* Finds a SET header command by its setting name.
37+
*
38+
* @param ast The root AST node.
39+
* @param settingName The name of the setting to find (e.g. "unmapped_fields").
40+
* @returns The matching SET header command, or undefined.
41+
*/
42+
export const findBySettingName = (
43+
ast: ESQLAstQueryExpression,
44+
settingName: string
45+
): ESQLAstSetHeaderCommand | undefined => {
46+
for (const cmd of list(ast)) {
47+
const assignment = cmd.args[0];
48+
if (isAssignment(assignment)) {
49+
const identifier = assignment.args[0] as ESQLSingleAstItem;
50+
if (identifier.name === settingName) {
51+
return cmd;
52+
}
53+
}
54+
}
55+
return undefined;
56+
};
57+
58+
/**
59+
* Updates the value of an existing SET setting. Returns the modified header
60+
* command, or `undefined` if no matching setting was found.
61+
*
62+
* @param ast The root AST node.
63+
* @param settingName The name of the setting to modify.
64+
* @param value The new string value for the setting.
65+
* @returns The modified SET header command, or undefined if not found.
66+
*/
67+
export const set = (
68+
ast: ESQLAstQueryExpression,
69+
settingName: string,
70+
value: string
71+
): ESQLAstSetHeaderCommand | undefined => {
72+
const cmd = findBySettingName(ast, settingName);
73+
if (!cmd) return undefined;
74+
75+
const assignment = cmd.args[0] as ESQLBinaryExpression<BinaryExpressionAssignmentOperator>;
76+
const newValue = Builder.expression.literal.string(value);
77+
assignment.args[1] = newValue;
78+
79+
return cmd;
80+
};
81+
82+
/**
83+
* Updates the value of an existing SET setting, or inserts a new SET header
84+
* command if the setting does not exist.
85+
*
86+
* @param ast The root AST node.
87+
* @param settingName The name of the setting (e.g. "unmapped_fields").
88+
* @param value The string value for the setting.
89+
* @returns The modified or newly created SET header command.
90+
*/
91+
export const upsert = (
92+
ast: ESQLAstQueryExpression,
93+
settingName: string,
94+
value: string
95+
): ESQLAstSetHeaderCommand => {
96+
const existing = set(ast, settingName, value);
97+
if (existing) return existing;
98+
99+
const identifier = Builder.identifier(settingName);
100+
const literal = Builder.expression.literal.string(value);
101+
const assignment = Builder.expression.func.binary('=', [identifier, literal]);
102+
const cmd = Builder.header.command.set([
103+
assignment as ESQLBinaryExpression<BinaryExpressionAssignmentOperator>,
104+
]);
105+
106+
if (!ast.header) {
107+
ast.header = [];
108+
}
109+
ast.header.push(cmd);
110+
111+
return cmd;
112+
};

0 commit comments

Comments
 (0)