From d25efb385c8464e3817ea7f7e8e2edcb9d24a650 Mon Sep 17 00:00:00 2001
From: Bob Paulin <bob@bobpaulin.com>
Date: Wed, 12 Jun 2024 18:21:33 -0500
Subject: [PATCH] NIFI-13394 Check candidate directory for python command

This closes #8961

Signed-off-by: David Handermann <exceptionfactory@apache.org>
---
 .../org/apache/nifi/py4j/PythonProcess.java   | 13 +++++++++--
 .../apache/nifi/py4j/PythonProcessTest.java   | 23 +++++++++++++++----
 2 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java b/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java
index cc805467c1..52fe2e8a7f 100644
--- a/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java
+++ b/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/main/java/org/apache/nifi/py4j/PythonProcess.java
@@ -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
diff --git a/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/test/java/org/apache/nifi/py4j/PythonProcessTest.java b/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/test/java/org/apache/nifi/py4j/PythonProcessTest.java
index 8797d66eba..8366f248f8 100644
--- a/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/test/java/org/apache/nifi/py4j/PythonProcessTest.java
+++ b/nifi-extension-bundles/nifi-py4j-bundle/nifi-py4j-bridge/src/test/java/org/apache/nifi/py4j/PythonProcessTest.java
@@ -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