|
5 | 5 |
|
6 | 6 | package org.opensearch.sql.legacy; |
7 | 7 |
|
| 8 | +import static org.opensearch.sql.common.setting.Settings.Key; |
| 9 | +import static org.opensearch.sql.legacy.TestUtils.getResponseBody; |
| 10 | +import static org.opensearch.sql.legacy.TestsConstants.PERSISTENT; |
| 11 | +import static org.opensearch.sql.legacy.TestsConstants.TRANSIENT; |
| 12 | + |
8 | 13 | import java.io.IOException; |
9 | 14 | import java.util.ArrayList; |
10 | 15 | import java.util.List; |
11 | 16 | import java.util.Map; |
| 17 | +import java.util.Objects; |
12 | 18 | import java.util.Optional; |
| 19 | +import org.apache.commons.lang3.StringUtils; |
13 | 20 | import org.apache.hc.client5.http.auth.AuthScope; |
14 | 21 | import org.apache.hc.client5.http.auth.UsernamePasswordCredentials; |
15 | 22 | import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider; |
|
29 | 36 | import org.json.JSONArray; |
30 | 37 | import org.json.JSONObject; |
31 | 38 | import org.junit.AfterClass; |
| 39 | +import org.junit.Assert; |
32 | 40 | import org.opensearch.client.Request; |
| 41 | +import org.opensearch.client.RequestOptions; |
33 | 42 | import org.opensearch.client.Response; |
34 | 43 | import org.opensearch.client.RestClient; |
35 | 44 | import org.opensearch.client.RestClientBuilder; |
36 | 45 | import org.opensearch.common.settings.Settings; |
37 | 46 | import org.opensearch.common.unit.TimeValue; |
38 | 47 | import org.opensearch.common.util.concurrent.ThreadContext; |
39 | 48 | import org.opensearch.common.util.io.IOUtils; |
| 49 | +import org.opensearch.common.xcontent.XContentFactory; |
| 50 | +import org.opensearch.core.rest.RestStatus; |
| 51 | +import org.opensearch.core.xcontent.XContentBuilder; |
40 | 52 | import org.opensearch.test.rest.OpenSearchRestTestCase; |
41 | 53 |
|
42 | 54 | /** |
@@ -135,7 +147,13 @@ public RestClient initClient(String clusterName) throws IOException { |
135 | 147 | int port = Integer.parseInt(stringUrl.substring(portSeparator + 1)); |
136 | 148 | hosts.add(buildHttpHost(host, port)); |
137 | 149 | } |
138 | | - return buildClient(restClientSettings(), hosts.toArray(new HttpHost[0])); |
| 150 | + Settings.Builder builder = Settings.builder(); |
| 151 | + if (System.getProperty("tests.rest.client_path_prefix") != null) { |
| 152 | + builder.put(CLIENT_PATH_PREFIX, System.getProperty("tests.rest.client_path_prefix")); |
| 153 | + } |
| 154 | + // Disable max compilations rate to avoid hitting compilations threshold during tests |
| 155 | + builder.put(Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE.getKeyValue(), "true"); |
| 156 | + return buildClient(builder.build(), hosts.toArray(new HttpHost[0])); |
139 | 157 | } |
140 | 158 |
|
141 | 159 | /** 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 { |
301 | 319 | connectionRequest.setJsonEntity(connectionSetting); |
302 | 320 | adminClient().performRequest(connectionRequest); |
303 | 321 | } |
| 322 | + |
| 323 | + /** |
| 324 | + * Increases the maximum script compilation rate for all script contexts for tests. This method |
| 325 | + * sets an unlimited compilation rate for each script context when the |
| 326 | + * script.disable_max_compilations_rate setting is not enabled, allowing tests to run without |
| 327 | + * hitting compilation rate limits. |
| 328 | + * |
| 329 | + * @throws IOException if there is an error retrieving cluster settings or updating them |
| 330 | + */ |
| 331 | + protected void increaseMaxCompilationsRate() throws IOException { |
| 332 | + // When script.disable_max_compilations_rate is set, custom context compilation rates cannot be |
| 333 | + // set |
| 334 | + if (!Objects.equals( |
| 335 | + getClusterSetting( |
| 336 | + org.opensearch.sql.common.setting.Settings.Key.SCRIPT_DISABLE_MAX_COMPILATIONS_RATE |
| 337 | + .getKeyValue(), |
| 338 | + "persistent"), |
| 339 | + "true")) { |
| 340 | + List<String> contexts = getScriptContexts(); |
| 341 | + for (String context : contexts) { |
| 342 | + String contextCompilationsRate = |
| 343 | + org.opensearch.sql.common.setting.Settings.Key |
| 344 | + .SCRIPT_CONTEXT_MAX_COMPILATIONS_RATE_PATTERN |
| 345 | + .getKeyValue() |
| 346 | + .replace("*", context); |
| 347 | + updateClusterSetting(contextCompilationsRate, "unlimited", true); |
| 348 | + } |
| 349 | + } |
| 350 | + } |
| 351 | + |
| 352 | + protected List<String> getScriptContexts() throws IOException { |
| 353 | + Request request = new Request("GET", "/_script_context"); |
| 354 | + String responseBody = executeRequest(request); |
| 355 | + JSONObject jsonResponse = new JSONObject(responseBody); |
| 356 | + JSONArray contexts = jsonResponse.getJSONArray("contexts"); |
| 357 | + List<String> contextNames = new ArrayList<>(); |
| 358 | + for (int i = 0; i < contexts.length(); i++) { |
| 359 | + JSONObject context = contexts.getJSONObject(i); |
| 360 | + String contextName = context.getString("name"); |
| 361 | + contextNames.add(contextName); |
| 362 | + } |
| 363 | + return contextNames; |
| 364 | + } |
| 365 | + |
| 366 | + protected static void updateClusterSetting(String settingKey, Object value) throws IOException { |
| 367 | + updateClusterSetting(settingKey, value, true); |
| 368 | + } |
| 369 | + |
| 370 | + protected static void updateClusterSetting(String settingKey, Object value, boolean persistent) |
| 371 | + throws IOException { |
| 372 | + String property = persistent ? PERSISTENT : TRANSIENT; |
| 373 | + XContentBuilder builder = |
| 374 | + XContentFactory.jsonBuilder() |
| 375 | + .startObject() |
| 376 | + .startObject(property) |
| 377 | + .field(settingKey, value) |
| 378 | + .endObject() |
| 379 | + .endObject(); |
| 380 | + Request request = new Request("PUT", "_cluster/settings"); |
| 381 | + request.setJsonEntity(builder.toString()); |
| 382 | + Response response = client().performRequest(request); |
| 383 | + Assert.assertEquals( |
| 384 | + RestStatus.OK, RestStatus.fromCode(response.getStatusLine().getStatusCode())); |
| 385 | + } |
| 386 | + |
| 387 | + protected static JSONObject getAllClusterSettings() throws IOException { |
| 388 | + Request request = new Request("GET", "/_cluster/settings?flat_settings&include_defaults"); |
| 389 | + RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder(); |
| 390 | + restOptionsBuilder.addHeader("Content-Type", "application/json"); |
| 391 | + request.setOptions(restOptionsBuilder); |
| 392 | + return new JSONObject(executeRequest(request)); |
| 393 | + } |
| 394 | + |
| 395 | + protected static String getClusterSetting(String settingPath, String type) throws IOException { |
| 396 | + JSONObject settings = getAllClusterSettings(); |
| 397 | + String value = settings.optJSONObject(type).optString(settingPath); |
| 398 | + if (StringUtils.isEmpty(value)) { |
| 399 | + return settings.optJSONObject("defaults").optString(settingPath); |
| 400 | + } else { |
| 401 | + return value; |
| 402 | + } |
| 403 | + } |
| 404 | + |
| 405 | + protected static String executeRequest(final Request request) throws IOException { |
| 406 | + return executeRequest(request, client()); |
| 407 | + } |
| 408 | + |
| 409 | + protected static String executeRequest(final Request request, RestClient client) |
| 410 | + throws IOException { |
| 411 | + Response response = client.performRequest(request); |
| 412 | + Assert.assertEquals(200, response.getStatusLine().getStatusCode()); |
| 413 | + return getResponseBody(response); |
| 414 | + } |
304 | 415 | } |
0 commit comments