Skip to content

Commit 1170bcd

Browse files
idegtiarenkoidegtiarenko
committed
Replace overloaded run with builder
1 parent 1ec4e0b commit 1170bcd

File tree

7 files changed

+120
-67
lines changed

7 files changed

+120
-67
lines changed

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/AbstractEsqlIntegTestCase.java‎

Lines changed: 3 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import org.elasticsearch.core.TimeValue;
2424
import org.elasticsearch.core.Tuple;
2525
import org.elasticsearch.health.node.selection.HealthNode;
26-
import org.elasticsearch.index.query.QueryBuilder;
2726
import org.elasticsearch.indices.breaker.CircuitBreakerService;
2827
import org.elasticsearch.indices.breaker.HierarchyCircuitBreakerService;
2928
import org.elasticsearch.plugins.Plugin;
@@ -43,6 +42,7 @@
4342

4443
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
4544
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
45+
import static org.elasticsearch.xpack.esql.action.EsqlQueryRequest.syncEsqlQueryRequest;
4646
import static org.hamcrest.Matchers.containsInAnyOrder;
4747
import static org.hamcrest.Matchers.equalTo;
4848

@@ -166,33 +166,10 @@ protected void setRequestCircuitBreakerLimit(ByteSizeValue limit) {
166166
}
167167

168168
protected final EsqlQueryResponse run(String esqlCommands) {
169-
return run(esqlCommands, randomPragmas());
169+
return run(syncEsqlQueryRequest().query(esqlCommands).pragmas(randomPragmas()));
170170
}
171171

