Skip to content

Commit

Permalink
[Backport 2.18] SQL pagination should work with the pretty parameter (
Browse files Browse the repository at this point in the history
#3106)

(cherry picked from commit 0c2e1da)

Signed-off-by: Lantao Jin <ltjin@amazon.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
  • Loading branch information
1 parent 3cd7702 commit 28ecac3
Show file tree
Hide file tree
Showing 3 changed files with 72 additions and 13 deletions.
20 changes: 20 additions & 0 deletions integ-test/src/test/java/org/opensearch/sql/jdbc/CursorIT.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -116,6 +117,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"));
}
}

Expand All @@ -134,6 +137,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"));
}
}

Expand All @@ -152,6 +157,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"));
}
}

Expand All @@ -170,6 +177,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"));
}
}

Expand All @@ -188,6 +197,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"));
}
}

Expand All @@ -206,6 +217,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"));
}
}

Expand All @@ -218,13 +231,20 @@ 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");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
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;
Expand Down Expand Up @@ -81,19 +82,21 @@ 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();

Predicate<String> supportedParams = Set.of(QUERY_PARAMS_FORMAT, QUERY_PARAMS_PRETTY)::contains;
boolean hasUnsupportedParams =
(!params.isEmpty())
&& params.keySet().stream().dropWhile(supportedParams).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 boolean isCursor() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,42 @@ public void should_support_cursor_request() {
() -> assertTrue(cursorRequest.isSupported()));
}

@Test
public void should_support_cursor_request_with_supported_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 =
Expand Down

0 comments on commit 28ecac3

Please sign in to comment.