Add sessionId parameters for create async query API

common/src/main/java/org/opensearch/sql/common/setting/Settings.java

@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.common.setting;
 
+import static org.opensearch.sql.common.setting.Settings.Key.SPARK_EXECUTION_SESSION_ENABLED;
+
 import com.google.common.base.Strings;
 import com.google.common.collect.ImmutableMap;
 import java.util.List;
@@ -36,7 +38,8 @@ public enum Key {
     METRICS_ROLLING_WINDOW("plugins.query.metrics.rolling_window"),
     METRICS_ROLLING_INTERVAL("plugins.query.metrics.rolling_interval"),
     SPARK_EXECUTION_ENGINE_CONFIG("plugins.query.executionengine.spark.config"),
-    CLUSTER_NAME("cluster.name");
+    CLUSTER_NAME("cluster.name"),
+    SPARK_EXECUTION_SESSION_ENABLED("plugins.query.executionengine.spark.session.enabled");
 
     @Getter private final String keyValue;
 
@@ -60,4 +63,9 @@ public static Optional<Key> of(String keyValue) {
   public abstract <T> T getSettingValue(Key key);
 
   public abstract List<?> getSettings();
+
+  /** Helper class */
+  public static boolean isSparkExecutionSessionEnabled(Settings settings) {
+    return settings.getSettingValue(SPARK_EXECUTION_SESSION_ENABLED);
+  }
 }

docs/user/admin/settings.rst

@@ -311,3 +311,39 @@ SQL query::
       "status": 400
     }
 