172-
protected final EsqlQueryResponse run(String esqlCommands, QueryPragmas pragmas) {
173-
return run(esqlCommands, pragmas, null);
174-
}
175-
176-
protected EsqlQueryResponse run(String esqlCommands, QueryPragmas pragmas, QueryBuilder filter) {
177-
return run(esqlCommands, pragmas, filter, null);
178-
}
179-
180-
protected EsqlQueryResponse run(String esqlCommands, QueryPragmas pragmas, QueryBuilder filter, Boolean allowPartialResults) {
181-
EsqlQueryRequest request = EsqlQueryRequest.syncEsqlQueryRequest();
182-
request.query(esqlCommands);
183-
if (pragmas != null) {
184-
request.pragmas(pragmas);
185-
}
186-
if (filter != null) {
187-
request.filter(filter);
188-
}
189-
if (allowPartialResults != null) {
190-
request.allowPartialResults(allowPartialResults);
191-
}
192-
return run(request);
193-
}
194-
195-
protected EsqlQueryResponse run(EsqlQueryRequest request) {
172+
public EsqlQueryResponse run(EsqlQueryRequest request) {
196173
try {
197174
return client().execute(EsqlQueryAction.INSTANCE, request).actionGet(30, TimeUnit.SECONDS);
198175
} catch (ElasticsearchTimeoutException e) {

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/EsqlActionIT.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@
8686
import static org.elasticsearch.test.MapMatcher.assertMap;
8787
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
8888
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
89+
import static org.elasticsearch.xpack.esql.action.EsqlQueryRequest.syncEsqlQueryRequest;
8990
import static org.hamcrest.Matchers.allOf;
9091
import static org.hamcrest.Matchers.anyOf;
9192
import static org.hamcrest.Matchers.contains;
@@ -1916,7 +1917,7 @@ public void testScriptField() throws Exception {
19161917
pragmaSettings.put("data_partitioning", "doc");
19171918
pragmas = new QueryPragmas(pragmaSettings.build());
19181919
}
1919-
try (EsqlQueryResponse resp = run("FROM test-script | SORT k1 | LIMIT " + numDocs, pragmas)) {
1920+
try (EsqlQueryResponse resp = run(syncEsqlQueryRequest().query("FROM test-script | SORT k1 | LIMIT " + numDocs).pragmas(pragmas))) {
19201921
List<Object> k1Column = Iterators.toList(resp.column(0));
19211922
assertThat(k1Column, equalTo(LongStream.range(0L, numDocs).boxed().toList()));
19221923
List<Object> k2Column = Iterators.toList(resp.column(1));

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/ManyShardsIT.java‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import java.util.concurrent.atomic.AtomicInteger;
4646
import java.util.concurrent.atomic.AtomicReference;
4747

48+
import static org.elasticsearch.xpack.esql.action.EsqlQueryRequest.syncEsqlQueryRequest;
4849
import static org.hamcrest.Matchers.equalTo;
4950
import static org.hamcrest.Matchers.instanceOf;
5051
import static org.hamcrest.Matchers.lessThanOrEqualTo;
@@ -116,7 +117,11 @@ public void testConcurrentQueries() throws Exception {
116117
.put("task_concurrency", between(1, 2))
117118
.put("exchange_concurrent_clients", between(1, 2));
118119
}
119-
try (var response = run("from test-* | stats count(user) by tags", new QueryPragmas(pragmas.build()))) {
120+
try (
121+
var response = run(
122+
syncEsqlQueryRequest().query("from test-* | stats count(user) by tags").pragmas(new QueryPragmas(pragmas.build()))
123+
)
124+
) {
120125
// do nothing
121126
} catch (Exception | AssertionError e) {
122127
logger.warn("Query failed with exception", e);
@@ -250,7 +255,7 @@ public void testLimitConcurrentShards() {
250255
mockSearchService.setOnPutContext(r -> counter.onNewContext());
251256
mockSearchService.setOnRemoveContext(r -> counter.onContextReleased());
252257
}
253-
run(q, pragmas).close();
258+
run(syncEsqlQueryRequest().query(q).pragmas(pragmas)).close();
254259
}
255260
} finally {
256261
for (SearchService searchService : searchServices) {
@@ -277,7 +282,7 @@ public void testCancelUnnecessaryRequests() {
277282
connection.sendRequest(requestId, action, request, options);
278283
});
279284

280-
var query = EsqlQueryRequest.syncEsqlQueryRequest();
285+
var query = syncEsqlQueryRequest();
281286
query.query("from test-* | LIMIT 1");
282287
query.pragmas(new QueryPragmas(Settings.builder().put(QueryPragmas.MAX_CONCURRENT_NODES_PER_CLUSTER.getKey(), 1).build()));
283288

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/TimeBasedIndicesIT.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
1919
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
20+
import static org.elasticsearch.xpack.esql.action.EsqlQueryRequest.syncEsqlQueryRequest;
2021
import static org.hamcrest.Matchers.hasSize;
2122

2223
public class TimeBasedIndicesIT extends AbstractEsqlIntegTestCase {
@@ -39,15 +40,15 @@ public void testFilter() {
3940
{
4041
String query = "FROM test | limit 1000";
4142
var filter = new RangeQueryBuilder("@timestamp").from(epoch - TimeValue.timeValueHours(3).millis()).to("now");
42-
try (var resp = run(query, null, filter)) {
43+
try (var resp = run(syncEsqlQueryRequest().query(query).filter(filter))) {
4344
List<List<Object>> values = getValuesList(resp);
4445
assertThat(values, hasSize(oldDocs));
4546
}
4647
}
4748
{
4849
String query = "FROM test | limit 1000";
4950
var filter = new RangeQueryBuilder("@timestamp").from("now").to(epoch + TimeValue.timeValueHours(3).millis());
50-
try (var resp = run(query, null, filter)) {
51+
try (var resp = run(syncEsqlQueryRequest().query(query).filter(filter))) {
5152
List<List<Object>> values = getValuesList(resp);
5253
assertThat(values, hasSize(newDocs));
5354
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/CanMatchIT.java‎

Lines changed: 90 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
3333
import static org.elasticsearch.xpack.esql.EsqlTestUtils.as;
3434
import static org.elasticsearch.xpack.esql.EsqlTestUtils.getValuesList;
35+
import static org.elasticsearch.xpack.esql.action.EsqlQueryRequest.syncEsqlQueryRequest;
3536
import static org.hamcrest.Matchers.containsString;
3637
import static org.hamcrest.Matchers.empty;
3738
import static org.hamcrest.Matchers.equalTo;
@@ -82,7 +83,13 @@ public void testCanMatch() {
8283
}
8384
);
8485
}
85-
try (EsqlQueryResponse resp = run("from events_*", randomPragmas(), new RangeQueryBuilder("@timestamp").gte("2023-01-01"))) {
86+
try (
87+
EsqlQueryResponse resp = run(
88+
syncEsqlQueryRequest().query("from events_*")
89+
.pragmas(randomPragmas())
90+
.filter(new RangeQueryBuilder("@timestamp").gte("2023-01-01"))
91+
)
92+
) {
8693
assertThat(getValuesList(resp), hasSize(4));
8794
assertThat(queriedIndices, equalTo(Set.of("events_2023")));
8895
queriedIndices.clear();
@@ -93,7 +100,13 @@ public void testCanMatch() {
93100
queriedIndices.clear();
94101
}
95102

96-
try (EsqlQueryResponse resp = run("from events_*", randomPragmas(), new RangeQueryBuilder("@timestamp").lt("2023-01-01"))) {
103+
try (
104+
EsqlQueryResponse resp = run(
105+
syncEsqlQueryRequest().query("from events_*")
106+
.pragmas(randomPragmas())
107+
.filter(new RangeQueryBuilder("@timestamp").lt("2023-01-01"))
108+
)
109+
) {
97110
assertThat(getValuesList(resp), hasSize(3));
98111
assertThat(queriedIndices, equalTo(Set.of("events_2022")));
99112
queriedIndices.clear();
@@ -106,9 +119,9 @@ public void testCanMatch() {
106119

107120
try (
108121
EsqlQueryResponse resp = run(
109-
"from events_*",
110-
randomPragmas(),
111-
new RangeQueryBuilder("@timestamp").gt("2022-01-01").lt("2023-12-31")
122+
syncEsqlQueryRequest().query("from events_*")
123+
.pragmas(randomPragmas())
124+
.filter(new RangeQueryBuilder("@timestamp").gt("2022-01-01").lt("2023-12-31"))
112125
)
113126
) {
114127
assertThat(getValuesList(resp), hasSize(7));
@@ -129,9 +142,9 @@ public void testCanMatch() {
129142

130143
try (
131144
EsqlQueryResponse resp = run(
132-
"from events_*",
133-
randomPragmas(),
134-
new RangeQueryBuilder("@timestamp").gt("2021-01-01").lt("2021-12-31")
145+
syncEsqlQueryRequest().query("from events_*")
146+
.pragmas(randomPragmas())
147+
.filter(new RangeQueryBuilder("@timestamp").gt("2021-01-01").lt("2021-12-31"))
135148
)
136149
) {
137150
assertThat(getValuesList(resp), hasSize(0));
@@ -181,62 +194,110 @@ public void testAliasFilters() {
181194
.addAlias("employees", "sales", new MatchQueryBuilder("dept", "sales"))
182195
);
183196
// employees index
184-
try (var resp = run("from employees | stats count(emp_no)", randomPragmas())) {
197+
try (var resp = run("from employees | stats count(emp_no)")) {
185198
assertThat(getValuesList(resp).get(0), equalTo(List.of(6L)));
186199
}
187-
try (var resp = run("from employees | stats avg(salary)", randomPragmas())) {
200+
try (var resp = run("from employees | stats avg(salary)")) {
188201
assertThat(getValuesList(resp).get(0), equalTo(List.of(26.95d)));
189202
}
190203

191-
try (var resp = run("from employees | stats count(emp_no)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
204+
try (
205+
var resp = run(
206+
syncEsqlQueryRequest().query("from employees | stats count(emp_no)")
207+
.pragmas(randomPragmas())
208+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
209+
)
210+
) {
192211
assertThat(getValuesList(resp).get(0), equalTo(List.of(4L)));
193212
}
194-
try (var resp = run("from employees | stats avg(salary)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
213+
try (
214+
var resp = run(
215+
syncEsqlQueryRequest().query("from employees | stats avg(salary)")
216+
.pragmas(randomPragmas())
217+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
218+
)
219+
) {
195220
assertThat(getValuesList(resp).get(0), equalTo(List.of(26.65d)));
196221
}
197222

198223
// match both employees index and engineers alias -> employees
199-
try (var resp = run("from e* | stats count(emp_no)", randomPragmas())) {
224+
try (var resp = run("from e* | stats count(emp_no)")) {
200225
assertThat(getValuesList(resp).get(0), equalTo(List.of(6L)));
201226
}
202-
try (var resp = run("from employees | stats avg(salary)", randomPragmas())) {
227+
try (var resp = run("from employees | stats avg(salary)")) {
203228
assertThat(getValuesList(resp).get(0), equalTo(List.of(26.95d)));
204229
}
205230

206-
try (var resp = run("from e* | stats count(emp_no)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
231+
try (
232+
var resp = run(
233+
syncEsqlQueryRequest().query("from e* | stats count(emp_no)")
234+
.pragmas(randomPragmas())
235+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
236+
)
237+
) {
207238
assertThat(getValuesList(resp).get(0), equalTo(List.of(4L)));
208239
}
209-
try (var resp = run("from e* | stats avg(salary)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
240+
try (
241+
var resp = run(
242+
syncEsqlQueryRequest().query("from e* | stats avg(salary)")
243+
.pragmas(randomPragmas())
244+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
245+
)
246+
) {
210247
assertThat(getValuesList(resp).get(0), equalTo(List.of(26.65d)));
211248
}
212249

213250
// engineers alias
214-
try (var resp = run("from engineer* | stats count(emp_no)", randomPragmas())) {
251+
try (var resp = run("from engineer* | stats count(emp_no)")) {
215252
assertThat(getValuesList(resp).get(0), equalTo(List.of(4L)));
216253
}
217-
try (var resp = run("from engineer* | stats avg(salary)", randomPragmas())) {
254+
try (var resp = run("from engineer* | stats avg(salary)")) {
218255
assertThat(getValuesList(resp).get(0), equalTo(List.of(26.65d)));
219256
}
220257

221-
try (var resp = run("from engineer* | stats count(emp_no)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
258+
try (
259+
var resp = run(
260+
syncEsqlQueryRequest().query("from engineer* | stats count(emp_no)")
261+
.pragmas(randomPragmas())
262+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
263+
)
264+
) {
222265
assertThat(getValuesList(resp).get(0), equalTo(List.of(3L)));
223266
}
224-
try (var resp = run("from engineer* | stats avg(salary)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
267+
try (
268+
var resp = run(
269+
syncEsqlQueryRequest().query("from engineer* | stats avg(salary)")
270+
.pragmas(randomPragmas())
271+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
272+
)
273+
) {
225274
assertThat(getValuesList(resp).get(0), equalTo(List.of(27.2d)));
226275
}
227276

228277
// sales alias
229-
try (var resp = run("from sales | stats count(emp_no)", randomPragmas())) {
278+
try (var resp = run("from sales | stats count(emp_no)")) {
230279
assertThat(getValuesList(resp).get(0), equalTo(List.of(2L)));
231280
}
232-
try (var resp = run("from sales | stats avg(salary)", randomPragmas())) {
281+
try (var resp = run("from sales | stats avg(salary)")) {
233282
assertThat(getValuesList(resp).get(0), equalTo(List.of(27.55d)));
234283
}
235284

236-
try (var resp = run("from sales | stats count(emp_no)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
285+
try (
286+
var resp = run(
287+
syncEsqlQueryRequest().query("from sales | stats count(emp_no)")
288+
.pragmas(randomPragmas())
289+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
290+
)
291+
) {
237292
assertThat(getValuesList(resp).get(0), equalTo(List.of(1L)));
238293
}
239-
try (var resp = run("from sales | stats avg(salary)", randomPragmas(), new RangeQueryBuilder("hired").lt("2012-04-30"))) {
294+
try (
295+
var resp = run(
296+
syncEsqlQueryRequest().query("from sales | stats avg(salary)")
297+
.pragmas(randomPragmas())
298+
.filter(new RangeQueryBuilder("hired").lt("2012-04-30"))
299+
)
300+
) {
240301
assertThat(getValuesList(resp).get(0), equalTo(List.of(25.0d)));
241302
}
242303
}
@@ -297,7 +358,11 @@ public void testFailOnUnavailableShards() throws Exception {
297358
containsString("index [logs] has no active shard copy"),
298359
() -> run("from * | KEEP timestamp,message")
299360
);
300-
try (EsqlQueryResponse resp = run("from events,logs | KEEP timestamp,message", null, null, true)) {
361+
try (
362+
EsqlQueryResponse resp = run(
363+
syncEsqlQueryRequest().query("from events,logs | KEEP timestamp,message").allowPartialResults(true)
364+
)
365+
) {
301366
assertThat(getValuesList(resp), hasSize(3));
302367
}
303368
}

‎x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/plugin/MatchOperatorIT.java‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import static org.elasticsearch.index.query.QueryBuilders.boolQuery;
2323
import static org.elasticsearch.index.query.QueryBuilders.matchQuery;
2424
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
25+
import static org.elasticsearch.xpack.esql.action.EsqlQueryRequest.syncEsqlQueryRequest;
2526
import static org.hamcrest.CoreMatchers.containsString;
2627

2728
//@TestLogging(value = "org.elasticsearch.xpack.esql:TRACE,org.elasticsearch.compute:TRACE", reason = "debug")
@@ -137,7 +138,7 @@ public void testWhereMatchWithScoring_AndRequestFilter() {
137138

138139
QueryBuilder filter = boolQuery().must(matchQuery("content", "brown"));
139140

140-
try (var resp = run(query, randomPragmas(), filter)) {
141+
try (var resp = run(syncEsqlQueryRequest().query(query).pragmas(randomPragmas()).filter(filter))) {
141142
assertColumnNames(resp.columns(), List.of("content", "_score"));
142143
assertColumnTypes(resp.columns(), List.of("text", "double"));
143144
assertValues(
@@ -160,7 +161,7 @@ public void testWhereMatchWithScoring_AndNoScoreRequestFilter() {
160161

161162
QueryBuilder filter = boolQuery().filter(matchQuery("content", "brown"));
162163

163-
try (var resp = run(query, randomPragmas(), filter)) {
164+
try (var resp = run(syncEsqlQueryRequest().query(query).pragmas(randomPragmas()).filter(filter))) {
164165
assertColumnNames(resp.columns(), List.of("content", "_score"));
165166
assertColumnTypes(resp.columns(), List.of("text", "double"));
166167
assertValues(
@@ -183,7 +184,7 @@ public void testWhereMatchWithScoring_And_MatchAllRequestFilter() {
183184

184185
QueryBuilder filter = QueryBuilders.matchAllQuery();
185186

186-
try (var resp = run(query, randomPragmas(), filter)) {
187+
try (var resp = run(syncEsqlQueryRequest().query(query).pragmas(randomPragmas()).filter(filter))) {
187188
assertColumnNames(resp.columns(), List.of("content", "_score"));
188189
assertColumnTypes(resp.columns(), List.of("text", "double"));
189190
assertValues(
@@ -205,7 +206,7 @@ public void testScoringOutsideQuery() {
205206

206207
QueryBuilder filter = boolQuery().must(matchQuery("content", "fox"));
207208

208-
try (var resp = run(query, randomPragmas(), filter)) {
209+
try (var resp = run(syncEsqlQueryRequest().query(query).pragmas(randomPragmas()).filter(filter))) {
209210
assertColumnNames(resp.columns(), List.of("content", "_score"));
210211
assertColumnTypes(resp.columns(), List.of("text", "double"));
211212
assertValues(
@@ -227,7 +228,7 @@ public void testScoring_Zero_OutsideQuery() {
227228

228229
QueryBuilder filter = boolQuery().filter(matchQuery("content", "fox"));
229230

230-
try (var resp = run(query, randomPragmas(), filter)) {
231+
try (var resp = run(syncEsqlQueryRequest().query(query).pragmas(randomPragmas()).filter(filter))) {
231232
assertColumnNames(resp.columns(), List.of("content", "_score"));
232233
assertColumnTypes(resp.columns(), List.of("text", "double"));
233234
assertValues(

0 commit comments

Comments
 (0)