Merge remote-tracking branch 'origin/main' into issues/4636

Signed-off-by: Yuanchun Shen <yuanchu@amazon.com>

Author: yuancu

common/build.gradle

@@ -45,7 +45,7 @@ dependencies {
     implementation "com.github.seancfoley:ipaddress:5.4.2"
 
     testImplementation group: 'junit', name: 'junit', version: '4.13.2'
-    testImplementation group: 'org.assertj', name: 'assertj-core', version: '3.9.1'
+    testImplementation group: 'org.assertj', name: 'assertj-core', version: '3.27.7'
     testImplementation group: 'com.google.guava', name: 'guava', version: "${guava_version}"
     testImplementation group: 'org.hamcrest', name: 'hamcrest-library', version: "${hamcrest_version}"
     testImplementation('org.junit.jupiter:junit-jupiter:5.9.3')

direct-query-core/src/main/java/org/opensearch/sql/prometheus/client/PrometheusClientImpl.java

@@ -26,6 +26,7 @@
 import org.json.JSONArray;
 import org.json.JSONException;
 import org.json.JSONObject;
+import org.opensearch.secure_sm.AccessController;
 import org.opensearch.sql.prometheus.exception.PrometheusClientException;
 import org.opensearch.sql.prometheus.model.MetricMetadata;
 
@@ -91,7 +92,7 @@ public JSONObject queryRange(
     Request request = new Request.Builder().url(queryUrl).build();
 
     logger.debug("Executing Prometheus request with headers: {}", request.headers().toString());
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
 
     logger.debug("Received Prometheus response for query_range: code={}", response);
 
@@ -126,7 +127,7 @@ public JSONObject query(String query, Long time, Integer limit, Integer timeout)
     Request request = new Request.Builder().url(queryUrl).build();
 
     logger.info("Executing Prometheus request with headers: {}", request.headers().toString());
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
 
     logger.info("Received Prometheus response for instant query: code={}", response);
     // Return the full response object, not just the data field
@@ -146,7 +147,7 @@ public List<String> getLabels(Map<String, String> queryParams) throws IOExceptio
             "%s/api/v1/labels%s", prometheusUri.toString().replaceAll("/$", ""), queryString);
     logger.debug("queryUrl: " + queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
     JSONObject jsonObject = readResponse(response);
     return toListOfLabels(jsonObject.getJSONArray("data"));
   }
@@ -161,7 +162,7 @@ public List<String> getLabel(String labelName, Map<String, String> queryParams)
             prometheusUri.toString().replaceAll("/$", ""), labelName, queryString);
     logger.debug("queryUrl: " + queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
     JSONObject jsonObject = readResponse(response);
     return toListOfLabels(jsonObject.getJSONArray("data"));
   }
@@ -175,7 +176,7 @@ public Map<String, List<MetricMetadata>> getAllMetrics(Map<String, String> query
             "%s/api/v1/metadata%s", prometheusUri.toString().replaceAll("/$", ""), queryString);
     logger.debug("queryUrl: " + queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
     JSONObject jsonObject = readResponse(response);
     TypeReference<HashMap<String, List<MetricMetadata>>> typeRef = new TypeReference<>() {};
     return new ObjectMapper().readValue(jsonObject.getJSONObject("data").toString(), typeRef);
@@ -194,7 +195,7 @@ public List<Map<String, String>> getSeries(Map<String, String> queryParams) thro
             "%s/api/v1/series%s", prometheusUri.toString().replaceAll("/$", ""), queryString);
     logger.debug("queryUrl: " + queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
     JSONObject jsonObject = readResponse(response);
     JSONArray dataArray = jsonObject.getJSONArray("data");
     return toListOfSeries(dataArray);
@@ -211,7 +212,7 @@ public JSONArray queryExemplars(String query, Long start, Long end) throws IOExc
             end);
     logger.debug("queryUrl: " + queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
     JSONObject jsonObject = readResponse(response);
     return jsonObject.getJSONArray("data");
   }
