CommandLine environment

3 files changed, 64 insertions, 5 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLine.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLine.java
index e00c1dabf0f..bda7590e38f 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLine.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLine.java
@@ -12,7 +12,9 @@ import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.TreeMap;
 import java.util.function.Predicate;
 import java.util.logging.Logger;
 import java.util.regex.Pattern;
@@ -47,6 +49,7 @@ public class CommandLine {
     public static final Duration DEFAULT_SIGKILL_GRACE_PERIOD = Duration.ofMinutes(30);
 
     private final List<String> arguments = new ArrayList<>();
+    private final TreeMap<String, String> environment = new TreeMap<>();
     private final TaskContext taskContext;
     private final ProcessFactory processFactory;
 
@@ -80,6 +83,25 @@ public class CommandLine {
         return add(arguments.split("\\s+"));
     }
 
+    /** Set an environment variable, overriding any existing. */
+    public CommandLine setEnvironmentVariable(String name, String value) {
+        if (name.indexOf('=') != -1) {
+            throw new IllegalArgumentException("name contains '=': " + name);
+        }
+        Objects.requireNonNull(value, "cannot set environment variable to null");
+
+        environment.put(name, value);
+        return this;
+    }
+
+    public CommandLine removeEnvironmentVariable(String name) {
+        if (name.indexOf('=') != -1) {
+            throw new IllegalArgumentException("name contains '=': " + name);
+        }
+        environment.put(name, null);
+        return this;
+    }
+
     /**
      * Execute a shell-like program in a child process:
      *  - the program is recorded and logged as modifying the system, but see executeSilently().
@@ -153,15 +175,30 @@ public class CommandLine {
     /** Returns a shell-like representation of the command. */
     @Override
     public String toString() {
-        String command = arguments.stream()
+        var command = new StringBuilder();
+
+        if (!environment.isEmpty()) {
+            // Pretend environment is propagated through the env program for display purposes
+            command.append(environment.entrySet().stream()
+                    .map(entry -> {
+                        if (entry.getValue() == null) {
+                            return "-u " + maybeEscapeArgument(entry.getKey());
+                        } else {
+                            return maybeEscapeArgument(entry.getKey() + "=" + entry.getValue());
+                        }
+                    })
+                    .collect(Collectors.joining(" ", "env ", " ")));
+        }
+
+        command.append(arguments.stream()
                 .map(CommandLine::maybeEscapeArgument)
-                .collect(Collectors.joining(" "));
+                .collect(Collectors.joining(" ")));
 
         // Note: Both of these cannot be confused with an argument since they would
         // require escaping.
-        command += redirectStderrToStdoutInsteadOfDiscard ? " 2>&1" : " 2>/dev/null";
+        command.append(redirectStderrToStdoutInsteadOfDiscard ? " 2>&1" : " 2>/dev/null");
 
-        return command;
+        return command.toString();
     }
 
 
@@ -254,6 +291,9 @@ public class CommandLine {
 
     public List<String> getArguments() { return Collections.unmodifiableList(arguments); }
 
+    /** Returns a copy of the environment overrides.  A null value means the environment variable should be removed. */
+    public TreeMap<String, String> getEnvironmentOverrides() { return new TreeMap<>(environment); }
+
     // Accessor fields necessary for classes in this package. Could be public if necessary.
     boolean getRedirectStderrToStdoutInsteadOfDiscard() { return redirectStderrToStdoutInsteadOfDiscard; }
     Predicate<Integer> getSuccessfulExitCodePredicate() { return successfulExitCodePredicate; }
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/ProcessFactoryImpl.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/ProcessFactoryImpl.java
index f4be8b7479f..65561793a09 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/ProcessFactoryImpl.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/task/util/process/ProcessFactoryImpl.java
@@ -3,7 +3,6 @@
 package com.yahoo.vespa.hosted.node.admin.task.util.process;
 
 import com.yahoo.jdisc.Timer;
-import java.util.logging.Level;
 
 import java.io.File;
 import java.io.IOException;
@@ -14,6 +13,7 @@ import java.nio.file.attribute.PosixFilePermission;
 import java.nio.file.attribute.PosixFilePermissions;
 import java.util.List;
 import java.util.Set;
+import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static com.yahoo.yolean.Exceptions.uncheck;
@@ -42,6 +42,14 @@ public class ProcessFactoryImpl implements ProcessFactory {
 
         ProcessBuilder processBuilder = new ProcessBuilder(arguments);
 
+        for (var entry : commandLine.getEnvironmentOverrides().entrySet()) {
+            if (entry.getValue() == null) {
+                processBuilder.environment().remove(entry.getKey());
+            } else {
+                processBuilder.environment().put(entry.getKey(), entry.getValue());
+            }
+        }
+
         if (commandLine.getRedirectStderrToStdoutInsteadOfDiscard()) {
             processBuilder.redirectErrorStream(true);
         } else {
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLineTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLineTest.java
index 6dd5087e8a0..9f601d53f73 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLineTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/task/util/process/CommandLineTest.java
@@ -161,4 +161,15 @@ public class CommandLineTest {
         terminal.verifyAllCommandsExecuted();
     }
 
+    @Test
+    public void testEnvironment() {
+        terminal.expectCommand("env k1=v1 -u k2 \"key 3=value 3\" programname 2>&1");
+        commandLine.add("programname")
+                .setEnvironmentVariable("key 3", "value 3")
+                .removeEnvironmentVariable("k2")
+                .setEnvironmentVariable("k1", "v1")
+                .execute();
+        terminal.verifyAllCommandsExecuted();
+    }
+
 }
\ No newline at end of file