From fdee707bb0c6673ac12e2a6b8d172278bfc0c582 Mon Sep 17 00:00:00 2001 From: Zac Blanco Date: Tue, 12 Mar 2024 22:08:40 -0700 Subject: [PATCH] Add JSON output format to CLI Cherry-pick of https://github.com/trinodb/trino/commit/daf5ce670ee70d14b8534c3e41a09c4a365b6fd8 Co-authored-by: Yuya Ebihara --- presto-cli/pom.xml | 5 ++ .../facebook/presto/cli/ClientOptions.java | 3 +- .../com/facebook/presto/cli/JsonPrinter.java | 71 ++++++++++++++++ .../java/com/facebook/presto/cli/Query.java | 2 + .../facebook/presto/cli/TestJsonPrinter.java | 82 +++++++++++++++++++ .../facebook/presto/cli/TestTsvPrinter.java | 2 +- 6 files changed, 163 insertions(+), 2 deletions(-) create mode 100644 presto-cli/src/main/java/com/facebook/presto/cli/JsonPrinter.java create mode 100644 presto-cli/src/test/java/com/facebook/presto/cli/TestJsonPrinter.java diff --git a/presto-cli/pom.xml b/presto-cli/pom.xml index 3693846212356..6c019fd931bcd 100644 --- a/presto-cli/pom.xml +++ b/presto-cli/pom.xml @@ -103,6 +103,11 @@ okhttp + + com.fasterxml.jackson.core + jackson-core + + org.testng diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java b/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java index 30ac745f3ea2e..c78548877f720 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java @@ -116,7 +116,7 @@ public class ClientOptions @Option(name = "--execute", title = "execute", description = "Execute specified statements and exit") public String execute; - @Option(name = "--output-format", title = "output-format", description = "Output format for batch mode [ALIGNED, VERTICAL, CSV, TSV, CSV_HEADER, TSV_HEADER, NULL] (default: CSV)") + @Option(name = "--output-format", title = "output-format", description = "Output format for batch mode [ALIGNED, VERTICAL, JSON, CSV, TSV, CSV_HEADER, TSV_HEADER, NULL] (default: CSV)") public OutputFormat outputFormat = OutputFormat.CSV; @Option(name = "--resource-estimate", title = "resource-estimate", description = "Resource estimate (property can be used multiple times; format is key=value)") @@ -157,6 +157,7 @@ public enum OutputFormat TSV, CSV_HEADER, TSV_HEADER, + JSON, NULL } diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/JsonPrinter.java b/presto-cli/src/main/java/com/facebook/presto/cli/JsonPrinter.java new file mode 100644 index 0000000000000..d055399fe9b19 --- /dev/null +++ b/presto-cli/src/main/java/com/facebook/presto/cli/JsonPrinter.java @@ -0,0 +1,71 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.cli; + +import com.fasterxml.jackson.core.JsonFactory; +import com.fasterxml.jackson.core.JsonGenerator; +import com.google.common.collect.ImmutableList; + +import java.io.IOException; +import java.io.Writer; +import java.util.List; + +import static com.facebook.presto.cli.AlignedTablePrinter.formatHexDump; +import static java.util.Objects.requireNonNull; + +public class JsonPrinter + implements OutputPrinter +{ + private final List fieldNames; + private final Writer writer; + + public JsonPrinter(List fieldNames, Writer writer) + { + this.fieldNames = ImmutableList.copyOf(requireNonNull(fieldNames, "fieldNames is null")); + this.writer = requireNonNull(writer, "writer is null"); + } + + @Override + public void printRows(List> rows, boolean complete) + throws IOException + { + JsonFactory jsonFactory = new JsonFactory(); + for (List row : rows) { + JsonGenerator jsonGenerator = jsonFactory.createGenerator(writer); + jsonGenerator.writeStartObject(); + for (int position = 0; position < row.size(); position++) { + String columnName = fieldNames.get(position); + jsonGenerator.writeObjectField(columnName, formatValue(row.get(position))); + } + jsonGenerator.writeEndObject(); + jsonGenerator.writeRaw('\n'); + jsonGenerator.flush(); + } + } + + @Override + public void finish() + throws IOException + { + writer.flush(); + } + + private static Object formatValue(Object o) + { + if (o instanceof byte[]) { + return formatHexDump((byte[]) o); + } + return o; + } +} diff --git a/presto-cli/src/main/java/com/facebook/presto/cli/Query.java b/presto-cli/src/main/java/com/facebook/presto/cli/Query.java index 83411e2a3ae5e..c68e62f21e6eb 100644 --- a/presto-cli/src/main/java/com/facebook/presto/cli/Query.java +++ b/presto-cli/src/main/java/com/facebook/presto/cli/Query.java @@ -334,6 +334,8 @@ private static OutputPrinter createOutputPrinter(OutputFormat format, Writer wri return new TsvPrinter(fieldNames, writer, false); case TSV_HEADER: return new TsvPrinter(fieldNames, writer, true); + case JSON: + return new JsonPrinter(fieldNames, writer); case NULL: return new NullPrinter(); } diff --git a/presto-cli/src/test/java/com/facebook/presto/cli/TestJsonPrinter.java b/presto-cli/src/test/java/com/facebook/presto/cli/TestJsonPrinter.java new file mode 100644 index 0000000000000..dddde5d00795c --- /dev/null +++ b/presto-cli/src/test/java/com/facebook/presto/cli/TestJsonPrinter.java @@ -0,0 +1,82 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.cli; + +import com.google.common.collect.ImmutableList; +import org.testng.annotations.Test; + +import java.io.IOException; +import java.io.StringWriter; +import java.util.List; + +import static com.facebook.presto.cli.TestAlignedTablePrinter.row; +import static com.facebook.presto.cli.TestAlignedTablePrinter.rows; +import static org.testng.Assert.assertEquals; + +public class TestJsonPrinter +{ + @Test + public void testJsonPrinting() + throws Exception + { + StringWriter writer = new StringWriter(); + List fieldNames = ImmutableList.of("first", "last", "quantity"); + OutputPrinter printer = new JsonPrinter(fieldNames, writer); + + printer.printRows(rows( + row("hello", "world", 123), + row("a", null, 4.5), + row("some long\ntext\tdone", "more\ntext", 4567), + row("bye", "done", -15)), + true); + printer.finish(); + + String expected = "" + + "{\"first\":\"hello\",\"last\":\"world\",\"quantity\":123}\n" + + "{\"first\":\"a\",\"last\":null,\"quantity\":4.5}\n" + + "{\"first\":\"some long\\ntext\\tdone\",\"last\":\"more\\ntext\",\"quantity\":4567}\n" + + "{\"first\":\"bye\",\"last\":\"done\",\"quantity\":-15}\n"; + + assertEquals(writer.getBuffer().toString(), expected); + } + + @Test + public void testJsonPrintingNoRows() + throws Exception + { + StringWriter writer = new StringWriter(); + List fieldNames = ImmutableList.of("first", "last"); + OutputPrinter printer = new JsonPrinter(fieldNames, writer); + + printer.finish(); + + assertEquals(writer.getBuffer().toString(), ""); + } + + @Test + public void testJsonVarbinaryPrinting() + throws IOException + { + StringWriter writer = new StringWriter(); + List fieldNames = ImmutableList.of("first", "last", "quantity"); + OutputPrinter printer = new JsonPrinter(fieldNames, writer); + + printer.printRows(rows(row("hello".getBytes(), null, 123)), true); + printer.finish(); + + String expected = "{\"first\":\"68 65 6c 6c 6f\",\"last\":null,\"quantity\":123}\n"; + + assertEquals(writer.getBuffer().toString(), expected); + } +} diff --git a/presto-cli/src/test/java/com/facebook/presto/cli/TestTsvPrinter.java b/presto-cli/src/test/java/com/facebook/presto/cli/TestTsvPrinter.java index 93475e3e5e477..c4a375b6bc603 100644 --- a/presto-cli/src/test/java/com/facebook/presto/cli/TestTsvPrinter.java +++ b/presto-cli/src/test/java/com/facebook/presto/cli/TestTsvPrinter.java @@ -89,7 +89,7 @@ public void testTsvPrintingNoHeader() } @Test - public void testCsvVarbinaryPrinting() + public void testTsvVarbinaryPrinting() throws IOException { StringWriter writer = new StringWriter();