Skip to content

Commit 63142e0

Browse files
setting(name) method and more readable tests
1 parent 12c2601 commit 63142e0

File tree

2 files changed

+89
-49
lines changed

2 files changed

+89
-49
lines changed

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/EsqlQuery.java

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,30 @@
77

88
package org.elasticsearch.xpack.esql.plan.logical;
99

10+
import org.elasticsearch.xpack.esql.core.expression.Alias;
11+
import org.elasticsearch.xpack.esql.core.expression.Expression;
12+
1013
import java.util.List;
1114

12-
public record EsqlQuery(LogicalPlan plan, List<QuerySettings> settings) {}
15+
public record EsqlQuery(LogicalPlan plan, List<QuerySettings> settings) {
16+
/**
17+
* Returns the expression corresponding to a setting value.
18+
* If the setting name appears multiple times (in one or more QuerySettings objects), this will return last occurrence.
19+
*
20+
* @param name the setting name
21+
*/
22+
public Expression setting(String name) {
23+
if (settings == null) {
24+
return null;
25+
}
26+
Expression result = null;
27+
for (QuerySettings setting : settings) {
28+
for (Alias field : setting.fields()) {
29+
if (field.name().equals(name)) {
30+
result = field.child();
31+
}
32+
}
33+
}
34+
return result;
35+
}
36+
}

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/analysis/ParsingTests.java

Lines changed: 64 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,6 @@
4848
import static org.hamcrest.Matchers.hasSize;
4949
import static org.hamcrest.Matchers.instanceOf;
5050
import static org.hamcrest.Matchers.is;
51-
import static org.hamcrest.Matchers.nullValue;
5251

5352
public class ParsingTests extends ESTestCase {
5453
private static final String INDEX_NAME = "test";
@@ -218,54 +217,44 @@ public void testSet() {
218217
assertThat(query.plan(), is(instanceOf(Row.class)));
219218
assertThat(query.settings().size(), is(1));
220219
assertThat(query.settings().get(0).fields().size(), is(1));
221-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
222-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("bar")));
220+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("bar"));
223221

224222
query = parse("SET foo = \"bar\", bar = 2; row a = 1 | eval x = 12", new QueryParams());
225223
assertThat(query.plan(), is(instanceOf(Eval.class)));
226224
assertThat(query.settings().size(), is(1));
227225
assertThat(query.settings().get(0).fields().size(), is(2));
228-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
229-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("bar")));
230-
assertThat(query.settings().get(0).fields().get(1).name(), is("bar"));
231-
assertThat(query.settings().get(0).fields().get(1).child().fold(FoldContext.small()), is(2));
226+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("bar"));
227+
checkSetting(query, 0, 1, "bar", 2);
232228
}
233229

234230
public void testSetWithTripleQuotes() {
235231
EsqlQuery query = parse("SET foo = \"\"\"bar\"baz\"\"\"; row a = 1", new QueryParams());
236232
assertThat(query.plan(), is(instanceOf(Row.class)));
237233
assertThat(query.settings().size(), is(1));
238234
assertThat(query.settings().get(0).fields().size(), is(1));
239-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
240-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("bar\"baz")));
235+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("bar\"baz"));
241236

242237
query = parse("SET foo = \"\"\"bar\"\"\"\"; row a = 1", new QueryParams());
243238
assertThat(query.plan(), is(instanceOf(Row.class)));
244239
assertThat(query.settings().size(), is(1));
245240
assertThat(query.settings().get(0).fields().size(), is(1));
246-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
247-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("bar\"")));
241+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("bar\""));
248242

249243
query = parse("SET foo = \"\"\"\"bar\"\"\"; row a = 1 | LIMIT 3", new QueryParams());
250244
assertThat(query.plan(), is(instanceOf(Limit.class)));
251245
assertThat(query.settings().size(), is(1));
252246
assertThat(query.settings().get(0).fields().size(), is(1));
253-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
254-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("\"bar")));
247+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("\"bar"));
255248
}
256249

