Skip to content

Commit abcd246

Browse files
ivanceaivancea
committed
Fixed serialization of functions with configuration
1 parent fb98bec commit abcd246

File tree

6 files changed

+88
-22
lines changed

6 files changed

+88
-22
lines changed

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/Configuration.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,8 @@ public boolean equals(Object o) {
266266
&& Objects.equals(locale, that.locale)
267267
&& Objects.equals(that.query, query)
268268
&& profile == that.profile
269-
&& tables.equals(that.tables);
269+
&& tables.equals(that.tables)
270+
&& activeEsqlFeatures.equals(that.activeEsqlFeatures);
270271
}
271272

272273
@Override
@@ -282,7 +283,8 @@ public int hashCode() {
282283
locale,
283284
query,
284285
profile,
285-
tables
286+
tables,
287+
activeEsqlFeatures
286288
);
287289
}
288290

@@ -304,6 +306,8 @@ public String toString() {
304306
+ profile
305307
+ ", tables="
306308
+ tables
309+
+ ", activeEsqlFeatures="
310+
+ activeEsqlFeatures
307311
+ '}';
308312
}
309313
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractConfigurationAggregationTestCase.java‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
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+
package org.elasticsearch.xpack.esql.expression.function;
9+
10+
import org.elasticsearch.common.settings.Settings;
11+
import org.elasticsearch.features.NodeFeature;
12+
import org.elasticsearch.xpack.esql.EsqlTestUtils;
13+
import org.elasticsearch.xpack.esql.core.expression.Expression;
14+
import org.elasticsearch.xpack.esql.core.tree.Source;
15+
import org.elasticsearch.xpack.esql.core.util.StringUtils;
16+
import org.elasticsearch.xpack.esql.plugin.EsqlFeatures;
17+
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
18+
import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
19+
import org.elasticsearch.xpack.esql.session.Configuration;
20+
21+
import java.util.List;
22+
import java.util.Map;
23+
import java.util.stream.Collectors;
24+
25+
import static org.elasticsearch.xpack.esql.SerializationTestUtils.assertSerialization;
26+
27+
public abstract class AbstractConfigurationAggregationTestCase extends AbstractAggregationTestCase {
28+
protected abstract Expression buildWithConfiguration(Source source, List<Expression> args, Configuration configuration);
29+
30+
@Override
31+
protected Expression build(Source source, List<Expression> args) {
32+
return buildWithConfiguration(source, args, EsqlTestUtils.TEST_CFG);
33+
}
34+
35+
static Configuration randomConfiguration() {
36+
// TODO: Randomize the query and maybe the pragmas.
37+
return new Configuration(
38+
randomZone(),
39+
randomLocale(random()),
40+
randomBoolean() ? null : randomAlphaOfLength(randomInt(64)),
41+
randomBoolean() ? null : randomAlphaOfLength(randomInt(64)),
42+
QueryPragmas.EMPTY,
43+
EsqlPlugin.QUERY_RESULT_TRUNCATION_MAX_SIZE.getDefault(Settings.EMPTY),
44+
EsqlPlugin.QUERY_RESULT_TRUNCATION_DEFAULT_SIZE.getDefault(Settings.EMPTY),
45+
StringUtils.EMPTY,
46+
randomBoolean(),
47+
Map.of(),
48+
System.nanoTime(),
49+
new EsqlFeatures().getFeatures().stream().map(NodeFeature::id).collect(Collectors.toSet())
50+
);
51+
}
52+
53+
public void testSerializationWithConfiguration() {
54+
Configuration config = randomConfiguration();
55+
Expression expr = buildWithConfiguration(testCase.getSource(), testCase.getDataAsFields(), config);
56+
57+
assertSerialization(expr, config);
58+
59+
Configuration differentConfig;
60+
do {
61+
differentConfig = randomConfiguration();
62+
} while (config.equals(differentConfig));
63+
64+
Expression differentExpr = buildWithConfiguration(testCase.getSource(), testCase.getDataAsFields(), differentConfig);
65+
assertFalse(expr.equals(differentExpr));
66+
}
67+
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java‎

Lines changed: 0 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -136,8 +136,6 @@ public abstract class AbstractFunctionTestCase extends ESTestCase {
136136

137137
private static EsqlFunctionRegistry functionRegistry = new EsqlFunctionRegistry().snapshotRegistry();
138138

139-
private Configuration config;
140-
141139
protected TestCaseSupplier.TestCase testCase;
142140

143141
/**
@@ -1353,13 +1351,4 @@ private static boolean shouldHideSignature(List<DataType> argTypes, DataType ret
13531351
}
13541352
return false;
13551353
}
1356-
1357-
@Before
1358-
public void initConfig() {
1359-
config = randomConfiguration("FROM test", Map.of());
1360-
}
1361-
1362-
protected Configuration configuration() {
1363-
return config;
1364-
}
13651354
}

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/AvgTests.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
import org.elasticsearch.xpack.esql.core.tree.Source;
1515
import org.elasticsearch.xpack.esql.core.type.DataType;
1616
import org.elasticsearch.xpack.esql.expression.function.AbstractAggregationTestCase;
17+
import org.elasticsearch.xpack.esql.expression.function.AbstractConfigurationAggregationTestCase;
1718
import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier;
1819
import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
20+
import org.elasticsearch.xpack.esql.session.Configuration;
1921

2022
import java.util.ArrayList;
2123
import java.util.List;
@@ -25,7 +27,7 @@
2527

2628
import static org.hamcrest.Matchers.equalTo;
2729

28-
public class AvgTests extends AbstractAggregationTestCase {
30+
public class AvgTests extends AbstractConfigurationAggregationTestCase {
2931
public AvgTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
3032
this.testCase = testCaseSupplier.get();
3133
}
@@ -57,8 +59,8 @@ public static Iterable<Object[]> parameters() {
5759
}
5860

5961
@Override
60-
protected Expression build(Source source, List<Expression> args) {
61-
return new Avg(source, args.get(0), configuration());
62+
protected Expression buildWithConfiguration(Source source, List<Expression> args, Configuration configuration) {
63+
return new Avg(source, args.get(0), configuration);
6264
}
6365

6466
private static TestCaseSupplier makeSupplier(TestCaseSupplier.TypedDataSupplier fieldSupplier) {

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/SumTests.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,10 @@
1515
import org.elasticsearch.xpack.esql.core.tree.Source;
1616
import org.elasticsearch.xpack.esql.core.type.DataType;
1717
import org.elasticsearch.xpack.esql.expression.function.AbstractAggregationTestCase;
18+
import org.elasticsearch.xpack.esql.expression.function.AbstractConfigurationAggregationTestCase;
1819
import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier;
1920
import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
21+
import org.elasticsearch.xpack.esql.session.Configuration;
2022

2123
import java.util.ArrayList;
2224
import java.util.List;
@@ -27,7 +29,7 @@
2729
import static org.elasticsearch.xpack.esql.core.type.DataType.UNSIGNED_LONG;
2830
import static org.hamcrest.Matchers.equalTo;
2931

30-
public class SumTests extends AbstractAggregationTestCase {
32+
public class SumTests extends AbstractConfigurationAggregationTestCase {
3133
public SumTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
3234
this.testCase = testCaseSupplier.get();
3335
}
@@ -81,8 +83,8 @@ public static Iterable<Object[]> parameters() {
8183
}
8284

8385
@Override
84-
protected Expression build(Source source, List<Expression> args) {
85-
return new Sum(source, args.get(0), configuration());
86+
protected Expression buildWithConfiguration(Source source, List<Expression> args, Configuration configuration) {
87+
return new Sum(source, args.get(0), configuration);
8688
}
8789

8890
private static TestCaseSupplier makeSupplier(TestCaseSupplier.TypedDataSupplier fieldSupplier) {

‎x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/aggregate/WeightedAvgTests.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
import org.elasticsearch.xpack.esql.core.tree.Source;
1515
import org.elasticsearch.xpack.esql.core.type.DataType;
1616
import org.elasticsearch.xpack.esql.expression.function.AbstractAggregationTestCase;
17+
import org.elasticsearch.xpack.esql.expression.function.AbstractConfigurationAggregationTestCase;
1718
import org.elasticsearch.xpack.esql.expression.function.MultiRowTestCaseSupplier;
1819
import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier;
20+
import org.elasticsearch.xpack.esql.session.Configuration;
1921

2022
import java.util.ArrayList;
2123
import java.util.List;
@@ -25,7 +27,7 @@
2527

2628
import static org.hamcrest.Matchers.equalTo;
2729

28-
public class WeightedAvgTests extends AbstractAggregationTestCase {
30+
public class WeightedAvgTests extends AbstractConfigurationAggregationTestCase {
2931
public WeightedAvgTests(@Name("TestCase") Supplier<TestCaseSupplier.TestCase> testCaseSupplier) {
3032
this.testCase = testCaseSupplier.get();
3133
}
@@ -94,8 +96,8 @@ public static Iterable<Object[]> parameters() {
9496
}
9597

9698
@Override
97-
protected Expression build(Source source, List<Expression> args) {
98-
return new WeightedAvg(source, args.get(0), args.get(1), configuration());
99+
protected Expression buildWithConfiguration(Source source, List<Expression> args, Configuration configuration) {
100+
return new WeightedAvg(source, args.get(0), args.get(1), configuration);
99101
}
100102

101103
private static TestCaseSupplier makeSupplier(

0 commit comments

Comments
 (0)