Well format the raw response when query parameter "pretty" enabled (o…

…pensearch-project#2727)

Signed-off-by: Lantao Jin <[email protected]>

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<String, String> 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());

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;

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 =

legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSQLQueryAction.java

@@ -172,7 +172,7 @@ private ResponseListener<QueryResponse> 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);
     }

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<QueryResult> {
+  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;
   }
 }

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<String> headers;
+  private final List<List<String>> data;
+
+  FlatResponseBase(QueryResult response, String separator) {
+    this.response = response;
+    this.separator = separator;
+    this.headers = getOriginalHeaders(response);
+    this.data = getOriginalData(response);
+  }
+
+  public String format() {
+    List<String> 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<String> getOriginalHeaders(QueryResult response) {
+    ImmutableList.Builder<String> headers = ImmutableList.builder();
+    response.columnNameTypes().forEach((column, type) -> headers.add(column));
+    List<String> result = headers.build();
+    return formatHeaders(result);
+  }
+
+  protected List<String> getDataLines() {
+    return data.stream().map(v -> String.join(separator, v)).collect(Collectors.toList());
+  }
+
+  private List<List<String>> getOriginalData(QueryResult response) {
+    ImmutableList.Builder<List<String>> dataLines = new ImmutableList.Builder<>();
+    response
+        .iterator()
+        .forEachRemaining(
+            row -> {
+              ImmutableList.Builder<String> 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<List<String>> result = dataLines.build();
+    return formatData(result);
+  }
+
+  protected List<String> formatHeaders(List<String> headers) {
+    return headers.stream()
+        .map(cell -> quoteIfRequired(separator, cell))
+        .collect(Collectors.toList());
+  }
+
+  protected List<List<String>> formatData(List<List<String>> lines) {
+    List<List<String>> result = new ArrayList<>();
+    for (List<String> 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;
+  }
+}