@@ -222,7 +223,7 @@ public JSONObject getAlerts() throws IOException {
         String.format("%s/api/v1/alerts", prometheusUri.toString().replaceAll("/$", ""));
     logger.debug("Making Prometheus alerts request: {}", queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
     JSONObject jsonObject = readResponse(response);
     return jsonObject.getJSONObject("data");
   }
@@ -235,7 +236,7 @@ public JSONObject getRules(Map<String, String> queryParams) throws IOException {
             "%s/api/v1/rules%s", prometheusUri.toString().replaceAll("/$", ""), queryString);
     logger.debug("Making Prometheus rules request: {}", queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.prometheusHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.prometheusHttpClient.newCall(request).execute());
     JSONObject jsonObject = readResponse(response);
     return jsonObject.getJSONObject("data");
   }
@@ -248,7 +249,7 @@ public JSONArray getAlertmanagerAlerts(Map<String, String> queryParams) throws I
 
     logger.debug("Making Alertmanager alerts request: {}", queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.alertmanagerHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.alertmanagerHttpClient.newCall(request).execute());
 
     return readAlertmanagerResponse(response);
   }
@@ -261,7 +262,7 @@ public JSONArray getAlertmanagerAlertGroups(Map<String, String> queryParams) thr
 
     logger.debug("Making Alertmanager alert groups request: {}", queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.alertmanagerHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.alertmanagerHttpClient.newCall(request).execute());
 
     return readAlertmanagerResponse(response);
   }
@@ -273,7 +274,7 @@ public JSONArray getAlertmanagerReceivers() throws IOException {
 
     logger.debug("Making Alertmanager receivers request: {}", queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.alertmanagerHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.alertmanagerHttpClient.newCall(request).execute());
 
     return readAlertmanagerResponse(response);
   }
@@ -285,7 +286,7 @@ public JSONArray getAlertmanagerSilences() throws IOException {
 
     logger.debug("Making Get Alertmanager silences request: {}", queryUrl);
     Request request = new Request.Builder().url(queryUrl).build();
-    Response response = this.alertmanagerHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.alertmanagerHttpClient.newCall(request).execute());
 
     return readAlertmanagerResponse(response);
   }
@@ -301,7 +302,7 @@ public String createAlertmanagerSilences(String silenceJson) throws IOException
             .header("Content-Type", "application/json")
             .post(RequestBody.create(silenceJson.getBytes(StandardCharsets.UTF_8)))
             .build();
