diff options
author | Håkon Hallingstad <hakon@verizonmedia.com> | 2021-02-16 13:31:32 +0100 |
---|---|---|
committer | Håkon Hallingstad <hakon@verizonmedia.com> | 2021-02-16 13:31:32 +0100 |
commit | b89171616df65727ff8dfc5d7e49b4f454be144e (patch) | |
tree | 643876837fce285556f1adbddd81e1a679d27b2d /node-admin | |
parent | 605a548159baba46523711ede9fc76924b1dd5df (diff) |
CommandLine environment
Diffstat (limited to 'node-admin')
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 |