Drop Index Implementation

Signed-off-by: Vamsi Manohar <[email protected]>

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

@@ -97,6 +97,7 @@
 import org.opensearch.sql.spark.client.EmrServerlessClientImplEMR;
 import org.opensearch.sql.spark.config.SparkExecutionEngineConfig;
 import org.opensearch.sql.spark.dispatcher.SparkQueryDispatcher;
+import org.opensearch.sql.spark.flint.FlintIndexMetadataReaderImpl;
 import org.opensearch.sql.spark.response.JobExecutionResponseReader;
 import org.opensearch.sql.spark.rest.RestAsyncQueryManagementAction;
 import org.opensearch.sql.spark.storage.SparkStorageFactory;
@@ -297,14 +298,15 @@ private DataSourceServiceImpl createDataSourceService() {
   private AsyncQueryExecutorService createAsyncQueryExecutorService() {
     AsyncQueryJobMetadataStorageService asyncQueryJobMetadataStorageService =
         new OpensearchAsyncQueryJobMetadataStorageService(client, clusterService);
-    EMRServerlessClient EMRServerlessClient = createEMRServerlessClient();
+    EMRServerlessClient emrServerlessClient = createEMRServerlessClient();
     JobExecutionResponseReader jobExecutionResponseReader = new JobExecutionResponseReader(client);
     SparkQueryDispatcher sparkQueryDispatcher =
         new SparkQueryDispatcher(
-            EMRServerlessClient,
+            emrServerlessClient,
             this.dataSourceService,
             new DataSourceUserAuthorizationHelperImpl(client),
-            jobExecutionResponseReader);
+            jobExecutionResponseReader,
+            new FlintIndexMetadataReaderImpl(client));
     return new AsyncQueryExecutorServiceImpl(
         asyncQueryJobMetadataStorageService, sparkQueryDispatcher, pluginSettings);
   }

spark/build.gradle

@@ -90,7 +90,8 @@ jacocoTestCoverageVerification {
                     'org.opensearch.sql.spark.transport.model.*',
                     'org.opensearch.sql.spark.asyncquery.model.*',
                     'org.opensearch.sql.spark.asyncquery.exceptions.*',
-                    'org.opensearch.sql.spark.dispatcher.model.*'
+                    'org.opensearch.sql.spark.dispatcher.model.*',
+                    'org.opensearch.sql.spark.flint.FlintIndexType'
             ]
             limit {
                 counter = 'LINE'

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

@@ -25,6 +25,7 @@
 import org.opensearch.sql.spark.config.SparkExecutionEngineConfig;
 import org.opensearch.sql.spark.dispatcher.SparkQueryDispatcher;
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest;
+import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse;
 import org.opensearch.sql.spark.functions.response.DefaultSparkSqlFunctionResponseHandle;
 import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest;
 import org.opensearch.sql.spark.rest.model.CreateAsyncQueryResponse;
@@ -64,7 +65,7 @@ public CreateAsyncQueryResponse createAsyncQuery(
                     SparkExecutionEngineConfig.toSparkExecutionEngineConfig(
                         sparkExecutionEngineConfigString));
     ClusterName clusterName = settings.getSettingValue(CLUSTER_NAME);
-    String jobId =
+    DispatchQueryResponse dispatchQueryResponse =
         sparkQueryDispatcher.dispatch(
             new DispatchQueryRequest(
                 sparkExecutionEngineConfig.getApplicationId(),
@@ -74,8 +75,11 @@ public CreateAsyncQueryResponse createAsyncQuery(
                 sparkExecutionEngineConfig.getExecutionRoleARN(),
                 clusterName.value()));
     asyncQueryJobMetadataStorageService.storeJobMetadata(
-        new AsyncQueryJobMetadata(jobId, sparkExecutionEngineConfig.getApplicationId()));
-    return new CreateAsyncQueryResponse(jobId);
+        new AsyncQueryJobMetadata(
+            sparkExecutionEngineConfig.getApplicationId(),
+            dispatchQueryResponse.getJobId(),
+            dispatchQueryResponse.isDropIndexQuery()));
+    return new CreateAsyncQueryResponse(dispatchQueryResponse.getJobId());
   }
 
   @Override
