NIFI-13394 Check candidate directory for python command

This closes #8961

Signed-off-by: David Handermann <exceptionfactory@apache.org>
This commit is contained in:
Bob Paulin 2024-06-12 18:21:33 -05:00 committed by exceptionfactory
parent 039cd2f18a
commit d25efb385c
No known key found for this signature in database
2 changed files with 30 additions and 6 deletions

View File

@ -17,6 +17,7 @@
package org.apache.nifi.py4j;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.nifi.logging.LogLevel;
import org.apache.nifi.py4j.client.JavaObjectBindings;
import org.apache.nifi.py4j.client.NiFiPythonGateway;
@ -302,14 +303,22 @@ public class PythonProcess {
} else if (virtualEnvDirectories.length == 1) {
commandExecutableDirectory = virtualEnvDirectories[0].getName();
} else {
// Default to bin directory for macOS and Linux
commandExecutableDirectory = "bin";
commandExecutableDirectory = findExecutableDirectory(pythonCmd, virtualEnvDirectories);
}
final File pythonCommandFile = new File(virtualEnvHome, commandExecutableDirectory + File.separator + pythonCmd);
return pythonCommandFile.getAbsolutePath();
}
String findExecutableDirectory(final String pythonCmd, final File[] virtualEnvDirectories) throws IOException {
// Check for python command.
return List.of(virtualEnvDirectories)
.stream()
.filter(file -> ArrayUtils.isNotEmpty(file.list((dir, name) -> name.startsWith(pythonCmd))))
.findFirst()
.orElseThrow(() -> new IOException("Failed to find Python command [%s]".formatted(pythonCmd)))
.getName();
}
private void setupEnvironment() throws IOException {
// Environment creation is only necessary if using PIP. Otherwise, the Process requires no outside dependencies, other than those

View File

@ -91,17 +91,32 @@ class PythonProcessTest {
}
@Test
void testResolvePythonCommandPreferBin() throws IOException {
void testResolvePythonCommandFindCommand() 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());
final File fakeWindowsPythonExe = new File(scriptsDir, PYTHON_CMD + ".exe");
assertTrue(fakeWindowsPythonExe.createNewFile());
when(pythonProcessConfig.getPythonCommand()).thenReturn(PYTHON_CMD);
final String result = this.pythonProcess.resolvePythonCommand();
final String expected = getExpectedBinaryPath(WINDOWS_SCRIPTS_DIR);
assertEquals(expected, result);
}
@Test
void testResolvePythonCommandFindCommandMissingPythonCmd() 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);
assertThrows(IOException.class, () -> this.pythonProcess.resolvePythonCommand());
}
@Test