diff --git a/nifi-nar-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java b/nifi-nar-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java index f34199b5bc..3cfb2b30f4 100644 --- a/nifi-nar-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java +++ b/nifi-nar-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java @@ -227,10 +227,7 @@ public class PythonProcess { private Process launchPythonProcess(final int listeningPort, final String authToken) throws IOException { final File pythonFrameworkDirectory = processConfig.getPythonFrameworkDirectory(); final File pythonApiDirectory = new File(pythonFrameworkDirectory.getParentFile(), "api"); - final File pythonCmdFile = new File(processConfig.getPythonCommand()); - final String pythonCmd = pythonCmdFile.getName(); - final File pythonCommandFile = new File(virtualEnvHome, "bin/" + pythonCmd); - final String pythonCommand = pythonCommandFile.getAbsolutePath(); + final String pythonCommand = resolvePythonCommand(); final File controllerPyFile = new File(pythonFrameworkDirectory, PYTHON_CONTROLLER_FILENAME); final ProcessBuilder processBuilder = new ProcessBuilder(); @@ -256,7 +253,7 @@ public class PythonProcess { processBuilder.environment().put("JAVA_PORT", String.valueOf(listeningPort)); processBuilder.environment().put("ENV_HOME", virtualEnvHome.getAbsolutePath()); processBuilder.environment().put("PYTHONPATH", pythonPath); - processBuilder.environment().put("PYTHON_CMD", pythonCommandFile.getAbsolutePath()); + processBuilder.environment().put("PYTHON_CMD", pythonCommand); processBuilder.environment().put("AUTH_TOKEN", authToken); // Redirect error stream to standard output stream @@ -267,6 +264,27 @@ public class PythonProcess { return processBuilder.start(); } + String resolvePythonCommand() throws IOException { + final File pythonCmdFile = new File(processConfig.getPythonCommand()); + final String pythonCmd = pythonCmdFile.getName(); + + // Find command directories according to standard Python venv conventions + final File[] virtualEnvDirectories = virtualEnvHome.listFiles((file, name) -> file.isDirectory() && (name.equals("bin") || name.equals("Scripts"))); + + final String commandExecutableDirectory; + if (virtualEnvDirectories == null || virtualEnvDirectories.length == 0) { + throw new IOException("Python binary directory could not be found in " + virtualEnvHome); + } else if( virtualEnvDirectories.length == 1) { + commandExecutableDirectory = virtualEnvDirectories[0].getName(); + } else { + // Default to bin directory for macOS and Linux + commandExecutableDirectory = "bin"; + } + + final File pythonCommandFile = new File(virtualEnvHome, commandExecutableDirectory + File.separator + pythonCmd); + return pythonCommandFile.getAbsolutePath(); + } + private void setupEnvironment() throws IOException { final File environmentCreationCompleteFile = new File(virtualEnvHome, "env-creation-complete.txt"); diff --git a/nifi-nar-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/test/java/org/apache/nifi/py4j/PythonProcessTest.java b/nifi-nar-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/test/java/org/apache/nifi/py4j/PythonProcessTest.java new file mode 100644 index 0000000000..4a1ad32638 --- /dev/null +++ b/nifi-nar-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/test/java/org/apache/nifi/py4j/PythonProcessTest.java @@ -0,0 +1,109 @@ +/* + * 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.nifi.py4j; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.when; + +import java.io.File; +import java.io.IOException; + +import org.apache.nifi.python.ControllerServiceTypeLookup; +import org.apache.nifi.python.PythonProcessConfig; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.io.CleanupMode; +import org.junit.jupiter.api.io.TempDir; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class PythonProcessTest { + + private static final String UNIX_BIN_DIR = "bin"; + + private static final String WINDOWS_SCRIPTS_DIR = "Scripts"; + + private static final String PYTHON_CMD = "python"; + + private PythonProcess pythonProcess; + + @TempDir(cleanup = CleanupMode.ON_SUCCESS) + private File virtualEnvHome; + + @Mock + private PythonProcessConfig pythonProcessConfig; + + @Mock + private ControllerServiceTypeLookup controllerServiceTypeLookup; + + @BeforeEach + public void setUp() { + this.pythonProcess = new PythonProcess(this.pythonProcessConfig, this.controllerServiceTypeLookup, virtualEnvHome, "Controller", "Controller"); + } + + @Test + void testResolvePythonCommandWindows() throws IOException { + final File scriptsDir = new File(virtualEnvHome, WINDOWS_SCRIPTS_DIR); + assertTrue(scriptsDir.mkdir()); + + when(pythonProcessConfig.getPythonCommand()).thenReturn(PYTHON_CMD); + final String result = this.pythonProcess.resolvePythonCommand(); + + final String expected = getExpectedBinaryPath(WINDOWS_SCRIPTS_DIR); + assertEquals(expected, result); + } + + @Test + void testResolvePythonCommandUnix() throws IOException { + final File binDir = new File(virtualEnvHome, UNIX_BIN_DIR); + assertTrue(binDir.mkdir()); + + when(pythonProcessConfig.getPythonCommand()).thenReturn(PYTHON_CMD); + final String result = this.pythonProcess.resolvePythonCommand(); + + final String expected = getExpectedBinaryPath(UNIX_BIN_DIR); + assertEquals(expected, result); + } + + @Test + void testResolvePythonCommandPreferBin() throws IOException { + final File binDir = new File(virtualEnvHome, UNIX_BIN_DIR); + assertTrue(binDir.mkdir()); + final File scriptsDir = new File(virtualEnvHome, WINDOWS_SCRIPTS_DIR); + assertTrue(scriptsDir.mkdir()); + + when(pythonProcessConfig.getPythonCommand()).thenReturn(PYTHON_CMD); + final String result = this.pythonProcess.resolvePythonCommand(); + + final String expected = getExpectedBinaryPath(UNIX_BIN_DIR); + assertEquals(expected, result); + } + + @Test + void testResolvePythonCommandNone() { + when(pythonProcessConfig.getPythonCommand()).thenReturn(PYTHON_CMD); + assertThrows(IOException.class, ()-> this.pythonProcess.resolvePythonCommand()); + } + + private String getExpectedBinaryPath(String binarySubDirectoryName) { + return this.virtualEnvHome.getAbsolutePath() + File.separator + binarySubDirectoryName + File.separator + PYTHON_CMD; + } +}