well formatted for raw response

Signed-off-by: Lantao Jin <ltjin@amazon.com>
LantaoJin · Jun 7, 2024 · 1de2a5d · 1de2a5d
1 parent c0a5123
commit 1de2a5d
Show file tree

Hide file tree

Showing 8 changed files with 159 additions and 36 deletions.
diff --git a/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java b/integ-test/src/test/java/org/opensearch/sql/legacy/SQLIntegTestCase.java
@@ -239,8 +239,12 @@ protected Request getSqlCursorCloseRequest(String cursorRequest) {
   }
 
   protected String executeQuery(String query, String requestType) {
+    return executeQuery(query, requestType, false);
+  }
+
+  protected String executeQuery(String query, String requestType, boolean pretty) {
     try {
-      String endpoint = "/_plugins/_sql?format=" + requestType;
+      String endpoint = "/_plugins/_sql?format=" + requestType + "&pretty=" + pretty;
       String requestBody = makeRequest(query);
 
       Request sqlRequest = new Request("POST", endpoint);

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
@@ -41,6 +41,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",
+            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 =

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
@@ -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);
     }

diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/CsvResponseFormatter.java
@@ -5,12 +5,13 @@
 
 package org.opensearch.sql.protocol.response.format;
 
+/** Response formatter to format response to csv format. */
 public class CsvResponseFormatter extends FlatResponseFormatter {
   public CsvResponseFormatter() {
-    super(",", true);
+    super(",", true, false);
   }
 
   public CsvResponseFormatter(boolean sanitize) {
-    super(",", sanitize);
+    super(",", sanitize, false);
   }
 }
diff --git a/...ocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java b/...ocol/src/main/java/org/opensearch/sql/protocol/response/format/FlatResponseFormatter.java
@@ -12,6 +12,7 @@
 import java.util.List;
 import java.util.Set;
 import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 import org.opensearch.sql.protocol.response.QueryResult;
@@ -25,10 +26,12 @@ public abstract class FlatResponseFormatter implements ResponseFormatter<QueryRe
   public static final String CONTENT_TYPE = "plain/text; charset=UTF-8";
 
   private boolean sanitize = false;
+  private boolean pretty = false;
 