-    Response response = this.alertmanagerHttpClient.newCall(request).execute();
+    Response response = AccessController.doPrivilegedChecked(() -> this.alertmanagerHttpClient.newCall(request).execute());
 
     if (response.isSuccessful()) {
       return Objects.requireNonNull(response.body()).string();

integ-test/build.gradle

@@ -406,7 +406,6 @@ task integTestWithSecurity(type: RestIntegTestTask) {
         configureSecurityPlugin(cluster)
     }
 
-    useJUnitPlatform()
     dependsOn ':opensearch-sql-plugin:bundlePlugin'
     testLogging {
         events "passed", "skipped", "failed"
@@ -447,11 +446,8 @@ task integTestWithSecurity(type: RestIntegTestTask) {
         jvmArgs '-agentlib:jdwp=transport=dt_socket,server=y,suspend=y,address=*:5005'
     }
 
-    // NOTE: this IT config discovers only junit5 (jupiter) tests.
-    // https://github.com/opensearch-project/sql/issues/1974
     filter {
-        includeTestsMatching 'org.opensearch.sql.security.CrossClusterSearchIT'
-        includeTestsMatching 'org.opensearch.sql.security.PPLPermissionsIT'
+        includeTestsMatching 'org.opensearch.sql.security.*'
     }
 }
 

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

@@ -311,7 +311,9 @@ protected static void configureHttpsClient(
    * cluster.
    */
   public void configureMultiClusters(String remote) throws IOException {
-    initRemoteClient(remote);
+    if (remoteClient == null) {
+      initRemoteClient(remote);
+    }
 
     Request connectionRequest = new Request("PUT", "_cluster/settings");
     String connectionSetting =

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

@@ -210,19 +210,6 @@ protected synchronized void loadIndex(Index index, RestClient client) throws IOE
       createIndexByRestClient(client, indexName, mapping);
       loadDataByRestClient(client, indexName, dataSet);
     }
-    // loadIndex() could directly return when isIndexExist()=true,
-    // e.g. the index is created in the cluster but data hasn't been flushed.
-    // We block loadIndex() until data loaded to resolve
-    // https://github.com/opensearch-project/sql/issues/4261
-    int countDown = 3; // 1500ms timeout
-    while (countDown != 0 && getDocCount(client, indexName) == 0) {
-      try {
-        Thread.sleep(500);
-        countDown--;
-      } catch (InterruptedException e) {
-        throw new IOException(e);
-      }
-    }
   }
 
   protected synchronized void loadIndex(Index index) throws IOException {

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

@@ -46,12 +46,29 @@ public class TestUtils {
    */
   public static void createIndexByRestClient(RestClient client, String indexName, String mapping) {
     Request request = new Request("PUT", "/" + indexName);
-    if (!isNullOrEmpty(mapping)) {
-      request.setJsonEntity(mapping);
-    }
+    JSONObject jsonObject = isNullOrEmpty(mapping) ? new JSONObject() : new JSONObject(mapping);
+    setZeroReplicas(jsonObject);
+    request.setJsonEntity(jsonObject.toString());
     performRequest(client, request);
   }
 
+  /**
+   * Sets number_of_replicas to 0 in the index settings. This makes multi-node behavior consistent
+   * (<a href="https://github.com/opensearch-project/sql/issues/4261">4261</a>) and prevents tests
+   * from hanging on single-node clusters when using wait_for_active_shards=all.
+   *
+   * @param jsonObject the index creation JSON object to modify
+   */
+  private static void setZeroReplicas(JSONObject jsonObject) {
+    JSONObject settings =
+        jsonObject.has("settings") ? jsonObject.getJSONObject("settings") : new JSONObject();
+    JSONObject indexSettings =
+        settings.has("index") ? settings.getJSONObject("index") : new JSONObject();
+    indexSettings.put("number_of_replicas", 0);
+    settings.put("index", indexSettings);
+    jsonObject.put("settings", settings);
+  }
+
   /**
    * https://github.com/elastic/elasticsearch/pull/49959<br>
    * Deprecate creation of dot-prefixed index names except for hidden and system indices. Create
@@ -99,7 +116,8 @@ public static boolean isIndexExist(RestClient client, String indexName) {
   public static void loadDataByRestClient(
       RestClient client, String indexName, String dataSetFilePath) throws IOException {
     Path path = Paths.get(getResourceFilePath(dataSetFilePath));
-    Request request = new Request("POST", "/" + indexName + "/_bulk?refresh=true");
+    Request request =
+        new Request("POST", "/" + indexName + "/_bulk?refresh=wait_for&wait_for_active_shards=all");
     request.setJsonEntity(new String(Files.readAllBytes(path)));
     performRequest(client, request);
   }

integ-test/src/test/java/org/opensearch/sql/security/CalciteCrossClusterSearchIT.java

@@ -5,9 +5,6 @@
 
 package org.opensearch.sql.security;
 
-import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_ACCOUNT;
-import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK;
-import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_DOG;
 import static org.opensearch.sql.util.MatcherUtils.columnName;
 import static org.opensearch.sql.util.MatcherUtils.rows;
 import static org.opensearch.sql.util.MatcherUtils.schema;
@@ -16,51 +13,21 @@
 import static org.opensearch.sql.util.MatcherUtils.verifySchema;
 
 import java.io.IOException;
-import lombok.SneakyThrows;
 import org.json.JSONObject;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Test;
-import org.opensearch.sql.ppl.PPLIntegTestCase;
+import org.junit.Test;
 
 /** Cross Cluster Search tests with Calcite enabled for enhanced fields features. */
-public class CalciteCrossClusterSearchIT extends PPLIntegTestCase {
-
-  static {
-    String[] clusterNames = System.getProperty("cluster.names").split(",");
-    var remote = "remoteCluster";
-    for (var cluster : clusterNames) {
-      if (cluster.startsWith("remote")) {
-        remote = cluster;
-        break;
-      }
-    }
-    REMOTE_CLUSTER = remote;
-  }
-
-  public static final String REMOTE_CLUSTER;
-  private static final String TEST_INDEX_ACCOUNT_REMOTE = REMOTE_CLUSTER + ":" + TEST_INDEX_ACCOUNT;
-  private static final String TEST_INDEX_DOG_REMOTE = REMOTE_CLUSTER + ":" + TEST_INDEX_DOG;
-  private static final String TEST_INDEX_BANK_REMOTE = REMOTE_CLUSTER + ":" + TEST_INDEX_BANK;
-  private static boolean initialized = false;
-
-  @SneakyThrows
-  @BeforeEach
-  public void initialize() {
-    if (!initialized) {
-      setUpIndices();
-      initialized = true;
-    }
-  }
+public class CalciteCrossClusterSearchIT extends CrossClusterTestBase {
 
   @Override
   protected void init() throws Exception {
-    configureMultiClusters(REMOTE_CLUSTER);
+    super.init();
     loadIndex(Index.BANK);
     loadIndex(Index.BANK, remoteClient());
-    loadIndex(Index.ACCOUNT);
-    loadIndex(Index.ACCOUNT, remoteClient());
     loadIndex(Index.DOG);
     loadIndex(Index.DOG, remoteClient());
+    loadIndex(Index.ACCOUNT);
+    loadIndex(Index.ACCOUNT, remoteClient());
     loadIndex(Index.TIME_TEST_DATA);
     loadIndex(Index.TIME_TEST_DATA, remoteClient());
     enableCalcite();
@@ -87,8 +54,8 @@ public void testCrossClusterFieldsWildcardPrefix() throws IOException {
   public void testCrossClusterFieldsWildcardSuffix() throws IOException {
     JSONObject result =
         executeQuery(String.format("search source=%s | fields *Name", TEST_INDEX_DOG_REMOTE));
-    verifyColumn(result, columnName("dog_name"), columnName("holdersName"));
-    verifySchema(result, schema("dog_name", "string"), schema("holdersName", "string"));
+    verifyColumn(result, columnName("holdersName"));
+    verifySchema(result, schema("holdersName", "string"));
   }
 
   @Test
@@ -165,7 +132,7 @@ public void testDefaultBinCrossCluster() throws IOException {
                 TEST_INDEX_ACCOUNT_REMOTE));
     verifySchema(result, schema("count()", null, "bigint"), schema("age", null, "string"));
 
-    verifyDataRows(result, rows(451L, "20-30"), rows(504L, "30-40"), rows(45L, "40-50"));
+    verifyDataRows(result, rows(451, "20.0-30.0"), rows(504L, "30.0-40.0"), rows(45L, "40.0-50.0"));
   }
 
   @Test
@@ -218,18 +185,18 @@ public void testRangeBinCrossCluster() throws IOException {
                 TEST_INDEX_ACCOUNT_REMOTE));
     verifySchema(result, schema("count()", null, "bigint"), schema("age", null, "string"));
 
-    verifyDataRows(result, rows(1000L, "0-100"));
+    verifyDataRows(result, rows(451, "20-30"), rows(504L, "30-40"), rows(45L, "40-50"));
   }
 
   @Test
   public void testTimeBinCrossCluster() throws IOException {
     // Time-based binning with span
     JSONObject result =
         executeQuery(
-            REMOTE_CLUSTER
-                + ":opensearch-sql_test_index_time_data"
-                + " | bin @timestamp span=1h"
-                + " | fields `@timestamp`, value | sort `@timestamp` | head 3");
+            String.format(
+                "source=%s | bin @timestamp span=1h | fields `@timestamp`, value | sort"
+                    + " `@timestamp` | head 3",
+                TEST_INDEX_TIME_DATA_REMOTE));
     verifySchema(result, schema("@timestamp", null, "timestamp"), schema("value", null, "int"));
 
     // With 1-hour spans
@@ -285,12 +252,27 @@ public void testCrossClusterRenameFullWildcard() throws IOException {
     JSONObject result =
         executeQuery(String.format("search source=%s | rename * as old_*", TEST_INDEX_DOG_REMOTE));
     verifyColumn(
-        result, columnName("old_dog_name"), columnName("old_holdersName"), columnName("old_age"));
+        result,
+        columnName("old_dog_name"),
+        columnName("old_holdersName"),
+        columnName("old_age"),
+        columnName("old__id"),
+        columnName("old__index"),
+        columnName("old__score"),
+        columnName("old__maxscore"),
+        columnName("old__sort"),
+        columnName("old__routing"));
     verifySchema(
         result,
         schema("old_dog_name", "string"),
         schema("old_holdersName", "string"),
-        schema("old_age", "bigint"));
+        schema("old_age", "bigint"),
+        schema("old__id", "string"),
+        schema("old__index", "string"),
+        schema("old__score", "float"),
+        schema("old__maxscore", "float"),
+        schema("old__sort", "bigint"),
+        schema("old__routing", "string"));
   }
 
   @Test