Another TCPH fix: Wait for all shards (#5106)

Swiddis · web-flow · commit 8073b4efbf68 · 2026-02-03T13:26:55.000-08:00
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
@@ -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 {
diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/TestUtils.java b/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);
   }
