[Backport 2.11] deprecated job-metadata-index

common/build.gradle

@@ -34,7 +34,7 @@ repositories {
 dependencies {
     api "org.antlr:antlr4-runtime:4.7.1"
     api group: 'com.google.guava', name: 'guava', version: '32.0.1-jre'
-    api group: 'org.apache.logging.log4j', name: 'log4j-core', version:'2.20.0'
+    api group: 'org.apache.logging.log4j', name: 'log4j-core', version:'2.21.0'
     api group: 'org.apache.commons', name: 'commons-lang3', version: '3.12.0'
     api group: 'com.squareup.okhttp3', name: 'okhttp', version: '4.9.3'
     implementation 'com.github.babbel:okhttp-aws-signer:1.0.2'

integ-test/build.gradle

@@ -167,7 +167,7 @@ dependencies {
     testImplementation group: 'org.opensearch.client', name: 'opensearch-rest-client', version: "${opensearch_version}"
     testImplementation group: 'org.opensearch.driver', name: 'opensearch-sql-jdbc', version: System.getProperty("jdbcDriverVersion", '1.2.0.0')
     testImplementation group: 'org.hamcrest', name: 'hamcrest', version: '2.1'
-    implementation group: 'org.apache.logging.log4j', name: 'log4j-core', version:'2.20.0'
+    implementation group: 'org.apache.logging.log4j', name: 'log4j-core', version:'2.21.0'
     testImplementation project(':opensearch-sql-plugin')
     testImplementation project(':legacy')
     testImplementation('org.junit.jupiter:junit-jupiter-api:5.6.2')

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

@@ -306,8 +306,9 @@ private DataSourceServiceImpl createDataSourceService() {
   private AsyncQueryExecutorService createAsyncQueryExecutorService(
       SparkExecutionEngineConfigSupplier sparkExecutionEngineConfigSupplier,
       SparkExecutionEngineConfig sparkExecutionEngineConfig) {
+    StateStore stateStore = new StateStore(client, clusterService);
     AsyncQueryJobMetadataStorageService asyncQueryJobMetadataStorageService =
-        new OpensearchAsyncQueryJobMetadataStorageService(client, clusterService);
+        new OpensearchAsyncQueryJobMetadataStorageService(stateStore);
     EMRServerlessClient emrServerlessClient =
         createEMRServerlessClient(sparkExecutionEngineConfig.getRegion());
     JobExecutionResponseReader jobExecutionResponseReader = new JobExecutionResponseReader(client);
@@ -319,8 +320,7 @@ private AsyncQueryExecutorService createAsyncQueryExecutorService(
             jobExecutionResponseReader,
             new FlintIndexMetadataReaderImpl(client),
             client,
-            new SessionManager(
-                new StateStore(client, clusterService), emrServerlessClient, pluginSettings));
+            new SessionManager(stateStore, emrServerlessClient, pluginSettings));
     return new AsyncQueryExecutorServiceImpl(
         asyncQueryJobMetadataStorageService,
         sparkQueryDispatcher,

ppl/build.gradle

@@ -49,7 +49,7 @@ dependencies {
     implementation "org.antlr:antlr4-runtime:4.7.1"
     implementation group: 'com.google.guava', name: 'guava', version: '32.0.1-jre'
     api group: 'org.json', name: 'json', version: '20231013'
-    implementation group: 'org.apache.logging.log4j', name: 'log4j-core', version:'2.20.0'
+    implementation group: 'org.apache.logging.log4j', name: 'log4j-core', version:'2.21.0'
     api project(':common')
     api project(':core')
     api project(':protocol')

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

@@ -69,13 +69,14 @@ public CreateAsyncQueryResponse createAsyncQuery(
                 createAsyncQueryRequest.getSessionId()));
     asyncQueryJobMetadataStorageService.storeJobMetadata(
         new AsyncQueryJobMetadata(
+            dispatchQueryResponse.getQueryId(),
             sparkExecutionEngineConfig.getApplicationId(),
             dispatchQueryResponse.getJobId(),
             dispatchQueryResponse.isDropIndexQuery(),
             dispatchQueryResponse.getResultIndex(),
             dispatchQueryResponse.getSessionId()));
     return new CreateAsyncQueryResponse(
-        dispatchQueryResponse.getJobId(), dispatchQueryResponse.getSessionId());
+        dispatchQueryResponse.getQueryId().getId(), dispatchQueryResponse.getSessionId());
   }
 
   @Override

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

@@ -7,166 +7,31 @@
 
 package org.opensearch.sql.spark.asyncquery;
 
-import java.io.IOException;
-import java.io.InputStream;
-import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
-import java.util.List;
+import static org.opensearch.sql.spark.execution.statestore.StateStore.createJobMetaData;
+
 import java.util.Optional;
-import org.apache.commons.io.IOUtils;
-import org.apache.logging.log4j.LogManager;
-import org.apache.logging.log4j.Logger;
-import org.opensearch.action.DocWriteRequest;
-import org.opensearch.action.DocWriteResponse;
-import org.opensearch.action.admin.indices.create.CreateIndexRequest;
-import org.opensearch.action.admin.indices.create.CreateIndexResponse;
-import org.opensearch.action.index.IndexRequest;
-import org.opensearch.action.index.IndexResponse;
-import org.opensearch.action.search.SearchRequest;
-import org.opensearch.action.search.SearchResponse;
-import org.opensearch.action.support.WriteRequest;
-import org.opensearch.client.Client;
-import org.opensearch.cluster.service.ClusterService;
-import org.opensearch.common.action.ActionFuture;
-import org.opensearch.common.util.concurrent.ThreadContext;
-import org.opensearch.common.xcontent.XContentType;
-import org.opensearch.index.query.QueryBuilder;
-import org.opensearch.index.query.QueryBuilders;
-import org.opensearch.search.SearchHit;
-import org.opensearch.search.builder.SearchSourceBuilder;
+import lombok.RequiredArgsConstructor;
+import org.opensearch.sql.spark.asyncquery.model.AsyncQueryId;
 import org.opensearch.sql.spark.asyncquery.model.AsyncQueryJobMetadata;
+import org.opensearch.sql.spark.execution.statestore.StateStore;
 
 /** Opensearch implementation of {@link AsyncQueryJobMetadataStorageService} */
+@RequiredArgsConstructor
 public class OpensearchAsyncQueryJobMetadataStorageService
     implements AsyncQueryJobMetadataStorageService {
 
-  public static final String JOB_METADATA_INDEX = ".ql-job-metadata";
-  private static final String JOB_METADATA_INDEX_MAPPING_FILE_NAME =
-      "job-metadata-index-mapping.yml";
-  private static final String JOB_METADATA_INDEX_SETTINGS_FILE_NAME =
-      "job-metadata-index-settings.yml";
-  private static final Logger LOG = LogManager.getLogger();
-  private final Client client;
-  private final ClusterService clusterService;
-
-  /**
-   * This class implements JobMetadataStorageService interface using OpenSearch as underlying
-   * storage.
-   *
-   * @param client opensearch NodeClient.
-   * @param clusterService ClusterService.
-   */
-  public OpensearchAsyncQueryJobMetadataStorageService(
-      Client client, ClusterService clusterService) {
-    this.client = client;
-    this.clusterService = clusterService;
-  }
+  private final StateStore stateStore;
 
   @Override
   public void storeJobMetadata(AsyncQueryJobMetadata asyncQueryJobMetadata) {
-    if (!this.clusterService.state().routingTable().hasIndex(JOB_METADATA_INDEX)) {
-      createJobMetadataIndex();
-    }
-    IndexRequest indexRequest = new IndexRequest(JOB_METADATA_INDEX);
-    indexRequest.id(asyncQueryJobMetadata.getJobId());
-    indexRequest.opType(DocWriteRequest.OpType.CREATE);
-    indexRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
-    ActionFuture<IndexResponse> indexResponseActionFuture;
-    IndexResponse indexResponse;
-    try (ThreadContext.StoredContext storedContext =
-        client.threadPool().getThreadContext().stashContext()) {
-      indexRequest.source(AsyncQueryJobMetadata.convertToXContent(asyncQueryJobMetadata));
-      indexResponseActionFuture = client.index(indexRequest);
-      indexResponse = indexResponseActionFuture.actionGet();
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
-
-    if (indexResponse.getResult().equals(DocWriteResponse.Result.CREATED)) {
-      LOG.debug("JobMetadata   : {}  successfully created", asyncQueryJobMetadata.getJobId());
-    } else {
-      throw new RuntimeException(
-          "Saving job metadata information failed with result : "
-              + indexResponse.getResult().getLowercase());
-    }
+    AsyncQueryId queryId = asyncQueryJobMetadata.getQueryId();
+    createJobMetaData(stateStore, queryId.getDataSourceName()).apply(asyncQueryJobMetadata);
   }
 
   @Override
-  public Optional<AsyncQueryJobMetadata> getJobMetadata(String jobId) {
-    if (!this.clusterService.state().routingTable().hasIndex(JOB_METADATA_INDEX)) {
-      createJobMetadataIndex();
-      return Optional.empty();
-    }
-    return searchInJobMetadataIndex(QueryBuilders.termQuery("jobId.keyword", jobId)).stream()
-        .findFirst();
-  }
-
-  private void createJobMetadataIndex() {
-    try {
-      InputStream mappingFileStream =
-          OpensearchAsyncQueryJobMetadataStorageService.class
-              .getClassLoader()
-              .getResourceAsStream(JOB_METADATA_INDEX_MAPPING_FILE_NAME);
-      InputStream settingsFileStream =
-          OpensearchAsyncQueryJobMetadataStorageService.class
-              .getClassLoader()
-              .getResourceAsStream(JOB_METADATA_INDEX_SETTINGS_FILE_NAME);
-      CreateIndexRequest createIndexRequest = new CreateIndexRequest(JOB_METADATA_INDEX);
-      createIndexRequest
-          .mapping(IOUtils.toString(mappingFileStream, StandardCharsets.UTF_8), XContentType.YAML)
-          .settings(
-              IOUtils.toString(settingsFileStream, StandardCharsets.UTF_8), XContentType.YAML);
-      ActionFuture<CreateIndexResponse> createIndexResponseActionFuture;
-      try (ThreadContext.StoredContext ignored =
-          client.threadPool().getThreadContext().stashContext()) {
-        createIndexResponseActionFuture = client.admin().indices().create(createIndexRequest);
-      }
-      CreateIndexResponse createIndexResponse = createIndexResponseActionFuture.actionGet();
-      if (createIndexResponse.isAcknowledged()) {
-        LOG.info("Index: {} creation Acknowledged", JOB_METADATA_INDEX);
-      } else {
-        throw new RuntimeException("Index creation is not acknowledged.");
-      }
-    } catch (Throwable e) {
-      throw new RuntimeException(
-          "Internal server error while creating"
-              + JOB_METADATA_INDEX
-              + " index:: "
-              + e.getMessage());
-    }
-  }
-
-  private List<AsyncQueryJobMetadata> searchInJobMetadataIndex(QueryBuilder query) {
-    SearchRequest searchRequest = new SearchRequest();
-    searchRequest.indices(JOB_METADATA_INDEX);
-    SearchSourceBuilder searchSourceBuilder = new SearchSourceBuilder();
-    searchSourceBuilder.query(query);
-    searchSourceBuilder.size(1);
-    searchRequest.source(searchSourceBuilder);
-    // https://github.com/opensearch-project/sql/issues/1801.
-    searchRequest.preference("_primary_first");
-    ActionFuture<SearchResponse> searchResponseActionFuture;
-    try (ThreadContext.StoredContext ignored =
-        client.threadPool().getThreadContext().stashContext()) {
-      searchResponseActionFuture = client.search(searchRequest);
-    }
-    SearchResponse searchResponse = searchResponseActionFuture.actionGet();
-    if (searchResponse.status().getStatus() != 200) {
-      throw new RuntimeException(
-          "Fetching job metadata information failed with status : " + searchResponse.status());
-    } else {
-      List<AsyncQueryJobMetadata> list = new ArrayList<>();
-      for (SearchHit searchHit : searchResponse.getHits().getHits()) {
-        String sourceAsString = searchHit.getSourceAsString();
-        AsyncQueryJobMetadata asyncQueryJobMetadata;
-        try {
-          asyncQueryJobMetadata = AsyncQueryJobMetadata.toJobMetadata(sourceAsString);
-        } catch (IOException e) {
-          throw new RuntimeException(e);
-        }
-        list.add(asyncQueryJobMetadata);
-      }
-      return list;
-    }
+  public Optional<AsyncQueryJobMetadata> getJobMetadata(String qid) {
+    AsyncQueryId queryId = new AsyncQueryId(qid);
+    return StateStore.getJobMetaData(stateStore, queryId.getDataSourceName())
+        .apply(queryId.docId());
   }
 }

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

@@ -0,0 +1,35 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.sql.spark.asyncquery.model;
+
+import static org.opensearch.sql.spark.utils.IDUtils.decode;
+import static org.opensearch.sql.spark.utils.IDUtils.encode;
+
+import lombok.Data;
+
+/** Async query id. */
+@Data
+public class AsyncQueryId {
+  private final String id;
+
+  public static AsyncQueryId newAsyncQueryId(String datasourceName) {
+    return new AsyncQueryId(encode(datasourceName));
+  }
+
+  public String getDataSourceName() {
+    return decode(id);
+  }
+
+  /** OpenSearch DocId. */
+  public String docId() {
+    return "qid" + id;
+  }
+
+  @Override
+  public String toString() {
+    return "asyncQueryId=" + id;
+  }
+}