Eliminate reliance on assert in Calcite for integration test

common/src/main/java/org/opensearch/sql/common/setting/Settings.java

@@ -61,6 +61,9 @@ public enum Key {
     SESSION_INACTIVITY_TIMEOUT_MILLIS(
         "plugins.query.executionengine.spark.session_inactivity_timeout_millis"),
 
+    SCRIPT_DISABLE_MAX_COMPILATIONS_RATE("script.disable_max_compilations_rate"),
+    SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN("script.context.*.max_compilations_rate"),
+
     /** Async query Settings * */
     ASYNC_QUERY_ENABLED("plugins.query.executionengine.async_query.enabled"),
     ASYNC_QUERY_EXTERNAL_SCHEDULER_ENABLED(

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

@@ -539,23 +539,22 @@ public RexNode visitInSubquery(InSubquery node, CalcitePlanContext context) {
     List<RexNode> 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

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

@@ -5,11 +5,18 @@
 
 package org.opensearch.sql.legacy;
 
+import static org.opensearch.sql.common.setting.Settings.Key;
+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,14 +36,19 @@
 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;
 import org.opensearch.common.settings.Settings;
 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.test.rest.OpenSearchRestTestCase;
 
 /**
@@ -135,7 +147,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(Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE.getKeyValue(), "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 +319,97 @@ 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(
+            org.opensearch.sql.common.setting.Settings.Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE
+                .getKeyValue(),
+            "persistent"),
+        "true")) {
+      List<String> contexts = getScriptContexts();
+      for (String context : contexts) {
+        String contextCompilationsRate =
+            org.opensearch.sql.common.setting.Settings.Key
+                .SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN
+                .getKeyValue()
+                .replace("*", context);
+        updateClusterSetting(contextCompilationsRate, "unlimited", true);
+      }
+    }
+  }
+
+  protected List<String> 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<String> 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);
+  }
 }

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 {

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;

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;
     }