-  public FlatResponseFormatter(String seperator, boolean sanitize) {
+  public FlatResponseFormatter(String seperator, boolean sanitize, boolean pretty) {
     this.INLINE_SEPARATOR = seperator;
     this.sanitize = sanitize;
+    this.pretty = pretty;
   }
 
   public String contentType() {
@@ -37,7 +40,7 @@ public String contentType() {
 
   @Override
   public String format(QueryResult response) {
-    FlatResult result = new FlatResult(response, sanitize);
+    FlatResult result = new FlatResult(response, sanitize, pretty);
     return result.getFlat();
   }
 
@@ -55,22 +58,54 @@ public String format(Throwable t) {
   static class FlatResult {
     private final QueryResult response;
     private final boolean sanitize;
+    private final boolean pretty;
+
+    private List<String> headers;
+    private List<List<String>> data;
+    private int[] maxWidths;
 
     public String getFlat() {
+      // Store headers and data in variables to avoid redundant calls
+      headers = getHeaders(response, sanitize);
+      data = getData(response, sanitize);
+
+      calculateMaxWidths();
+
       List<String> headersAndData = new ArrayList<>();
-      headersAndData.add(getHeaderLine(response, sanitize));
-      headersAndData.addAll(getDataLines(response, sanitize));
+      headersAndData.add(getHeaderLine());
+      headersAndData.addAll(getDataLines());
       return String.join(INTERLINE_SEPARATOR, headersAndData);
     }
 
-    private String getHeaderLine(QueryResult response, boolean sanitize) {
-      List<String> headers = getHeaders(response, sanitize);
-      return String.join(INLINE_SEPARATOR, headers);
+    private void calculateMaxWidths() {
+      int columns = headers.size();
+      maxWidths = new int[columns];
+
+      for (int i = 0; i < columns; i++) {
+        int maxWidth = headers.get(i).length();
+        for (List<String> row : data) {
+          maxWidth = Math.max(maxWidth, row.get(i).length());
+        }
+        maxWidths[i] = maxWidth;
+      }
     }
 
-    private List<String> getDataLines(QueryResult response, boolean sanitize) {
-      List<List<String>> data = getData(response, sanitize);
-      return data.stream().map(v -> String.join(INLINE_SEPARATOR, v)).collect(Collectors.toList());
+    private List<String> getDataLines() {
+      if (pretty) {
+        return data.stream().map(this::prettyFormatLine).collect(Collectors.toList());
+      } else {
+        return data.stream()
+            .map(v -> String.join(INLINE_SEPARATOR, v))
+            .collect(Collectors.toList());
+      }
+    }
+
+    private String getHeaderLine() {
+      if (pretty) {
+        return prettyFormatLine(headers);
+      } else {
+        return String.join(INLINE_SEPARATOR, headers);
+      }
     }
 
     private List<String> getHeaders(QueryResult response, boolean sanitize) {
@@ -145,5 +180,15 @@ private String quoteIfRequired(String separator, String cell) {
     private boolean isStartWithSensitiveChar(String cell) {
       return SENSITIVE_CHAR.stream().anyMatch(cell::startsWith);
     }
+
+    private String prettyFormatLine(List<String> line) {
+      return IntStream.range(0, line.size())
+          .mapToObj(i -> padRight(line.get(i), maxWidths[i]))
+          .collect(Collectors.joining(INLINE_SEPARATOR));
+    }
+
+    private String padRight(String s, int n) {
+      return String.format("%-" + n + "s", s);
+    }
   }
 }
diff --git a/protocol/src/main/java/org/opensearch/sql/protocol/response/format/RawResponseFormatter.java b/protocol/src/main/java/org/opensearch/sql/protocol/response/format/RawResponseFormatter.java
@@ -5,9 +5,19 @@
 
 package org.opensearch.sql.protocol.response.format;
 
-/** Response formatter to format response to csv or raw format. */
+/** Response formatter to format response to raw format. */
 public class RawResponseFormatter extends FlatResponseFormatter {
   public RawResponseFormatter() {
-    super("|", false);
+    super("|", false, false);
+  }
+
+  /**
+   * Create a raw response formatter with pretty parameter.
+   *
+   * @param pretty if true, display the columns with proper padding. Tracks the maximum width for
+   *     each column to ensure proper formatting.
+   */
+  public RawResponseFormatter(boolean pretty) {
+    super("|", false, pretty);
   }
 }
diff --git a/...l/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java b/...l/src/test/java/org/opensearch/sql/protocol/response/format/RawResponseFormatterTest.java
@@ -25,7 +25,6 @@
 
 /** Unit test for {@link FlatResponseFormatter}. */
 public class RawResponseFormatterTest {
-  private FlatResponseFormatter rawFormatter = new RawResponseFormatter();
 
   @Test
   void formatResponse() {
@@ -40,8 +39,10 @@ void formatResponse() {
             Arrays.asList(
                 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), rawFormatter.format(response));
+    String expected = "name|age%n" + "John|20%n" + "Smith|30";
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty = "name |age%n" + "John |20 %n" + "Smith|30 ";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }
 
   @Test
@@ -67,8 +68,11 @@ void sanitizeHeaders() {
                         "Seattle",
                         "@age",
                         20))));
-    String expected = "=firstname|+lastname|-city|@age%nJohn|Smith|Seattle|20";
-    assertEquals(format(expected), rawFormatter.format(response));
+    String expected = "=firstname|+lastname|-city|@age%n" + "John|Smith|Seattle|20";
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty =
+        "=firstname|+lastname|-city  |@age%n" + "John      |Smith    |Seattle|20  ";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }
 
   @Test
@@ -94,7 +98,16 @@ void sanitizeData() {
             + "-Seattle%n"
             + "@Seattle%n"
             + "Seattle=";
-    assertEquals(format(expected), rawFormatter.format(response));
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty =
+        "city    %n"
+            + "Seattle %n"
+            + "=Seattle%n"
+            + "+Seattle%n"
+            + "-Seattle%n"
+            + "@Seattle%n"
+            + "Seattle=";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }
 
   @Test
@@ -108,21 +121,23 @@ void quoteIfRequired() {
         new QueryResult(
             schema,
             Arrays.asList(tupleValue(ImmutableMap.of("na|me", "John|Smith", "||age", "30|||"))));
-    String expected = "\"na|me\"|\"||age\"%n\"John|Smith\"|\"30|||\"";
-    assertEquals(format(expected), rawFormatter.format(response));
+    String expected = "\"na|me\"|\"||age\"%n" + "\"John|Smith\"|\"30|||\"";
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty = "\"na|me\"     |\"||age\"%n" + "\"John|Smith\"|\"30|||\"";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().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, rawFormatter.format(t));
+    assertEquals(expected, getRawFormatter().format(t));
+    assertEquals(expected, getRawFormatterPretty().format(t));
   }
 
   @Test
   void escapeSanitize() {
-    FlatResponseFormatter escapeFormatter = new RawResponseFormatter();
     ExecutionEngine.Schema schema =
         new ExecutionEngine.Schema(
             ImmutableList.of(new ExecutionEngine.Schema.Column("city", "city", STRING)));
@@ -132,8 +147,10 @@ void escapeSanitize() {
             Arrays.asList(
                 tupleValue(ImmutableMap.of("city", "=Seattle")),
                 tupleValue(ImmutableMap.of("city", "||Seattle"))));
-    String expected = "city%n=Seattle%n\"||Seattle\"";
-    assertEquals(format(expected), escapeFormatter.format(response));
+    String expected = "city%n" + "=Seattle%n" + "\"||Seattle\"";
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty = "city       %n" + "=Seattle   %n" + "\"||Seattle\"";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }
 
   @Test
@@ -147,13 +164,14 @@ void senstiveCharater() {
             Arrays.asList(
                 tupleValue(ImmutableMap.of("city", "@Seattle")),
                 tupleValue(ImmutableMap.of("city", "++Seattle"))));
-    String expected = "city%n@Seattle%n++Seattle";
-    assertEquals(format(expected), rawFormatter.format(response));
+    String expected = "city%n" + "@Seattle%n" + "++Seattle";
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty = "city     %n" + "@Seattle %n" + "++Seattle";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }
 
   @Test
   void senstiveCharaterWithSanitize() {
-    FlatResponseFormatter testFormater = new RawResponseFormatter();
     ExecutionEngine.Schema schema =
         new ExecutionEngine.Schema(
             ImmutableList.of(new ExecutionEngine.Schema.Column("city", "city", STRING)));
@@ -163,8 +181,10 @@ void senstiveCharaterWithSanitize() {
             Arrays.asList(
                 tupleValue(ImmutableMap.of("city", "@Seattle")),
                 tupleValue(ImmutableMap.of("city", "++Seattle|||"))));
-    String expected = "city%n@Seattle%n\"++Seattle|||\"";
-    assertEquals(format(expected), testFormater.format(response));
+    String expected = "city%n" + "@Seattle%n" + "\"++Seattle|||\"";
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty = "city          %n" + "@Seattle      %n" + "\"++Seattle|||\"";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }
 
   @Test
@@ -183,12 +203,23 @@ void replaceNullValues() {
                     ImmutableMap.of("firstname", LITERAL_NULL, "city", stringValue("Seattle"))),
                 ExprTupleValue.fromExprValueMap(
                     ImmutableMap.of("firstname", stringValue("John"), "city", LITERAL_MISSING))));
