refactor

Signed-off-by: Lantao Jin <ltjin@amazon.com>

integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java

@@ -239,16 +239,17 @@ protected Request getSqlCursorCloseRequest(String cursorRequest) {
   }
 
   protected String executeQuery(String query, String requestType) {
-    return executeQuery(query, requestType, false);
+    return executeQuery(query, requestType, Map.of());
   }
 
-  protected String executeQuery(String query, String requestType, boolean pretty) {
+  protected String executeQuery(String query, String requestType, Map<String, String> params) {
     try {
-      String endpoint = "/_plugins/_sql?format=" + requestType + "&pretty=" + pretty;
+      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;
@@ -48,7 +49,7 @@ public void rawFormatPrettyWithPipeFieldTest() {
             String.format(
                 Locale.ROOT, "SELECT firstname, lastname FROM %s", TEST_INDEX_BANK_RAW_SANITIZE),
             "raw",
-            true);
+            Map.of("pretty", "true"));
     assertEquals(
         StringUtils.format(
             "firstname    |lastname         %n"

protocol/src/main/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatter.java

@@ -5,13 +5,40 @@
 
 package org.opensearch.sql.protocol.response.format;
 
+import org.opensearch.sql.protocol.response.QueryResult;
+
 /** Response formatter to format response to csv format. */
-public class CsvResponseFormatter extends FlatResponseFormatter {
+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, false);
+    this(",", true);
   }
 
   public CsvResponseFormatter(boolean sanitize) {
-    super(",", sanitize, false);
+    this(",", sanitize);
+  }
+
+  public CsvResponseFormatter(String separator, boolean sanitize) {
+    this.separator = separator;
+    this.sanitize = sanitize;
+  }
+
+  @Override
+  public String format(QueryResult response) {
+    Sanitizer sanitizer = new Sanitizer(response, separator, sanitize);
+    return sanitizer.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/FlatResponse.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 FlatResponse {
+  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;
+
+  FlatResponse(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;
+  }
+}