[HUDI-398]Add set env for spark launcher

hudi-cli/src/main/java/org/apache/hudi/cli/commands/SetSparkEnvCommand.java

@@ -0,0 +1,46 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hudi.cli.commands;
+
+import org.springframework.shell.core.CommandMarker;
+import org.springframework.shell.core.annotation.CliCommand;
+import org.springframework.shell.core.annotation.CliOption;
+import org.springframework.stereotype.Component;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * CLI command to set env.
+ */
+@Component
+public class SetSparkEnvCommand implements CommandMarker {
+
+  public static Map env = new HashMap<String, String>();
+
+  @CliCommand(value = "set", help = "Set spark launcher env to cli")
+  public void setEnv(@CliOption(key = {"conf"}, help = "Env config to be set") final String confMap)
+      throws IllegalArgumentException {
+    String[] map = confMap.split("=");
+    if (map.length != 2) {
+      throw new IllegalArgumentException("Illegal set parameter, please use like [set SPARK_HOME=/usr/etc/spark]'");
+    }
+    env.put(map[0].trim(), map[1].trim());
+  }
+}

hudi-cli/src/main/java/org/apache/hudi/cli/utils/SparkUtil.java

@@ -19,6 +19,7 @@
 package org.apache.hudi.cli.utils;
 
 import org.apache.hudi.HoodieWriteClient;
+import org.apache.hudi.cli.commands.SetSparkEnvCommand;
 import org.apache.hudi.cli.commands.SparkMain;
 import org.apache.hudi.common.util.FSUtils;
 import org.apache.hudi.common.util.StringUtils;
@@ -30,6 +31,7 @@
 
 import java.io.File;
 import java.net.URISyntaxException;
+import java.util.Map;
 
 /**
  * Utility functions dealing with Spark.
@@ -45,13 +47,13 @@ public class SparkUtil {
   public static SparkLauncher initLauncher(String propertiesFile) throws URISyntaxException {
     String currentJar = new File(SparkUtil.class.getProtectionDomain().getCodeSource().getLocation().toURI().getPath())
         .getAbsolutePath();
+    Map<String, String> env = SetSparkEnvCommand.env;
     SparkLauncher sparkLauncher =
-        new SparkLauncher().setAppResource(currentJar).setMainClass(SparkMain.class.getName());
+        new SparkLauncher(env).setAppResource(currentJar).setMainClass(SparkMain.class.getName());
 
     if (!StringUtils.isNullOrEmpty(propertiesFile)) {
       sparkLauncher.setPropertiesFile(propertiesFile);
     }
-
     File libDirectory = new File(new File(currentJar).getParent(), "lib");
     for (String library : libDirectory.list()) {
       sparkLauncher.addJar(new File(libDirectory, library).getAbsolutePath());