From be9dd71dda17c4e0123ba2956c6ac1b9d088d907 Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Mon, 3 Jun 2024 13:19:34 +0800 Subject: [PATCH] well formatted for raw response Signed-off-by: Lantao Jin --- .../sql/legacy/SQLIntegTestCase.java | 6 +- .../org/opensearch/sql/sql/RawFormatIT.java | 19 +++++ .../sql/legacy/plugin/RestSQLQueryAction.java | 2 +- .../response/format/CsvResponseFormatter.java | 5 +- .../format/FlatResponseFormatter.java | 65 ++++++++++++++--- .../response/format/RawResponseFormatter.java | 14 +++- .../format/RawResponseFormatterTest.java | 71 +++++++++++++------ .../sql/sql/domain/SQLQueryRequest.java | 13 ++++ 8 files changed, 159 insertions(+), 36 deletions(-) diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java index 058182f123..38d44df24a 100644 --- a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java +++ b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java @@ -239,8 +239,12 @@ protected Request getSqlCursorCloseRequest(String cursorRequest) { } protected String executeQuery(String query, String requestType) { + return executeQuery(query, requestType, false); + } + + protected String executeQuery(String query, String requestType, boolean pretty) { try { - String endpoint = "/_plugins/_sql?format=" + requestType; + String endpoint = "/_plugins/_sql?format=" + requestType + "&pretty=" + pretty; String requestBody = makeRequest(query); Request sqlRequest = new Request("POST", endpoint); diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/RawFormatIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/RawFormatIT.java index 9d2861ce98..ba29672623 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/RawFormatIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/RawFormatIT.java @@ -41,6 +41,25 @@ public void rawFormatWithPipeFieldTest() { result); } + @Test + public void rawFormatPrettyWithPipeFieldTest() { + String result = + executeQuery( + String.format( + Locale.ROOT, "SELECT firstname, lastname FROM %s", TEST_INDEX_BANK_RAW_SANITIZE), + "raw", + true); + assertEquals( + StringUtils.format( + "firstname |lastname %n" + + "+Amber JOHnny|Duke Willmington+%n" + + "-Hattie |Bond- %n" + + "=Nanette |Bates= %n" + + "@Dale |Adams@ %n" + + "@Elinor |\"Ratliff|||\" %n"), + result); + } + @Test public void contentHeaderTest() throws IOException { String query = diff --git a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java index 309a7c9c2a..4440219f1b 100644 --- a/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java +++ b/legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java @@ -172,7 +172,7 @@ private ResponseListener createQueryResponseListener( } else if (format.equals(Format.CSV)) { formatter = new CsvResponseFormatter(request.sanitize()); } else if (format.equals(Format.RAW)) { - formatter = new RawResponseFormatter(); + formatter = new RawResponseFormatter(request.pretty()); } else { formatter = new JdbcResponseFormatter(PRETTY); } diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatter.java index a61b54b258..5abe5f92e2 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatter.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatter.java @@ -5,12 +5,13 @@ package org.opensearch.sql.protocol.response.format; +/** Response formatter to format response to csv format. */ public class CsvResponseFormatter extends FlatResponseFormatter { public CsvResponseFormatter() { - super(",", true); + super(",", true, false); } public CsvResponseFormatter(boolean sanitize) { - super(",", sanitize); + super(",", sanitize, false); } } diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java index 8c67d524b8..456c9fa510 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java @@ -12,6 +12,7 @@ import java.util.List; import java.util.Set; import java.util.stream.Collectors; +import java.util.stream.IntStream; import lombok.Getter; import lombok.RequiredArgsConstructor; import org.opensearch.sql.protocol.response.QueryResult; @@ -25,10 +26,12 @@ public abstract class FlatResponseFormatter implements ResponseFormatter headers; + private List> data; + private int[] maxWidths; public String getFlat() { + // Store headers and data in variables to avoid redundant calls + headers = getHeaders(response, sanitize); + data = getData(response, sanitize); + + calculateMaxWidths(); + List headersAndData = new ArrayList<>(); - headersAndData.add(getHeaderLine(response, sanitize)); - headersAndData.addAll(getDataLines(response, sanitize)); + headersAndData.add(getHeaderLine()); + headersAndData.addAll(getDataLines()); return String.join(INTERLINE_SEPARATOR, headersAndData); } - private String getHeaderLine(QueryResult response, boolean sanitize) { - List headers = getHeaders(response, sanitize); - return String.join(INLINE_SEPARATOR, headers); + private void calculateMaxWidths() { + int columns = headers.size(); + maxWidths = new int[columns]; + + for (int i = 0; i < columns; i++) { + int maxWidth = headers.get(i).length(); + for (List row : data) { + maxWidth = Math.max(maxWidth, row.get(i).length()); + } + maxWidths[i] = maxWidth; + } } - private List getDataLines(QueryResult response, boolean sanitize) { - List> data = getData(response, sanitize); - return data.stream().map(v -> String.join(INLINE_SEPARATOR, v)).collect(Collectors.toList()); + private List getDataLines() { + if (pretty) { + return data.stream().map(this::prettyFormatLine).collect(Collectors.toList()); + } else { + return data.stream() + .map(v -> String.join(INLINE_SEPARATOR, v)) + .collect(Collectors.toList()); + } + } + + private String getHeaderLine() { + if (pretty) { + return prettyFormatLine(headers); + } else { + return String.join(INLINE_SEPARATOR, headers); + } } private List getHeaders(QueryResult response, boolean sanitize) { @@ -145,5 +180,15 @@ private String quoteIfRequired(String separator, String cell) { private boolean isStartWithSensitiveChar(String cell) { return SENSITIVE_CHAR.stream().anyMatch(cell::startsWith); } + + private String prettyFormatLine(List line) { + return IntStream.range(0, line.size()) + .mapToObj(i -> padRight(line.get(i), maxWidths[i])) + .collect(Collectors.joining(INLINE_SEPARATOR)); + } + + private String padRight(String s, int n) { + return String.format("%-" + n + "s", s); + } } } diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/RawResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/RawResponseFormatter.java index 3b64be7062..a6924422b6 100644 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/RawResponseFormatter.java +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/RawResponseFormatter.java @@ -5,9 +5,19 @@ package org.opensearch.sql.protocol.response.format; -/** Response formatter to format response to csv or raw format. */ +/** Response formatter to format response to raw format. */ public class RawResponseFormatter extends FlatResponseFormatter { public RawResponseFormatter() { - super("|", false); + super("|", false, false); + } + + /** + * Create a raw response formatter with pretty parameter. + * + * @param pretty if true, display the columns with proper padding. Tracks the maximum width for + * each column to ensure proper formatting. + */ + public RawResponseFormatter(boolean pretty) { + super("|", false, pretty); } } diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java index 65111bd3b9..f837b8d274 100644 --- a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java +++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java @@ -25,7 +25,6 @@ /** Unit test for {@link FlatResponseFormatter}. */ public class RawResponseFormatterTest { - private FlatResponseFormatter rawFormatter = new RawResponseFormatter(); @Test void formatResponse() { @@ -40,8 +39,10 @@ void formatResponse() { Arrays.asList( tupleValue(ImmutableMap.of("name", "John", "age", 20)), tupleValue(ImmutableMap.of("name", "Smith", "age", 30)))); - String expected = "name|age%nJohn|20%nSmith|30"; - assertEquals(format(expected), rawFormatter.format(response)); + String expected = "name|age%n" + "John|20%n" + "Smith|30"; + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = "name |age%n" + "John |20 %n" + "Smith|30 "; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test @@ -67,8 +68,11 @@ void sanitizeHeaders() { "Seattle", "@age", 20)))); - String expected = "=firstname|+lastname|-city|@age%nJohn|Smith|Seattle|20"; - assertEquals(format(expected), rawFormatter.format(response)); + String expected = "=firstname|+lastname|-city|@age%n" + "John|Smith|Seattle|20"; + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = + "=firstname|+lastname|-city |@age%n" + "John |Smith |Seattle|20 "; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test @@ -94,7 +98,16 @@ void sanitizeData() { + "-Seattle%n" + "@Seattle%n" + "Seattle="; - assertEquals(format(expected), rawFormatter.format(response)); + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = + "city %n" + + "Seattle %n" + + "=Seattle%n" + + "+Seattle%n" + + "-Seattle%n" + + "@Seattle%n" + + "Seattle="; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test @@ -108,8 +121,10 @@ void quoteIfRequired() { new QueryResult( schema, Arrays.asList(tupleValue(ImmutableMap.of("na|me", "John|Smith", "||age", "30|||")))); - String expected = "\"na|me\"|\"||age\"%n\"John|Smith\"|\"30|||\""; - assertEquals(format(expected), rawFormatter.format(response)); + String expected = "\"na|me\"|\"||age\"%n" + "\"John|Smith\"|\"30|||\""; + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = "\"na|me\" |\"||age\"\n" + "\"John|Smith\"|\"30|||\""; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test @@ -117,12 +132,12 @@ void formatError() { Throwable t = new RuntimeException("This is an exception"); String expected = "{\n \"type\": \"RuntimeException\",\n \"reason\": \"This is an exception\"\n}"; - assertEquals(expected, rawFormatter.format(t)); + assertEquals(expected, getRawFormatter().format(t)); + assertEquals(expected, getRawFormatterPretty().format(t)); } @Test void escapeSanitize() { - FlatResponseFormatter escapeFormatter = new RawResponseFormatter(); ExecutionEngine.Schema schema = new ExecutionEngine.Schema( ImmutableList.of(new ExecutionEngine.Schema.Column("city", "city", STRING))); @@ -132,8 +147,10 @@ void escapeSanitize() { Arrays.asList( tupleValue(ImmutableMap.of("city", "=Seattle")), tupleValue(ImmutableMap.of("city", "||Seattle")))); - String expected = "city%n=Seattle%n\"||Seattle\""; - assertEquals(format(expected), escapeFormatter.format(response)); + String expected = "city%n" + "=Seattle%n" + "\"||Seattle\""; + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = "city %n" + "=Seattle %n" + "\"||Seattle\""; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test @@ -147,13 +164,14 @@ void senstiveCharater() { Arrays.asList( tupleValue(ImmutableMap.of("city", "@Seattle")), tupleValue(ImmutableMap.of("city", "++Seattle")))); - String expected = "city%n@Seattle%n++Seattle"; - assertEquals(format(expected), rawFormatter.format(response)); + String expected = "city%n" + "@Seattle%n" + "++Seattle"; + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = "city %n" + "@Seattle %n" + "++Seattle"; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test void senstiveCharaterWithSanitize() { - FlatResponseFormatter testFormater = new RawResponseFormatter(); ExecutionEngine.Schema schema = new ExecutionEngine.Schema( ImmutableList.of(new ExecutionEngine.Schema.Column("city", "city", STRING))); @@ -163,8 +181,10 @@ void senstiveCharaterWithSanitize() { Arrays.asList( tupleValue(ImmutableMap.of("city", "@Seattle")), tupleValue(ImmutableMap.of("city", "++Seattle|||")))); - String expected = "city%n@Seattle%n\"++Seattle|||\""; - assertEquals(format(expected), testFormater.format(response)); + String expected = "city%n" + "@Seattle%n" + "\"++Seattle|||\""; + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = "city %n" + "@Seattle %n" + "\"++Seattle|||\""; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test @@ -183,12 +203,23 @@ void replaceNullValues() { ImmutableMap.of("firstname", LITERAL_NULL, "city", stringValue("Seattle"))), ExprTupleValue.fromExprValueMap( ImmutableMap.of("firstname", stringValue("John"), "city", LITERAL_MISSING)))); - String expected = "name|city%nJohn|Seattle%n|Seattle%nJohn|"; - assertEquals(format(expected), rawFormatter.format(response)); + String expected = "name|city%n" + "John|Seattle%n" + "|Seattle%n" + "John|"; + assertEquals(format(expected), getRawFormatter().format(response)); + String expectedPretty = "name|city %n" + "John|Seattle%n" + " |Seattle%n" + "John| "; + assertEquals(format(expectedPretty), getRawFormatterPretty().format(response)); } @Test void testContentType() { - assertEquals(rawFormatter.contentType(), CONTENT_TYPE); + assertEquals(getRawFormatter().contentType(), CONTENT_TYPE); + assertEquals(getRawFormatterPretty().contentType(), CONTENT_TYPE); + } + + private FlatResponseFormatter getRawFormatter() { + return new RawResponseFormatter(); + } + + private FlatResponseFormatter getRawFormatterPretty() { + return new RawResponseFormatter(true); } } diff --git a/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java b/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java index 4e902cb67d..79ad9bf317 100644 --- a/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java +++ b/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java @@ -29,6 +29,7 @@ public class SQLQueryRequest { Set.of("query", "fetch_size", "parameters", QUERY_FIELD_CURSOR); private static final String QUERY_PARAMS_FORMAT = "format"; private static final String QUERY_PARAMS_SANITIZE = "sanitize"; + private static final String QUERY_PARAMS_PRETTY = "pretty"; /** JSON payload in REST request. */ private final JSONObject jsonContent; @@ -49,6 +50,10 @@ public class SQLQueryRequest { @Accessors(fluent = true) private boolean sanitize = true; + @Getter + @Accessors(fluent = true) + private boolean pretty = false; + private String cursor; /** Constructor of SQLQueryRequest that passes request params. */ @@ -64,6 +69,7 @@ public SQLQueryRequest( this.params = params; this.format = getFormat(params); this.sanitize = shouldSanitize(params); + this.pretty = shouldPretty(params); this.cursor = cursor; } @@ -148,4 +154,11 @@ private boolean shouldSanitize(Map params) { } return true; } + + private boolean shouldPretty(Map params) { + if (params.containsKey(QUERY_PARAMS_PRETTY)) { + return Boolean.parseBoolean(params.get(QUERY_PARAMS_PRETTY)); + } + return false; + } }