Handle EMR Exceptions in FlintCancelJob Operation

Signed-off-by: Vamsi Manohar <[email protected]>
vmmusings · Mar 25, 2024 · 3ad9827 · 3ad9827
1 parent 85dae6f
commit 3ad9827
Show file tree

Hide file tree

Showing 14 changed files with 1,173 additions and 903 deletions.
diff --git a/integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java b/integ-test/src/test/java/org/opensearch/sql/datasource/DataSourceAPIsIT.java
@@ -11,7 +11,10 @@
 import static org.opensearch.sql.datasources.utils.XContentParserUtils.DESCRIPTION_FIELD;
 import static org.opensearch.sql.datasources.utils.XContentParserUtils.NAME_FIELD;
 import static org.opensearch.sql.datasources.utils.XContentParserUtils.STATUS_FIELD;
+import static org.opensearch.sql.legacy.TestUtils.createIndexByRestClient;
 import static org.opensearch.sql.legacy.TestUtils.getResponseBody;
+import static org.opensearch.sql.legacy.TestUtils.isIndexExist;
+import static org.opensearch.sql.legacy.TestUtils.loadDataByRestClient;
 
 import com.google.common.collect.ImmutableMap;
 import com.google.gson.Gson;
@@ -37,11 +40,6 @@
 
 public class DataSourceAPIsIT extends PPLIntegTestCase {
 
-  @Override
-  protected void init() throws Exception {
-    loadIndex(Index.DATASOURCES);
-  }
-
   @After
   public void cleanUp() throws IOException {
     wipeAllClusterSettings();
@@ -397,6 +395,16 @@ public void patchDataSourceAPITest() {
   @SneakyThrows
   @Test
   public void testOldDataSourceModelLoadingThroughGetDataSourcesAPI() {
+    Index index = Index.DATASOURCES;
+    String indexName = index.getName();
+    String mapping = index.getMapping();
+    String dataSet = index.getDataSet();
+    if (isIndexExist(client(), indexName)) {
+      createIndexByRestClient(client(), indexName, mapping);
+    }
+    loadDataByRestClient(client(), indexName, dataSet);
+    // waiting for loaded indices.
+    Thread.sleep(1000);
     // get datasource to validate the creation.
     Request getRequest = getFetchDataSourceRequest(null);
     Response getResponse = client().performRequest(getRequest);

diff --git a/spark/src/main/java/org/opensearch/sql/spark/client/EMRServerlessClient.java b/spark/src/main/java/org/opensearch/sql/spark/client/EMRServerlessClient.java
@@ -41,5 +41,6 @@ public interface EMRServerlessClient {
    * @param jobId jobId.
    * @return {@link CancelJobRunResult}
    */
-  CancelJobRunResult cancelJobRun(String applicationId, String jobId);
+  CancelJobRunResult cancelJobRun(
+      String applicationId, String jobId, boolean allowExceptionPropagation);
 }
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
@@ -32,7 +32,7 @@ public class EmrServerlessClientImpl implements EMRServerlessClient {
 
   private static final int MAX_JOB_NAME_LENGTH = 255;
 
-  private static final String GENERIC_INTERNAL_SERVER_ERROR_MESSAGE = "Internal Server Error.";
+  public static final String GENERIC_INTERNAL_SERVER_ERROR_MESSAGE = "Internal Server Error.";
 
   public EmrServerlessClientImpl(AWSEMRServerless emrServerless) {
     this.emrServerless = emrServerless;
@@ -98,7 +98,8 @@ public GetJobRunResult getJobRunResult(String applicationId, String jobId) {
   }
 
   @Override
-  public CancelJobRunResult cancelJobRun(String applicationId, String jobId) {
+  public CancelJobRunResult cancelJobRun(
+      String applicationId, String jobId, boolean allowExceptionPropagation) {
     CancelJobRunRequest cancelJobRunRequest =
         new CancelJobRunRequest().withJobRunId(jobId).withApplicationId(applicationId);
     CancelJobRunResult cancelJobRunResult =
@@ -108,10 +109,14 @@ public CancelJobRunResult cancelJobRun(String applicationId, String jobId) {
                   try {
                     return emrServerless.cancelJobRun(cancelJobRunRequest);
                   } catch (Throwable t) {
-                    logger.error("Error while making cancel job request to emr:", t);
-                    MetricUtils.incrementNumericalMetric(
-                        MetricName.EMR_CANCEL_JOB_REQUEST_FAILURE_COUNT);
-                    throw new RuntimeException(GENERIC_INTERNAL_SERVER_ERROR_MESSAGE);
+                    if (allowExceptionPropagation) {
+                      throw t;
+                    } else {
+                      logger.error("Error while making cancel job request to emr:", t);
+                      MetricUtils.incrementNumericalMetric(
+                          MetricName.EMR_CANCEL_JOB_REQUEST_FAILURE_COUNT);
+                      throw new RuntimeException(GENERIC_INTERNAL_SERVER_ERROR_MESSAGE);
+                    }
                   }
                 });
     logger.info(String.format("Job : %s cancelled", cancelJobRunResult.getJobRunId()));

diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/BatchQueryHandler.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/BatchQueryHandler.java
@@ -58,7 +58,7 @@ protected JSONObject getResponseFromExecutor(AsyncQueryJobMetadata asyncQueryJob
   @Override
   public String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata) {
     emrServerlessClient.cancelJobRun(
-        asyncQueryJobMetadata.getApplicationId(), asyncQueryJobMetadata.getJobId());
+        asyncQueryJobMetadata.getApplicationId(), asyncQueryJobMetadata.getJobId(), false);
     return asyncQueryJobMetadata.getQueryId().getId();
   }
 

diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/session/InteractiveSession.java b/spark/src/main/java/org/opensearch/sql/spark/execution/session/InteractiveSession.java
@@ -80,7 +80,8 @@ public void close() {
     if (model.isEmpty()) {
       throw new IllegalStateException("session does not exist. " + sessionModel.getSessionId());
     } else {
-      serverlessClient.cancelJobRun(sessionModel.getApplicationId(), sessionModel.getJobId());
+      serverlessClient.cancelJobRun(
+          sessionModel.getApplicationId(), sessionModel.getJobId(), false);
     }
   }
 

diff --git a/spark/src/main/java/org/opensearch/sql/spark/flint/operation/FlintIndexOp.java b/spark/src/main/java/org/opensearch/sql/spark/flint/operation/FlintIndexOp.java
@@ -5,10 +5,12 @@
 
 package org.opensearch.sql.spark.flint.operation;
 
+import static org.opensearch.sql.spark.client.EmrServerlessClientImpl.GENERIC_INTERNAL_SERVER_ERROR_MESSAGE;
 import static org.opensearch.sql.spark.execution.statestore.StateStore.deleteFlintIndexState;
 import static org.opensearch.sql.spark.execution.statestore.StateStore.getFlintIndexState;
 import static org.opensearch.sql.spark.execution.statestore.StateStore.updateFlintIndexState;
 
+import com.amazonaws.services.emrserverless.model.ValidationException;
 import java.util.Locale;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
@@ -145,11 +147,18 @@ public void cancelStreamingJob(
     String jobId = flintIndexStateModel.getJobId();
     try {
       emrServerlessClient.cancelJobRun(
-          flintIndexStateModel.getApplicationId(), flintIndexStateModel.getJobId());
-    } catch (IllegalArgumentException e) {
-      // handle job does not exist case.
+          flintIndexStateModel.getApplicationId(), flintIndexStateModel.getJobId(), true);
+    } catch (ValidationException e) {
+      // Exception when the job is not in cancellable state and already in terminal state.
+      if (e.getMessage().contains("Job run is not in a cancellable state")) {
+        LOG.error(e);
+        return;
+      } else {
+        throw new RuntimeException(GENERIC_INTERNAL_SERVER_ERROR_MESSAGE);
+      }
+    } catch (Exception e) {
       LOG.error(e);
-      return;
+      throw new RuntimeException(GENERIC_INTERNAL_SERVER_ERROR_MESSAGE);
     }
 
     // pull job state until timeout or cancelled.

diff --git a/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java b/spark/src/test/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceSpec.java
@@ -243,7 +243,8 @@ public GetJobRunResult getJobRunResult(String applicationId, String jobId) {
     }
 
     @Override
-    public CancelJobRunResult cancelJobRun(String applicationId, String jobId) {
+    public CancelJobRunResult cancelJobRun(
+        String applicationId, String jobId, boolean allowExceptionPropagation) {
       cancelJobRunCalled++;
       return new CancelJobRunResult().withJobRunId(jobId);
     }