SQL pagination should work with the pretty parameter

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

Hide file tree

Showing 2 changed files with 55 additions and 13 deletions.
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
@@ -6,10 +6,12 @@
 package org.opensearch.sql.sql.domain;
 
 import java.util.Collections;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.function.Predicate;
 import java.util.stream.Stream;
 import lombok.EqualsAndHashCode;
 import lombok.Getter;
@@ -29,6 +31,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;
@@ -79,21 +82,24 @@ public SQLQueryRequest(
    * @return true if supported.
    */
   public boolean isSupported() {
-    var noCursor = !isCursor();
-    var noQuery = query == null;
-    var noUnsupportedParams =
-        params.isEmpty() || (params.size() == 1 && params.containsKey(QUERY_PARAMS_FORMAT));
-    var noContent = jsonContent == null || jsonContent.isEmpty();
-
-    return ((!noCursor
-                && noQuery
-                && noUnsupportedParams
-                && noContent) // if cursor is given, but other things
-            || (noCursor && !noQuery)) // or if cursor is not given, but query
-        && isOnlySupportedFieldInPayload() // and request has supported fields only
-        && isSupportedFormat(); // and request is in supported format
+    boolean hasCursor = isCursor();
+    boolean hasQuery = query != null;
+    boolean hasContent = jsonContent != null && !jsonContent.isEmpty();
+    boolean hasUnsupportedParams =
+        (!params.isEmpty())
+            && params.keySet().stream().dropWhile(builtinSupported).findAny().isPresent();
+
+    boolean validCursor = hasCursor && !hasQuery && !hasUnsupportedParams && !hasContent;
+    boolean validQuery = !hasCursor && hasQuery;
+
+    return (validCursor || validQuery) // It's a valid cursor or a valid query
+        && isOnlySupportedFieldInPayload() // and request must contain supported fields only
+        && isSupportedFormat(); // and request must be a supported format
   }
 
+  private final Predicate<String> builtinSupported =
+      List.of(QUERY_PARAMS_FORMAT, QUERY_PARAMS_PRETTY)::contains;
+
   private boolean isCursor() {
     return cursor != null && !cursor.isEmpty();
   }

diff --git a/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java b/sql/src/test/java/org/opensearch/sql/sql/domain/SQLQueryRequestTest.java
@@ -106,6 +106,42 @@ public void should_support_cursor_request() {
         () -> assertTrue(cursorRequest.isSupported()));
   }
 
+  @Test
+  public void should_support_cursor_request_with_builtin_parameters() {
+    SQLQueryRequest fetchSizeRequest =
+        SQLQueryRequestBuilder.request("SELECT 1")
+            .jsonContent("{\"query\": \"SELECT 1\", \"fetch_size\": 5}")
+            .build();
+
+    SQLQueryRequest cursorRequest =
+        SQLQueryRequestBuilder.request(null)
+            .cursor("abcdefgh...")
+            .params(Map.of("format", "csv", "pretty", "true"))
+            .build();
+
+    assertAll(
+        () -> assertTrue(fetchSizeRequest.isSupported()),
+        () -> assertTrue(cursorRequest.isSupported()));
+  }
+
+  @Test
+  public void should_not_support_cursor_request_with_unsupported_parameters() {
+    SQLQueryRequest fetchSizeRequest =
+        SQLQueryRequestBuilder.request("SELECT 1")
+            .jsonContent("{\"query\": \"SELECT 1\", \"fetch_size\": 5}")
+            .build();
+
+    SQLQueryRequest cursorRequest =
+        SQLQueryRequestBuilder.request(null)
+            .cursor("abcdefgh...")
+            .params(Map.of("one", "two"))
+            .build();
+
+    assertAll(
+        () -> assertTrue(fetchSizeRequest.isSupported()),
+        () -> assertFalse(cursorRequest.isSupported()));
+  }
+
   @Test
   public void should_support_cursor_close_request() {
     SQLQueryRequest closeRequest =