mirror of
https://github.com/google/nomulus.git
synced 2025-05-13 16:07:15 +02:00
Make OutputEncapsulator a CommandRunner
This is in preparation for having other "command changing things" like redirecting to file and maybe variable substitutions in the arguments. "On the way" added a RUNNING "some_command" "--some_flag" "some_value" to the output encapsulator so that if we run multiple commands, we know what command was called where. ------------- Created by MOE: https://github.com/google/moe MOE_MIGRATED_REVID=219837452
This commit is contained in:
guyben
jianglai
parent
f59005ad35
commit
9ce07db38a
2 changed files with 107 additions and 70 deletions
|
|
@ -28,6 +28,8 @@ import com.google.common.annotations.VisibleForTesting;
|
|||
import com.google.common.base.Ascii;
|
||||
import com.google.common.collect.ImmutableList;
|
||||
import com.google.common.collect.ImmutableTable;
|
||||
import com.google.common.escape.Escaper;
|
||||
import com.google.common.escape.Escapers;
|
||||
import google.registry.util.Clock;
|
||||
import google.registry.util.SystemClock;
|
||||
import java.io.BufferedReader;
|
||||
|
|
@ -70,10 +72,25 @@ public class ShellCommand implements Command {
|
|||
private static final String ALERT_COLOR = "\u001b[1;41;97m"; // red background
|
||||
private static final Duration IDLE_THRESHOLD = Duration.standardHours(1);
|
||||
private static final String SUCCESS = "SUCCESS";
|
||||
private static final String FAILURE = "FAILURE ";
|
||||
private static final String FAILURE = "FAILURE";
|
||||
private static final String RUNNING = "RUNNING";
|
||||
private static final Escaper STRING_ESCAPER =
|
||||
Escapers.builder()
|
||||
.addEscape('\\', "\\\\")
|
||||
.addEscape('"', "\\\"")
|
||||
.addEscape('\n', "\\n")
|
||||
.addEscape('\r', "\\r")
|
||||
.addEscape('\t', "\\t")
|
||||
.build();
|
||||
|
||||
/**
|
||||
* The runner we received in the constructor.
|
||||
*
|
||||
* <p>We might want to update this runner based on flags (e.g. --encapsulate_output), but these
|
||||
* flags aren't available in the constructor so we have to do it in the {@link #run} function.
|
||||
*/
|
||||
private final CommandRunner originalRunner;
|
||||
|
||||
private final CommandRunner runner;
|
||||
private final BufferedReader lineReader;
|
||||
private final ConsoleReader consoleReader;
|
||||
private final Clock clock;
|
||||
|
|
@ -96,7 +113,7 @@ public class ShellCommand implements Command {
|
|||
boolean encapsulateOutput = false;
|
||||
|
||||
public ShellCommand(CommandRunner runner) throws IOException {
|
||||
this.runner = runner;
|
||||
this.originalRunner = runner;
|
||||
InputStream in = System.in;
|
||||
if (System.console() != null) {
|
||||
consoleReader = new ConsoleReader();
|
||||
|
|
@ -114,7 +131,7 @@ public class ShellCommand implements Command {
|
|||
|
||||
@VisibleForTesting
|
||||
ShellCommand(BufferedReader bufferedReader, Clock clock, CommandRunner runner) {
|
||||
this.runner = runner;
|
||||
this.originalRunner = runner;
|
||||
this.lineReader = bufferedReader;
|
||||
this.clock = clock;
|
||||
this.consoleReader = null;
|
||||
|
|
@ -145,42 +162,11 @@ public class ShellCommand implements Command {
|
|||
return this;
|
||||
}
|
||||
|
||||
private static class OutputEncapsulator {
|
||||
private PrintStream orgStdout;
|
||||
private PrintStream orgStderr;
|
||||
private static class OutputEncapsulator implements CommandRunner {
|
||||
private final CommandRunner runner;
|
||||
|
||||
private EncapsulatingOutputStream encapsulatedOutputStream = null;
|
||||
private EncapsulatingOutputStream encapsulatedErrorStream = null;
|
||||
|
||||
private Exception error;
|
||||
|
||||
private OutputEncapsulator() {
|
||||
orgStdout = System.out;
|
||||
orgStderr = System.err;
|
||||
encapsulatedOutputStream = new EncapsulatingOutputStream(System.out, "out: ");
|
||||
encapsulatedErrorStream = new EncapsulatingOutputStream(System.out, "err: ");
|
||||
System.setOut(new PrintStream(encapsulatedOutputStream));
|
||||
System.setErr(new PrintStream(encapsulatedErrorStream));
|
||||
}
|
||||
|
||||
void setError(Exception e) {
|
||||
error = e;
|
||||
}
|
||||
|
||||
private void restoreOriginalStreams() {
|
||||
try {
|
||||
encapsulatedOutputStream.dumpLastLine();
|
||||
encapsulatedErrorStream.dumpLastLine();
|
||||
System.setOut(orgStdout);
|
||||
System.setErr(orgStderr);
|
||||
if (error != null) {
|
||||
emitFailure(error);
|
||||
} else {
|
||||
emitSuccess();
|
||||
}
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException(e);
|
||||
}
|
||||
private OutputEncapsulator(CommandRunner runner) {
|
||||
this.runner = runner;
|
||||
}
|
||||
|
||||
/**
|
||||
|
|
@ -188,7 +174,7 @@ public class ShellCommand implements Command {
|
|||
*
|
||||
* <p>Dumps the last line of output prior to doing this.
|
||||
*/
|
||||
private void emitSuccess() {
|
||||
private static void emitSuccess() {
|
||||
System.out.println(SUCCESS);
|
||||
System.out.flush();
|
||||
}
|
||||
|
|
@ -198,23 +184,45 @@ public class ShellCommand implements Command {
|
|||
*
|
||||
* <p>Dumps the last line of output prior to doing this.
|
||||
*/
|
||||
private void emitFailure(Throwable e) {
|
||||
System.out.println(
|
||||
FAILURE
|
||||
+ e.getClass().getName()
|
||||
+ " "
|
||||
+ e.getMessage().replace("\\", "\\\\").replace("\n", "\\n"));
|
||||
private static void emitFailure(Throwable e) {
|
||||
System.out.format(
|
||||
"%s %s %s\n", FAILURE, e.getClass().getName(), STRING_ESCAPER.escape(e.getMessage()));
|
||||
System.out.flush();
|
||||
}
|
||||
|
||||
private static void emitArguments(String[] args) {
|
||||
System.out.print(RUNNING);
|
||||
Arrays.stream(args).forEach(arg -> System.out.format(" \"%s\"", STRING_ESCAPER.escape(arg)));
|
||||
System.out.println();
|
||||
System.out.flush();
|
||||
}
|
||||
|
||||
private void encapsulatedRun(String[] args) throws Exception {
|
||||
PrintStream orgOut = System.out;
|
||||
PrintStream orgErr = System.err;
|
||||
try (PrintStream newOut =
|
||||
new PrintStream(new EncapsulatingOutputStream(System.out, "out: "));
|
||||
PrintStream newErr =
|
||||
new PrintStream(new EncapsulatingOutputStream(System.out, "err: "))) {
|
||||
System.setOut(newOut);
|
||||
System.setErr(newErr);
|
||||
runner.run(args);
|
||||
} finally {
|
||||
System.setOut(orgOut);
|
||||
System.setErr(orgErr);
|
||||
}
|
||||
}
|
||||
|
||||
/** Run "func" with output encapsulation. */
|
||||
static void run(CommandRunner runner, String[] args) {
|
||||
OutputEncapsulator encapsulator = new OutputEncapsulator();
|
||||
@Override
|
||||
public void run(String[] args) {
|
||||
|
||||
try {
|
||||
runner.run(args);
|
||||
emitArguments(args);
|
||||
encapsulatedRun(args);
|
||||
emitSuccess();
|
||||
} catch (Exception e) {
|
||||
encapsulator.setError(e);
|
||||
} finally {
|
||||
encapsulator.restoreOriginalStreams();
|
||||
emitFailure(e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -222,6 +230,10 @@ public class ShellCommand implements Command {
|
|||
/** Run the shell until the user presses "Ctrl-D". */
|
||||
@Override
|
||||
public void run() {
|
||||
// Wrap standard output and error if requested. We have to do so here in run because the flags
|
||||
// haven't been processed in the constructor.
|
||||
CommandRunner runner =
|
||||
encapsulateOutput ? new OutputEncapsulator(originalRunner) : originalRunner;
|
||||
// On Production we want to be extra careful - to prevent accidental use.
|
||||
boolean beExtraCareful = (RegistryToolEnvironment.get() == RegistryToolEnvironment.PRODUCTION);
|
||||
setPrompt(RegistryToolEnvironment.get(), beExtraCareful);
|
||||
|
|
@ -242,16 +254,11 @@ public class ShellCommand implements Command {
|
|||
continue;
|
||||
}
|
||||
|
||||
// Wrap standard output and error if requested. We have to do so here in run because the flags
|
||||
// haven't been processed in the constructor.
|
||||
if (encapsulateOutput) {
|
||||
OutputEncapsulator.run(runner, lineArgs);
|
||||
} else {
|
||||
try {
|
||||
runner.run(lineArgs);
|
||||
} catch (Exception e) {
|
||||
System.err.println("Got an exception:\n" + e);
|
||||
}
|
||||
try {
|
||||
runner.run(lineArgs);
|
||||
} catch (Exception e) {
|
||||
System.err.println("Got an exception:\n" + e);
|
||||
e.printStackTrace();
|
||||
}
|
||||
}
|
||||
if (!encapsulateOutput) {
|
||||
|
|
@ -566,6 +573,14 @@ public class ShellCommand implements Command {
|
|||
@Override
|
||||
public void flush() throws IOException {
|
||||
dumpLastLine();
|
||||
super.flush();
|
||||
}
|
||||
|
||||
@Override
|
||||
public void close() throws IOException {
|
||||
dumpLastLine();
|
||||
// We do NOT want to call super.close as that would close the original outputStream
|
||||
// (System.out)
|
||||
}
|
||||
|
||||
/** Dump the accumulated last line of output, if there was one. */
|
||||
|
|
|
|||
|
|
@ -253,10 +253,11 @@ public class ShellCommandTest {
|
|||
@Test
|
||||
public void testEncapsulatedOutputStream_basicFuncionality() {
|
||||
ByteArrayOutputStream backing = new ByteArrayOutputStream();
|
||||
PrintStream out = new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "));
|
||||
out.println("first line");
|
||||
out.print("second line\ntrailing data");
|
||||
out.flush();
|
||||
try (PrintStream out =
|
||||
new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "))) {
|
||||
out.println("first line");
|
||||
out.print("second line\ntrailing data");
|
||||
}
|
||||
assertThat(backing.toString())
|
||||
.isEqualTo("out: first line\nout: second line\nout: trailing data\n");
|
||||
}
|
||||
|
|
@ -264,8 +265,8 @@ public class ShellCommandTest {
|
|||
@Test
|
||||
public void testEncapsulatedOutputStream_emptyStream() {
|
||||
ByteArrayOutputStream backing = new ByteArrayOutputStream();
|
||||
PrintStream out = new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "));
|
||||
out.flush();
|
||||
try (PrintStream out =
|
||||
new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: "))) {}
|
||||
assertThat(backing.toString()).isEqualTo("");
|
||||
}
|
||||
|
||||
|
|
@ -288,10 +289,31 @@ public class ShellCommandTest {
|
|||
assertThat(stderr.toString()).isEmpty();
|
||||
assertThat(stdout.toString())
|
||||
.isEqualTo(
|
||||
"out: first line\nerr: second line\nerr: surprise!\nout: fragmented line\n"
|
||||
"RUNNING \"command1\"\n"
|
||||
+ "out: first line\nerr: second line\nerr: surprise!\nout: fragmented line\n"
|
||||
+ "SUCCESS\n");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testEncapsulatedOutput_throws() throws Exception {
|
||||
RegistryToolEnvironment.ALPHA.setup();
|
||||
captureOutput();
|
||||
ShellCommand shellCommand =
|
||||
new ShellCommand(
|
||||
args -> {
|
||||
System.out.println("first line");
|
||||
throw new Exception("some error!");
|
||||
});
|
||||
shellCommand.encapsulateOutput = true;
|
||||
shellCommand.run();
|
||||
assertThat(stderr.toString()).isEmpty();
|
||||
assertThat(stdout.toString())
|
||||
.isEqualTo(
|
||||
"RUNNING \"command1\"\n"
|
||||
+ "out: first line\n"
|
||||
+ "FAILURE java.lang.Exception some error!\n");
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testEncapsulatedOutput_noCommand() throws Exception {
|
||||
captureOutput();
|
||||
|
|
@ -307,7 +329,7 @@ public class ShellCommandTest {
|
|||
shellCommand.run();
|
||||
assertThat(stderr.toString()).isEmpty();
|
||||
assertThat(stdout.toString())
|
||||
.isEqualTo("out: first line\nSUCCESS\n");
|
||||
.isEqualTo("RUNNING \"do\" \"something\"\nout: first line\nSUCCESS\n");
|
||||
}
|
||||
|
||||
void captureOutput() {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue