Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Backport 2.18] [Backport 2.x] SQL pagination should work with the pretty parameter #3104

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading