Add flaky retry on CalcitePPLTpchIT

integ-test/src/test/java/org/opensearch/sql/calcite/tpch/CalcitePPLTpchIT.java

@@ -18,7 +18,9 @@
 import org.junit.Ignore;
 import org.junit.Test;
 import org.opensearch.sql.ppl.PPLIntegTestCase;
+import org.opensearch.sql.util.Retry;
 
+@Retry
 public class CalcitePPLTpchIT extends PPLIntegTestCase {
 
   @Override
@@ -142,7 +144,6 @@ public void testQ3() throws IOException {
   // TODO: Aggregation push down has a hard-coded limit of 1000 buckets for output, so this query
   // will not return the correct results with aggregation push down and it's unstable
   @Ignore
-  @Test
   public void testQ4() throws IOException {
     String ppl = sanitize(loadFromFile("tpch/queries/q4.ppl"));
     JSONObject actual = executeQuery(ppl);
@@ -173,6 +174,7 @@ public void testQ6() throws IOException {
     verifyDataRows(actual, rows(77949.9186));
   }
 
+  @Test
   public void testQ7() throws IOException {
     String ppl = sanitize(loadFromFile("tpch/queries/q7.ppl"));
     JSONObject actual = executeQuery(ppl);
@@ -185,6 +187,7 @@ public void testQ7() throws IOException {
     verifyNumOfRows(actual, 0);
   }
 
+  @Test
   public void testQ8() throws IOException {
     String ppl = sanitize(loadFromFile("tpch/queries/q8.ppl"));
     JSONObject actual = executeQuery(ppl);

integ-test/src/test/java/org/opensearch/sql/ppl/PPLIntegTestCase.java

@@ -21,6 +21,7 @@
 import org.json.JSONException;
 import org.json.JSONObject;
 import org.junit.Assert;
+import org.junit.Rule;
 import org.opensearch.client.Request;
 import org.opensearch.client.RequestOptions;
 import org.opensearch.client.Response;
@@ -29,10 +30,12 @@
 import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.common.setting.Settings.Key;
 import org.opensearch.sql.legacy.SQLIntegTestCase;
+import org.opensearch.sql.util.RetryProcessor;
 
 /** OpenSearch Rest integration test base for PPL testing. */
 public abstract class PPLIntegTestCase extends SQLIntegTestCase {
   private static final Logger LOG = LogManager.getLogger();
+  @Rule public final RetryProcessor retryProcessor = new RetryProcessor();
 
   @Override
   protected void init() throws Exception {

integ-test/src/test/java/org/opensearch/sql/util/Retry.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.util;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/**
+ * Retry annotation to indicate a test should be retried when exception happens. The default retry
+ * count is 3. You can specify the retry count by passing a value to the annotation. For
+ * example: @Retry(5)
+ */
+@Target({ElementType.TYPE, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+public @interface Retry {
+  int value() default 3;
+}

integ-test/src/test/java/org/opensearch/sql/util/RetryProcessor.java

@@ -0,0 +1,42 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.util;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.junit.rules.TestWatcher;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+/** Retry processor to retry a test when exception happens. Retry a test by adding @Retry. */
+public class RetryProcessor extends TestWatcher {
+  private static final Logger LOG = LogManager.getLogger();
+
+  @Override
+  public Statement apply(Statement base, Description description) {
+    Retry retry = description.getAnnotation(Retry.class);
+    if (retry == null) {
+      return base;
+    }
+    return new Statement() {
+      @Override
+      public void evaluate() throws Throwable {
+        Throwable lastException = null;
+        for (int i = 0; i < retry.value(); i++) {
+          try {
+            base.evaluate();
+            return;
+          } catch (Throwable t) {
+            lastException = t;
+            LOG.info("Retrying {} {} times", description.getDisplayName(), (i + 1));
+          }
+        }
+        assert lastException != null;
+        throw lastException;
+      }
+    };
+  }
+}