From aec482577143a327a6dc94e073a0a6cec709a7ed Mon Sep 17 00:00:00 2001 From: Lantao Jin Date: Tue, 9 Jul 2024 23:19:44 +0800 Subject: [PATCH] Well format the raw response when query parameter "pretty" enabled (#2727) Signed-off-by: Lantao Jin --- .../sql/legacy/SQLIntegTestCase.java | 5 + .../org/opensearch/sql/sql/CsvFormatIT.java | 2 +- .../org/opensearch/sql/sql/RawFormatIT.java | 22 ++- .../sql/legacy/plugin/RestSQLQueryAction.java | 2 +- .../response/format/CsvResponseFormatter.java | 39 ++++- .../response/format/FlatResponseBase.java | 89 +++++++++++ .../format/FlatResponseFormatter.java | 149 ------------------ .../format/FlatResponseWithPrettifier.java | 53 +++++++ .../format/FlatResponseWithSanitizer.java | 53 +++++++ .../response/format/RawResponseFormatter.java | 48 +++++- .../format/CsvResponseFormatterTest.java | 2 +- .../format/RawResponseFormatterTest.java | 75 ++++++--- .../sql/sql/domain/SQLQueryRequest.java | 13 ++ .../opensearch/sql/sql/SQLServiceTest.java | 39 +++++ 14 files changed, 410 insertions(+), 181 deletions(-) create mode 100644 protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java delete mode 100644 protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java create mode 100644 protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseWithPrettifier.java create mode 100644 protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseWithSanitizer.java 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 8a0ad563a6..4c70a291d1 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,12 +239,17 @@ protected Request getSqlCursorCloseRequest(String cursorRequest) { } protected String executeQuery(String query, String requestType) { + return executeQuery(query, requestType, Map.of()); + } + + protected String executeQuery(String query, String requestType, Map params) { try { String endpoint = "/_plugins/_sql?format=" + requestType; String requestBody = makeRequest(query); Request sqlRequest = new Request("POST", endpoint); sqlRequest.setJsonEntity(requestBody); + sqlRequest.addParameters(params); Response response = client().performRequest(sqlRequest); Assert.assertEquals(200, response.getStatusLine().getStatusCode()); diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/CsvFormatIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/CsvFormatIT.java index 330268c0e4..d481e0ad49 100644 --- a/integ-test/src/test/java/org/opensearch/sql/sql/CsvFormatIT.java +++ b/integ-test/src/test/java/org/opensearch/sql/sql/CsvFormatIT.java @@ -6,7 +6,7 @@ package org.opensearch.sql.sql; import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK_CSV_SANITIZE; -import static org.opensearch.sql.protocol.response.format.FlatResponseFormatter.CONTENT_TYPE; +import static org.opensearch.sql.protocol.response.format.CsvResponseFormatter.CONTENT_TYPE; import java.io.IOException; import java.util.Locale; 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..d0a5a37db3 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 @@ -6,10 +6,11 @@ package org.opensearch.sql.sql; import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK_RAW_SANITIZE; -import static org.opensearch.sql.protocol.response.format.FlatResponseFormatter.CONTENT_TYPE; +import static org.opensearch.sql.protocol.response.format.RawResponseFormatter.CONTENT_TYPE; import java.io.IOException; import java.util.Locale; +import java.util.Map; import org.junit.Test; import org.opensearch.client.Request; import org.opensearch.client.Response; @@ -41,6 +42,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", + Map.of("pretty", "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..f5899e88e0 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,45 @@ package org.opensearch.sql.protocol.response.format; -public class CsvResponseFormatter extends FlatResponseFormatter { +import org.opensearch.sql.protocol.response.QueryResult; + +/** Response formatter to format response to csv format. */ +public class CsvResponseFormatter implements ResponseFormatter { + public static final String CONTENT_TYPE = "plain/text; charset=UTF-8"; + private final String separator; + private final boolean sanitize; + public CsvResponseFormatter() { - super(",", true); + this(",", true); } public CsvResponseFormatter(boolean sanitize) { - super(",", sanitize); + this(",", sanitize); + } + + public CsvResponseFormatter(String separator, boolean sanitize) { + this.separator = separator; + this.sanitize = sanitize; + } + + @Override + public String format(QueryResult response) { + FlatResponseBase flatResponse; + if (sanitize) { + flatResponse = new FlatResponseWithSanitizer(response, separator); + } else { + flatResponse = new FlatResponseBase(response, separator); + } + return flatResponse.format(); + } + + @Override + public String format(Throwable t) { + return ErrorFormatter.prettyFormat(t); + } + + @Override + public String contentType() { + return CONTENT_TYPE; } } diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java new file mode 100644 index 0000000000..98e79a4048 --- /dev/null +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseBase.java @@ -0,0 +1,89 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.protocol.response.format; + +import com.google.common.collect.ImmutableList; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; +import lombok.Getter; +import org.opensearch.sql.protocol.response.QueryResult; + +@Getter +public class FlatResponseBase { + protected static final String INTERLINE_SEPARATOR = System.lineSeparator(); + + private final QueryResult response; + protected final String separator; + + private final List headers; + private final List> data; + + FlatResponseBase(QueryResult response, String separator) { + this.response = response; + this.separator = separator; + this.headers = getOriginalHeaders(response); + this.data = getOriginalData(response); + } + + public String format() { + List headersAndData = new ArrayList<>(); + headersAndData.add(getHeaderLine()); + headersAndData.addAll(getDataLines()); + return String.join(INTERLINE_SEPARATOR, headersAndData); + } + + protected String getHeaderLine() { + return String.join(separator, headers); + } + + private List getOriginalHeaders(QueryResult response) { + ImmutableList.Builder headers = ImmutableList.builder(); + response.columnNameTypes().forEach((column, type) -> headers.add(column)); + List result = headers.build(); + return formatHeaders(result); + } + + protected List getDataLines() { + return data.stream().map(v -> String.join(separator, v)).collect(Collectors.toList()); + } + + private List> getOriginalData(QueryResult response) { + ImmutableList.Builder> dataLines = new ImmutableList.Builder<>(); + response + .iterator() + .forEachRemaining( + row -> { + ImmutableList.Builder line = new ImmutableList.Builder<>(); + // replace null values with empty string + Arrays.asList(row).forEach(val -> line.add(val == null ? "" : val.toString())); + dataLines.add(line.build()); + }); + List> result = dataLines.build(); + return formatData(result); + } + + protected List formatHeaders(List headers) { + return headers.stream() + .map(cell -> quoteIfRequired(separator, cell)) + .collect(Collectors.toList()); + } + + protected List> formatData(List> lines) { + List> result = new ArrayList<>(); + for (List line : lines) { + result.add( + line.stream().map(cell -> quoteIfRequired(separator, cell)).collect(Collectors.toList())); + } + return result; + } + + protected String quoteIfRequired(String separator, String cell) { + final String quote = "\""; + return cell.contains(separator) ? quote + cell.replaceAll("\"", "\"\"") + quote : cell; + } +} 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 deleted file mode 100644 index 8c67d524b8..0000000000 --- a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java +++ /dev/null @@ -1,149 +0,0 @@ -/* - * Copyright OpenSearch Contributors - * SPDX-License-Identifier: Apache-2.0 - */ - -package org.opensearch.sql.protocol.response.format; - -import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableSet; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Set; -import java.util.stream.Collectors; -import lombok.Getter; -import lombok.RequiredArgsConstructor; -import org.opensearch.sql.protocol.response.QueryResult; - -@RequiredArgsConstructor -public abstract class FlatResponseFormatter implements ResponseFormatter { - private static String INLINE_SEPARATOR = ","; - private static final String INTERLINE_SEPARATOR = System.lineSeparator(); - private static final Set SENSITIVE_CHAR = ImmutableSet.of("=", "+", "-", "@"); - - public static final String CONTENT_TYPE = "plain/text; charset=UTF-8"; - - private boolean sanitize = false; - - public FlatResponseFormatter(String seperator, boolean sanitize) { - this.INLINE_SEPARATOR = seperator; - this.sanitize = sanitize; - } - - public String contentType() { - return CONTENT_TYPE; - } - - @Override - public String format(QueryResult response) { - FlatResult result = new FlatResult(response, sanitize); - return result.getFlat(); - } - - @Override - public String format(Throwable t) { - return ErrorFormatter.prettyFormat(t); - } - - /** - * Sanitize methods are migrated from legacy CSV result. Sanitize both headers and data lines by: - * 1) Second double quote entire cell if any comma is found. - */ - @Getter - @RequiredArgsConstructor - static class FlatResult { - private final QueryResult response; - private final boolean sanitize; - - public String getFlat() { - List headersAndData = new ArrayList<>(); - headersAndData.add(getHeaderLine(response, sanitize)); - headersAndData.addAll(getDataLines(response, sanitize)); - 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 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 getHeaders(QueryResult response, boolean sanitize) { - ImmutableList.Builder headers = ImmutableList.builder(); - response.columnNameTypes().forEach((column, type) -> headers.add(column)); - List result = headers.build(); - return sanitizeHeaders(result); - } - - private List> getData(QueryResult response, boolean sanitize) { - ImmutableList.Builder> dataLines = new ImmutableList.Builder<>(); - response - .iterator() - .forEachRemaining( - row -> { - ImmutableList.Builder line = new ImmutableList.Builder<>(); - // replace null values with empty string - Arrays.asList(row).forEach(val -> line.add(val == null ? "" : val.toString())); - dataLines.add(line.build()); - }); - List> result = dataLines.build(); - return sanitizeData(result); - } - - /** Sanitize headers because OpenSearch allows special character present in field names. */ - private List sanitizeHeaders(List headers) { - if (sanitize) { - return headers.stream() - .map(this::sanitizeCell) - .map(cell -> quoteIfRequired(INLINE_SEPARATOR, cell)) - .collect(Collectors.toList()); - } else { - return headers.stream() - .map(cell -> quoteIfRequired(INLINE_SEPARATOR, cell)) - .collect(Collectors.toList()); - } - } - - private List> sanitizeData(List> lines) { - List> result = new ArrayList<>(); - if (sanitize) { - for (List line : lines) { - result.add( - line.stream() - .map(this::sanitizeCell) - .map(cell -> quoteIfRequired(INLINE_SEPARATOR, cell)) - .collect(Collectors.toList())); - } - } else { - for (List line : lines) { - result.add( - line.stream() - .map(cell -> quoteIfRequired(INLINE_SEPARATOR, cell)) - .collect(Collectors.toList())); - } - } - return result; - } - - private String sanitizeCell(String cell) { - if (isStartWithSensitiveChar(cell)) { - return "'" + cell; - } - return cell; - } - - private String quoteIfRequired(String separator, String cell) { - final String quote = "\""; - return cell.contains(separator) ? quote + cell.replaceAll("\"", "\"\"") + quote : cell; - } - - private boolean isStartWithSensitiveChar(String cell) { - return SENSITIVE_CHAR.stream().anyMatch(cell::startsWith); - } - } -} diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseWithPrettifier.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseWithPrettifier.java new file mode 100644 index 0000000000..0e2527fbfe --- /dev/null +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseWithPrettifier.java @@ -0,0 +1,53 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.sql.protocol.response.format; + +import java.util.List; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.opensearch.sql.protocol.response.QueryResult; + +public class FlatResponseWithPrettifier extends FlatResponseBase { + private int[] maxWidths; + + FlatResponseWithPrettifier(QueryResult response, String inlineSeparator) { + super(response, inlineSeparator); + calculateMaxWidths(); + } + + private void calculateMaxWidths() { + int columns = getHeaders().size(); + maxWidths = new int[columns]; + + for (int i = 0; i < columns; i++) { + int maxWidth = getHeaders().get(i).length(); + for (List row : getData()) { + maxWidth = Math.max(maxWidth, row.get(i).length()); + } + maxWidths[i] = maxWidth; + } + } + + @Override + protected List getDataLines() { + return getData().stream().map(this::prettyFormatLine).collect(Collectors.toList()); + } + + @Override + protected String getHeaderLine() { + return prettyFormatLine(getHeaders()); + } + + private String prettyFormatLine(List line) { + return IntStream.range(0, line.size()) + .mapToObj(i -> padRight(line.get(i), maxWidths[i])) + .collect(Collectors.joining(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/FlatResponseWithSanitizer.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseWithSanitizer.java new file mode 100644 index 0000000000..69a29c5cce --- /dev/null +++ b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseWithSanitizer.java @@ -0,0 +1,53 @@ +package org.opensearch.sql.protocol.response.format; + +import com.google.common.collect.ImmutableSet; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; +import org.opensearch.sql.protocol.response.QueryResult; + +/** + * Sanitize methods are migrated from legacy CSV result. Sanitize both headers and data lines by: 1) + * Second double quote entire cell if any comma is found. + */ +public class FlatResponseWithSanitizer extends FlatResponseBase { + private static final Set SENSITIVE_CHAR = ImmutableSet.of("=", "+", "-", "@"); + + FlatResponseWithSanitizer(QueryResult response, String inlineSeparator) { + super(response, inlineSeparator); + } + + /** Sanitize headers because OpenSearch allows special character present in field names. */ + @Override + protected List formatHeaders(List headers) { + return headers.stream() + .map(this::sanitizeCell) + .map(cell -> quoteIfRequired(separator, cell)) + .collect(Collectors.toList()); + } + + @Override + protected List> formatData(List> lines) { + List> result = new ArrayList<>(); + for (List line : lines) { + result.add( + line.stream() + .map(this::sanitizeCell) + .map(cell -> quoteIfRequired(separator, cell)) + .collect(Collectors.toList())); + } + return result; + } + + private String sanitizeCell(String cell) { + if (isStartWithSensitiveChar(cell)) { + return "'" + cell; + } + return cell; + } + + private boolean isStartWithSensitiveChar(String cell) { + return SENSITIVE_CHAR.stream().anyMatch(cell::startsWith); + } +} 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..061c2fcb8f 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,51 @@ package org.opensearch.sql.protocol.response.format; -/** Response formatter to format response to csv or raw format. */ -public class RawResponseFormatter extends FlatResponseFormatter { +import org.opensearch.sql.protocol.response.QueryResult; + +/** Response formatter to format response to raw format. */ +public class RawResponseFormatter implements ResponseFormatter { + public static final String CONTENT_TYPE = "plain/text; charset=UTF-8"; + private final String separator; + private final boolean pretty; + public RawResponseFormatter() { - super("|", false); + this("|", false); + } + + public RawResponseFormatter(boolean pretty) { + this("|", pretty); + } + + /** + * Create a raw response formatter with separator and 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(String separator, boolean pretty) { + this.separator = separator; + this.pretty = pretty; + } + + @Override + public String format(QueryResult response) { + FlatResponseBase flatResponse; + if (pretty) { + flatResponse = new FlatResponseWithPrettifier(response, separator); + } else { + flatResponse = new FlatResponseBase(response, separator); + } + return flatResponse.format(); + } + + @Override + public String format(Throwable t) { + return ErrorFormatter.prettyFormat(t); + } + + @Override + public String contentType() { + return CONTENT_TYPE; } } diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java index d27ac72373..554d5ae8bb 100644 --- a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java +++ b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatterTest.java @@ -13,7 +13,7 @@ import static org.opensearch.sql.data.model.ExprValueUtils.tupleValue; import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.protocol.response.format.FlatResponseFormatter.CONTENT_TYPE; +import static org.opensearch.sql.protocol.response.format.CsvResponseFormatter.CONTENT_TYPE; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; 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..fd057437a0 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 @@ -13,7 +13,7 @@ import static org.opensearch.sql.data.model.ExprValueUtils.tupleValue; import static org.opensearch.sql.data.type.ExprCoreType.INTEGER; import static org.opensearch.sql.data.type.ExprCoreType.STRING; -import static org.opensearch.sql.protocol.response.format.FlatResponseFormatter.CONTENT_TYPE; +import static org.opensearch.sql.protocol.response.format.RawResponseFormatter.CONTENT_TYPE; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -23,9 +23,8 @@ import org.opensearch.sql.executor.ExecutionEngine; import org.opensearch.sql.protocol.response.QueryResult; -/** Unit test for {@link FlatResponseFormatter}. */ +/** Unit test for {@link RawResponseFormatter}. */ 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 RawResponseFormatter getRawFormatter() { + return new RawResponseFormatter(); + } + + private RawResponseFormatter 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; + } } diff --git a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java index 8cb2994dc3..b124757069 100644 --- a/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java +++ b/sql/src/test/java/org/opensearch/sql/sql/SQLServiceTest.java @@ -127,6 +127,45 @@ public void onFailure(Exception e) { }); } + @Test + public void can_execute_raw_format_request() { + sqlService.execute( + new SQLQueryRequest(new JSONObject(), "SELECT 123", QUERY, "raw"), + new ResponseListener() { + @Override + public void onResponse(QueryResponse response) { + assertNotNull(response); + } + + @Override + public void onFailure(Exception e) { + fail(e); + } + }); + } + + @Test + public void can_execute_pretty_raw_format_request() { + sqlService.execute( + new SQLQueryRequest( + new JSONObject(), + "SELECT 123", + QUERY, + Map.of("format", "jdbc", "pretty", "true"), + "n:cursor"), + new ResponseListener() { + @Override + public void onResponse(QueryResponse response) { + assertNotNull(response); + } + + @Override + public void onFailure(Exception e) { + fail(e); + } + }); + } + @Test public void can_explain_sql_query() { doAnswer(