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 index 4e902cb67d..1cd09a103b 100644 --- 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 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 index 2b64b13b35..c92af990c2 100644 --- 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 =