Refactor StateStore

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

@@ -83,6 +83,7 @@
 import org.opensearch.sql.spark.cluster.ClusterManagerEventListener;
 import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.flint.FlintIndexMetadataServiceImpl;
+import org.opensearch.sql.spark.flint.FlintIndexStateModelService;
 import org.opensearch.sql.spark.rest.RestAsyncQueryManagementAction;
 import org.opensearch.sql.spark.storage.SparkStorageFactory;
 import org.opensearch.sql.spark.transport.TransportCancelAsyncQueryRequestAction;
@@ -229,7 +230,8 @@ public Collection<Object> createComponents(
             dataSourceService,
             injector.getInstance(FlintIndexMetadataServiceImpl.class),
             injector.getInstance(StateStore.class),
-            injector.getInstance(EMRServerlessClientFactory.class));
+            injector.getInstance(EMRServerlessClientFactory.class),
+            injector.getInstance(FlintIndexStateModelService.class));
     return ImmutableList.of(
         dataSourceService,
         injector.getInstance(AsyncQueryExecutorService.class),

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

@@ -7,7 +7,7 @@
 
 package org.opensearch.sql.spark.asyncquery;
 
-import static org.opensearch.sql.spark.execution.statestore.StateStore.createJobMetaData;
+import static org.opensearch.sql.spark.execution.statestore.StateStore.DATASOURCE_TO_REQUEST_INDEX;
 
 import java.util.Optional;
 import lombok.RequiredArgsConstructor;
@@ -17,6 +17,7 @@
 import org.opensearch.sql.spark.asyncquery.model.AsyncQueryId;
 import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata;
 import org.opensearch.sql.spark.execution.statestore.StateStore;
+import org.opensearch.sql.spark.execution.xcontent.AsyncQueryJobMetadataXContentSerializer;
 
 /** Opensearch implementation of {@link AsyncQueryJobMetadataStorageService} */
 @RequiredArgsConstructor
@@ -31,15 +32,22 @@ public class OpensearchAsyncQueryJobMetadataStorageService
   @Override
   public void storeJobMetadata(AsyncQueryJobMetadata asyncQueryJobMetadata) {
     AsyncQueryId queryId = asyncQueryJobMetadata.getQueryId();
-    createJobMetaData(stateStore, queryId.getDataSourceName()).apply(asyncQueryJobMetadata);
+    stateStore.create(
+        asyncQueryJobMetadata,
+        AsyncQueryJobMetadata::copy,
+        DATASOURCE_TO_REQUEST_INDEX.apply(queryId.getDataSourceName()));
   }
 
   @Override
   public Optional<AsyncQueryJobMetadata> getJobMetadata(String qid) {
     try {
       AsyncQueryId queryId = new AsyncQueryId(qid);
-      return StateStore.getJobMetaData(stateStore, queryId.getDataSourceName())
-          .apply(queryId.docId());
+      AsyncQueryJobMetadataXContentSerializer asyncQueryJobMetadataXContentSerializer =
+          new AsyncQueryJobMetadataXContentSerializer();
+      return stateStore.get(
+          queryId.docId(),
+          asyncQueryJobMetadataXContentSerializer::fromXContent,
+          DATASOURCE_TO_REQUEST_INDEX.apply(queryId.getDataSourceName()));
     } catch (Exception e) {
       LOGGER.error("Error while fetching the job metadata.", e);
       throw new AsyncQueryNotFoundException(String.format("Invalid QueryId: %s", qid));

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

@@ -7,19 +7,9 @@
 
 package org.opensearch.sql.spark.asyncquery.model;
 
-import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
-import static org.opensearch.sql.spark.execution.session.SessionModel.DATASOURCE_NAME;
-import static org.opensearch.sql.spark.execution.statement.StatementModel.QUERY_ID;
-
 import com.google.gson.Gson;
-import java.io.IOException;
-import java.util.Locale;
 import lombok.Data;
 import lombok.EqualsAndHashCode;
-import lombok.SneakyThrows;
-import org.opensearch.core.common.Strings;
-import org.opensearch.core.xcontent.XContentBuilder;
-import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.index.seqno.SequenceNumbers;
 import org.opensearch.sql.spark.dispatcher.model.JobType;
 import org.opensearch.sql.spark.execution.statestore.StateModel;
@@ -134,29 +124,6 @@ public String toString() {
     return new Gson().toJson(this);
   }
 
-  /**
-   * Converts JobMetadata to XContentBuilder.
-   *
-   * @return XContentBuilder {@link XContentBuilder}
-   * @throws Exception Exception.
-   */
-  @Override
-  public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
-    builder
-        .startObject()
-        .field(QUERY_ID, queryId.getId())
-        .field("type", TYPE_JOBMETA)
-        .field("jobId", jobId)
-        .field("applicationId", applicationId)
-        .field("resultIndex", resultIndex)
-        .field("sessionId", sessionId)
-        .field(DATASOURCE_NAME, datasourceName)
-        .field(JOB_TYPE, jobType.getText().toLowerCase(Locale.ROOT))
-        .field(INDEX_NAME, indexName)
-        .endObject();
-    return builder;
-  }
-
   /** copy builder. update seqNo and primaryTerm */
   public static AsyncQueryJobMetadata copy(
       AsyncQueryJobMetadata copy, long seqNo, long primaryTerm) {
@@ -173,72 +140,6 @@ public static AsyncQueryJobMetadata copy(
         primaryTerm);
   }
 
-  /**
-   * Convert xcontent parser to JobMetadata.
-   *
-   * @param parser parser.
-   * @return JobMetadata {@link AsyncQueryJobMetadata}
-   * @throws IOException IOException.
-   */
-  @SneakyThrows
-  public static AsyncQueryJobMetadata fromXContent(
-      XContentParser parser, long seqNo, long primaryTerm) {
-    AsyncQueryId queryId = null;
-    String jobId = null;
-    String applicationId = null;
-    String resultIndex = null;
-    String sessionId = null;
-    String datasourceName = null;
-    String jobTypeStr = null;
-    String indexName = null;
-    ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
-    while (!XContentParser.Token.END_OBJECT.equals(parser.nextToken())) {
-      String fieldName = parser.currentName();
-      parser.nextToken();
-      switch (fieldName) {
-        case QUERY_ID:
-          queryId = new AsyncQueryId(parser.textOrNull());
-          break;
-        case "jobId":
-          jobId = parser.textOrNull();
-          break;
-        case "applicationId":
-          applicationId = parser.textOrNull();
-          break;
-        case "resultIndex":
-          resultIndex = parser.textOrNull();
-          break;
-        case "sessionId":
-          sessionId = parser.textOrNull();
-          break;
-        case DATASOURCE_NAME:
-          datasourceName = parser.textOrNull();
-        case JOB_TYPE:
-          jobTypeStr = parser.textOrNull();
-        case INDEX_NAME:
-          indexName = parser.textOrNull();
-        case "type":
-          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(
-        queryId,
-        applicationId,
-        jobId,
-        resultIndex,
-        sessionId,
-        datasourceName,
-        Strings.isNullOrEmpty(jobTypeStr) ? null : JobType.fromString(jobTypeStr),
-        indexName,
-        seqNo,
-        primaryTerm);
-  }
-
   @Override
   public String getId() {
     return queryId.docId();

spark/src/main/java/org/opensearch/sql/spark/cluster/ClusterManagerEventListener.java

@@ -24,6 +24,7 @@
 import org.opensearch.sql.spark.client.EMRServerlessClientFactory;
 import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.flint.FlintIndexMetadataService;
+import org.opensearch.sql.spark.flint.FlintIndexStateModelService;
 import org.opensearch.threadpool.Scheduler.Cancellable;
 import org.opensearch.threadpool.ThreadPool;
 
@@ -39,6 +40,8 @@ public class ClusterManagerEventListener implements LocalNodeClusterManagerListe
   private FlintIndexMetadataService flintIndexMetadataService;
   private StateStore stateStore;
   private EMRServerlessClientFactory emrServerlessClientFactory;
+
+  private FlintIndexStateModelService flintIndexStateModelService;
   private Duration sessionTtlDuration;
   private Duration resultTtlDuration;
   private TimeValue streamingJobHouseKeepingInterval;
@@ -57,7 +60,8 @@ public ClusterManagerEventListener(
       DataSourceService dataSourceService,
       FlintIndexMetadataService flintIndexMetadataService,
       StateStore stateStore,
-      EMRServerlessClientFactory emrServerlessClientFactory) {
+      EMRServerlessClientFactory emrServerlessClientFactory,
+      FlintIndexStateModelService flintIndexStateModelService) {
     this.clusterService = clusterService;
     this.threadPool = threadPool;
     this.client = client;
@@ -70,7 +74,7 @@ public ClusterManagerEventListener(
     this.sessionTtlDuration = toDuration(sessionTtl.get(settings));
     this.resultTtlDuration = toDuration(resultTtl.get(settings));
     this.streamingJobHouseKeepingInterval = streamingJobHouseKeepingInterval.get(settings);
-
+    this.flintIndexStateModelService = flintIndexStateModelService;
     clusterService
         .getClusterSettings()
         .addSettingsUpdateConsumer(
@@ -153,7 +157,7 @@ private void initializeStreamingJobHouseKeeperCron() {
             new FlintStreamingJobHouseKeeperTask(
                 dataSourceService,
                 flintIndexMetadataService,
-                stateStore,
+                flintIndexStateModelService,
                 emrServerlessClientFactory),
             streamingJobHouseKeepingInterval,
             executorName());

spark/src/main/java/org/opensearch/sql/spark/cluster/FlintStreamingJobHouseKeeperTask.java

@@ -19,9 +19,9 @@
 import org.opensearch.sql.legacy.metrics.Metrics;
 import org.opensearch.sql.spark.client.EMRServerlessClientFactory;
 import org.opensearch.sql.spark.dispatcher.model.FlintIndexOptions;
-import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.flint.FlintIndexMetadata;
 import org.opensearch.sql.spark.flint.FlintIndexMetadataService;
+import org.opensearch.sql.spark.flint.FlintIndexStateModelService;
 import org.opensearch.sql.spark.flint.operation.FlintIndexOpAlter;
 import org.opensearch.sql.spark.flint.operation.FlintIndexOpDrop;
 
@@ -31,7 +31,7 @@ public class FlintStreamingJobHouseKeeperTask implements Runnable {
 
   private final DataSourceService dataSourceService;
   private final FlintIndexMetadataService flintIndexMetadataService;
-  private final StateStore stateStore;
+  private final FlintIndexStateModelService flintIndexStateModelService;
   private final EMRServerlessClientFactory emrServerlessClientFactory;
 
   private static final Logger LOGGER = LogManager.getLogger(FlintStreamingJobHouseKeeperTask.class);
@@ -96,7 +96,8 @@ private void dropAutoRefreshIndex(
     // When the datasource is deleted. Possibly Replace with VACUUM Operation.
     LOGGER.info("Attempting to drop auto refresh index: {}", autoRefreshIndex);
     FlintIndexOpDrop flintIndexOpDrop =
-        new FlintIndexOpDrop(stateStore, datasourceName, emrServerlessClientFactory.getClient());
+        new FlintIndexOpDrop(
+            flintIndexStateModelService, datasourceName, emrServerlessClientFactory.getClient());
     flintIndexOpDrop.apply(flintIndexMetadata);
     LOGGER.info("Successfully dropped index: {}", autoRefreshIndex);
   }
@@ -109,7 +110,7 @@ private void alterAutoRefreshIndex(
     FlintIndexOpAlter flintIndexOpAlter =
         new FlintIndexOpAlter(
             flintIndexOptions,
-            stateStore,
+            flintIndexStateModelService,
             datasourceName,
             emrServerlessClientFactory.getClient(),
             flintIndexMetadataService);

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

@@ -7,7 +7,6 @@
 
 import static org.opensearch.sql.spark.data.constants.SparkConstants.ERROR_FIELD;
 import static org.opensearch.sql.spark.data.constants.SparkConstants.STATUS_FIELD;
-import static org.opensearch.sql.spark.execution.statestore.StateStore.createIndexDMLResult;
 
 import com.amazonaws.services.emrserverless.model.JobRunState;
 import java.util.Map;
@@ -27,9 +26,10 @@
 import org.opensearch.sql.spark.dispatcher.model.IndexDMLResult;
 import org.opensearch.sql.spark.dispatcher.model.IndexQueryDetails;
 import org.opensearch.sql.spark.execution.statement.StatementState;
-import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.flint.FlintIndexMetadata;
 import org.opensearch.sql.spark.flint.FlintIndexMetadataService;
+import org.opensearch.sql.spark.flint.FlintIndexStateModelService;
+import org.opensearch.sql.spark.flint.IndexDMLResultStorageService;
 import org.opensearch.sql.spark.flint.operation.FlintIndexOp;
 import org.opensearch.sql.spark.flint.operation.FlintIndexOpAlter;
 import org.opensearch.sql.spark.flint.operation.FlintIndexOpDrop;
@@ -51,7 +51,8 @@ public class IndexDMLHandler extends AsyncQueryHandler {
 
   private final FlintIndexMetadataService flintIndexMetadataService;
 
-  private final StateStore stateStore;
+  private final FlintIndexStateModelService flintIndexStateModelService;
+  private final IndexDMLResultStorageService indexDMLResultStorageService;
 
   private final Client client;
 
@@ -106,7 +107,7 @@ private AsyncQueryId storeIndexDMLResult(
             dispatchQueryRequest.getDatasource(),
             System.currentTimeMillis() - startTime,
             System.currentTimeMillis());
-    createIndexDMLResult(stateStore, dataSourceMetadata.getResultIndex()).apply(indexDMLResult);
+    indexDMLResultStorageService.createIndexDMLResult(indexDMLResult, dataSourceMetadata.getName());
     return asyncQueryId;
   }
 
@@ -118,22 +119,25 @@ private void executeIndexOp(
       case DROP:
         FlintIndexOp dropOp =
             new FlintIndexOpDrop(
-                stateStore, dispatchQueryRequest.getDatasource(), emrServerlessClient);
+                flintIndexStateModelService,
+                dispatchQueryRequest.getDatasource(),
+                emrServerlessClient);
         dropOp.apply(indexMetadata);
         break;
       case ALTER:
         FlintIndexOpAlter flintIndexOpAlter =
             new FlintIndexOpAlter(
                 indexQueryDetails.getFlintIndexOptions(),
-                stateStore,
+                flintIndexStateModelService,
                 dispatchQueryRequest.getDatasource(),
                 emrServerlessClient,
                 flintIndexMetadataService);
         flintIndexOpAlter.apply(indexMetadata);
         break;
       case VACUUM:
         FlintIndexOp indexVacuumOp =
-            new FlintIndexOpVacuum(stateStore, dispatchQueryRequest.getDatasource(), client);
+            new FlintIndexOpVacuum(
+                flintIndexStateModelService, dispatchQueryRequest.getDatasource(), client);
         indexVacuumOp.apply(indexMetadata);
         break;
       default:

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

@@ -13,9 +13,9 @@
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryRequest;
 import org.opensearch.sql.spark.dispatcher.model.DispatchQueryResponse;
 import org.opensearch.sql.spark.dispatcher.model.JobType;
-import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.flint.FlintIndexMetadata;
 import org.opensearch.sql.spark.flint.FlintIndexMetadataService;
+import org.opensearch.sql.spark.flint.FlintIndexStateModelService;
 import org.opensearch.sql.spark.flint.operation.FlintIndexOp;
 import org.opensearch.sql.spark.flint.operation.FlintIndexOpCancel;
 import org.opensearch.sql.spark.leasemanager.LeaseManager;
@@ -25,18 +25,18 @@
 public class RefreshQueryHandler extends BatchQueryHandler {
 
   private final FlintIndexMetadataService flintIndexMetadataService;
-  private final StateStore stateStore;
+  private final FlintIndexStateModelService flintIndexStateModelService;
   private final EMRServerlessClient emrServerlessClient;
 
   public RefreshQueryHandler(
       EMRServerlessClient emrServerlessClient,
       JobExecutionResponseReader jobExecutionResponseReader,
       FlintIndexMetadataService flintIndexMetadataService,
-      StateStore stateStore,
+      FlintIndexStateModelService flintIndexStateModelService,
       LeaseManager leaseManager) {
     super(emrServerlessClient, jobExecutionResponseReader, leaseManager);
     this.flintIndexMetadataService = flintIndexMetadataService;
-    this.stateStore = stateStore;
+    this.flintIndexStateModelService = flintIndexStateModelService;
     this.emrServerlessClient = emrServerlessClient;
   }
 
@@ -52,7 +52,7 @@ public String cancelJob(AsyncQueryJobMetadata asyncQueryJobMetadata) {
     }
     FlintIndexMetadata indexMetadata = indexMetadataMap.get(asyncQueryJobMetadata.getIndexName());
     FlintIndexOp jobCancelOp =
-        new FlintIndexOpCancel(stateStore, datasourceName, emrServerlessClient);
+        new FlintIndexOpCancel(flintIndexStateModelService, datasourceName, emrServerlessClient);
     jobCancelOp.apply(indexMetadata);
     return asyncQueryJobMetadata.getQueryId().getId();
   }