diff --git a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java index f545f224d04..fec5d2dfa1d 100644 --- a/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java +++ b/core/src/main/java/org/opensearch/sql/calcite/CalciteRexNodeVisitor.java @@ -539,23 +539,22 @@ public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) { List nodes = node.getChild().stream().map(child -> analyze(child, context)).toList(); UnresolvedPlan subquery = node.getQuery(); RelNode subqueryRel = resolveSubqueryPlan(subquery, context); - try { - return context.relBuilder.in(subqueryRel, nodes); - // TODO - // The {@link org.apache.calcite.tools.RelBuilder#in(RexNode,java.util.function.Function)} - // only support one expression. Change to follow code after calcite fixed. - // return context.relBuilder.in( - // nodes.getFirst(), - // b -> { - // RelNode subqueryRel = subquery.accept(planVisitor, context); - // b.build(); - // return subqueryRel; - // }); - } catch (AssertionError e) { + if (subqueryRel.getRowType().getFieldCount() != nodes.size()) { throw new SemanticCheckException( "The number of columns in the left hand side of an IN subquery does not match the number" + " of columns in the output of subquery"); } + // TODO + // The {@link org.apache.calcite.tools.RelBuilder#in(RexNode,java.util.function.Function)} + // only support one expression. Change to follow code after calcite fixed. + // return context.relBuilder.in( + // nodes.getFirst(), + // b -> { + // RelNode subqueryRel = subquery.accept(planVisitor, context); + // b.build(); + // return subqueryRel; + // }); + return context.relBuilder.in(subqueryRel, nodes); } @Override diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/OpenSearchSQLRestTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/OpenSearchSQLRestTestCase.java index ced69d54a07..91cf5dc943c 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/OpenSearchSQLRestTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/OpenSearchSQLRestTestCase.java @@ -5,11 +5,17 @@ package org.opensearch.sql.legacy; +import static org.opensearch.sql.legacy.TestUtils.getResponseBody; +import static org.opensearch.sql.legacy.TestsConstants.PERSISTENT; +import static org.opensearch.sql.legacy.TestsConstants.TRANSIENT; + import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; +import org.apache.commons.lang3.StringUtils; import org.apache.hc.client5.http.auth.AuthScope; import org.apache.hc.client5.http.auth.UsernamePasswordCredentials; import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider; @@ -29,7 +35,9 @@ import org.json.JSONArray; import org.json.JSONObject; import org.junit.AfterClass; +import org.junit.Assert; import org.opensearch.client.Request; +import org.opensearch.client.RequestOptions; import org.opensearch.client.Response; import org.opensearch.client.RestClient; import org.opensearch.client.RestClientBuilder; @@ -37,6 +45,10 @@ import org.opensearch.common.unit.TimeValue; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.common.util.io.IOUtils; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.core.rest.RestStatus; +import org.opensearch.core.xcontent.XContentBuilder; +import org.opensearch.script.ScriptService; import org.opensearch.test.rest.OpenSearchRestTestCase; /** @@ -63,7 +75,8 @@ public abstract class OpenSearchSQLRestTestCase extends OpenSearchRestTestCase { + " }" + "}" + "}"; - + private static final String SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN = + "script.context.*.max_compilations_rate"; private static RestClient remoteClient; /** @@ -135,7 +148,13 @@ public RestClient initClient(String clusterName) throws IOException { int port = Integer.parseInt(stringUrl.substring(portSeparator + 1)); hosts.add(buildHttpHost(host, port)); } - return buildClient(restClientSettings(), hosts.toArray(new HttpHost[0])); + Settings.Builder builder = Settings.builder(); + if (System.getProperty("tests.rest.client_path_prefix") != null) { + builder.put(CLIENT_PATH_PREFIX, System.getProperty("tests.rest.client_path_prefix")); + } + // Disable max compilations rate to avoid hitting compilations threshold during tests + builder.put(ScriptService.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE_SETTING.getKey(), "true"); + return buildClient(builder.build(), hosts.toArray(new HttpHost[0])); } /** 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 { connectionRequest.setJsonEntity(connectionSetting); adminClient().performRequest(connectionRequest); } + + /** + * Increases the maximum script compilation rate for all script contexts for tests. This method + * sets an unlimited compilation rate for each script context when the + * script.disable_max_compilations_rate setting is not enabled, allowing tests to run without + * hitting compilation rate limits. + * + * @throws IOException if there is an error retrieving cluster settings or updating them + */ + protected void increaseMaxCompilationsRate() throws IOException { + // When script.disable_max_compilations_rate is set, custom context compilation rates cannot be + // set + if (!Objects.equals( + getClusterSetting( + ScriptService.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE_SETTING.getKey(), "persistent"), + "true")) { + List contexts = getScriptContexts(); + for (String context : contexts) { + String contextCompilationsRate = + SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN.replace("*", context); + updateClusterSetting(contextCompilationsRate, "unlimited", true); + } + } + } + + protected List getScriptContexts() throws IOException { + Request request = new Request("GET", "/_script_context"); + String responseBody = executeRequest(request); + JSONObject jsonResponse = new JSONObject(responseBody); + JSONArray contexts = jsonResponse.getJSONArray("contexts"); + List contextNames = new ArrayList<>(); + for (int i = 0; i < contexts.length(); i++) { + JSONObject context = contexts.getJSONObject(i); + String contextName = context.getString("name"); + contextNames.add(contextName); + } + return contextNames; + } + + protected static void updateClusterSetting(String settingKey, Object value) throws IOException { + updateClusterSetting(settingKey, value, true); + } + + protected static void updateClusterSetting(String settingKey, Object value, boolean persistent) + throws IOException { + String property = persistent ? PERSISTENT : TRANSIENT; + XContentBuilder builder = + XContentFactory.jsonBuilder() + .startObject() + .startObject(property) + .field(settingKey, value) + .endObject() + .endObject(); + Request request = new Request("PUT", "_cluster/settings"); + request.setJsonEntity(builder.toString()); + Response response = client().performRequest(request); + Assert.assertEquals( + RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + } + + protected static JSONObject getAllClusterSettings() throws IOException { + Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults"); + RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder(); + restOptionsBuilder.addHeader("Content-Type", "application/json"); + request.setOptions(restOptionsBuilder); + return new JSONObject(executeRequest(request)); + } + + protected static String getClusterSetting(String settingPath, String type) throws IOException { + JSONObject settings = getAllClusterSettings(); + String value = settings.optJSONObject(type).optString(settingPath); + if (StringUtils.isEmpty(value)) { + return settings.optJSONObject("defaults").optString(settingPath); + } else { + return value; + } + } + + protected static String executeRequest(final Request request) throws IOException { + return executeRequest(request, client()); + } + + protected static String executeRequest(final Request request, RestClient client) + throws IOException { + Response response = client.performRequest(request); + Assert.assertEquals(200, response.getStatusLine().getStatusCode()); + return getResponseBody(response); + } } diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java index 3d53b966683..a91e7581b8f 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/RestIntegTestCase.java @@ -26,8 +26,6 @@ import static org.opensearch.sql.legacy.TestUtils.getWeblogsIndexMapping; import static org.opensearch.sql.legacy.TestUtils.isIndexExist; import static org.opensearch.sql.legacy.TestUtils.loadDataByRestClient; -import static org.opensearch.sql.legacy.TestsConstants.PERSISTENT; -import static org.opensearch.sql.legacy.TestsConstants.TRANSIENT; import java.io.IOException; import java.nio.file.Files; @@ -39,13 +37,7 @@ import javax.management.remote.JMXConnectorFactory; import javax.management.remote.JMXServiceURL; import org.junit.AfterClass; -import org.junit.Assert; import org.junit.Before; -import org.opensearch.client.Request; -import org.opensearch.client.Response; -import org.opensearch.common.xcontent.XContentFactory; -import org.opensearch.core.rest.RestStatus; -import org.opensearch.core.xcontent.XContentBuilder; /** * @@ -148,27 +140,8 @@ protected synchronized void loadIndex(Index index) throws IOException { } /** Provide for each test to load test index, data and other setup work */ - protected void init() throws Exception {} - - protected static void updateClusterSetting(String settingKey, Object value) throws IOException { - updateClusterSetting(settingKey, value, true); - } - - protected static void updateClusterSetting(String settingKey, Object value, boolean persistent) - throws IOException { - String property = persistent ? PERSISTENT : TRANSIENT; - XContentBuilder builder = - XContentFactory.jsonBuilder() - .startObject() - .startObject(property) - .field(settingKey, value) - .endObject() - .endObject(); - Request request = new Request("PUT", "_cluster/settings"); - request.setJsonEntity(builder.toString()); - Response response = client().performRequest(request); - Assert.assertEquals( - RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); + protected void init() throws Exception { + increaseMaxCompilationsRate(); } protected static void wipeAllClusterSettings() throws IOException { diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java index e951cbc678b..a6afc2d0ddd 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java @@ -221,6 +221,7 @@ protected void resetMaxResultWindow(String indexName) throws IOException { /** Provide for each test to load test index, data and other setup work */ protected void init() throws Exception { disableCalcite(); + increaseMaxCompilationsRate(); } /** @@ -397,17 +398,6 @@ private String executeRequest(final String requestBody, final boolean isExplainQ return executeRequest(sqlRequest); } - protected static String executeRequest(final Request request, RestClient client) - throws IOException { - Response response = client.performRequest(request); - Assert.assertEquals(200, response.getStatusLine().getStatusCode()); - return getResponseBody(response); - } - - protected static String executeRequest(final Request request) throws IOException { - return executeRequest(request, client()); - } - protected JSONObject executeQueryWithGetRequest(final String sqlQuery) throws IOException { final Request request = buildGetEndpointRequest(sqlQuery); @@ -444,24 +434,6 @@ protected static JSONObject updateClusterSettings(ClusterSetting setting) throws return updateClusterSettings(setting, client()); } - protected static JSONObject getAllClusterSettings() throws IOException { - Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults"); - RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder(); - restOptionsBuilder.addHeader("Content-Type", "application/json"); - request.setOptions(restOptionsBuilder); - return new JSONObject(executeRequest(request)); - } - - protected static String getClusterSetting(String settingPath, String type) throws IOException { - JSONObject settings = getAllClusterSettings(); - String value = settings.optJSONObject(type).optString(settingPath); - if (StringUtils.isEmpty(value)) { - return settings.optJSONObject("defaults").optString(settingPath); - } else { - return value; - } - } - protected static class ClusterSetting { private final String type; private final String name; diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/CorrectnessTestBase.java b/integ-test/src/test/java/org/opensearch/sql/sql/CorrectnessTestBase.java index 33c9c0687f1..ce50a54d84c 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/CorrectnessTestBase.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/CorrectnessTestBase.java @@ -35,6 +35,7 @@ public abstract class CorrectnessTestBase extends RestIntegTestCase { @Override protected void init() throws Exception { + super.init(); if (runner != null) { return; }