Handle EMRS exception as 400 (#2612)

Signed-off-by: Louis Chu <[email protected]>
opensearch-project · Apr 18, 2024 · 0d8341f · 0d8341f
1 parent 204c7da
commit 0d8341f
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 1 deletion.
diff --git a/spark/src/main/java/org/opensearch/sql/spark/client/EmrServerlessClientImpl.java b/spark/src/main/java/org/opensearch/sql/spark/client/EmrServerlessClientImpl.java
@@ -17,6 +17,7 @@
 import com.amazonaws.services.emrserverless.model.SparkSubmit;
 import com.amazonaws.services.emrserverless.model.StartJobRunRequest;
 import com.amazonaws.services.emrserverless.model.StartJobRunResult;
+import com.amazonaws.services.emrserverless.model.ValidationException;
 import java.security.AccessController;
 import java.security.PrivilegedAction;
 import org.apache.commons.lang3.StringUtils;
@@ -69,6 +70,11 @@ public String startJobRun(StartJobRequest startJobRequest) {
                     logger.error("Error while making start job request to emr:", t);
                     MetricUtils.incrementNumericalMetric(
                         MetricName.EMR_START_JOB_REQUEST_FAILURE_COUNT);
+                    if (t instanceof ValidationException) {
+                      throw new IllegalArgumentException(
+                          "The input fails to satisfy the constraints specified by AWS EMR"
+                              + " Serverless.");
+                    }
                     throw new RuntimeException(GENERIC_INTERNAL_SERVER_ERROR_MESSAGE);
                   }
                 });

diff --git a/spark/src/test/java/org/opensearch/sql/spark/client/EmrServerlessClientImplTest.java b/spark/src/test/java/org/opensearch/sql/spark/client/EmrServerlessClientImplTest.java
@@ -21,6 +21,7 @@
 import static org.opensearch.sql.spark.constants.TestConstants.SPARK_SUBMIT_PARAMETERS;
 
 import com.amazonaws.services.emrserverless.AWSEMRServerless;
+import com.amazonaws.services.emrserverless.model.AWSEMRServerlessException;
 import com.amazonaws.services.emrserverless.model.CancelJobRunResult;
 import com.amazonaws.services.emrserverless.model.GetJobRunResult;
 import com.amazonaws.services.emrserverless.model.JobRun;
@@ -97,7 +98,9 @@ void testStartJobRun() {
 
   @Test
   void testStartJobRunWithErrorMetric() {
-    doThrow(new ValidationException("Couldn't start job")).when(emrServerless).startJobRun(any());
+    doThrow(new AWSEMRServerlessException("Couldn't start job"))
+        .when(emrServerless)
+        .startJobRun(any());
     EmrServerlessClientImpl emrServerlessClient = new EmrServerlessClientImpl(emrServerless);
     RuntimeException runtimeException =
         Assertions.assertThrows(
@@ -224,4 +227,33 @@ void testStartJobRunWithLongJobName() {
     StartJobRunRequest startJobRunRequest = startJobRunRequestArgumentCaptor.getValue();
     Assertions.assertEquals(255, startJobRunRequest.getName().length());
   }
+
+  @Test
+  void testStartJobRunThrowsValidationException() {
+    when(emrServerless.startJobRun(any())).thenThrow(new ValidationException("Unmatched quote"));
+    EmrServerlessClientImpl emrServerlessClient = new EmrServerlessClientImpl(emrServerless);
+
+    IllegalArgumentException exception =
+        Assertions.assertThrows(
+            IllegalArgumentException.class,
+            () ->
+                emrServerlessClient.startJobRun(
+                    new StartJobRequest(
+                        EMRS_JOB_NAME,
+                        EMRS_APPLICATION_ID,
+                        EMRS_EXECUTION_ROLE,
+                        SPARK_SUBMIT_PARAMETERS,
+                        new HashMap<>(),
+                        false,
+                        DEFAULT_RESULT_INDEX)),
+            "Expected ValidationException to be thrown");
+
+    // Verify that the message in the exception is correct
+    Assertions.assertEquals(
+        "The input fails to satisfy the constraints specified by AWS EMR Serverless.",
+        exception.getMessage());
+
+    // Optionally verify that no job run is started
+    verify(emrServerless, times(1)).startJobRun(any());
+  }
 }