257250
public void testSetArrays() {
258251
EsqlQuery query = parse("SET foo = [\"bar\", \"baz\"], bar = [1, 2, 3]; row a = 1", new QueryParams());
259252
assertThat(query.plan(), is(instanceOf(Row.class)));
260253
assertThat(query.settings().size(), is(1));
254+
261255
assertThat(query.settings().get(0).fields().size(), is(2));
262-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
263-
assertThat(
264-
query.settings().get(0).fields().get(0).child().fold(FoldContext.small()),
265-
is(List.of(BytesRefs.toBytesRef("bar"), BytesRefs.toBytesRef("baz")))
266-
);
267-
assertThat(query.settings().get(0).fields().get(1).name(), is("bar"));
268-
assertThat(query.settings().get(0).fields().get(1).child().fold(FoldContext.small()), is(List.of(1, 2, 3)));
256+
checkSetting(query, 0, 0, "foo", List.of(BytesRefs.toBytesRef("bar"), BytesRefs.toBytesRef("baz")));
257+
checkSetting(query, 0, 1, "bar", List.of(1, 2, 3));
269258
}
270259

271260
public void testMultipleSet() {
@@ -275,20 +264,16 @@ public void testMultipleSet() {
275264
);
276265
assertThat(query.plan(), is(instanceOf(Row.class)));
277266
assertThat(query.settings().size(), is(2));
267+
278268
assertThat(query.settings().get(0).fields().size(), is(2));
279-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
280-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("bar")));
281-
assertThat(query.settings().get(0).fields().get(1).name(), is("bar"));
282-
assertThat(query.settings().get(0).fields().get(1).child().fold(FoldContext.small()), is(2));
269+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("bar"), BytesRefs.toBytesRef("baz"));
270+
checkSetting(query, 0, 1, "bar", 2);
271+
283272
assertThat(query.settings().get(1).fields().size(), is(4));
284-
assertThat(query.settings().get(1).fields().get(0).name(), is("foo"));
285-
assertThat(query.settings().get(1).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("baz")));
286-
assertThat(query.settings().get(1).fields().get(1).name(), is("x"));
287-
assertThat(query.settings().get(1).fields().get(1).child().fold(FoldContext.small()), is(3.5));
288-
assertThat(query.settings().get(1).fields().get(2).name(), is("y"));
289-
assertThat(query.settings().get(1).fields().get(2).child().fold(FoldContext.small()), is(false));
290-
assertThat(query.settings().get(1).fields().get(3).name(), is("z"));
291-
assertThat(query.settings().get(1).fields().get(3).child().fold(FoldContext.small()), is(nullValue()));
273+
checkSetting(query, 1, 0, "foo", BytesRefs.toBytesRef("baz"));
274+
checkSetting(query, 1, 1, "x", 3.5);
275+
checkSetting(query, 1, 2, "y", false);
276+
checkSetting(query, 1, 3, "z", null);
292277
}
293278

294279
public void testSetWithNamedParams() {
@@ -303,16 +288,14 @@ public void testSetWithNamedParams() {
303288
);
304289
assertThat(query.plan(), is(instanceOf(Row.class)));
305290
assertThat(query.settings().size(), is(2));
291+
306292
assertThat(query.settings().get(0).fields().size(), is(2));
307-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
308-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("bar")));
309-
assertThat(query.settings().get(0).fields().get(1).name(), is("bar"));
310-
assertThat(query.settings().get(0).fields().get(1).child().fold(FoldContext.small()), is(2));
293+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("bar"), BytesRefs.toBytesRef("baz"));
294+
checkSetting(query, 0, 1, "bar", 2);
295+
311296
assertThat(query.settings().get(1).fields().size(), is(2));
312-
assertThat(query.settings().get(1).fields().get(0).name(), is("foo"));
313-
assertThat(query.settings().get(1).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("baz")));
314-
assertThat(query.settings().get(1).fields().get(1).name(), is("x"));
315-
assertThat(query.settings().get(1).fields().get(1).child().fold(FoldContext.small()), is(3.5));
297+
checkSetting(query, 1, 0, "foo", BytesRefs.toBytesRef("baz"));
298+
checkSetting(query, 1, 1, "x", 3.5);
316299
}
317300

