Include cluster name to spark env

Signed-off-by: Louis Chu <[email protected]>
opensearch-project · Jan 5, 2024 · c9bb54d · c9bb54d
1 parent e8ca07a
commit c9bb54d
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParameters.java b/spark/src/main/java/org/opensearch/sql/spark/asyncquery/model/SparkSubmitParameters.java
@@ -59,7 +59,8 @@ private Builder() {
       config.put(SPARK_JAR_REPOSITORIES_KEY, AWS_SNAPSHOT_REPOSITORY);
       config.put(SPARK_DRIVER_ENV_JAVA_HOME_KEY, JAVA_HOME_LOCATION);
       config.put(SPARK_EXECUTOR_ENV_JAVA_HOME_KEY, JAVA_HOME_LOCATION);
-      config.put(FLINT_INDEX_STORE_CLUSTER_NAME_KEY, FLINT_DEFAULT_CLUSTER_NAME);
+      config.put(SPARK_DRIVER_ENV_FLINT_CLUSTER_NAME_KEY, FLINT_DEFAULT_CLUSTER_NAME);
+      config.put(SPARK_EXECUTOR_ENV_FLINT_CLUSTER_NAME_KEY, FLINT_DEFAULT_CLUSTER_NAME);
       config.put(FLINT_INDEX_STORE_HOST_KEY, FLINT_DEFAULT_HOST);
       config.put(FLINT_INDEX_STORE_PORT_KEY, FLINT_DEFAULT_PORT);
       config.put(FLINT_INDEX_STORE_SCHEME_KEY, FLINT_DEFAULT_SCHEME);
@@ -79,7 +80,8 @@ public Builder className(String className) {
     }
 
     public Builder clusterName(String clusterName) {
-      config.put(FLINT_INDEX_STORE_CLUSTER_NAME_KEY, clusterName);
+      config.put(SPARK_DRIVER_ENV_FLINT_CLUSTER_NAME_KEY, clusterName);
+      config.put(SPARK_EXECUTOR_ENV_FLINT_CLUSTER_NAME_KEY, clusterName);
       return this;
     }
 

diff --git a/spark/src/main/java/org/opensearch/sql/spark/data/constants/SparkConstants.java b/spark/src/main/java/org/opensearch/sql/spark/data/constants/SparkConstants.java
@@ -46,7 +46,10 @@ public class SparkConstants {
   public static final String SPARK_DRIVER_ENV_JAVA_HOME_KEY =
       "spark.emr-serverless.driverEnv.JAVA_HOME";
   public static final String SPARK_EXECUTOR_ENV_JAVA_HOME_KEY = "spark.executorEnv.JAVA_HOME";
-  public static final String FLINT_INDEX_STORE_CLUSTER_NAME_KEY = "spark.flint.clusterName";
+  public static final String SPARK_DRIVER_ENV_FLINT_CLUSTER_NAME_KEY =
+      "spark.emr-serverless.driverEnv.FLINT_CLUSTER_NAME";
+  public static final String SPARK_EXECUTOR_ENV_FLINT_CLUSTER_NAME_KEY =
+      "spark.executorEnv.FLINT_CLUSTER_NAME";
   public static final String FLINT_INDEX_STORE_HOST_KEY = "spark.datasource.flint.host";
   public static final String FLINT_INDEX_STORE_PORT_KEY = "spark.datasource.flint.port";
   public static final String FLINT_INDEX_STORE_SCHEME_KEY = "spark.datasource.flint.scheme";

diff --git a/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java b/spark/src/test/java/org/opensearch/sql/spark/dispatcher/SparkQueryDispatcherTest.java
@@ -1052,7 +1052,8 @@ private String constructExpectedSparkSubmitParameterString(
                + "  --conf"
                + " spark.emr-serverless.driverEnv.JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto.x86_64/"
                + "  --conf spark.executorEnv.JAVA_HOME=/usr/lib/jvm/java-17-amazon-corretto.x86_64/"
-               + "  --conf spark.flint.clusterName=TEST_CLUSTER  --conf"
+               + "  --conf spark.emr-serverless.driverEnv.FLINT_CLUSTER_NAME=TEST_CLUSTER  --conf"
+               + " spark.executorEnv.FLINT_CLUSTER_NAME=TEST_CLUSTER  --conf"
                + " spark.datasource.flint.host=search-flint-dp-benchmark-cf5crj5mj2kfzvgwdeynkxnefy.eu-west-1.es.amazonaws.com"
                + "  --conf spark.datasource.flint.port=-1  --conf"
                + " spark.datasource.flint.scheme=https  --conf spark.datasource.flint.auth="