-    String expected = "name|city%nJohn|Seattle%n|Seattle%nJohn|";
-    assertEquals(format(expected), rawFormatter.format(response));
+    String expected = "name|city%n" + "John|Seattle%n" + "|Seattle%n" + "John|";
+    assertEquals(format(expected), getRawFormatter().format(response));
+    String expectedPretty = "name|city   %n" + "John|Seattle%n" + "    |Seattle%n" + "John|       ";
+    assertEquals(format(expectedPretty), getRawFormatterPretty().format(response));
   }
 
   @Test
   void testContentType() {
-    assertEquals(rawFormatter.contentType(), CONTENT_TYPE);
+    assertEquals(getRawFormatter().contentType(), CONTENT_TYPE);
+    assertEquals(getRawFormatterPretty().contentType(), CONTENT_TYPE);
+  }
+
+  private FlatResponseFormatter getRawFormatter() {
+    return new RawResponseFormatter();
+  }
+
+  private FlatResponseFormatter getRawFormatterPretty() {
+    return new RawResponseFormatter(true);
   }
 }
diff --git a/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java b/sql/src/main/java/org/opensearch/sql/sql/domain/SQLQueryRequest.java
@@ -29,6 +29,7 @@ public class SQLQueryRequest {
       Set.of("query", "fetch_size", "parameters", QUERY_FIELD_CURSOR);
   private static final String QUERY_PARAMS_FORMAT = "format";
   private static final String QUERY_PARAMS_SANITIZE = "sanitize";
+  private static final String QUERY_PARAMS_PRETTY = "pretty";
 
   /** JSON payload in REST request. */
   private final JSONObject jsonContent;
@@ -49,6 +50,10 @@ public class SQLQueryRequest {
   @Accessors(fluent = true)
   private boolean sanitize = true;
 
+  @Getter
+  @Accessors(fluent = true)
+  private boolean pretty = false;
+
   private String cursor;
 
   /** Constructor of SQLQueryRequest that passes request params. */
@@ -64,6 +69,7 @@ public SQLQueryRequest(
     this.params = params;
     this.format = getFormat(params);
     this.sanitize = shouldSanitize(params);
+    this.pretty = shouldPretty(params);
     this.cursor = cursor;
   }
 
@@ -148,4 +154,11 @@ private boolean shouldSanitize(Map<String, String> params) {
     }
     return true;
   }
+
+  private boolean shouldPretty(Map<String, String> params) {
+    if (params.containsKey(QUERY_PARAMS_PRETTY)) {
+      return Boolean.parseBoolean(params.get(QUERY_PARAMS_PRETTY));
+    }
+    return false;
+  }
 }