318301
public void testSetWithPositionalParams() {
@@ -329,16 +312,48 @@ public void testSetWithPositionalParams() {
329312
assertThat(query.plan(), is(instanceOf(Row.class)));
330313
assertThat(((Row) query.plan()).fields().get(0).child().fold(FoldContext.small()), is(8));
331314
assertThat(query.settings().size(), is(2));
315+
332316
assertThat(query.settings().get(0).fields().size(), is(2));
333-
assertThat(query.settings().get(0).fields().get(0).name(), is("foo"));
334-
assertThat(query.settings().get(0).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("bar")));
335-
assertThat(query.settings().get(0).fields().get(1).name(), is("bar"));
336-
assertThat(query.settings().get(0).fields().get(1).child().fold(FoldContext.small()), is(2));
317+
checkSetting(query, 0, 0, "foo", BytesRefs.toBytesRef("bar"), BytesRefs.toBytesRef("baz"));
318+
checkSetting(query, 0, 1, "bar", 2);
319+
337320
assertThat(query.settings().get(1).fields().size(), is(2));
338-
assertThat(query.settings().get(1).fields().get(0).name(), is("foo"));
339-
assertThat(query.settings().get(1).fields().get(0).child().fold(FoldContext.small()), is(BytesRefs.toBytesRef("baz")));
340-
assertThat(query.settings().get(1).fields().get(1).name(), is("x"));
341-
assertThat(query.settings().get(1).fields().get(1).child().fold(FoldContext.small()), is(3.5));
321+
checkSetting(query, 1, 0, "foo", BytesRefs.toBytesRef("baz"));
322+
checkSetting(query, 1, 1, "x", 3.5);
323+
}
324+
325+
/**
326+
* @param query the query
327+
* @param group the order of the corresponding SET statement
328+
* @param position the position within the SET statement
329+
* @param name the setting name
330+
* @param value the setting value as it appears in the query at that position
331+
*/
332+
private void checkSetting(EsqlQuery query, int group, int position, String name, Object value) {
333+
checkSetting(query, group, position, name, value, value);
334+
}
335+
336+
/**
337+
* @param query the query
338+
* @param group the order of the corresponding SET statement
339+
* @param position the position within the SET statement
340+
* @param name the setting name
341+
* @param value the setting value as it appears in the query at that position
342+
* @param maskingValue the final value you'll obtain if you use query.setting(name).
343+
* It could be different from value in case of name collisions in the query
344+
*/
345+
private void checkSetting(EsqlQuery query, int group, int position, String name, Object value, Object maskingValue) {
346+
assertThat(settingName(query, group, position), is(name));
347+
assertThat(settingValue(query, group, position), is(value));
348+
assertThat(query.setting(name).fold(FoldContext.small()), is(maskingValue));
349+
}
350+
351+
private String settingName(EsqlQuery query, int group, int position) {
352+
return query.settings().get(group).fields().get(position).name();
353+
}
354+
355+
private Object settingValue(EsqlQuery query, int group, int position) {
356+
return query.settings().get(group).fields().get(position).child().fold(FoldContext.small());
342357
}
343358

344359
private String error(String query, QueryParams params) {
@@ -359,4 +374,5 @@ private String error(String query) {
359374
private static IndexResolution loadIndexResolution(String name) {
360375
return IndexResolution.valid(new EsIndex(INDEX_NAME, LoadMapping.loadMapping(name)));
361376
}
377+
362378
}

0 commit comments

Comments
 (0)