diff --git a/plugin/sql/sql-proto/src/test/java/org/elasticsearch/xpack/sql/plugin/SqlQueryResponseTests.java b/plugin/sql/sql-proto/src/test/java/org/elasticsearch/xpack/sql/plugin/SqlQueryResponseTests.java index 47918809f19..42c08bb0914 100644 --- a/plugin/sql/sql-proto/src/test/java/org/elasticsearch/xpack/sql/plugin/SqlQueryResponseTests.java +++ b/plugin/sql/sql-proto/src/test/java/org/elasticsearch/xpack/sql/plugin/SqlQueryResponseTests.java @@ -37,7 +37,7 @@ public class SqlQueryResponseTests extends AbstractStreamableXContentTestCase columns = null; diff --git a/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/RestSqlQueryAction.java b/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/RestSqlQueryAction.java index 57e730f262f..b668e431f16 100644 --- a/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/RestSqlQueryAction.java +++ b/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/RestSqlQueryAction.java @@ -55,41 +55,34 @@ public class RestSqlQueryAction extends BaseRestHandler { } }); } - // The client accepts plain text + + // The client accepts a text format + TextFormat text = TextFormat.fromMediaTypeOrFormat(request); long startNanos = System.nanoTime(); return channel -> client.execute(SqlQueryAction.INSTANCE, sqlRequest, new RestResponseListener(channel) { @Override public RestResponse buildResponse(SqlQueryResponse response) throws Exception { - final String data; - final CliFormatter formatter; - if (sqlRequest.cursor().equals("") == false) { - formatter = ((CliFormatterCursor) Cursor.decodeFromString(sqlRequest.cursor())).getCliFormatter(); - data = formatter.formatWithoutHeader(response); - } else { - formatter = new CliFormatter(response); - data = formatter.formatWithHeader(response); - } + Cursor cursor = Cursor.decodeFromString(sqlRequest.cursor()); + final String data = text.format(cursor, request, response); - return buildTextResponse(CliFormatterCursor.wrap(Cursor.decodeFromString(response.cursor()), formatter), - System.nanoTime() - startNanos, data); + RestResponse restResponse = new BytesRestResponse(RestStatus.OK, text.contentType(request), + data.getBytes(StandardCharsets.UTF_8)); + + Cursor responseCursor = text.wrapCursor(cursor, response); + + if (responseCursor != Cursor.EMPTY) { + restResponse.addHeader("Cursor", Cursor.encodeToString(Version.CURRENT, responseCursor)); + } + restResponse.addHeader("Took-nanos", Long.toString(System.nanoTime() - startNanos)); + + return restResponse; } }); } - private RestResponse buildTextResponse(Cursor responseCursor, long tookNanos, String data) { - RestResponse restResponse = new BytesRestResponse(RestStatus.OK, BytesRestResponse.TEXT_CONTENT_TYPE, - data.getBytes(StandardCharsets.UTF_8)); - if (responseCursor != Cursor.EMPTY) { - restResponse.addHeader("Cursor", Cursor.encodeToString(Version.CURRENT, responseCursor)); - } - restResponse.addHeader("Took-nanos", Long.toString(tookNanos)); - return restResponse; - } - @Override public String getName() { return "xpack_sql_action"; } -} - +} \ No newline at end of file diff --git a/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TextFormat.java b/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TextFormat.java new file mode 100644 index 00000000000..4de4244345c --- /dev/null +++ b/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/TextFormat.java @@ -0,0 +1,305 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ +package org.elasticsearch.xpack.sql.plugin; + +import org.elasticsearch.common.Strings; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.xpack.sql.session.Cursor; +import org.elasticsearch.xpack.sql.util.StringUtils; + +import java.util.List; +import java.util.Locale; +import java.util.Objects; +import java.util.function.Function; + +/** + * Templating class for displaying SQL responses in text formats. + */ + +// TODO are we sure toString is correct here? What about dates that come back as longs. +// Tracked by https://github.com/elastic/x-pack-elasticsearch/issues/3081 +enum TextFormat { + + /** + * Default text writer. + * + * The implementation is a bit weird since state needs to be passed around, namely the formatter + * since it is initialized based on the first page of data. + * To avoid leaking the formatter, it gets discovered again in the wrapping method to attach it + * to the next cursor and so on. + */ + PLAIN_TEXT() { + @Override + String format(Cursor cursor, RestRequest request, SqlQueryResponse response) { + final CliFormatter formatter; + if (cursor instanceof CliFormatterCursor) { + formatter = ((CliFormatterCursor) cursor).getCliFormatter(); + return formatter.formatWithoutHeader(response); + } else { + formatter = new CliFormatter(response); + return formatter.formatWithHeader(response); + } + } + + @Override + Cursor wrapCursor(Cursor oldCursor, SqlQueryResponse response) { + CliFormatter formatter = (oldCursor instanceof CliFormatterCursor) ? + ((CliFormatterCursor) oldCursor).getCliFormatter() : new CliFormatter(response); + return CliFormatterCursor.wrap(super.wrapCursor(oldCursor, response), formatter); + } + + @Override + String shortName() { + return "txt"; + } + + @Override + String contentType() { + return "text/plain"; + } + + @Override + protected String delimiter() { + throw new UnsupportedOperationException(); + } + + @Override + protected String eol() { + throw new UnsupportedOperationException(); + } + }, + + /** + * Comma Separated Values implementation. + * + * Based on: + * https://tools.ietf.org/html/rfc4180 + * https://www.iana.org/assignments/media-types/text/csv + * https://www.w3.org/TR/sparql11-results-csv-tsv/ + * + */ + CSV() { + + @Override + protected String delimiter() { + return ","; + } + + @Override + protected String eol() { + //LFCR + return "\r\n"; + } + + @Override + String shortName() { + return "csv"; + } + + @Override + String contentType() { + return "text/csv"; + } + + @Override + String contentType(RestRequest request) { + return contentType() + "; charset=utf-8; header=" + (hasHeader(request) ? "present" : "absent"); + } + + @Override + String maybeEscape(String value) { + boolean needsEscaping = false; + + for (int i = 0; i < value.length(); i++) { + char c = value.charAt(i); + if (c == '"' || c == ',' || c == '\n' || c == '\r') { + needsEscaping = true; + break; + } + } + + if (needsEscaping) { + StringBuilder sb = new StringBuilder(); + + sb.append('"'); + for (int i = 0; i < value.length(); i++) { + char c = value.charAt(i); + if (value.charAt(i) == '"') { + sb.append('"'); + } + sb.append(c); + } + sb.append('"'); + value = sb.toString(); + } + return value; + } + + @Override + boolean hasHeader(RestRequest request) { + String header = request.param("header"); + if (header == null) { + List values = request.getAllHeaderValues("Accept"); + if (values != null) { + // header is a parameter specified by ; so try breaking it down + for (String value : values) { + String[] params = Strings.tokenizeToStringArray(value, ";"); + for (String param : params) { + if (param.toLowerCase(Locale.ROOT).equals("header=absent")) { + return false; + } + } + } + } + return true; + } else { + return !header.toLowerCase(Locale.ROOT).equals("absent"); + } + } + }, + + TSV() { + @Override + protected String delimiter() { + return "\t"; + } + + @Override + protected String eol() { + // only CR + return "\n"; + } + + @Override + String shortName() { + return "tsv"; + } + + @Override + String contentType() { + return "text/tab-separated-values"; + } + + @Override + String contentType(RestRequest request) { + return contentType() + "; charset=utf-8"; + } + + @Override + String maybeEscape(String value) { + StringBuilder sb = new StringBuilder(); + + for (int i = 0; i < value.length(); i++) { + char c = value.charAt(i); + switch (c) { + case '\n' : + sb.append("\\n"); + break; + case '\t' : + sb.append("\\t"); + break; + default: + sb.append(c); + } + } + + return sb.toString(); + } + }; + + + String format(Cursor cursor, RestRequest request, SqlQueryResponse response) { + StringBuilder sb = new StringBuilder(); + + boolean header = hasHeader(request); + + if (header) { + row(sb, response.columns(), ColumnInfo::name); + } + + for (List row : response.rows()) { + row(sb, row, f -> Objects.toString(f, StringUtils.EMPTY)); + } + + return sb.toString(); + } + + boolean hasHeader(RestRequest request) { + return true; + } + + Cursor wrapCursor(Cursor oldCursor, SqlQueryResponse response) { + return Cursor.decodeFromString(response.cursor()); + } + + static TextFormat fromMediaTypeOrFormat(RestRequest request) { + String format = request.param("format", request.header("Accept")); + + if (format == null) { + return PLAIN_TEXT; + } + + for (TextFormat text : values()) { + String contentType = text.contentType(); + if (contentType.equalsIgnoreCase(format) + || format.toLowerCase(Locale.ROOT).startsWith(contentType + ";") + || text.shortName().equalsIgnoreCase(format)) { + return text; + } + } + + return PLAIN_TEXT; + } + + /** + * Short name typically used by format parameter. + * Can differ from the IANA mime type. + */ + abstract String shortName(); + + + /** + * Formal IANA mime type. + */ + abstract String contentType(); + + /** + * Content type depending on the request. + * Might be used by some formatters (like CSV) to specify certain metadata like + * whether the header is returned or not. + */ + String contentType(RestRequest request) { + return contentType(); + } + + // utility method for consuming a row. + void row(StringBuilder sb, List row, Function toString) { + for (int i = 0; i < row.size(); i++) { + sb.append(maybeEscape(toString.apply(row.get(i)))); + if (i < row.size() - 1) { + sb.append(delimiter()); + } + } + sb.append(eol()); + } + + /** + * Delimiter between fields + */ + protected abstract String delimiter(); + + /** + * String indicating end-of-line or row. + */ + protected abstract String eol(); + + /** + * Method used for escaping (if needed) a given value. + */ + String maybeEscape(String value) { + return value; + } +} \ No newline at end of file diff --git a/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/plugin/TextFormatTests.java b/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/plugin/TextFormatTests.java new file mode 100644 index 00000000000..50f117d2665 --- /dev/null +++ b/plugin/sql/src/test/java/org/elasticsearch/xpack/sql/plugin/TextFormatTests.java @@ -0,0 +1,165 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License; + * you may not use this file except in compliance with the Elastic License. + */ +package org.elasticsearch.xpack.sql.plugin; + + +import org.elasticsearch.common.xcontent.NamedXContentRegistry; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.test.rest.FakeRestRequest; + +import java.util.ArrayList; +import java.util.List; + +import static java.util.Arrays.asList; +import static java.util.Collections.emptyList; +import static java.util.Collections.singletonList; +import static java.util.Collections.singletonMap; +import static org.elasticsearch.xpack.sql.plugin.TextFormat.CSV; +import static org.elasticsearch.xpack.sql.plugin.TextFormat.TSV; +import static org.hamcrest.CoreMatchers.is; + +public class TextFormatTests extends ESTestCase { + + public void testPlainTextDetection() { + TextFormat text = TextFormat.fromMediaTypeOrFormat(withHeader("Accept", "text/plain")); + assertThat(text, is(TextFormat.PLAIN_TEXT)); + } + + public void testTextFallbackDetection() { + TextFormat text = TextFormat.fromMediaTypeOrFormat(withHeader("Accept", "text/*")); + assertThat(text, is(TextFormat.PLAIN_TEXT)); + } + + public void testTextFallbackNoHeader() { + assertThat(TextFormat.fromMediaTypeOrFormat(req()), is(TextFormat.PLAIN_TEXT)); + } + + public void testCsvDetection() { + TextFormat text = TextFormat.fromMediaTypeOrFormat(withHeader("Accept", "text/csv")); + assertThat(text, is(CSV)); + } + + public void testTsvDetection() { + TextFormat text = TextFormat.fromMediaTypeOrFormat(withHeader("Accept", "text/tab-separated-values")); + assertThat(text, is(TSV)); + } + + public void testCsvContentType() { + assertEquals("text/csv; charset=utf-8; header=present", CSV.contentType(req())); + } + + public void testCsvContentTypeWithoutHeader() { + assertEquals("text/csv; charset=utf-8; header=absent", CSV.contentType(reqNoHeader())); + } + + public void testTsvContentType() { + assertEquals("text/tab-separated-values; charset=utf-8", TSV.contentType(req())); + } + + public void testCsvEscaping() { + assertEquals("string", CSV.maybeEscape("string")); + assertEquals("", CSV.maybeEscape("")); + assertEquals("\"\"\"\"", CSV.maybeEscape("\"")); + assertEquals("\"\"\",\"\"\"", CSV.maybeEscape("\",\"")); + assertEquals("\"\"\"quo\"\"ted\"\"\"", CSV.maybeEscape("\"quo\"ted\"")); + } + + public void testTsvEscaping() { + assertEquals("string", TSV.maybeEscape("string")); + assertEquals("", TSV.maybeEscape("")); + assertEquals("\"", TSV.maybeEscape("\"")); + assertEquals("\\t", TSV.maybeEscape("\t")); + assertEquals("\\n\"\\t", TSV.maybeEscape("\n\"\t")); + } + + public void testCsvFormatWithEmptyData() { + String text = CSV.format(null, req(), emptyData()); + assertEquals("name\r\n", text); + } + + public void testTsvFormatWithEmptyData() { + String text = TSV.format(null, req(), emptyData()); + assertEquals("name\n", text); + } + + public void testCsvFormatWithRegularData() { + String text = CSV.format(null, req(), regularData()); + assertEquals("string,number\r\n" + + "Along The River Bank,708\r\n" + + "Mind Train,280\r\n", + text); + } + + public void testTsvFormatWithRegularData() { + String text = TSV.format(null, req(), regularData()); + assertEquals("string\tnumber\n" + + "Along The River Bank\t708\n" + + "Mind Train\t280\n", + text); + } + + public void testCsvFormatWithEscapedData() { + String text = CSV.format(null, req(), escapedData()); + assertEquals("first,\"\"\"special\"\"\"\r\n" + + "normal,\"\"\"quo\"\"ted\"\",\n\"\r\n" + + "commas,\"a,b,c,\n,d,e,\t\n\"\r\n" + , text); + } + + public void testTsvFormatWithEscapedData() { + String text = TSV.format(null, req(), escapedData()); + assertEquals("first\t\"special\"\n" + + "normal\t\"quo\"ted\",\\n\n" + + "commas\ta,b,c,\\n,d,e,\\t\\n\n" + , text); + } + + private static SqlQueryResponse emptyData() { + return new SqlQueryResponse(null, singletonList(new ColumnInfo("index", "name", "keyword")), emptyList()); + } + + private static SqlQueryResponse regularData() { + // headers + List headers = new ArrayList<>(); + headers.add(new ColumnInfo("index", "string", "keyword")); + headers.add(new ColumnInfo("index", "number", "integer")); + + // values + List> values = new ArrayList<>(); + values.add(asList("Along The River Bank", 11 * 60 + 48)); + values.add(asList("Mind Train", 4 * 60 + 40)); + + return new SqlQueryResponse(null, headers, values); + } + + private static SqlQueryResponse escapedData() { + // headers + List headers = new ArrayList<>(); + headers.add(new ColumnInfo("index", "first", "keyword")); + headers.add(new ColumnInfo("index", "\"special\"", "keyword")); + + // values + List> values = new ArrayList<>(); + values.add(asList("normal", "\"quo\"ted\",\n")); + values.add(asList("commas", "a,b,c,\n,d,e,\t\n")); + + return new SqlQueryResponse(null, headers, values); + } + + private static RestRequest req() { + return new FakeRestRequest(); + } + + private static RestRequest reqNoHeader() { + return new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY).withParams(singletonMap("header", "absent")).build(); + } + + + private static RestRequest withHeader(String key, String value) { + return new FakeRestRequest.Builder(NamedXContentRegistry.EMPTY).withHeaders(singletonMap(key, singletonList(value))).build(); + } +} \ No newline at end of file diff --git a/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/rest/RestSqlTestCase.java b/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/rest/RestSqlTestCase.java index c038e058b2f..8860758cf66 100644 --- a/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/rest/RestSqlTestCase.java +++ b/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/rest/RestSqlTestCase.java @@ -5,9 +5,13 @@ */ package org.elasticsearch.xpack.qa.sql.rest; +import com.fasterxml.jackson.core.io.JsonStringEncoder; + +import org.apache.http.Header; import org.apache.http.HttpEntity; import org.apache.http.entity.ContentType; import org.apache.http.entity.StringEntity; +import org.apache.http.message.BasicHeader; import org.elasticsearch.client.Response; import org.elasticsearch.client.ResponseException; import org.elasticsearch.common.CheckedSupplier; @@ -489,18 +493,97 @@ public abstract class RestSqlTestCase extends ESRestTestCase implements ErrorsTe assertEquals(0, getNumberOfSearchContexts("test")); } + // CSV/TSV tests + + private static String toJson(String value) { + return "\"" + new String(JsonStringEncoder.getInstance().quoteAsString(value)) + "\""; + } + + public void testDefaultQueryInCSV() throws IOException { + index("{\"name\":" + toJson("first") + ", \"number\" : 1 }", + "{\"name\":" + toJson("second\t") + ", \"number\": 2 }", + "{\"name\":" + toJson("\"third,\"") + ", \"number\": 3 }"); + + String expected = + "name,number\r\n" + + "first,1\r\n" + + "second\t,2\r\n" + + "\"\"\"third,\"\"\",3\r\n"; + + String query = "SELECT * FROM test ORDER BY number"; + Tuple response = runSqlAsText(query, "text/csv"); + assertEquals(expected, response.v1()); + + response = runSqlAsTextFormat(query, "csv"); + assertEquals(expected, response.v1()); + } + + public void testQueryWithoutHeaderInCSV() throws IOException { + index("{\"name\":" + toJson("first") + ", \"number\" : 1 }", + "{\"name\":" + toJson("second\t") + ", \"number\": 2 }", + "{\"name\":" + toJson("\"third,\"") + ", \"number\": 3 }"); + + String expected = + "first,1\r\n" + + "second\t,2\r\n" + + "\"\"\"third,\"\"\",3\r\n"; + + String query = "SELECT * FROM test ORDER BY number"; + Tuple response = runSqlAsText(query, "text/csv; header=absent"); + assertEquals(expected, response.v1()); + } + + public void testQueryInTSV() throws IOException { + index("{\"name\":" + toJson("first") + ", \"number\" : 1 }", + "{\"name\":" + toJson("second\t") + ", \"number\": 2 }", + "{\"name\":" + toJson("\"third,\"") + ", \"number\": 3 }"); + + String expected = + "name\tnumber\n" + + "first\t1\n" + + "second\\t\t2\n" + + "\"third,\"\t3\n"; + + String query = "SELECT * FROM test ORDER BY number"; + Tuple response = runSqlAsText(query, "text/tab-separated-values"); + assertEquals(expected, response.v1()); + response = runSqlAsTextFormat(query, "tsv"); + assertEquals(expected, response.v1()); + } + private Tuple runSqlAsText(String sql) throws IOException { return runSqlAsText("", new StringEntity("{\"query\":\"" + sql + "\"}", ContentType.APPLICATION_JSON)); } - private Tuple runSqlAsText(String suffix, HttpEntity sql) throws IOException { - Response response = client().performRequest("POST", "/_xpack/sql" + suffix, singletonMap("error_trace", "true"), sql); + private Tuple runSqlAsText(String sql, String acceptHeader) throws IOException { + return runSqlAsText("", new StringEntity("{\"query\":\"" + sql + "\"}", ContentType.APPLICATION_JSON), + new BasicHeader("Accept", acceptHeader)); + } + + private Tuple runSqlAsText(String suffix, HttpEntity sql, Header... headers) throws IOException { + Response response = client().performRequest("POST", "/_xpack/sql" + suffix, singletonMap("error_trace", "true"), sql, headers); return new Tuple<>( Streams.copyToString(new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8)), response.getHeader("Cursor") ); } + private Tuple runSqlAsTextFormat(String sql, String format) throws IOException { + StringEntity entity = new StringEntity("{\"query\":\"" + sql + "\"}", ContentType.APPLICATION_JSON); + + Map params = new HashMap<>(); + params.put("error_trace", "true"); + params.put("format", format); + + Response response = client().performRequest("POST", "/_xpack/sql", params, entity); + return new Tuple<>( + Streams.copyToString(new InputStreamReader(response.getEntity().getContent(), StandardCharsets.UTF_8)), + response.getHeader("Cursor") + ); + + } + + private void assertResponse(Map expected, Map actual) { if (false == expected.equals(actual)) { NotEqualMessageBuilder message = new NotEqualMessageBuilder(); @@ -542,5 +625,4 @@ public abstract class RestSqlTestCase extends ESRestTestCase implements ErrorsTe public static String randomMode() { return randomFrom("", "jdbc", "plain"); } - -} +} \ No newline at end of file