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

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

@@ -508,6 +508,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

@@ -15,6 +15,7 @@
 import org.elasticsearch.client.ResponseException;
 import org.elasticsearch.common.CheckedSupplier;
 import org.elasticsearch.common.Strings;
+import org.elasticsearch.common.bytes.BytesArray;
 import org.elasticsearch.common.collect.Tuple;
 import org.elasticsearch.common.io.Streams;
 import org.elasticsearch.common.xcontent.XContentHelper;
@@ -414,6 +415,85 @@ 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()) {
+            String actualJson = new BytesArray(content.readAllBytes()).utf8ToString();
+            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);
                 }