Skip to content

Commit e8873ce

Browse files
dnhatndnhatn
committed
Add allow_partial_results cluster settings in ES|QL
1 parent b46dd1c commit e8873ce

File tree

11 files changed

+249
-6
lines changed

11 files changed

+249
-6
lines changed

‎test/external-modules/failing-field/build.gradle

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
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", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
apply plugin: 'elasticsearch.internal-java-rest-test'
11+
12+
13+
tasks.named('javaRestTest') {
14+
usesDefaultDistribution()
15+
it.onlyIf("snapshot build") { buildParams.isSnapshotBuild() }
16+
}
17+
18+
dependencies {
19+
api project(':test:framework')
20+
}
21+
22+
esplugin {
23+
description = 'A test module that includes runtime fields which throw exceptions when accessed'
24+
classname ='org.elasticsearch.test.FailingFieldPlugin'
25+
}

‎test/external-modules/failing-field/src/javaRestTest/java/org/elasticsearch/test/failingfield/EsqlPartialResultsIT.java

+140
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
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", the "GNU Affero General Public License v3.0 only", and the "Server Side
5+
* Public License v 1"; you may not use this file except in compliance with, at
6+
* your election, the "Elastic License 2.0", the "GNU Affero General Public
7+
* License v3.0 only", or the "Server Side Public License, v 1".
8+
*/
9+
10+
package org.elasticsearch.test.failingfield;
11+
12+
import org.apache.http.util.EntityUtils;
13+
import org.elasticsearch.client.Request;
14+
import org.elasticsearch.client.Response;
15+
import org.elasticsearch.client.ResponseException;
16+
import org.elasticsearch.common.settings.Settings;
17+
import org.elasticsearch.test.cluster.ElasticsearchCluster;
18+
import org.elasticsearch.test.cluster.local.distribution.DistributionType;
19+
import org.elasticsearch.test.rest.ESRestTestCase;
20+
import org.junit.ClassRule;
21+
22+
import java.util.HashSet;
23+
import java.util.List;
24+
import java.util.Map;
25+
import java.util.Set;
26+
27+
import static org.hamcrest.Matchers.containsString;
28+
import static org.hamcrest.Matchers.equalTo;
29+
import static org.hamcrest.Matchers.hasSize;
30+
31+
public class EsqlPartialResultsIT extends ESRestTestCase {
32+
@ClassRule
33+
public static ElasticsearchCluster cluster = ElasticsearchCluster.local()
34+
.distribution(DistributionType.DEFAULT)
35+
.module("test-failing-field")
36+
.setting("xpack.security.enabled", "false")
37+
.setting("xpack.license.self_generated.type", "trial")
38+
.setting("esql.query.allow_partial_results", "true")
39+
.build();
40+
41+
@Override
42+
protected String getTestRestCluster() {
43+
return cluster.getHttpAddresses();
44+
}
45+
46+
public Set<String> populateIndices() throws Exception {
47+
int nextId = 0;
48+
{
49+
createIndex("failing-index", Settings.EMPTY, """
50+
{
51+
"runtime": {
52+
"fail_me": {
53+
"type": "long",
54+
"script": {
55+
"source": "",
56+
"lang": "failing_field"
57+
}
58+
}
59+
},
60+
"properties": {
61+
"v": {
62+
"type": "long"
63+
}
64+
}
65+
}
66+
""");
67+
int numDocs = between(1, 50);
68+
for (int i = 0; i < numDocs; i++) {
69+
String id = Integer.toString(nextId++);
70+
Request doc = new Request("PUT", "failing-index/_doc/" + id);
71+
doc.setJsonEntity("{\"v\": " + id + "}");
72+
client().performRequest(doc);
73+
}
74+
75+
}
76+
Set<String> okIds = new HashSet<>();
77+
{
78+
createIndex("ok-index", Settings.EMPTY, """
79+
{
80+
"properties": {
81+
"v": {
82+
"type": "long"
83+
}
84+
}
85+
}
86+
""");
87+
int numDocs = between(1, 50);
88+
for (int i = 0; i < numDocs; i++) {
89+
String id = Integer.toString(nextId++);
90+
okIds.add(id);
91+
Request doc = new Request("PUT", "ok-index/_doc/" + id);
92+
doc.setJsonEntity("{\"v\": " + id + "}");
93+
client().performRequest(doc);
94+
}
95+
}
96+
refresh(client(), "failing-index,ok-index");
97+
return okIds;
98+
}
99+
100+
public void testPartialResult() throws Exception {
101+
Set<String> okIds = populateIndices();
102+
String query = """
103+
{
104+
"query": "FROM ok-index,failing-index | LIMIT 100 | KEEP fail_me,v"
105+
}
106+
""";
107+
// allow_partial_results = true
108+
{
109+
Request request = new Request("POST", "/_query");
110+
request.setJsonEntity(query);
111+
if (randomBoolean()) {
112+
request.addParameter("allow_partial_results", "true");
113+
}
114+
Response resp = client().performRequest(request);
115+
Map<String, Object> results = entityAsMap(resp);
116+
assertThat(results.get("is_partial"), equalTo(true));
117+
List<?> columns = (List<?>) results.get("columns");
118+
assertThat(columns, equalTo(List.of(
119+
Map.of("name", "fail_me", "type", "long"),
120+
Map.of("name", "v", "type", "long")
121+
)));
122+
List<?> values = (List<?>) results.get("values");
123+
assertThat(values, hasSize(okIds.size()));
124+
}
125+
// allow_partial_results = false
126+
{
127+
Request request = new Request("POST", "/_query");
128+
request.setJsonEntity("""
129+
{
130+
"query": "FROM ok-index,failing-index | LIMIT 100"
131+
}
132+
""");
133+
request.addParameter("allow_partial_results", "false");
134+
var error = expectThrows(ResponseException.class, () -> client().performRequest(request));
135+
Response resp = error.getResponse();
136+
assertThat(resp.getStatusLine().getStatusCode(), equalTo(500));
137+
assertThat(EntityUtils.toString(resp.getEntity()), containsString("Accessing failing field"));
138+
}
139+
}
140+
}

