Redefine Drop Index as logical delete

plugin/src/main/java/org/opensearch/sql/plugin/SQLPlugin.java

@@ -320,7 +320,8 @@ private AsyncQueryExecutorService createAsyncQueryExecutorService(
             jobExecutionResponseReader,
             new FlintIndexMetadataReaderImpl(client),
             client,
-            new SessionManager(stateStore, emrServerlessClient, pluginSettings));
+            new SessionManager(stateStore, emrServerlessClient, pluginSettings),
+            stateStore);
     return new AsyncQueryExecutorServiceImpl(
         asyncQueryJobMetadataStorageService,
         sparkQueryDispatcher,

spark/build.gradle

@@ -122,7 +122,8 @@ jacocoTestCoverageVerification {
                     // ignore because XContext IOException
                     'org.opensearch.sql.spark.execution.statestore.StateStore',
                     'org.opensearch.sql.spark.execution.session.SessionModel',
-                    'org.opensearch.sql.spark.execution.statement.StatementModel'
+                    'org.opensearch.sql.spark.execution.statement.StatementModel',
+                    'org.opensearch.sql.spark.flint.FlintIndexStateModel'
             ]
             limit {
                 counter = 'LINE'

spark/src/main/java/org/opensearch/sql/spark/asyncquery/AsyncQueryExecutorServiceImpl.java

@@ -72,7 +72,6 @@ public CreateAsyncQueryResponse createAsyncQuery(
             dispatchQueryResponse.getQueryId(),
             sparkExecutionEngineConfig.getApplicationId(),
             dispatchQueryResponse.getJobId(),
-            dispatchQueryResponse.isDropIndexQuery(),
             dispatchQueryResponse.getResultIndex(),
             dispatchQueryResponse.getSessionId()));
     return new CreateAsyncQueryResponse(

spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryJobMetadata.java

@@ -29,7 +29,6 @@ public class AsyncQueryJobMetadata extends StateModel {
   private final AsyncQueryId queryId;
   private final String applicationId;
   private final String jobId;
-  private final boolean isDropIndexQuery;
   private final String resultIndex;
   // optional sessionId.
   private final String sessionId;
@@ -43,7 +42,6 @@ public AsyncQueryJobMetadata(
         queryId,
         applicationId,
         jobId,
-        false,
         resultIndex,
         null,
         SequenceNumbers.UNASSIGNED_SEQ_NO,
@@ -54,14 +52,12 @@ public AsyncQueryJobMetadata(
       AsyncQueryId queryId,
       String applicationId,
       String jobId,
-      boolean isDropIndexQuery,
       String resultIndex,
       String sessionId) {
     this(
         queryId,
         applicationId,
         jobId,
-        isDropIndexQuery,
         resultIndex,
         sessionId,
         SequenceNumbers.UNASSIGNED_SEQ_NO,
@@ -72,15 +68,13 @@ public AsyncQueryJobMetadata(
       AsyncQueryId queryId,
       String applicationId,
       String jobId,
-      boolean isDropIndexQuery,
       String resultIndex,
       String sessionId,
       long seqNo,
       long primaryTerm) {
     this.queryId = queryId;
     this.applicationId = applicationId;
     this.jobId = jobId;
-    this.isDropIndexQuery = isDropIndexQuery;
     this.resultIndex = resultIndex;
     this.sessionId = sessionId;
     this.seqNo = seqNo;
@@ -106,7 +100,6 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         .field("type", TYPE_JOBMETA)
         .field("jobId", jobId)
         .field("applicationId", applicationId)
-        .field("isDropIndexQuery", isDropIndexQuery)
         .field("resultIndex", resultIndex)
         .field("sessionId", sessionId)
         .endObject();
@@ -120,7 +113,6 @@ public static AsyncQueryJobMetadata copy(
         copy.getQueryId(),
         copy.getApplicationId(),
         copy.getJobId(),
-        copy.isDropIndexQuery(),
         copy.getResultIndex(),
         copy.getSessionId(),
         seqNo,
@@ -176,14 +168,7 @@ public static AsyncQueryJobMetadata fromXContent(
       throw new IllegalArgumentException("jobId and applicationId are required fields.");
     }
     return new AsyncQueryJobMetadata(
-        queryId,
-        applicationId,
-        jobId,
-        isDropIndexQuery,
-        resultIndex,
-        sessionId,
-        seqNo,
-        primaryTerm);
+        queryId, applicationId, jobId, resultIndex, sessionId, seqNo, primaryTerm);
   }
 
   @Override

spark/src/main/java/org/opensearch/sql/spark/dispatcher/AsyncQueryHandler.java

@@ -17,11 +17,6 @@
 public abstract class AsyncQueryHandler {
 
   public JSONObject getQueryResponse(AsyncQueryJobMetadata asyncQueryJobMetadata) {
-    if (asyncQueryJobMetadata.isDropIndexQuery()) {
-      return SparkQueryDispatcher.DropIndexResult.fromJobId(asyncQueryJobMetadata.getJobId())
-          .result();
-    }
-
     JSONObject result = getResponseFromResultIndex(asyncQueryJobMetadata);
     if (result.has(DATA_FIELD)) {
       JSONObject items = result.getJSONObject(DATA_FIELD);
@@ -45,5 +40,5 @@ protected abstract JSONObject getResponseFromResultIndex(
   protected abstract JSONObject getResponseFromExecutor(
       AsyncQueryJobMetadata asyncQueryJobMetadata);
 
-  abstract String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata);
+  public abstract String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata);
 }

spark/src/main/java/org/opensearch/sql/spark/dispatcher/IndexDMLHandler.java

@@ -0,0 +1,119 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.spark.dispatcher;
+
+import static org.opensearch.sql.spark.data.constants.SparkConstants.SPARK_RESPONSE_BUFFER_INDEX_NAME;
+import static org.opensearch.sql.spark.execution.statestore.StateStore.createIndexDMLResult;
+
+import com.amazonaws.services.emrserverless.model.JobRunState;
+import java.util.Optional;
+import lombok.RequiredArgsConstructor;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.json.JSONObject;
+import org.opensearch.client.Client;
+import org.opensearch.sql.datasource.DataSourceService;
+import org.opensearch.sql.datasource.model.DataSourceMetadata;
+import org.opensearch.sql.datasources.auth.DataSourceUserAuthorizationHelperImpl;
+import org.opensearch.sql.spark.asyncquery.model.AsyncQueryId;
+import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata;
+import org.opensearch.sql.spark.client.EMRServerlessClient;
+import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest;
+import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse;
+import org.opensearch.sql.spark.dispatcher.model.IndexDMLResult;
+import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails;
+import org.opensearch.sql.spark.execution.statestore.StateStore;
+import org.opensearch.sql.spark.flint.FlintIndexMetadata;
+import org.opensearch.sql.spark.flint.FlintIndexMetadataReader;
+import org.opensearch.sql.spark.flint.operation.FlintIndexOp;
+import org.opensearch.sql.spark.flint.operation.FlintIndexOpCancel;
+import org.opensearch.sql.spark.flint.operation.FlintIndexOpDelete;
+import org.opensearch.sql.spark.response.JobExecutionResponseReader;
+
+/** Handle Index DML query. includes * DROP * ALT? */
+@RequiredArgsConstructor
+public class IndexDMLHandler extends AsyncQueryHandler {
+  private static final Logger LOG = LogManager.getLogger();
+
+  public static final String DROP_INDEX_JOB_ID = "dropIndexJobId";
+
+  private final EMRServerlessClient emrServerlessClient;
+
+  private final DataSourceService dataSourceService;
+
+  private final DataSourceUserAuthorizationHelperImpl dataSourceUserAuthorizationHelper;
+
+  private final JobExecutionResponseReader jobExecutionResponseReader;
+
+  private final FlintIndexMetadataReader flintIndexMetadataReader;
+
+  private final Client client;
+
+  private final StateStore stateStore;
+
+  public static boolean isIndexDMLQuery(String jobId) {
+    return DROP_INDEX_JOB_ID.equalsIgnoreCase(jobId);
+  }
+
+  public DispatchQueryResponse handle(
+      DispatchQueryRequest dispatchQueryRequest, IndexQueryDetails indexDetails) {
+    DataSourceMetadata dataSourceMetadata =
+        dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource());
+    dataSourceUserAuthorizationHelper.authorizeDataSource(dataSourceMetadata);
+    FlintIndexMetadata indexMetadata = flintIndexMetadataReader.getFlintIndexMetadata(indexDetails);
+    // if index is created without auto refresh. there is no job to cancel.
+    String status = JobRunState.FAILED.toString();
+    String error = "";
+    long startTime = 0L;
+    try {
+      FlintIndexOp jobCancelOp =
+          new FlintIndexOpCancel(
+              stateStore, dispatchQueryRequest.getDatasource(), emrServerlessClient);
+      jobCancelOp.apply(indexMetadata);
+
+      FlintIndexOp indexDeleteOp =
+          new FlintIndexOpDelete(stateStore, dispatchQueryRequest.getDatasource());
+      indexDeleteOp.apply(indexMetadata);
+      status = JobRunState.SUCCESS.toString();
+    } catch (Exception e) {
+      error = e.getMessage();
+      LOG.error(e);
+    }
+
+    AsyncQueryId asyncQueryId = AsyncQueryId.newAsyncQueryId(dataSourceMetadata.getName());
+    IndexDMLResult indexDMLResult =
+        new IndexDMLResult(
+            asyncQueryId.getId(),
+            status,
+            error,
+            dispatchQueryRequest.getDatasource(),
+            System.currentTimeMillis() - startTime,
+            System.currentTimeMillis());
+    String resultIndex =
+        Optional.ofNullable(dataSourceMetadata.getResultIndex())
+            .orElse(SPARK_RESPONSE_BUFFER_INDEX_NAME);
+    createIndexDMLResult(stateStore, resultIndex).apply(indexDMLResult);
+
+    return new DispatchQueryResponse(asyncQueryId, DROP_INDEX_JOB_ID, resultIndex, null);
+  }
+
+  @Override
+  protected JSONObject getResponseFromResultIndex(AsyncQueryJobMetadata asyncQueryJobMetadata) {
+    String queryId = asyncQueryJobMetadata.getQueryId().getId();
+    return jobExecutionResponseReader.getResultWithQueryId(
+        queryId, asyncQueryJobMetadata.getResultIndex());
+  }
+
+  @Override
+  protected JSONObject getResponseFromExecutor(AsyncQueryJobMetadata asyncQueryJobMetadata) {
+    throw new IllegalStateException("[BUG] can't fetch result of index DML query form server");
+  }
+
+  @Override
+  public String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata) {
+    throw new IllegalArgumentException("can't cancel index DML query");
+  }
+}