Fix CSV/RAW outputting wrong format (#279)

matthewryanwells · web-flow · commit 294dd8f51948 · 2023-06-26T09:36:09.000-07:00
* Fixed bug where CSV/RAW outputs as JSON rather than plain text

Signed-off-by: Matthew Wells &lt;matthew.wells@improving.com&gt;
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
@@ -7,10 +7,14 @@
 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 java.io.IOException;
 import java.util.Locale;
+
 import org.junit.Test;
+import org.opensearch.client.Request;
+import org.opensearch.client.Response;
 import org.opensearch.sql.common.utils.StringUtils;
 import org.opensearch.sql.legacy.SQLIntegTestCase;
 
@@ -49,4 +53,17 @@ public void escapeSanitizeTest() {
             + "\",Elinor\",\"Ratliff,,,\"%n"),
         result);
   }
+
+  @Test
+  public void contentHeaderTest() throws IOException {
+    String query = String.format(Locale.ROOT, "SELECT firstname, lastname FROM %s", TEST_INDEX_BANK_CSV_SANITIZE);
+    String requestBody = makeRequest(query);
+
+    Request sqlRequest = new Request("POST", "/_plugins/_sql?format=csv");
+    sqlRequest.setJsonEntity(requestBody);
+
+    Response response = client().performRequest(sqlRequest);
+
+    assertEquals(response.getEntity().getContentType(), CONTENT_TYPE);
+  }
 }
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
@@ -6,11 +6,15 @@
 
 package org.opensearch.sql.sql;
 
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK_CSV_SANITIZE;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK_RAW_SANITIZE;
+import static org.opensearch.sql.protocol.response.format.FlatResponseFormatter.CONTENT_TYPE;
 
 import java.io.IOException;
 import java.util.Locale;
 import org.junit.Test;
+import org.opensearch.client.Request;
+import org.opensearch.client.Response;
 import org.opensearch.sql.common.utils.StringUtils;
 import org.opensearch.sql.legacy.SQLIntegTestCase;
 
@@ -35,4 +39,16 @@ public void rawFormatWithPipeFieldTest() {
         result);
   }
 
+  @Test
+  public void contentHeaderTest() throws IOException {
+    String query = String.format(Locale.ROOT, "SELECT firstname, lastname FROM %s", TEST_INDEX_BANK_RAW_SANITIZE);
+    String requestBody = makeRequest(query);
+
+    Request sqlRequest = new Request("POST", "/_plugins/_sql?format=raw");
+    sqlRequest.setJsonEntity(requestBody);
+
+    Response response = client().performRequest(sqlRequest);
+
+    assertEquals(response.getEntity().getContentType(), CONTENT_TYPE);
+  }
 }
diff --git a/integ-test/src/test/java/org/opensearch/sql/sql/SimpleQueryStringIT.java b/integ-test/src/test/java/org/opensearch/sql/sql/SimpleQueryStringIT.java
@@ -5,11 +5,17 @@
 
 package org.opensearch.sql.sql;
 
+import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK_CSV_SANITIZE;
 import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BEER;
+import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.CONTENT_TYPE;
 
 import java.io.IOException;
+import java.util.Locale;
+
 import org.json.JSONObject;
 import org.junit.Test;
+import org.opensearch.client.Request;
+import org.opensearch.client.Response;
 import org.opensearch.sql.legacy.SQLIntegTestCase;
 
 public class SimpleQueryStringIT extends SQLIntegTestCase {
@@ -61,4 +67,18 @@ public void verify_wildcard_test() throws IOException {
     var result = new JSONObject(executeQuery(query, "jdbc"));
     assertEquals(10, result.getInt("total"));
   }
+
+  @Test
+  public void contentHeaderTest() throws IOException {
+    String query = "SELECT Id FROM " + TEST_INDEX_BEER
+            + " WHERE simple_query_string([\\\"Tags\\\" ^ 1.5, Title, 'Body' 4.2], 'taste')";
+    String requestBody = makeRequest(query);
+
+    Request sqlRequest = new Request("POST", "/_plugins/_sql");
+    sqlRequest.setJsonEntity(requestBody);
+
+    Response response = client().performRequest(sqlRequest);
+
+    assertEquals(response.getEntity().getContentType(), CONTENT_TYPE);
+  }
 }
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
@@ -141,15 +141,16 @@ public void onFailure(Exception e) {
 
   private ResponseListener<ExplainResponse> createExplainResponseListener(
       RestChannel channel, BiConsumer<RestChannel, Exception> errorHandler) {
-    return new ResponseListener<ExplainResponse>() {
+    return new ResponseListener<>() {
       @Override
       public void onResponse(ExplainResponse response) {
-        sendResponse(channel, OK, new JsonResponseFormatter<ExplainResponse>(PRETTY) {
+        JsonResponseFormatter<ExplainResponse> formatter = new JsonResponseFormatter<>(PRETTY) {
           @Override
           protected Object buildJsonObject(ExplainResponse response) {
             return response;
           }
-        }.format(response));
+        };
+        sendResponse(channel, OK, formatter.format(response), formatter.contentType());
       }
 
       @Override
@@ -180,7 +181,7 @@ private ResponseListener<QueryResponse> createQueryResponseListener(
       public void onResponse(QueryResponse response) {
         sendResponse(channel, OK,
             formatter.format(new QueryResult(response.getSchema(), response.getResults(),
-                response.getCursor())));
+                response.getCursor())), formatter.contentType());
       }
 
       @Override
@@ -190,9 +191,9 @@ public void onFailure(Exception e) {
     };
   }
 
-  private void sendResponse(RestChannel channel, RestStatus status, String content) {
+  private void sendResponse(RestChannel channel, RestStatus status, String content, String contentType) {
     channel.sendResponse(new BytesRestResponse(
-        status, "application/json; charset=UTF-8", content));
+        status, contentType, content));
   }
 
   private static void logAndPublishMetrics(Exception e) {
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
@@ -23,13 +23,19 @@ public abstract class FlatResponseFormatter implements ResponseFormatter<QueryRe
   private static final String INTERLINE_SEPARATOR = System.lineSeparator();
   private static final Set<String> 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);
diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JsonResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/JsonResponseFormatter.java
@@ -36,6 +36,8 @@ public enum Style {
    */
   private final Style style;
 
+  public static final String CONTENT_TYPE = "application/json; charset=UTF-8";
+
   @Override
   public String format(R response) {
     return jsonify(buildJsonObject(response));
@@ -47,6 +49,10 @@ public String format(Throwable t) {
         (style == PRETTY) ? prettyFormat(t) : compactFormat(t));
   }
 
+  public String contentType() {
+    return CONTENT_TYPE;
+  }
+
   /**
    * Build JSON object to generate response json string.
    *
diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/ResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/ResponseFormatter.java
@@ -27,4 +27,11 @@ public interface ResponseFormatter<R> {
    */
   String format(Throwable t);
 
+  /**
+   * Getter for the content type header of the response.
+   *
+   * @return string
+   */
+  String contentType();
+
 }
diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CommandResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/CommandResponseFormatterTest.java
@@ -10,6 +10,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.JsonResponseFormatter.CONTENT_TYPE;
 import static org.opensearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY;
 
 import com.google.common.collect.ImmutableList;
@@ -56,4 +57,10 @@ public void formats_error_as_default_formatter() {
     assertEquals(new JdbcResponseFormatter(PRETTY).format(exception),
         new CommandResponseFormatter().format(exception));
   }
+
+  @Test
+  void testContentType() {
+    var formatter = new CommandResponseFormatter();
+    assertEquals(formatter.contentType(), 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
@@ -14,6 +14,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 com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -130,4 +131,9 @@ void replaceNullValues() {
     assertEquals(format(expected), formatter.format(response));
   }
 
+  @Test
+  void testContentType() {
+    assertEquals(formatter.contentType(), CONTENT_TYPE);
+  }
+
 }
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
@@ -14,6 +14,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 com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -27,7 +28,7 @@
  * Unit test for {@link FlatResponseFormatter}.
  */
 public class RawResponseFormatterTest {
-  private FlatResponseFormatter rawFormater = new RawResponseFormatter();
+  private FlatResponseFormatter rawFormatter = new RawResponseFormatter();
 
   @Test
   void formatResponse() {
@@ -38,7 +39,7 @@ void formatResponse() {
         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), rawFormater.format(response));
+    assertEquals(format(expected), rawFormatter.format(response));
   }
 
   @Test
@@ -53,7 +54,7 @@ void sanitizeHeaders() {
             "=firstname", "John", "+lastname", "Smith", "-city", "Seattle", "@age", 20))));
     String expected = "=firstname|+lastname|-city|@age%n"
         + "John|Smith|Seattle|20";
-    assertEquals(format(expected), rawFormater.format(response));
+    assertEquals(format(expected), rawFormatter.format(response));
   }
 
   @Test
@@ -74,7 +75,7 @@ void sanitizeData() {
         + "-Seattle%n"
         + "@Seattle%n"
         + "Seattle=";
-    assertEquals(format(expected), rawFormater.format(response));
+    assertEquals(format(expected), rawFormatter.format(response));
   }
 
   @Test
@@ -86,15 +87,15 @@ void quoteIfRequired() {
             tupleValue(ImmutableMap.of("na|me", "John|Smith", "||age", "30|||"))));
     String expected = "\"na|me\"|\"||age\"%n"
             + "\"John|Smith\"|\"30|||\"";
-    assertEquals(format(expected), rawFormater.format(response));
+    assertEquals(format(expected), rawFormatter.format(response));
   }
 
   @Test
   void formatError() {
     Throwable t = new RuntimeException("This is an exception");
     String expected =
         "{\n  \"type\": \"RuntimeException\",\n  \"reason\": \"This is an exception\"\n}";
-    assertEquals(expected, rawFormater.format(t));
+    assertEquals(expected, rawFormatter.format(t));
   }
 
   @Test
@@ -121,7 +122,7 @@ void senstiveCharater() {
     String expected = "city%n"
             + "@Seattle%n"
             + "++Seattle";
-    assertEquals(format(expected), rawFormater.format(response));
+    assertEquals(format(expected), rawFormatter.format(response));
   }
 
   @Test
@@ -153,7 +154,12 @@ void replaceNullValues() {
         + "John|Seattle%n"
         + "|Seattle%n"
         + "John|";
-    assertEquals(format(expected), rawFormater.format(response));
+    assertEquals(format(expected), rawFormatter.format(response));
+  }
+
+  @Test
+  void testContentType() {
+    assertEquals(rawFormatter.contentType(), CONTENT_TYPE);
   }
 
 }
diff --git a/protocol/src/test/java/org/opensearch/sql/protocol/response/format/VisualizationResponseFormatterTest.java b/protocol/src/test/java/org/opensearch/sql/protocol/response/format/VisualizationResponseFormatterTest.java
@@ -11,6 +11,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.JsonResponseFormatter.CONTENT_TYPE;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
@@ -187,4 +188,10 @@ private static void assertJsonEquals(String expected, String actual) {
         JsonParser.parseString(expected),
         JsonParser.parseString(actual));
   }
+
+  @Test
+  void testContentType() {
+    var formatter = new CommandResponseFormatter();
+    assertEquals(formatter.contentType(), CONTENT_TYPE);
+  }
 }