‎x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/esql/action/EsqlQueryRequestBuilder.java

+2
Original file line numberDiff line numberDiff line change
@@ -39,4 +39,6 @@ public final ActionType<Response> action() {
3939

4040
public abstract EsqlQueryRequestBuilder<Request, Response> filter(QueryBuilder filter);
4141

42+
public abstract EsqlQueryRequestBuilder<Request, Response> allowPartialResults(boolean allowPartialResults);
43+
4244
}

‎x-pack/plugin/esql/qa/server/src/main/java/org/elasticsearch/xpack/esql/qa/rest/RestEsqlTestCase.java

+9
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,7 @@ public static class RequestObjectBuilder {
127127
private Boolean includeCCSMetadata = null;
128128

129129
private CheckedConsumer<XContentBuilder, IOException> filter;
130+
private Boolean allPartialResults = null;
130131

131132
public RequestObjectBuilder() throws IOException {
132133
this(randomFrom(XContentType.values()));
@@ -204,6 +205,11 @@ public RequestObjectBuilder filter(CheckedConsumer<XContentBuilder, IOException>
204205
return this;
205206
}
206207

208+
public RequestObjectBuilder allPartialResults(boolean allPartialResults) {
209+
this.allPartialResults = allPartialResults;
210+
return this;
211+
}
212+
207213
public RequestObjectBuilder build() throws IOException {
208214
if (isBuilt == false) {
209215
if (tables != null) {
@@ -1151,6 +1157,9 @@ static Request prepareRequestWithOptions(RequestObjectBuilder requestObject, Mod
11511157
requestObject.build();
11521158
Request request = prepareRequest(mode);
11531159
String mediaType = attachBody(requestObject, request);
1160+
if (requestObject.allPartialResults != null) {
1161+
request.addParameter("allow_partial_results", String.valueOf(requestObject.allPartialResults));
1162+
}
11541163

11551164
RequestOptions.Builder options = request.getOptions().toBuilder();
11561165
options.setWarningsHandler(WarningsHandler.PERMISSIVE); // We assert the warnings ourselves

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

+44
Original file line numberDiff line numberDiff line change
@@ -17,13 +17,15 @@
1717
import org.elasticsearch.xcontent.XContentBuilder;
1818
import org.elasticsearch.xcontent.json.JsonXContent;
1919
import org.elasticsearch.xpack.esql.EsqlTestUtils;
20+
import org.elasticsearch.xpack.esql.plugin.EsqlPlugin;
2021

2122
import java.util.ArrayList;
2223
import java.util.Collection;
2324
import java.util.HashSet;
2425
import java.util.List;
2526
import java.util.Set;
2627

28+
import static org.elasticsearch.test.hamcrest.ElasticsearchAssertions.assertAcked;
2729
import static org.hamcrest.Matchers.equalTo;
2830
import static org.hamcrest.Matchers.in;
2931
import static org.hamcrest.Matchers.lessThanOrEqualTo;
@@ -122,4 +124,46 @@ public void testPartialResults() throws Exception {
122124
}
123125
}
124126
}
127+
128+
public void testDefaultPartialResults() throws Exception {
129+
Set<String> okIds = populateIndices();
130+
assertAcked(
131+
client().admin()
132+
.cluster()
133+
.prepareUpdateSettings(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS)
134+
.setPersistentSettings(Settings.builder().put(EsqlPlugin.QUERY_ALLOW_PARTIAL_RESULTS.getKey(), true))
135+
);
136+
try {
137+
// allow_partial_results = default
138+
{
139+
EsqlQueryRequest request = new EsqlQueryRequest();
140+
request.query("FROM fail,ok | LIMIT 100");
141+
request.pragmas(randomPragmas());
142+
if (randomBoolean()) {
143+
request.allowPartialResults(true);
144+
}
145+
try (EsqlQueryResponse resp = run(request)) {
146+
assertTrue(resp.isPartial());
147+
List<List<Object>> rows = EsqlTestUtils.getValuesList(resp);
148+
assertThat(rows.size(), lessThanOrEqualTo(okIds.size()));
149+
}
150+
}
151+
// allow_partial_results = false
152+
{
153+
EsqlQueryRequest request = new EsqlQueryRequest();
154+
request.query("FROM fail,ok | LIMIT 100");
155+
request.pragmas(randomPragmas());
156+
request.allowPartialResults(false);
157+
IllegalStateException e = expectThrows(IllegalStateException.class, () -> run(request).close());
158+
assertThat(e.getMessage(), equalTo("Accessing failing field"));
159+
}
160+
} finally {
161+
assertAcked(
162+
client().admin()
163+
.cluster()
164+
.prepareUpdateSettings(TimeValue.THIRTY_SECONDS, TimeValue.THIRTY_SECONDS)
165+
.setPersistentSettings(Settings.builder().putNull(EsqlPlugin.QUERY_ALLOW_PARTIAL_RESULTS.getKey()))
166+
);
167+
}
168+
}
125169
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequest.java

+4-3
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public class EsqlQueryRequest extends org.elasticsearch.xpack.core.esql.action.E
5252
private boolean keepOnCompletion;
5353
private boolean onSnapshotBuild = Build.current().isSnapshot();
5454
private boolean acceptedPragmaRisks = false;
55-
private boolean allowPartialResults = false;
55+
private Boolean allowPartialResults = null;
5656

5757
/**
5858
* "Tables" provided in the request for use with things like {@code LOOKUP}.
@@ -232,12 +232,13 @@ public Map<String, Map<String, Column>> tables() {
232232
return tables;
233233
}
234234

235-
public boolean allowPartialResults() {
235+
public Boolean allowPartialResults() {
236236
return allowPartialResults;
237237
}
238238

239-
public void allowPartialResults(boolean allowPartialResults) {
239+
public EsqlQueryRequest allowPartialResults(boolean allowPartialResults) {
240240
this.allowPartialResults = allowPartialResults;
241+
return this;
241242
}
242243

243244
@Override

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlQueryRequestBuilder.java

+6
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,12 @@ public EsqlQueryRequestBuilder keepOnCompletion(boolean keepOnCompletion) {
6666
return this;
6767
}
6868

69+
@Override
70+
public EsqlQueryRequestBuilder allowPartialResults(boolean allowPartialResults) {
71+
request.allowPartialResults(allowPartialResults);
72+
return this;
73+
}
74+
6975
static { // plumb access from x-pack core
7076
SharedSecrets.setEsqlQueryRequestBuilderAccess(EsqlQueryRequestBuilder::newSyncEsqlQueryRequestBuilder);
7177
}

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/RequestXContent.java

-2
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,6 @@ String fields() {
8585
static final ParseField WAIT_FOR_COMPLETION_TIMEOUT = new ParseField("wait_for_completion_timeout");
8686
static final ParseField KEEP_ALIVE = new ParseField("keep_alive");
8787
static final ParseField KEEP_ON_COMPLETION = new ParseField("keep_on_completion");
88-
static final ParseField ALLOW_PARTIAL_RESULTS = new ParseField("allow_partial_results");
8988

9089
private static final ObjectParser<EsqlQueryRequest, Void> SYNC_PARSER = objectParserSync(EsqlQueryRequest::syncEsqlQueryRequest);
9190
private static final ObjectParser<EsqlQueryRequest, Void> ASYNC_PARSER = objectParserAsync(EsqlQueryRequest::asyncEsqlQueryRequest);
@@ -115,7 +114,6 @@ private static void objectParserCommon(ObjectParser<EsqlQueryRequest, ?> parser)
115114
parser.declareString((request, localeTag) -> request.locale(Locale.forLanguageTag(localeTag)), LOCALE_FIELD);
116115
parser.declareBoolean(EsqlQueryRequest::profile, PROFILE_FIELD);
117116
parser.declareField((p, r, c) -> new ParseTables(r, p).parseTables(), TABLES_FIELD, ObjectParser.ValueType.OBJECT);
118-
parser.declareBoolean(EsqlQueryRequest::allowPartialResults, ALLOW_PARTIAL_RESULTS);
119117
}
120118

121119
private static ObjectParser<EsqlQueryRequest, Void> objectParserSync(Supplier<EsqlQueryRequest> supplier) {

‎x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/RestEsqlQueryAction.java

+4
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,10 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
5151
}
5252

5353
protected static RestChannelConsumer restChannelConsumer(EsqlQueryRequest esqlRequest, RestRequest request, NodeClient client) {
54+
final Boolean partialResults = request.paramAsBoolean("allow_partial_results", null);
55+
if (partialResults != null) {
56+
esqlRequest.allowPartialResults(partialResults);
57+
}
5458
LOGGER.debug("Beginning execution of ESQL query.\nQuery string: [{}]", esqlRequest.query());
5559

5660
return channel -> {

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

+8-1
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,13 @@ public class EsqlPlugin extends Plugin implements ActionPlugin {
102102
Setting.Property.Dynamic
103103
);
104104

105+
public static final Setting<Boolean> QUERY_ALLOW_PARTIAL_RESULTS = Setting.boolSetting(
106+
"esql.query.allow_partial_results",
107+
false,
108+
Setting.Property.NodeScope,
109+
Setting.Property.Dynamic
110+
);
111+
105112
@Override
106113
public Collection<?> createComponents(PluginServices services) {
107114
CircuitBreaker circuitBreaker = services.indicesService().getBigArrays().breakerService().getBreaker("request");
@@ -151,7 +158,7 @@ protected XPackLicenseState getLicenseState() {
151158
*/
152159
@Override
153160
public List<Setting<?>> getSettings() {
154-
return List.of(QUERY_RESULT_TRUNCATION_DEFAULT_SIZE, QUERY_RESULT_TRUNCATION_MAX_SIZE);
161+
return List.of(QUERY_RESULT_TRUNCATION_DEFAULT_SIZE, QUERY_RESULT_TRUNCATION_MAX_SIZE, QUERY_ALLOW_PARTIAL_RESULTS);
155162
}
156163

157164
@Override

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

+7
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ public class TransportEsqlQueryAction extends HandledTransportAction<EsqlQueryRe
8080
private final RemoteClusterService remoteClusterService;
8181
private final UsageService usageService;
8282
private final TransportActionServices services;
83+
private volatile boolean defaultAllowPartialResults;
8384

8485
@Inject
8586
@SuppressWarnings("this-escape")
@@ -158,6 +159,9 @@ public TransportEsqlQueryAction(
158159
indexNameExpressionResolver,
159160
usageService
160161
);
162+
defaultAllowPartialResults = EsqlPlugin.QUERY_ALLOW_PARTIAL_RESULTS.get(clusterService.getSettings());
163+
clusterService.getClusterSettings()
164+
.addSettingsUpdateConsumer(EsqlPlugin.QUERY_ALLOW_PARTIAL_RESULTS, v -> defaultAllowPartialResults = v);
161165
}
162166

163167
@Override
@@ -194,6 +198,9 @@ public void execute(EsqlQueryRequest request, EsqlQueryTask task, ActionListener
194198
}
195199

196200
private void innerExecute(Task task, EsqlQueryRequest request, ActionListener<EsqlQueryResponse> listener) {
201+
if (request.allowPartialResults() == null) {
202+
request.allowPartialResults(defaultAllowPartialResults);
203+
}
197204
Configuration configuration = new Configuration(
198205
ZoneOffset.UTC,
199206
request.locale() != null ? request.locale() : Locale.US,

0 commit comments

Comments
 (0)