+plugins.query.executionengine.spark.session.enabled
+===================================================
+
+Description
+-----------
+
+By default, execution engine is executed in job mode. You can enable session mode by this setting.
+
+1. The default value is false.
+2. This setting is node scope.
+3. This setting can be updated dynamically.
+
+You can update the setting with a new value like this.
+
+SQL query::
+
+    sh$ curl -sS -H 'Content-Type: application/json' -X PUT localhost:9200/_plugins/_query/settings \
+    ... -d '{"transient":{"plugins.query.executionengine.spark.session.enabled":"false"}}'
+    {
+      "acknowledged": true,
+      "persistent": {},
+      "transient": {
+        "plugins": {
+          "query": {
+            "executionengine": {
+              "spark": {
+                "session": {
+                  "enabled": "false"
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+

docs/user/interfaces/asyncqueryinterface.rst

@@ -62,6 +62,49 @@ Sample Response::
       "queryId": "00fd796ut1a7eg0q"
     }
 
+Execute query in session
+------------------------
+
+if plugins.query.executionengine.spark.session.enabled is set to true, session based execution is enabled. Under the hood, all queries submitted to the same session will be executed in the same SparkContext. Session is auto closed if not query submission in 10 minutes.
+
+Async query response include ``sessionId`` indicate the query is executed in session.
+
+Sample Request::
+
+    curl --location 'http://localhost:9200/_plugins/_async_query' \
+    --header 'Content-Type: application/json' \
+    --data '{
+        "datasource" : "my_glue",
+        "lang" : "sql",
+        "query" : "select * from my_glue.default.http_logs limit 10"
+    }'
+
+Sample Response::
+
+    {
+      "queryId": "HlbM61kX6MDkAktO",
+      "sessionId": "1Giy65ZnzNlmsPAm"
+    }
+
+User could reuse the session by using ``sessionId`` query parameters.
+
+Sample Request::
+
+    curl --location 'http://localhost:9200/_plugins/_async_query?sessionId=1Giy65ZnzNlmsPAm' \
+    --header 'Content-Type: application/json' \
+    --data '{
+        "datasource" : "my_glue",
+        "lang" : "sql",
+        "query" : "select * from my_glue.default.http_logs limit 10"
+    }'
+
+Sample Response::
+
+    {
+      "queryId": "7GC4mHhftiTejvxN",
+      "sessionId": "1Giy65ZnzNlmsPAm"
+    }
+
 
 Async Query Result API
 ======================================

opensearch/src/main/java/org/opensearch/sql/opensearch/setting/OpenSearchSettings.java

@@ -135,6 +135,13 @@ public class OpenSearchSettings extends Settings {
           Setting.Property.NodeScope,
           Setting.Property.Dynamic);
 
+  public static final Setting<?> SPARK_EXECUTION_SESSION_ENABLED_SETTING =
+      Setting.boolSetting(
+          Key.SPARK_EXECUTION_SESSION_ENABLED.getKeyValue(),
+          false,
+          Setting.Property.NodeScope,
+          Setting.Property.Dynamic);
+
   /** Construct OpenSearchSetting. The OpenSearchSetting must be singleton. */
   @SuppressWarnings("unchecked")
   public OpenSearchSettings(ClusterSettings clusterSettings) {
@@ -205,6 +212,12 @@ public OpenSearchSettings(ClusterSettings clusterSettings) {
         Key.SPARK_EXECUTION_ENGINE_CONFIG,
         SPARK_EXECUTION_ENGINE_CONFIG,
         new Updater(Key.SPARK_EXECUTION_ENGINE_CONFIG));
+    register(
+        settingBuilder,
+        clusterSettings,
+        Key.SPARK_EXECUTION_SESSION_ENABLED,
+        SPARK_EXECUTION_SESSION_ENABLED_SETTING,
+        new Updater(Key.SPARK_EXECUTION_SESSION_ENABLED));
     registerNonDynamicSettings(
         settingBuilder, clusterSettings, Key.CLUSTER_NAME, ClusterName.CLUSTER_NAME_SETTING);
     defaultSettings = settingBuilder.build();
@@ -270,6 +283,7 @@ public static List<Setting<?>> pluginSettings() {
         .add(METRICS_ROLLING_INTERVAL_SETTING)
         .add(DATASOURCE_URI_HOSTS_DENY_LIST)
         .add(SPARK_EXECUTION_ENGINE_CONFIG)
+        .add(SPARK_EXECUTION_SESSION_ENABLED_SETTING)
         .build();
   }
 

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

@@ -7,6 +7,7 @@
 
 import static org.opensearch.sql.common.setting.Settings.Key.SPARK_EXECUTION_ENGINE_CONFIG;
 import static org.opensearch.sql.datasource.model.DataSourceMetadata.defaultOpenSearchDataSourceMetadata;
+import static org.opensearch.sql.spark.data.constants.SparkConstants.SPARK_REQUEST_BUFFER_INDEX_NAME;
 
 import com.amazonaws.auth.DefaultAWSCredentialsProviderChain;
 import com.amazonaws.services.emrserverless.AWSEMRServerless;
@@ -99,6 +100,8 @@
 import org.opensearch.sql.spark.config.SparkExecutionEngineConfigSupplier;
 import org.opensearch.sql.spark.config.SparkExecutionEngineConfigSupplierImpl;
 import org.opensearch.sql.spark.dispatcher.SparkQueryDispatcher;
+import org.opensearch.sql.spark.execution.session.SessionManager;
+import org.opensearch.sql.spark.execution.statestore.StateStore;
 import org.opensearch.sql.spark.flint.FlintIndexMetadataReaderImpl;
 import org.opensearch.sql.spark.response.JobExecutionResponseReader;
 import org.opensearch.sql.spark.rest.RestAsyncQueryManagementAction;
@@ -318,7 +321,11 @@ private AsyncQueryExecutorService createAsyncQueryExecutorService(
             new DataSourceUserAuthorizationHelperImpl(client),
             jobExecutionResponseReader,
             new FlintIndexMetadataReaderImpl(client),
-            client);
+            client,
+            new SessionManager(
+                new StateStore(SPARK_REQUEST_BUFFER_INDEX_NAME, client),
+                emrServerlessClient,
+                pluginSettings));
     return new AsyncQueryExecutorServiceImpl(
         asyncQueryJobMetadataStorageService,
         sparkQueryDispatcher,

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

@@ -65,14 +65,17 @@ public CreateAsyncQueryResponse createAsyncQuery(
                 createAsyncQueryRequest.getLang(),
                 sparkExecutionEngineConfig.getExecutionRoleARN(),
                 sparkExecutionEngineConfig.getClusterName(),
-                sparkExecutionEngineConfig.getSparkSubmitParameters()));
+                sparkExecutionEngineConfig.getSparkSubmitParameters(),
+                createAsyncQueryRequest.getSessionId()));
     asyncQueryJobMetadataStorageService.storeJobMetadata(
         new AsyncQueryJobMetadata(
             sparkExecutionEngineConfig.getApplicationId(),
             dispatchQueryResponse.getJobId(),
             dispatchQueryResponse.isDropIndexQuery(),
-            dispatchQueryResponse.getResultIndex()));
-    return new CreateAsyncQueryResponse(dispatchQueryResponse.getJobId());
+            dispatchQueryResponse.getResultIndex(),
+            dispatchQueryResponse.getSessionId()));
+    return new CreateAsyncQueryResponse(
+        dispatchQueryResponse.getJobId(), dispatchQueryResponse.getSessionId());
   }
 
   @Override
