diff --git a/java/google/registry/tools/ShellCommand.java b/java/google/registry/tools/ShellCommand.java index 0726d2a3d..f7ae7edd9 100644 --- a/java/google/registry/tools/ShellCommand.java +++ b/java/google/registry/tools/ShellCommand.java @@ -17,6 +17,7 @@ package google.registry.tools; import static com.google.common.base.StandardSystemProperty.USER_HOME; import static com.google.common.collect.ImmutableList.toImmutableList; import static java.nio.charset.StandardCharsets.US_ASCII; +import static java.nio.charset.StandardCharsets.UTF_8; import com.beust.jcommander.JCommander; import com.beust.jcommander.Parameter; @@ -27,13 +28,18 @@ 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.SourceCodeEscapers; import google.registry.util.Clock; import google.registry.util.SystemClock; import java.io.BufferedReader; +import java.io.ByteArrayOutputStream; import java.io.File; +import java.io.FilterOutputStream; import java.io.IOException; import java.io.InputStream; import java.io.InputStreamReader; +import java.io.OutputStream; +import java.io.PrintStream; import java.io.StreamTokenizer; import java.io.StringReader; import java.util.Arrays; @@ -64,6 +70,9 @@ public class ShellCommand implements Command { private static final String NON_ALERT_COLOR = "\u001b[32m"; // green foreground 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 final CommandRunner runner; private final BufferedReader lineReader; @@ -77,6 +86,16 @@ public class ShellCommand implements Command { + "Will instead warn you and require re-running the command.") boolean dontExitOnIdle = false; + @Parameter( + names = {"--encapsulate_output"}, + description = + "Encapsulate command standard output and error by combining the two streams to standard " + + "output and inserting a prefix ('out:' or 'err:') at the beginning of every line " + + "of normal output and adding a line consisting of either 'SUCCESS' or " + + "'FAILURE ' at the end of the output for a " + + "command, allowing the output to be easily parsed by wrapper scripts.") + boolean encapsulateOutput = false; + public ShellCommand(CommandRunner runner) throws IOException { this.runner = runner; InputStream in = System.in; @@ -136,6 +155,23 @@ public class ShellCommand implements Command { String line; DateTime lastTime = clock.nowUtc(); while ((line = getLine()) != null) { + PrintStream orgStdout = null; + PrintStream orgStderr = null; + EncapsulatingOutputStream encapsulatedOutputStream = null; + EncapsulatingOutputStream encapsulatedErrorStream = null; + + + // 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) { + 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)); + } + // Make sure we're not idle for too long. Only relevant when we're "extra careful" if (!dontExitOnIdle && beExtraCareful @@ -149,13 +185,32 @@ public class ShellCommand implements Command { if (lineArgs.length == 0) { continue; } + Exception lastError = null; try { runner.run(lineArgs); } catch (Exception e) { + lastError = e; System.err.println("Got an exception:\n" + e); } + try { + if (encapsulatedOutputStream != null) { + encapsulatedOutputStream.dumpLastLine(); + encapsulatedErrorStream.dumpLastLine(); + System.setOut(orgStdout); + System.setErr(orgStderr); + if (lastError == null) { + emitSuccess(); + } else { + emitFailure(lastError); + } + } + } catch (IOException e) { + throw new RuntimeException(e); + } + } + if (!encapsulateOutput) { + System.err.println(); } - System.err.println(); } private String getLine() { @@ -189,6 +244,26 @@ public class ShellCommand implements Command { return resultBuilder.build().toArray(new String[0]); } + /** + * Emit a success command separator. + * + *

Dumps the last line of output prior to doing this. + */ + private void emitSuccess() throws IOException { + System.out.println(SUCCESS); + System.out.flush(); + } + + /** + * Emit a failure message obtained from the throwable. + * + *

Dumps the last line of output prior to doing this. + */ + private void emitFailure(Throwable e) throws IOException { + System.out.println(FAILURE + e.getClass().getName() + " " + + SourceCodeEscapers.javaCharEscaper().escape(e.getMessage())); + } + @VisibleForTesting static class JCommanderCompletor implements Completor { @@ -448,4 +523,51 @@ public class ShellCommand implements Command { .orElse(DEFAULT_PARAM_DOC); } } + + /** + * Encapsulate output according to the protocol described in the documentation for the + * --encapsulate_output flag. + */ + @VisibleForTesting + static class EncapsulatingOutputStream extends FilterOutputStream { + + private final byte[] prefix; + private final ByteArrayOutputStream lastLine = new ByteArrayOutputStream(); + + // Flag to keep track of whether the last character written was a newline. We initialize this + // to "true" because we always want the first line of output to be escaped with a leading space. + boolean lastWasNewline = true; + + EncapsulatingOutputStream(OutputStream out, String identifier) { + super(out); + this.prefix = identifier.getBytes(UTF_8); + } + + @Override + public void write(int b) throws IOException { + lastLine.write(b); + if (b == '\n') { + out.write(prefix); + lastLine.writeTo(out); + out.flush(); + lastLine.reset(); + } + } + + @Override + public void flush() throws IOException { + dumpLastLine(); + } + + /** Dump the accumulated last line of output, if there was one. */ + public void dumpLastLine() throws IOException { + if (lastLine.size() > 0) { + out.write(prefix); + lastLine.writeTo(out); + out.write('\n'); + out.flush(); + lastLine.reset(); + } + } + } } diff --git a/javatests/google/registry/tools/ShellCommandTest.java b/javatests/google/registry/tools/ShellCommandTest.java index d9e3a826b..be9520912 100644 --- a/javatests/google/registry/tools/ShellCommandTest.java +++ b/javatests/google/registry/tools/ShellCommandTest.java @@ -16,6 +16,7 @@ package google.registry.tools; import static com.google.common.truth.Truth.assertThat; import static google.registry.testing.JUnitBackports.assertThrows; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -28,12 +29,17 @@ import com.google.common.collect.ImmutableMap; import google.registry.testing.FakeClock; import google.registry.tools.ShellCommand.JCommanderCompletor; import java.io.BufferedReader; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.IOException; +import java.io.PrintStream; import java.util.ArrayDeque; import java.util.ArrayList; import java.util.List; import org.joda.time.DateTime; import org.joda.time.Duration; +import org.junit.After; +import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -44,8 +50,23 @@ public class ShellCommandTest { CommandRunner cli = mock(CommandRunner.class); FakeClock clock = new FakeClock(DateTime.parse("2000-01-01TZ")); + PrintStream orgStdout; + PrintStream orgStderr; + public ShellCommandTest() {} + @Before + public void setUp() { + orgStdout = System.out; + orgStderr = System.err; + } + + @After + public void tearDown() { + System.setOut(orgStdout); + System.setErr(orgStderr); + } + @Test public void testParsing() { assertThat(ShellCommand.parseCommand("foo bar 123 baz+ // comment \"string data\"")) @@ -226,6 +247,55 @@ public class ShellCommandTest { "testCommand --xorg ", 0, "", "Flag documentation: test organization\n (PRIVATE, PUBLIC)"); } + @Test + public void testEncapsulatedOutputStream_basicFuncionality() throws Exception { + 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(); + assertThat(backing.toString()) + .isEqualTo("out: first line\nout: second line\nout: trailing data\n"); + } + + @Test + public void testEncapsulatedOutputStream_emptyStream() throws Exception { + ByteArrayOutputStream backing = new ByteArrayOutputStream(); + PrintStream out = new PrintStream(new ShellCommand.EncapsulatingOutputStream(backing, "out: ")); + out.flush(); + assertThat(backing.toString()).isEqualTo(""); + } + + @Test + public void testEncapsulatedOutput_command() throws Exception { + RegistryToolEnvironment.ALPHA.setup(); + + // capture output (have to do this before the shell command is created) + ByteArrayOutputStream stdout = new ByteArrayOutputStream(); + ByteArrayOutputStream stderr = new ByteArrayOutputStream(); + System.setOut(new PrintStream(stdout)); + System.setErr(new PrintStream(stderr)); + System.setIn(new ByteArrayInputStream("command1\n".getBytes(UTF_8))); + + ShellCommand shellCommand = + new ShellCommand( + args -> { + System.out.println("first line"); + System.err.println("second line"); + System.out.print("fragmented "); + System.err.println("surprise!"); + System.out.println("line"); + }); + shellCommand.encapsulateOutput = true; + + shellCommand.run(); + assertThat(stderr.toString()).isEmpty(); + assertThat(stdout.toString()) + .isEqualTo( + "out: first line\nerr: second line\nerr: surprise!\nout: fragmented line\n" + + "SUCCESS\n"); + } + @Parameters(commandDescription = "Test command") static class TestCommand implements Command { enum OrgType {