@@ -84,9 +88,7 @@ public AsyncQueryExecutionResponse getAsyncQueryResults(String queryId) {
     Optional<AsyncQueryJobMetadata> jobMetadata =
         asyncQueryJobMetadataStorageService.getJobMetadata(queryId);
     if (jobMetadata.isPresent()) {
-      JSONObject jsonObject =
-          sparkQueryDispatcher.getQueryResponse(
-              jobMetadata.get().getApplicationId(), jobMetadata.get().getJobId());
+      JSONObject jsonObject = sparkQueryDispatcher.getQueryResponse(jobMetadata.get());
       if (JobRunState.SUCCESS.toString().equals(jsonObject.getString("status"))) {
         DefaultSparkSqlFunctionResponseHandle sparkSqlFunctionResponseHandle =
             new DefaultSparkSqlFunctionResponseHandle(jsonObject);
@@ -108,8 +110,7 @@ public String cancelQuery(String queryId) {
     Optional<AsyncQueryJobMetadata> asyncQueryJobMetadata =
         asyncQueryJobMetadataStorageService.getJobMetadata(queryId);
     if (asyncQueryJobMetadata.isPresent()) {
-      return sparkQueryDispatcher.cancelJob(
-          asyncQueryJobMetadata.get().getApplicationId(), queryId);
+      return sparkQueryDispatcher.cancelJob(asyncQueryJobMetadata.get());
     }
     throw new AsyncQueryNotFoundException(String.format("QueryId: %s not found", queryId));
   }

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

@@ -13,6 +13,7 @@
 import java.io.IOException;
 import lombok.AllArgsConstructor;
 import lombok.Data;
+import lombok.EqualsAndHashCode;
 import org.opensearch.common.xcontent.XContentFactory;
 import org.opensearch.common.xcontent.XContentType;
 import org.opensearch.core.xcontent.DeprecationHandler;
@@ -23,9 +24,17 @@
 /** This class models all the metadata required for a job. */
 @Data
 @AllArgsConstructor
+@EqualsAndHashCode
 public class AsyncQueryJobMetadata {
-  private String jobId;
   private String applicationId;
+  private String jobId;
+  private boolean isDropIndexQuery;
+
+  public AsyncQueryJobMetadata(String applicationId, String jobId) {
+    this.applicationId = applicationId;
+    this.jobId = jobId;
+    this.isDropIndexQuery = false;
+  }
 
   @Override
   public String toString() {
@@ -44,6 +53,7 @@ public static XContentBuilder convertToXContent(AsyncQueryJobMetadata metadata)
     builder.startObject();
     builder.field("jobId", metadata.getJobId());
     builder.field("applicationId", metadata.getApplicationId());
+    builder.field("isDropIndexQuery", metadata.isDropIndexQuery());
     builder.endObject();
     return builder;
   }
@@ -77,6 +87,7 @@ public static AsyncQueryJobMetadata toJobMetadata(String json) throws IOExceptio
   public static AsyncQueryJobMetadata toJobMetadata(XContentParser parser) throws IOException {
     String jobId = null;
     String applicationId = null;
+    boolean isDropIndexQuery = false;
     ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
     while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
       String fieldName = parser.currentName();
@@ -88,13 +99,16 @@ public static AsyncQueryJobMetadata toJobMetadata(XContentParser parser) throws
         case "applicationId":
           applicationId = parser.textOrNull();
           break;
+        case "isDropIndexQuery":
+          isDropIndexQuery = parser.booleanValue();
+          break;
         default:
           throw new IllegalArgumentException("Unknown field: " + fieldName);
       }
     }
     if (jobId == null || applicationId == null) {
       throw new IllegalArgumentException("jobId and applicationId are required fields.");
     }
-    return new AsyncQueryJobMetadata(jobId, applicationId);
+    return new AsyncQueryJobMetadata(applicationId, jobId, isDropIndexQuery);
   }
 }

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

@@ -11,15 +11,19 @@
 import java.util.HashMap;
 import java.util.Map;
 import lombok.AllArgsConstructor;
+import org.apache.commons.lang3.RandomStringUtils;
 import org.json.JSONObject;
 import org.opensearch.sql.datasource.DataSourceService;
 import org.opensearch.sql.datasources.auth.DataSourceUserAuthorizationHelperImpl;
+import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata;
 import org.opensearch.sql.spark.asyncquery.model.SparkSubmitParameters;
 import org.opensearch.sql.spark.client.EMRServerlessClient;
 import org.opensearch.sql.spark.client.StartJobRequest;
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest;
+import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse;
 import org.opensearch.sql.spark.dispatcher.model.FullyQualifiedTableName;
 import org.opensearch.sql.spark.dispatcher.model.IndexDetails;
+import org.opensearch.sql.spark.flint.FlintIndexMetadataReader;
 import org.opensearch.sql.spark.response.JobExecutionResponseReader;
 import org.opensearch.sql.spark.rest.model.LangType;
 import org.opensearch.sql.spark.utils.SQLQueryUtils;
@@ -42,49 +46,64 @@ public class SparkQueryDispatcher {
 
   private JobExecutionResponseReader jobExecutionResponseReader;
 
-  public String dispatch(DispatchQueryRequest dispatchQueryRequest) {
-    return emrServerlessClient.startJobRun(getStartJobRequest(dispatchQueryRequest));
+  private FlintIndexMetadataReader flintIndexMetadataReader;
+
+  public DispatchQueryResponse dispatch(DispatchQueryRequest dispatchQueryRequest) {
+    if (LangType.SQL.equals(dispatchQueryRequest.getLangType())) {
+      return handleSQLQuery(dispatchQueryRequest);
+    } else {
+      // Since we don't need any extra handling for PPL, we are treating it as normal dispatch
+      // Query.
+      return handleNonIndexQuery(dispatchQueryRequest);
+    }
   }
 
   // TODO : Fetch from Result Index and then make call to EMR Serverless.
-  public JSONObject getQueryResponse(String applicationId, String queryId) {
-    GetJobRunResult getJobRunResult = emrServerlessClient.getJobRunResult(applicationId, queryId);
+  public JSONObject getQueryResponse(AsyncQueryJobMetadata asyncQueryJobMetadata) {
+    GetJobRunResult getJobRunResult =
+        emrServerlessClient.getJobRunResult(
+            asyncQueryJobMetadata.getApplicationId(), asyncQueryJobMetadata.getJobId());
     JSONObject result = new JSONObject();
     if (getJobRunResult.getJobRun().getState().equals(JobRunState.SUCCESS.toString())) {
-      result = jobExecutionResponseReader.getResultFromOpensearchIndex(queryId);
+      result =
+          jobExecutionResponseReader.getResultFromOpensearchIndex(asyncQueryJobMetadata.getJobId());
     }
     result.put("status", getJobRunResult.getJobRun().getState());
     return result;
   }
 
-  public String cancelJob(String applicationId, String jobId) {
-    CancelJobRunResult cancelJobRunResult = emrServerlessClient.cancelJobRun(applicationId, jobId);
+  public String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata) {
+    CancelJobRunResult cancelJobRunResult =
+        emrServerlessClient.cancelJobRun(
+            asyncQueryJobMetadata.getApplicationId(), asyncQueryJobMetadata.getJobId());
     return cancelJobRunResult.getJobRunId();
   }
 
-  // we currently don't support index queries in PPL language.
-  // so we are treating all of them as non-index queries which don't require any kind of query
-  // parsing.
-  private StartJobRequest getStartJobRequest(DispatchQueryRequest dispatchQueryRequest) {
-    if (LangType.SQL.equals(dispatchQueryRequest.getLangType())) {
-      if (SQLQueryUtils.isIndexQuery(dispatchQueryRequest.getQuery()))
-        return getStartJobRequestForIndexRequest(dispatchQueryRequest);
-      else {
-        return getStartJobRequestForNonIndexQueries(dispatchQueryRequest);
+  private DispatchQueryResponse handleSQLQuery(DispatchQueryRequest dispatchQueryRequest) {
+    if (SQLQueryUtils.isIndexQuery(dispatchQueryRequest.getQuery())) {
+      IndexDetails indexDetails =
+          SQLQueryUtils.extractIndexDetails(dispatchQueryRequest.getQuery());
+      if (indexDetails.isDropIndex()) {
+        return handleDropIndexQuery(dispatchQueryRequest, indexDetails);
+      } else {
+        return handleIndexQuery(dispatchQueryRequest, indexDetails);
       }
     } else {
-      return getStartJobRequestForNonIndexQueries(dispatchQueryRequest);
+      return handleNonIndexQuery(dispatchQueryRequest);
     }
   }
 
-  private StartJobRequest getStartJobRequestForNonIndexQueries(
-      DispatchQueryRequest dispatchQueryRequest) {
-    StartJobRequest startJobRequest;
+  private DispatchQueryResponse handleIndexQuery(
+      DispatchQueryRequest dispatchQueryRequest, IndexDetails indexDetails) {
+    FullyQualifiedTableName fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName();
     dataSourceUserAuthorizationHelper.authorizeDataSource(
         this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource()));
-    String jobName = dispatchQueryRequest.getClusterName() + ":" + "non-index-query";
+    String jobName = dispatchQueryRequest.getClusterName() + ":" + "index-query";
     Map<String, String> tags = getDefaultTagsForJobSubmission(dispatchQueryRequest);
-    startJobRequest =
+    tags.put(INDEX_TAG_KEY, indexDetails.getIndexName());
+    tags.put(TABLE_TAG_KEY, fullyQualifiedTableName.getTableName());
+    tags.put(SCHEMA_TAG_KEY, fullyQualifiedTableName.getSchemaName());
+    StartJobRequest startJobRequest =
         new StartJobRequest(
             dispatchQueryRequest.getQuery(),
             jobName,
@@ -94,26 +113,21 @@ private StartJobRequest getStartJobRequestForNonIndexQueries(
                 .dataSource(
                     dataSourceService.getRawDataSourceMetadata(
                         dispatchQueryRequest.getDatasource()))
+                .structuredStreaming(indexDetails.getAutoRefresh())
                 .build()
                 .toString(),
             tags,
-            false);
-    return startJobRequest;
+            indexDetails.getAutoRefresh());
+    String jobId = emrServerlessClient.startJobRun(startJobRequest);
+    return new DispatchQueryResponse(jobId, false);
   }
 
-  private StartJobRequest getStartJobRequestForIndexRequest(
-      DispatchQueryRequest dispatchQueryRequest) {
-    StartJobRequest startJobRequest;
-    IndexDetails indexDetails = SQLQueryUtils.extractIndexDetails(dispatchQueryRequest.getQuery());
-    FullyQualifiedTableName fullyQualifiedTableName = indexDetails.getFullyQualifiedTableName();
+  private DispatchQueryResponse handleNonIndexQuery(DispatchQueryRequest dispatchQueryRequest) {
     dataSourceUserAuthorizationHelper.authorizeDataSource(
         this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource()));
-    String jobName = dispatchQueryRequest.getClusterName() + ":" + "index-query";
+    String jobName = dispatchQueryRequest.getClusterName() + ":" + "non-index-query";
     Map<String, String> tags = getDefaultTagsForJobSubmission(dispatchQueryRequest);
-    tags.put(INDEX_TAG_KEY, indexDetails.getIndexName());
-    tags.put(TABLE_TAG_KEY, fullyQualifiedTableName.getTableName());
-    tags.put(SCHEMA_TAG_KEY, fullyQualifiedTableName.getSchemaName());
-    startJobRequest =
+    StartJobRequest startJobRequest =
         new StartJobRequest(
             dispatchQueryRequest.getQuery(),
             jobName,
@@ -123,12 +137,22 @@ private StartJobRequest getStartJobRequestForIndexRequest(
                 .dataSource(
                     dataSourceService.getRawDataSourceMetadata(
                         dispatchQueryRequest.getDatasource()))
-                .structuredStreaming(indexDetails.getAutoRefresh())
                 .build()
                 .toString(),
             tags,
-            indexDetails.getAutoRefresh());
-    return startJobRequest;
+            false);
+    String jobId = emrServerlessClient.startJobRun(startJobRequest);
+    return new DispatchQueryResponse(jobId, false);
+  }
+
+  private DispatchQueryResponse handleDropIndexQuery(
+      DispatchQueryRequest dispatchQueryRequest, IndexDetails indexDetails) {
+    dataSourceUserAuthorizationHelper.authorizeDataSource(
+        this.dataSourceService.getRawDataSourceMetadata(dispatchQueryRequest.getDatasource()));
+    String jobId = flintIndexMetadataReader.getJobIdFromFlintIndexMetadata(indexDetails);
+    emrServerlessClient.cancelJobRun(dispatchQueryRequest.getApplicationId(), jobId);
+    String dropIndexDummyJobId = RandomStringUtils.randomAlphanumeric(16);
+    return new DispatchQueryResponse(dropIndexDummyJobId, true);
   }
 
   private static Map<String, String> getDefaultTagsForJobSubmission(

spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/DispatchQueryResponse.java

@@ -0,0 +1,11 @@
+package org.opensearch.sql.spark.dispatcher.model;
+
+import lombok.AllArgsConstructor;
+import lombok.Data;
+
+@Data
+@AllArgsConstructor
+public class DispatchQueryResponse {
+  private String jobId;
+  private boolean isDropIndexQuery;
+}

spark/src/main/java/org/opensearch/sql/spark/dispatcher/model/IndexDetails.java

@@ -5,13 +5,22 @@
 
 package org.opensearch.sql.spark.dispatcher.model;
 
+import lombok.AllArgsConstructor;
 import lombok.Data;
+import lombok.EqualsAndHashCode;
+import lombok.NoArgsConstructor;
+import org.opensearch.sql.spark.flint.FlintIndexType;
 
 /** Index details in an async query. */
 @Data
+@AllArgsConstructor
+@NoArgsConstructor
+@EqualsAndHashCode
 public class IndexDetails {
   private String indexName;
   private FullyQualifiedTableName fullyQualifiedTableName;
   // by default, auto_refresh = false;
   private Boolean autoRefresh = false;
+  private boolean isDropIndex;
+  private FlintIndexType indexType;
 }

spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataReader.java

@@ -0,0 +1,15 @@
+package org.opensearch.sql.spark.flint;
+
+import org.opensearch.sql.spark.dispatcher.model.IndexDetails;
+
+/** Interface for FlintIndexMetadataReader */
+public interface FlintIndexMetadataReader {
+
+  /**
+   * Given Index details, get the streaming job Id.
+   *
+   * @param indexDetails indexDetails.
+   * @return jobId.
+   */
+  String getJobIdFromFlintIndexMetadata(IndexDetails indexDetails);
+}