From a1cc2fd11ac97b2bace6977f946c0bc0480cc620 Mon Sep 17 00:00:00 2001
From: Wenjun Ruan <wenjun@apache.org>
Date: Wed, 26 Jan 2022 11:09:57 +0800
Subject: [PATCH] [task] Python task support PYTHON_HOME (#8196)

---
 .../plugin/task/python/PythonTask.java        | 31 +++++++++++++--
 .../plugin/task/python/PythonTaskTest.java    | 38 +++++++++++++++++++
 2 files changed, 66 insertions(+), 3 deletions(-)
 create mode 100644 dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/test/java/org/apache/dolphinscheduler/plugin/task/python/PythonTaskTest.java

diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/main/java/org/apache/dolphinscheduler/plugin/task/python/PythonTask.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/main/java/org/apache/dolphinscheduler/plugin/task/python/PythonTask.java
index 382f1778b9..6d924780d9 100644
--- a/dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/main/java/org/apache/dolphinscheduler/plugin/task/python/PythonTask.java
+++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/main/java/org/apache/dolphinscheduler/plugin/task/python/PythonTask.java
@@ -31,6 +31,7 @@ import org.apache.dolphinscheduler.spi.task.request.TaskRequest;
 import org.apache.dolphinscheduler.spi.utils.JSONUtils;
 
 import org.apache.commons.io.FileUtils;
+import org.apache.commons.lang.StringUtils;
 
 import java.io.File;
 import java.io.IOException;
@@ -40,6 +41,8 @@ import java.nio.file.Paths;
 import java.util.HashMap;
 import java.util.Map;
 
+import com.google.common.base.Preconditions;
+
 /**
  * python task
  */
@@ -57,6 +60,10 @@ public class PythonTask extends AbstractTaskExecutor {
 
     private TaskRequest taskRequest;
 
+    private final String PYTHON_HOME = "PYTHON_HOME";
+
+    private final String DEFAULT_PYTHON_VERSION = "python";
+
     /**
      * constructor
      *
@@ -103,7 +110,7 @@ public class PythonTask extends AbstractTaskExecutor {
 
             // create this file
             createPythonCommandFileIfNotExists(pythonScriptContent,pythonScriptFile);
-            String command = "python " + pythonScriptFile;
+            String command = buildPythonExecuteCommand(pythonScriptFile);
 
             TaskResponse taskResponse = shellCommandExecutor.run(command);
             setExitStatusCode(taskResponse.getExitStatusCode());
@@ -182,8 +189,8 @@ public class PythonTask extends AbstractTaskExecutor {
 
             // write data to file
             FileUtils.writeStringToFile(new File(pythonScriptFile),
-                sb.toString(),
-                StandardCharsets.UTF_8);
+                    sb.toString(),
+                    StandardCharsets.UTF_8);
         }
     }
 
@@ -220,4 +227,22 @@ public class PythonTask extends AbstractTaskExecutor {
         return rawPythonScript;
     }
 
+    /**
+     * Build the python task command.
+     * If user have set the 'PYTHON_HOME' environment, we will use the 'PYTHON_HOME',
+     * if not, we will default use python.
+     *
+     * @param pythonFile Python file, cannot be empty.
+     * @return Python execute command, e.g. 'python test.py'.
+     */
+    private String buildPythonExecuteCommand(String pythonFile) {
+        Preconditions.checkNotNull(pythonFile, "Python file cannot be null");
+
+        String pythonHome = System.getenv(PYTHON_HOME);
+        if (StringUtils.isEmpty(pythonHome)) {
+            return DEFAULT_PYTHON_VERSION + " " + pythonFile;
+        }
+        return pythonHome + " " + pythonFile;
+    }
+
 }
diff --git a/dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/test/java/org/apache/dolphinscheduler/plugin/task/python/PythonTaskTest.java b/dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/test/java/org/apache/dolphinscheduler/plugin/task/python/PythonTaskTest.java
new file mode 100644
index 0000000000..08c5e0d8cf
--- /dev/null
+++ b/dolphinscheduler-task-plugin/dolphinscheduler-task-python/src/test/java/org/apache/dolphinscheduler/plugin/task/python/PythonTaskTest.java
@@ -0,0 +1,38 @@
+/*
+ * 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.dolphinscheduler.plugin.task.python;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.powermock.reflect.Whitebox;
+
+public class PythonTaskTest {
+
+    @Test
+    public void buildPythonExecuteCommand() throws Exception {
+        PythonTask pythonTask = createPythonTask();
+        String methodName = "buildPythonExecuteCommand";
+        String pythonFile = "test.py";
+        String result1 = Whitebox.invokeMethod(pythonTask, methodName, pythonFile);
+        Assert.assertEquals("python test.py", result1);
+    }
+
+    private PythonTask createPythonTask() {
+        return new PythonTask(null);
+    }
+}
\ No newline at end of file