SQL: add pretty printing to JSON format (#43756) (#44220)

(cherry picked from commit cbd9d4c)

Author: astefan

server/src/main/java/org/elasticsearch/rest/AbstractRestChannel.java

@@ -84,9 +84,22 @@ public XContentBuilder newErrorBuilder() throws IOException {
      */
     @Override
     public XContentBuilder newBuilder(@Nullable XContentType requestContentType, boolean useFiltering) throws IOException {
+        return newBuilder(requestContentType, null, useFiltering);
+    }
+
+    /**
+     * Creates a new {@link XContentBuilder} for a response to be sent using this channel. The builder's type can be sent as a parameter,
+     * through {@code responseContentType} or it can fallback to {@link #newBuilder(XContentType, boolean)} logic if the sent type value
+     * is {@code null}.
+     */
+    @Override
+    public XContentBuilder newBuilder(@Nullable XContentType requestContentType, @Nullable XContentType responseContentType,
+            boolean useFiltering) throws IOException {
+        if (responseContentType == null) {
+            responseContentType = XContentType.fromMediaTypeOrFormat(format);
+        }
         // try to determine the response content type from the media type or the format query string parameter, with the format parameter
         // taking precedence over the Accept header
-        XContentType responseContentType = XContentType.fromMediaTypeOrFormat(format);
         if (responseContentType == null) {
             if (requestContentType != null) {
                 // if there was a parsed content-type for the incoming request use that since no format was specified using the query

server/src/main/java/org/elasticsearch/rest/RestChannel.java

@@ -36,6 +36,9 @@ public interface RestChannel {
     XContentBuilder newErrorBuilder() throws IOException;
 
     XContentBuilder newBuilder(@Nullable XContentType xContentType, boolean useFiltering) throws IOException;
+    
+    XContentBuilder newBuilder(@Nullable XContentType xContentType, @Nullable XContentType responseContentType,
+            boolean useFiltering) throws IOException;
 
     BytesStreamOutput bytesOutput();
 

server/src/main/java/org/elasticsearch/rest/RestController.java

@@ -509,6 +509,12 @@ public XContentBuilder newBuilder(@Nullable XContentType xContentType, boolean u
             return delegate.newBuilder(xContentType, useFiltering);
         }
 
+        @Override
+        public XContentBuilder newBuilder(XContentType xContentType, XContentType responseContentType, boolean useFiltering)
+                throws IOException {
+            return delegate.newBuilder(xContentType, responseContentType, useFiltering);
+        }
+
         @Override
         public BytesStreamOutput bytesOutput() {
             return delegate.bytesOutput();

x-pack/plugin/sql/qa/src/main/java/org/elasticsearch/xpack/sql/qa/rest/RestSqlTestCase.java

@@ -6,6 +6,7 @@
 package org.elasticsearch.xpack.sql.qa.rest;
 
 import com.fasterxml.jackson.core.io.JsonStringEncoder;
+
 import org.apache.http.HttpEntity;
 import org.apache.http.entity.ContentType;
 import org.apache.http.entity.StringEntity;
@@ -25,6 +26,7 @@
 import org.elasticsearch.xpack.sql.qa.ErrorsTestCase;
 import org.hamcrest.Matcher;
 
+import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.InputStreamReader;
@@ -414,6 +416,91 @@ protected Map<String, Object> runSql(HttpEntity sql, String suffix) throws IOExc
         }
     }
 
+    public void testPrettyPrintingEnabled() throws IOException {
+        boolean columnar = randomBoolean();
+        String expected = "";
+        if (columnar) {
+            expected = "{\n" + 
+                    "  \"columns\" : [\n" + 
+                    "    {\n" + 
+                    "      \"name\" : \"test1\",\n" + 
+                    "      \"type\" : \"text\"\n" + 
+                    "    }\n" + 
+                    "  ],\n" + 
+                    "  \"values\" : [\n" + 
+                    "    [\n" + 
+                    "      \"test1\",\n" + 
+                    "      \"test2\"\n" + 
+                    "    ]\n" + 
+                    "  ]\n" + 
+                    "}\n";
+        } else {
+            expected = "{\n" + 
+                    "  \"columns\" : [\n" + 
+                    "    {\n" + 
+                    "      \"name\" : \"test1\",\n" + 
+                    "      \"type\" : \"text\"\n" + 
+                    "    }\n" + 
+                    "  ],\n" + 
+                    "  \"rows\" : [\n" + 
+                    "    [\n" + 
+                    "      \"test1\"\n" + 
+                    "    ],\n" + 
+                    "    [\n" + 
+                    "      \"test2\"\n" + 
+                    "    ]\n" + 
+                    "  ]\n" + 
+                    "}\n";
+        }
+        executeAndAssertPrettyPrinting(expected, "true", columnar);
+    }
+    
+    public void testPrettyPrintingDisabled() throws IOException {
+        boolean columnar = randomBoolean();
+        String expected = "";
+        if (columnar) {
+            expected = "{\"columns\":[{\"name\":\"test1\",\"type\":\"text\"}],\"values\":[[\"test1\",\"test2\"]]}";
+        } else {
+            expected = "{\"columns\":[{\"name\":\"test1\",\"type\":\"text\"}],\"rows\":[[\"test1\"],[\"test2\"]]}";
+        }
+        executeAndAssertPrettyPrinting(expected, randomFrom("false", null), columnar);
+    }
+    
+    private void executeAndAssertPrettyPrinting(String expectedJson, String prettyParameter, boolean columnar)
+            throws IOException {
+        index("{\"test1\":\"test1\"}",
+              "{\"test1\":\"test2\"}");
+
+        Request request = new Request("POST", SQL_QUERY_REST_ENDPOINT);
+        if (prettyParameter != null) {
+            request.addParameter("pretty", prettyParameter);
+        }
+        if (randomBoolean()) {
+            // We default to JSON but we force it randomly for extra coverage
+            request.addParameter("format", "json");
+        }
+        if (randomBoolean()) {
+            // JSON is the default but randomly set it sometime for extra coverage
+            RequestOptions.Builder options = request.getOptions().toBuilder();
+            options.addHeader("Accept", randomFrom("*/*", "application/json"));
+            request.setOptions(options);
+        }
+        request.setEntity(new StringEntity("{\"query\":\"SELECT * FROM test\"" + mode("plain") + columnarParameter(columnar) + "}",
+                  ContentType.APPLICATION_JSON));
+        
+        Response response = client().performRequest(request);
+        try (InputStream content = response.getEntity().getContent()) {
+            ByteArrayOutputStream result = new ByteArrayOutputStream();
+            byte[] buffer = new byte[1024];
+            int length;
+            while ((length = content.read(buffer)) != -1) {
+                result.write(buffer, 0, length);
+            }
+            String actualJson = result.toString("UTF-8");
+            assertEquals(expectedJson, actualJson);
+        }
+    }
+
     public void testBasicTranslateQuery() throws IOException {
         index("{\"test\":\"test\"}", "{\"test\":\"test\"}");
 

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/plugin/RestSqlQueryAction.java

@@ -95,7 +95,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
             return channel -> client.execute(SqlQueryAction.INSTANCE, sqlRequest, new RestResponseListener<SqlQueryResponse>(channel) {
                 @Override
                 public RestResponse buildResponse(SqlQueryResponse response) throws Exception {
-                    XContentBuilder builder = XContentBuilder.builder(xContentType.xContent());
+                    XContentBuilder builder = channel.newBuilder(request.getXContentType(), xContentType, true);
                     response.toXContent(builder, request);
                     return new BytesRestResponse(RestStatus.OK, builder);
                 }