add async-query-core module

Signed-off-by: Vamsi Manohar <[email protected]>
vamsi-amazon · Apr 18, 2024 · 53bce1c · 53bce1c
1 parent 5978512
commit 53bce1c
Show file tree

Hide file tree

Showing 31 changed files with 257 additions and 217 deletions.
diff --git a/settings.gradle b/settings.gradle
@@ -21,4 +21,4 @@ include 'prometheus'
 include 'benchmarks'
 include 'datasources'
 include 'spark'
-
+include 'async-query-core'
diff --git a/spark/build.gradle b/spark/build.gradle
@@ -8,6 +8,7 @@ plugins {
     id "io.freefair.lombok"
     id 'jacoco'
     id 'antlr'
+    id 'com.github.johnrengelman.shadow'
 }
 
 repositories {
@@ -75,6 +76,15 @@ dependencies {
     testImplementation project(':opensearch')
 }
 
+shadowJar {
+    archiveBaseName.set('async-query-core')
+    archiveVersion.set('1.0.0')  // Set the desired version
+    archiveClassifier.set('all')
+
+    from sourceSets.main.output
+    configurations = [project.configurations.runtimeClasspath]
+}
+
 test {
     useJUnitPlatform {
         includeEngines("junit-jupiter")

diff --git a/spark/src/main/antlr/SqlBaseLexer.g4 b/spark/src/main/antlr/SqlBaseLexer.g4
@@ -182,6 +182,7 @@ ELSE: 'ELSE';
 END: 'END';
 ESCAPE: 'ESCAPE';
 ESCAPED: 'ESCAPED';
+EVOLUTION: 'EVOLUTION';
 EXCEPT: 'EXCEPT';
 EXCHANGE: 'EXCHANGE';
 EXCLUDE: 'EXCLUDE';

diff --git a/spark/src/main/antlr/SqlBaseParser.g4 b/spark/src/main/antlr/SqlBaseParser.g4
@@ -480,7 +480,7 @@ dmlStatementNoWith
     | fromClause multiInsertQueryBody+                                             #multiInsertQuery
     | DELETE FROM identifierReference tableAlias whereClause?                      #deleteFromTable
     | UPDATE identifierReference tableAlias setClause whereClause?                 #updateTable
-    | MERGE INTO target=identifierReference targetAlias=tableAlias
+    | MERGE (WITH SCHEMA EVOLUTION)? INTO target=identifierReference targetAlias=tableAlias
         USING (source=identifierReference |
           LEFT_PAREN sourceQuery=query RIGHT_PAREN) sourceAlias=tableAlias
         ON mergeCondition=booleanExpression
@@ -1399,6 +1399,7 @@ ansiNonReserved
     | DOUBLE
     | DROP
     | ESCAPED
+    | EVOLUTION
     | EXCHANGE
     | EXCLUDE
     | EXISTS
@@ -1715,6 +1716,7 @@ nonReserved
     | END
     | ESCAPE
     | ESCAPED
+    | EVOLUTION
     | EXCHANGE
     | EXCLUDE
     | EXECUTE

diff --git a/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryJobMetadata.java b/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/AsyncQueryJobMetadata.java
@@ -25,7 +25,7 @@ public class AsyncQueryJobMetadata extends StateModel {
   private final AsyncQueryId queryId;
   private final String applicationId;
   private final String jobId;
-  private final String resultIndex;
+  private final String resultLocation;
   // optional sessionId.
   private final String sessionId;
   // since 2.13
@@ -41,12 +41,12 @@ public class AsyncQueryJobMetadata extends StateModel {
   @EqualsAndHashCode.Exclude private final long primaryTerm;
 
   public AsyncQueryJobMetadata(
-      AsyncQueryId queryId, String applicationId, String jobId, String resultIndex) {
+      AsyncQueryId queryId, String applicationId, String jobId, String resultLocation) {
     this(
         queryId,
         applicationId,
         jobId,
-        resultIndex,
+        resultLocation,
         null,
         null,
         JobType.INTERACTIVE,
@@ -59,13 +59,13 @@ public AsyncQueryJobMetadata(
       AsyncQueryId queryId,
       String applicationId,
       String jobId,
-      String resultIndex,
+      String resultLocation,
       String sessionId) {
     this(
         queryId,
         applicationId,
         jobId,
-        resultIndex,
+        resultLocation,
         sessionId,
         null,
         JobType.INTERACTIVE,
@@ -78,7 +78,7 @@ public AsyncQueryJobMetadata(
       AsyncQueryId queryId,
       String applicationId,
       String jobId,
-      String resultIndex,
+      String resultLocation,
       String sessionId,
       String datasourceName,
       JobType jobType,
@@ -87,7 +87,7 @@ public AsyncQueryJobMetadata(
         queryId,
         applicationId,
         jobId,
-        resultIndex,
+        resultLocation,
         sessionId,
         datasourceName,
         jobType,
@@ -100,7 +100,7 @@ public AsyncQueryJobMetadata(
       AsyncQueryId queryId,
       String applicationId,
       String jobId,
-      String resultIndex,
+      String resultLocation,
       String sessionId,
       String datasourceName,
       JobType jobType,
@@ -110,7 +110,7 @@ public AsyncQueryJobMetadata(
     this.queryId = queryId;
     this.applicationId = applicationId;
     this.jobId = jobId;
-    this.resultIndex = resultIndex;
+    this.resultLocation = resultLocation;
     this.sessionId = sessionId;
     this.datasourceName = datasourceName;
     this.jobType = jobType;
@@ -131,7 +131,7 @@ public static AsyncQueryJobMetadata copy(
         copy.getQueryId(),
         copy.getApplicationId(),
         copy.getJobId(),
-        copy.getResultIndex(),
+        copy.getResultLocation(),
         copy.getSessionId(),
         copy.datasourceName,
         copy.jobType,

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
@@ -38,8 +38,8 @@ public class BatchQueryHandler extends AsyncQueryHandler {
   protected JSONObject getResponseFromResultIndex(AsyncQueryJobMetadata asyncQueryJobMetadata) {
     // either empty json when the result is not available or data with status
     // Fetch from Result Index
-    return jobExecutionResponseReader.getResultFromOpensearchIndex(
-        asyncQueryJobMetadata.getJobId(), asyncQueryJobMetadata.getResultIndex());
+    return jobExecutionResponseReader.getResultWithJobId(
+        asyncQueryJobMetadata.getJobId(), asyncQueryJobMetadata.getResultLocation());
   }
 
   @Override

diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/IndexDMLHandler.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/IndexDMLHandler.java
@@ -15,7 +15,6 @@
 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.model.DataSourceMetadata;
 import org.opensearch.sql.spark.asyncquery.model.AsyncQueryId;
 import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata;
@@ -54,8 +53,6 @@ public class IndexDMLHandler extends AsyncQueryHandler {
   private final FlintIndexStateModelService flintIndexStateModelService;
   private final IndexDMLResultStorageService indexDMLResultStorageService;
 
-  private final Client client;
-
   public static boolean isIndexDMLQuery(String jobId) {
     return DROP_INDEX_JOB_ID.equalsIgnoreCase(jobId) || DML_QUERY_JOB_ID.equalsIgnoreCase(jobId);
   }
@@ -137,7 +134,9 @@ private void executeIndexOp(
       case VACUUM:
         FlintIndexOp indexVacuumOp =
             new FlintIndexOpVacuum(
-                flintIndexStateModelService, dispatchQueryRequest.getDatasource(), client);
+                flintIndexStateModelService,
+                dispatchQueryRequest.getDatasource(),
+                flintIndexMetadataService);
         indexVacuumOp.apply(indexMetadata);
         break;
       default:
@@ -163,7 +162,7 @@ private FlintIndexMetadata getFlintIndexMetadata(IndexQueryDetails indexDetails)
   protected JSONObject getResponseFromResultIndex(AsyncQueryJobMetadata asyncQueryJobMetadata) {
     String queryId = asyncQueryJobMetadata.getQueryId().getId();
     return jobExecutionResponseReader.getResultWithQueryId(
-        queryId, asyncQueryJobMetadata.getResultIndex());
+        queryId, asyncQueryJobMetadata.getResultLocation());
   }
 
   @Override

diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/InteractiveQueryHandler.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/InteractiveQueryHandler.java
@@ -45,7 +45,7 @@ public class InteractiveQueryHandler extends AsyncQueryHandler {
   protected JSONObject getResponseFromResultIndex(AsyncQueryJobMetadata asyncQueryJobMetadata) {
     String queryId = asyncQueryJobMetadata.getQueryId().getId();
     return jobExecutionResponseReader.getResultWithQueryId(
-        queryId, asyncQueryJobMetadata.getResultIndex());
+        queryId, asyncQueryJobMetadata.getResultLocation());
   }
 
   @Override

diff --git a/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java b/spark/src/main/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcher.java
@@ -9,7 +9,6 @@
 import java.util.Map;
 import lombok.AllArgsConstructor;
 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.spark.asyncquery.model.AsyncQueryId;
@@ -48,8 +47,6 @@ public class SparkQueryDispatcher {
 
   private FlintIndexMetadataService flintIndexMetadataService;
 
-  private Client client;
-
   private SessionManager sessionManager;
 
   private LeaseManager leaseManager;
@@ -166,8 +163,7 @@ private IndexDMLHandler createIndexDMLHandler(EMRServerlessClient emrServerlessC
         jobExecutionResponseReader,
         flintIndexMetadataService,
         flintIndexStateModelService,
-        indexDMLResultStorageService,
-        client);
+        indexDMLResultStorageService);
   }
 
   // TODO: Revisit this logic.

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
@@ -24,7 +24,6 @@
 import org.opensearch.sql.spark.execution.statement.StatementId;
 import org.opensearch.sql.spark.execution.statement.StatementStorageService;
 import org.opensearch.sql.spark.execution.statestore.SessionStorageService;
-import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.rest.model.LangType;
 import org.opensearch.sql.spark.utils.TimeProvider;
 
@@ -41,7 +40,6 @@ public class InteractiveSession implements Session {
   public static final String SESSION_ID_TAG_KEY = "sid";
 
   private final SessionId sessionId;
-  private final StateStore stateStore;
   private final StatementStorageService statementStorageService;
   private final SessionStorageService sessionStorageService;
   private final EMRServerlessClient serverlessClient;
@@ -103,7 +101,6 @@ public StatementId submit(QueryRequest request) {
                 .sessionId(sessionId)
                 .applicationId(sessionModel.getApplicationId())
                 .jobId(sessionModel.getJobId())
-                .stateStore(stateStore)
                 .statementStorageService(statementStorageService)
                 .statementId(statementId)
                 .langType(LangType.SQL)
@@ -139,7 +136,6 @@ public Optional<Statement> get(StatementId stID) {
                     .langType(model.getLangType())
                     .query(model.getQuery())
                     .queryId(model.getQueryId())
-                    .stateStore(stateStore)
                     .statementStorageService(statementStorageService)
                     .statementModel(model)
                     .build());

diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/session/SessionConfigSupplier.java b/spark/src/main/java/org/opensearch/sql/spark/execution/session/SessionConfigSupplier.java
@@ -0,0 +1,6 @@
+package org.opensearch.sql.spark.execution.session;
+
+public interface SessionConfigSupplier {
+
+  Long getSessionInactivityTimeoutMillis();
+}
diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/session/SessionManager.java b/spark/src/main/java/org/opensearch/sql/spark/execution/session/SessionManager.java
@@ -5,15 +5,12 @@
 
 package org.opensearch.sql.spark.execution.session;
 
-import static org.opensearch.sql.common.setting.Settings.Key.SESSION_INACTIVITY_TIMEOUT_MILLIS;
 import static org.opensearch.sql.spark.execution.session.SessionId.newSessionId;
 
 import java.util.Optional;
-import org.opensearch.sql.common.setting.Settings;
 import org.opensearch.sql.spark.client.EMRServerlessClientFactory;
 import org.opensearch.sql.spark.execution.statement.StatementStorageService;
 import org.opensearch.sql.spark.execution.statestore.SessionStorageService;
-import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.utils.RealTimeProvider;
 
 /**
@@ -22,31 +19,26 @@
  * <p>todo. add Session cache and Session sweeper.
  */
 public class SessionManager {
-  private final StateStore stateStore;
   private final StatementStorageService statementStorageService;
-
   private final SessionStorageService sessionStorageService;
   private final EMRServerlessClientFactory emrServerlessClientFactory;
-  private Settings settings;
+  private SessionConfigSupplier sessionConfigSupplier;
 
   public SessionManager(
-      StateStore stateStore,
       StatementStorageService statementStorageService,
       SessionStorageService sessionStorageService,
       EMRServerlessClientFactory emrServerlessClientFactory,
-      Settings settings) {
-    this.stateStore = stateStore;
+      SessionConfigSupplier sessionConfigSupplier) {
     this.statementStorageService = statementStorageService;
     this.sessionStorageService = sessionStorageService;
     this.emrServerlessClientFactory = emrServerlessClientFactory;
-    this.settings = settings;
+    this.sessionConfigSupplier = sessionConfigSupplier;
   }
 
   public Session createSession(CreateSessionRequest request) {
     InteractiveSession session =
         InteractiveSession.builder()
             .sessionId(newSessionId(request.getDatasourceName()))
-            .stateStore(stateStore)
             .statementStorageService(statementStorageService)
             .sessionStorageService(sessionStorageService)
             .serverlessClient(emrServerlessClientFactory.getClient())
@@ -80,13 +72,12 @@ public Optional<Session> getSession(SessionId sid, String dataSourceName) {
       InteractiveSession session =
           InteractiveSession.builder()
               .sessionId(sid)
-              .stateStore(stateStore)
               .statementStorageService(statementStorageService)
               .sessionStorageService(sessionStorageService)
               .serverlessClient(emrServerlessClientFactory.getClient())
               .sessionModel(model.get())
               .sessionInactivityTimeoutMilli(
-                  settings.getSettingValue(SESSION_INACTIVITY_TIMEOUT_MILLIS))
+                  sessionConfigSupplier.getSessionInactivityTimeoutMillis())
               .timeProvider(new RealTimeProvider())
               .build();
       return Optional.ofNullable(session);

diff --git a/spark/src/main/java/org/opensearch/sql/spark/execution/statement/Statement.java b/spark/src/main/java/org/opensearch/sql/spark/execution/statement/Statement.java
@@ -15,7 +15,6 @@
 import org.opensearch.index.engine.DocumentMissingException;
 import org.opensearch.index.engine.VersionConflictEngineException;
 import org.opensearch.sql.spark.execution.session.SessionId;
-import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.rest.model.LangType;
 
 /** Statement represent query to execute in session. One statement map to one session. */
@@ -32,7 +31,6 @@ public class Statement {
   private final String datasourceName;
   private final String query;
   private final String queryId;
-  private final StateStore stateStore;
   private final StatementStorageService statementStorageService;
 
   @Setter private StatementModel statementModel;

diff --git a/.../org/opensearch/sql/spark/execution/xcontent/AsyncQueryJobMetadataXContentSerializer.java b/.../org/opensearch/sql/spark/execution/xcontent/AsyncQueryJobMetadataXContentSerializer.java
@@ -35,7 +35,7 @@ public XContentBuilder toXContent(AsyncQueryJobMetadata jobMetadata, ToXContent.
         .field("type", TYPE_JOBMETA)
         .field("jobId", jobMetadata.getJobId())
         .field("applicationId", jobMetadata.getApplicationId())
-        .field("resultIndex", jobMetadata.getResultIndex())
+        .field("resultIndex", jobMetadata.getResultLocation())
         .field("sessionId", jobMetadata.getSessionId())
         .field(DATASOURCE_NAME, jobMetadata.getDatasourceName())
         .field(JOB_TYPE, jobMetadata.getJobType().getText().toLowerCase(Locale.ROOT))

diff --git a/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataService.java b/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataService.java
@@ -27,4 +27,11 @@ public interface FlintIndexMetadataService {
    * @param flintIndexOptions flintIndexOptions.
    */
   void updateIndexToManualRefresh(String indexName, FlintIndexOptions flintIndexOptions);
+
+  /**
+   * Deletes FlintIndex.
+   *
+   * @param indexName indexName.
+   */
+  void deleteFlintIndex(String indexName);
 }
diff --git a/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataServiceImpl.java b/spark/src/main/java/org/opensearch/sql/spark/flint/FlintIndexMetadataServiceImpl.java
@@ -31,7 +31,9 @@
 import org.apache.commons.lang3.StringUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import org.opensearch.action.admin.indices.delete.DeleteIndexRequest;
 import org.opensearch.action.admin.indices.mapping.get.GetMappingsResponse;
+import org.opensearch.action.support.master.AcknowledgedResponse;
 import org.opensearch.client.Client;
 import org.opensearch.sql.spark.dispatcher.model.FlintIndexOptions;
 
@@ -87,6 +89,14 @@ public void updateIndexToManualRefresh(String indexName, FlintIndexOptions flint
     client.admin().indices().preparePutMapping(indexName).setSource(flintMetadataMap).get();
   }
 
+  @Override
+  public void deleteFlintIndex(String indexName) {
+    LOGGER.info("Vacuuming Flint index {}", indexName);
+    DeleteIndexRequest request = new DeleteIndexRequest().indices(indexName);
+    AcknowledgedResponse response = client.admin().indices().delete(request).actionGet();
+    LOGGER.info("OpenSearch index delete result: {}", response.isAcknowledged());
+  }
+
   private void validateFlintIndexOptions(
       String kind, Map<String, Object> existingOptions, Map<String, String> newOptions) {
     if ((newOptions.containsKey(INCREMENTAL_REFRESH)