Skip to content

Commit 569bf94

Browse files
yuancuyuancu
authored
Eliminate reliance on assert in Calcite for integration test (#4016)
* Move num-of-column check of in subquery ahead from RexSubQuery.java#L78 because assert is disabled in production Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Increase script.context.filter.max_compilations_rate for SQLCorrectnessIT Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Check script.disable_max_compilations_rate before setting context-specific compilations rate Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Refactor: remove some methods in tests to upper level to reduce duplication Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> * Replace plugin-level setting strings with private test-specific ones Signed-off-by: Yuanchun Shen <yuanchu@amazon.com> --------- Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>
1 parent 1977083 commit 569bf94

File tree

5 files changed

+125
-73
lines changed

5 files changed

+125
-73
lines changed

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java‎

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -539,23 +539,22 @@ public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) {
539539
List<RexNode> nodes = node.getChild().stream().map(child -> analyze(child, context)).toList();
540540
UnresolvedPlan subquery = node.getQuery();
541541
RelNode subqueryRel = resolveSubqueryPlan(subquery, context);
542-
try {
543-
return context.relBuilder.in(subqueryRel, nodes);
544-
// TODO
545-
// The {@link org.apache.calcite.tools.RelBuilder#in(RexNode,java.util.function.Function)}
546-
// only support one expression. Change to follow code after calcite fixed.
547-
// return context.relBuilder.in(
548-
// nodes.getFirst(),
549-
// b -> {
550-
// RelNode subqueryRel = subquery.accept(planVisitor, context);
551-
// b.build();
552-
// return subqueryRel;
553-
// });
554-
} catch (AssertionError e) {
542+
if (subqueryRel.getRowType().getFieldCount() != nodes.size()) {
555543
throw new SemanticCheckException(
556544
"The number of columns in the left hand side of an IN subquery does not match the number"
557545
+ " of columns in the output of subquery");
558546
}
547+
// TODO
548+
// The {@link org.apache.calcite.tools.RelBuilder#in(RexNode,java.util.function.Function)}
549+
// only support one expression. Change to follow code after calcite fixed.
550+
// return context.relBuilder.in(
551+
// nodes.getFirst(),
552+
// b -> {
553+
// RelNode subqueryRel = subquery.accept(planVisitor, context);
554+
// b.build();
555+
// return subqueryRel;
556+
// });
557+
return context.relBuilder.in(subqueryRel, nodes);
559558
}
560559

561560
@Override

‎integ-test/src/test/java/org/opensearch/sql/legacy/OpenSearchSQLRestTestCase.java‎

Lines changed: 109 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,11 +5,17 @@
55

66
package org.opensearch.sql.legacy;
77

8+
import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
9+
import static org.opensearch.sql.legacy.TestsConstants.PERSISTENT;
10+
import static org.opensearch.sql.legacy.TestsConstants.TRANSIENT;
11+
812
import java.io.IOException;
913
import java.util.ArrayList;
1014
import java.util.List;
1115
import java.util.Map;
16+
import java.util.Objects;
1217
import java.util.Optional;
18+
import org.apache.commons.lang3.StringUtils;
1319
import org.apache.hc.client5.http.auth.AuthScope;
1420
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
1521
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
@@ -29,14 +35,20 @@
2935
import org.json.JSONArray;
3036
import org.json.JSONObject;
3137
import org.junit.AfterClass;
38+
import org.junit.Assert;
3239
import org.opensearch.client.Request;
40+
import org.opensearch.client.RequestOptions;
3341
import org.opensearch.client.Response;
3442
import org.opensearch.client.RestClient;
3543
import org.opensearch.client.RestClientBuilder;
3644
import org.opensearch.common.settings.Settings;
3745
import org.opensearch.common.unit.TimeValue;
3846
import org.opensearch.common.util.concurrent.ThreadContext;
3947
import org.opensearch.common.util.io.IOUtils;
48+
import org.opensearch.common.xcontent.XContentFactory;
49+
import org.opensearch.core.rest.RestStatus;
50+
import org.opensearch.core.xcontent.XContentBuilder;
51+
import org.opensearch.script.ScriptService;
4052
import org.opensearch.test.rest.OpenSearchRestTestCase;
4153

4254
/**
@@ -63,7 +75,8 @@ public abstract class OpenSearchSQLRestTestCase extends OpenSearchRestTestCase {
6375
+ " }"
6476
+ "}"
6577
+ "}";
66-
78+
private static final String SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN =
79+
"script.context.*.max_compilations_rate";
6780
private static RestClient remoteClient;
6881

6982
/**
@@ -135,7 +148,13 @@ public RestClient initClient(String clusterName) throws IOException {
135148
int port = Integer.parseInt(stringUrl.substring(portSeparator + 1));
136149
hosts.add(buildHttpHost(host, port));
137150
}
138-
return buildClient(restClientSettings(), hosts.toArray(new HttpHost[0]));
151+
Settings.Builder builder = Settings.builder();
152+
if (System.getProperty("tests.rest.client_path_prefix") != null) {
153+
builder.put(CLIENT_PATH_PREFIX, System.getProperty("tests.rest.client_path_prefix"));
154+
}
155+
// Disable max compilations rate to avoid hitting compilations threshold during tests
156+
builder.put(ScriptService.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE_SETTING.getKey(), "true");
157+
return buildClient(builder.build(), hosts.toArray(new HttpHost[0]));
139158
}
140159

141160
/** Get a comma delimited list of [host:port] to which to send REST requests. */
@@ -301,4 +320,92 @@ public void configureMultiClusters(String remote) throws IOException {
301320
connectionRequest.setJsonEntity(connectionSetting);
302321
adminClient().performRequest(connectionRequest);
303322
}
323+
324+
/**
325+
* Increases the maximum script compilation rate for all script contexts for tests. This method
326+
* sets an unlimited compilation rate for each script context when the
327+
* script.disable_max_compilations_rate setting is not enabled, allowing tests to run without
328+
* hitting compilation rate limits.
329+
*
330+
* @throws IOException if there is an error retrieving cluster settings or updating them
331+
*/
332+
protected void increaseMaxCompilationsRate() throws IOException {
333+
// When script.disable_max_compilations_rate is set, custom context compilation rates cannot be
334+
// set
335+
if (!Objects.equals(
336+
getClusterSetting(
337+
ScriptService.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE_SETTING.getKey(), "persistent"),
338+
"true")) {
339+
List<String> contexts = getScriptContexts();
340+
for (String context : contexts) {
341+
String contextCompilationsRate =
342+
SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN.replace("*", context);
343+
updateClusterSetting(contextCompilationsRate, "unlimited", true);
344+
}
345+
}
346+
}
347+
348+
protected List<String> getScriptContexts() throws IOException {
349+
Request request = new Request("GET", "/_script_context");
350+
String responseBody = executeRequest(request);
351+
JSONObject jsonResponse = new JSONObject(responseBody);
352+
JSONArray contexts = jsonResponse.getJSONArray("contexts");
353+
List<String> contextNames = new ArrayList<>();
354+
for (int i = 0; i < contexts.length(); i++) {
355+
JSONObject context = contexts.getJSONObject(i);
356+
String contextName = context.getString("name");
357+
contextNames.add(contextName);
358+
}
359+
return contextNames;
360+
}
361+
362+
protected static void updateClusterSetting(String settingKey, Object value) throws IOException {
363+
updateClusterSetting(settingKey, value, true);
364+
}
365+
366+
protected static void updateClusterSetting(String settingKey, Object value, boolean persistent)
367+
throws IOException {
368+
String property = persistent ? PERSISTENT : TRANSIENT;
369+
XContentBuilder builder =
370+
XContentFactory.jsonBuilder()
371+
.startObject()
372+
.startObject(property)
373+
.field(settingKey, value)
374+
.endObject()
375+
.endObject();
376+
Request request = new Request("PUT", "_cluster/settings");
377+
request.setJsonEntity(builder.toString());
378+
Response response = client().performRequest(request);
379+
Assert.assertEquals(
380+
RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
381+
}
382+
383+
protected static JSONObject getAllClusterSettings() throws IOException {
384+
Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults");
385+
RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
386+
restOptionsBuilder.addHeader("Content-Type", "application/json");
387+
request.setOptions(restOptionsBuilder);
388+
return new JSONObject(executeRequest(request));
389+
}
390+
391+
protected static String getClusterSetting(String settingPath, String type) throws IOException {
392+
JSONObject settings = getAllClusterSettings();
393+
String value = settings.optJSONObject(type).optString(settingPath);
394+
if (StringUtils.isEmpty(value)) {
395+
return settings.optJSONObject("defaults").optString(settingPath);
396+
} else {
397+
return value;
398+
}
399+
}
400+
401+
protected static String executeRequest(final Request request) throws IOException {
402+
return executeRequest(request, client());
403+
}
404+
405+
protected static String executeRequest(final Request request, RestClient client)
406+
throws IOException {
407+
Response response = client.performRequest(request);
408+
Assert.assertEquals(200, response.getStatusLine().getStatusCode());
409+
return getResponseBody(response);
410+
}
304411
}

‎integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java‎

Lines changed: 2 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,6 @@
2626
import static org.opensearch.sql.legacy.TestUtils.getWeblogsIndexMapping;
2727
import static org.opensearch.sql.legacy.TestUtils.isIndexExist;
2828
import static org.opensearch.sql.legacy.TestUtils.loadDataByRestClient;
29-
import static org.opensearch.sql.legacy.TestsConstants.PERSISTENT;
30-
import static org.opensearch.sql.legacy.TestsConstants.TRANSIENT;
3129

3230
import java.io.IOException;
3331
import java.nio.file.Files;
@@ -39,13 +37,7 @@
3937
import javax.management.remote.JMXConnectorFactory;
4038
import javax.management.remote.JMXServiceURL;
4139
import org.junit.AfterClass;
42-
import org.junit.Assert;
4340
import org.junit.Before;
44-
import org.opensearch.client.Request;
45-
import org.opensearch.client.Response;
46-
import org.opensearch.common.xcontent.XContentFactory;
47-
import org.opensearch.core.rest.RestStatus;
48-
import org.opensearch.core.xcontent.XContentBuilder;
4941

5042
/**
5143
*
@@ -148,27 +140,8 @@ protected synchronized void loadIndex(Index index) throws IOException {
148140
}
149141

150142
/** Provide for each test to load test index, data and other setup work */
151-
protected void init() throws Exception {}
152-
153-
protected static void updateClusterSetting(String settingKey, Object value) throws IOException {
154-
updateClusterSetting(settingKey, value, true);
155-
}
156-
157-
protected static void updateClusterSetting(String settingKey, Object value, boolean persistent)
158-
throws IOException {
159-
String property = persistent ? PERSISTENT : TRANSIENT;
160-
XContentBuilder builder =
161-
XContentFactory.jsonBuilder()
162-
.startObject()
163-
.startObject(property)
164-
.field(settingKey, value)
165-
.endObject()
166-
.endObject();
167-
Request request = new Request("PUT", "_cluster/settings");
168-
request.setJsonEntity(builder.toString());
169-
Response response = client().performRequest(request);
170-
Assert.assertEquals(
171-
RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode()));
143+
protected void init() throws Exception {
144+
increaseMaxCompilationsRate();
172145
}
173146

174147
protected static void wipeAllClusterSettings() throws IOException {

‎integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java‎

Lines changed: 1 addition & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,7 @@ protected void resetMaxResultWindow(String indexName) throws IOException {
221221
/** Provide for each test to load test index, data and other setup work */
222222
protected void init() throws Exception {
223223
disableCalcite();
224+
increaseMaxCompilationsRate();
224225
}
225226

226227
/**
@@ -397,17 +398,6 @@ private String executeRequest(final String requestBody, final boolean isExplainQ
397398
return executeRequest(sqlRequest);
398399
}
399400

400-
protected static String executeRequest(final Request request, RestClient client)
401-
throws IOException {
402-
Response response = client.performRequest(request);
403-
Assert.assertEquals(200, response.getStatusLine().getStatusCode());
404-
return getResponseBody(response);
405-
}
406-
407-
protected static String executeRequest(final Request request) throws IOException {
408-
return executeRequest(request, client());
409-
}
410-
411401
protected JSONObject executeQueryWithGetRequest(final String sqlQuery) throws IOException {
412402

413403
final Request request = buildGetEndpointRequest(sqlQuery);
@@ -444,24 +434,6 @@ protected static JSONObject updateClusterSettings(ClusterSetting setting) throws
444434
return updateClusterSettings(setting, client());
445435
}
446436

447-
protected static JSONObject getAllClusterSettings() throws IOException {
448-
Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults");
449-
RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
450-
restOptionsBuilder.addHeader("Content-Type", "application/json");
451-
request.setOptions(restOptionsBuilder);
452-
return new JSONObject(executeRequest(request));
453-
}
454-
455-
protected static String getClusterSetting(String settingPath, String type) throws IOException {
456-
JSONObject settings = getAllClusterSettings();
457-
String value = settings.optJSONObject(type).optString(settingPath);
458-
if (StringUtils.isEmpty(value)) {
459-
return settings.optJSONObject("defaults").optString(settingPath);
460-
} else {
461-
return value;
462-
}
463-
}
464-
465437
protected static class ClusterSetting {
466438
private final String type;
467439
private final String name;

‎integ-test/src/test/java/org/opensearch/sql/sql/CorrectnessTestBase.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ public abstract class CorrectnessTestBase extends RestIntegTestCase {
3535

3636
@Override
3737
protected void init() throws Exception {
38+
super.init();
3839
if (runner != null) {
3940
return;
4041
}

0 commit comments

Comments
 (0)