address comments

Signed-off-by: Lantao Jin <[email protected]>
opensearch-project · Jun 25, 2024 · c25aca6 · c25aca6
1 parent bb29410
commit c25aca6
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/integ-test/src/test/java/org/opensearch/sql/jdbc/CursorIT.java b/integ-test/src/test/java/org/opensearch/sql/jdbc/CursorIT.java
@@ -22,6 +22,7 @@
 import java.sql.ResultSet;
 import java.sql.Statement;
 import java.util.List;
+import java.util.Map;
 import javax.annotation.Nullable;
 import lombok.SneakyThrows;
 import org.json.JSONObject;
@@ -115,6 +116,8 @@ public void select_all_no_cursor() {
 
       var restResponse = executeRestQuery(query, null);
       assertEquals(rows, restResponse.getInt("total"));
+      var restPrettyResponse = executeRestQuery(query, null, Map.of("pretty", "true"));
+      assertEquals(rows, restPrettyResponse.getInt("total"));
     }
   }
 
@@ -133,6 +136,8 @@ public void select_count_all_no_cursor() {
 
       var restResponse = executeRestQuery(query, null);
       assertEquals(rows, restResponse.getInt("total"));
+      var restPrettyResponse = executeRestQuery(query, null, Map.of("pretty", "true"));
+      assertEquals(rows, restPrettyResponse.getInt("total"));
     }
   }
 
@@ -151,6 +156,8 @@ public void select_all_small_table_big_cursor() {
 
       var restResponse = executeRestQuery(query, null);
       assertEquals(rows, restResponse.getInt("total"));
+      var restPrettyResponse = executeRestQuery(query, null, Map.of("pretty", "true"));
+      assertEquals(rows, restPrettyResponse.getInt("total"));
     }
   }
 
@@ -169,6 +176,8 @@ public void select_all_small_table_small_cursor() {
 
       var restResponse = executeRestQuery(query, null);
       assertEquals(rows, restResponse.getInt("total"));
+      var restPrettyResponse = executeRestQuery(query, null, Map.of("pretty", "true"));
+      assertEquals(rows, restPrettyResponse.getInt("total"));
     }
   }
 
@@ -187,6 +196,8 @@ public void select_all_big_table_small_cursor() {
 
       var restResponse = executeRestQuery(query, null);
       assertEquals(rows, restResponse.getInt("total"));
+      var restPrettyResponse = executeRestQuery(query, null, Map.of("pretty", "true"));
+      assertEquals(rows, restPrettyResponse.getInt("total"));
     }
   }
 
@@ -205,6 +216,8 @@ public void select_all_big_table_big_cursor() {
 
       var restResponse = executeRestQuery(query, null);
       assertEquals(rows, restResponse.getInt("total"));
+      var restPrettyResponse = executeRestQuery(query, null, Map.of("pretty", "true"));
+      assertEquals(rows, restPrettyResponse.getInt("total"));
     }
   }
 
@@ -217,13 +230,19 @@ private static String getConnectionString() {
 
   @SneakyThrows
   protected JSONObject executeRestQuery(String query, @Nullable Integer fetch_size) {
+    return executeRestQuery(query, fetch_size, Map.of());
+  }
+
+  @SneakyThrows
+  protected JSONObject executeRestQuery(String query, @Nullable Integer fetch_size, Map<String, String> params) {
     Request request = new Request("POST", QUERY_API_ENDPOINT);
     if (fetch_size != null) {
       request.setJsonEntity(
           String.format("{ \"query\": \"%s\", \"fetch_size\": %d }", query, fetch_size));
     } else {
       request.setJsonEntity(String.format("{ \"query\": \"%s\" }", query));
     }
+    request.addParameters(params);
 
     RequestOptions.Builder restOptionsBuilder = RequestOptions.DEFAULT.toBuilder();
     restOptionsBuilder.addHeader("Content-Type", "application/json");

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,7 +6,6 @@
 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;
@@ -86,7 +85,7 @@ public boolean isSupported() {
     boolean hasQuery = query != null;
     boolean hasContent = jsonContent != null && !jsonContent.isEmpty();
 
-    Predicate<String> supportedParams = List.of(QUERY_PARAMS_FORMAT, QUERY_PARAMS_PRETTY)::contains;
+    Predicate<String> supportedParams = Set.of(QUERY_PARAMS_FORMAT, QUERY_PARAMS_PRETTY)::contains;
     boolean hasUnsupportedParams =
         (!params.isEmpty())
             && params.keySet().stream().dropWhile(supportedParams).findAny().isPresent();