@@ -81,6 +84,7 @@ public AsyncQueryExecutionResponse getAsyncQueryResults(String queryId) {
     Optional<AsyncQueryJobMetadata> jobMetadata =
         asyncQueryJobMetadataStorageService.getJobMetadata(queryId);
     if (jobMetadata.isPresent()) {
+      String sessionId = jobMetadata.get().getSessionId();
       JSONObject jsonObject = sparkQueryDispatcher.getQueryResponse(jobMetadata.get());
       if (JobRunState.SUCCESS.toString().equals(jsonObject.getString(STATUS_FIELD))) {
         DefaultSparkSqlFunctionResponseHandle sparkSqlFunctionResponseHandle =
@@ -90,13 +94,18 @@ public AsyncQueryExecutionResponse getAsyncQueryResults(String queryId) {
           result.add(sparkSqlFunctionResponseHandle.next());
         }
         return new AsyncQueryExecutionResponse(
-            JobRunState.SUCCESS.toString(), sparkSqlFunctionResponseHandle.schema(), result, null);
+            JobRunState.SUCCESS.toString(),
+            sparkSqlFunctionResponseHandle.schema(),
+            result,
+            null,
+            sessionId);
       } else {
         return new AsyncQueryExecutionResponse(
             jsonObject.optString(STATUS_FIELD, JobRunState.FAILED.toString()),
             null,
             null,
-            jsonObject.optString(ERROR_FIELD, ""));
+            jsonObject.optString(ERROR_FIELD, ""),
+            sessionId);
       }
     }
     throw new AsyncQueryNotFoundException(String.format("QueryId: %s not found", queryId));

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

@@ -19,4 +19,5 @@ public class AsyncQueryExecutionResponse {
   private final ExecutionEngine.Schema schema;
   private final List<ExprValue> results;
   private final String error;
+  private final String sessionId;
 }

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

@@ -30,12 +30,15 @@
   private String jobId;
   private boolean isDropIndexQuery;
   private String resultIndex;
+  // optional sessionId.
+  private String sessionId;
 
   public AsyncQueryJobMetadata(String applicationId, String jobId, String resultIndex) {
     this.applicationId = applicationId;
     this.jobId = jobId;
     this.isDropIndexQuery = false;
     this.resultIndex = resultIndex;
+    this.sessionId = null;
   }
 
   @Override
@@ -57,6 +60,7 @@
     builder.field("applicationId", metadata.getApplicationId());
     builder.field("isDropIndexQuery", metadata.isDropIndexQuery());
     builder.field("resultIndex", metadata.getResultIndex());
+    builder.field("sessionId", metadata.getSessionId());
     builder.endObject();
     return builder;
   }
@@ -92,6 +96,7 @@
     String applicationId = null;
     boolean isDropIndexQuery = false;
     String resultIndex = null;
+    String sessionId = null;
     ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
     while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
       String fieldName = parser.currentName();
@@ -109,13 +114,17 @@
         case "resultIndex":
           resultIndex = parser.textOrNull();
           break;
+        case "sessionId":
+          sessionId = parser.textOrNull();
+          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(applicationId, jobId, isDropIndexQuery, resultIndex);
+    return new AsyncQueryJobMetadata(
+        applicationId, jobId, isDropIndexQuery, resultIndex, sessionId);
   }
 }

spark/src/main/java/org/opensearch/sql/spark/data/constants/SparkConstants.java

@@ -21,6 +21,7 @@ public class SparkConstants {
   public static final String SPARK_SQL_APPLICATION_JAR =
       "file:///home/hadoop/.ivy2/jars/org.opensearch_opensearch-spark-sql-application_2.12-0.1.0-SNAPSHOT.jar";
   public static final String SPARK_RESPONSE_BUFFER_INDEX_NAME = ".query_execution_result";
+  public static final String SPARK_REQUEST_BUFFER_INDEX_NAME = ".query_execution_request";
   // TODO should be replaced with mvn jar.
   public static final String FLINT_INTEGRATION_JAR =
       "s3://spark-datasource/flint-spark-integration-assembly-0.1.0-SNAPSHOT.jar";