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.x] Well format the raw response when query parameter "pretty" enabled #2829

Merged
Merged
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
Original file line number Diff line number Diff line change
Expand Up @@ -244,12 +244,17 @@ protected Request getSqlCursorCloseRequest(String cursorRequest) {
}

protected String executeQuery(String query, String requestType) {
return executeQuery(query, requestType, Map.of());
}

protected String executeQuery(String query, String requestType, Map<String, String> params) {
try {
String endpoint = "/_plugins/_sql?format=" + requestType;
String requestBody = makeRequest(query);

Request sqlRequest = new Request("POST", endpoint);
sqlRequest.setJsonEntity(requestBody);
sqlRequest.addParameters(params);

Response response = client().performRequest(sqlRequest);
Assert.assertEquals(200, response.getStatusLine().getStatusCode());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
package org.opensearch.sql.sql;

import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK_CSV_SANITIZE;
import static org.opensearch.sql.protocol.response.format.FlatResponseFormatter.CONTENT_TYPE;
import static org.opensearch.sql.protocol.response.format.CsvResponseFormatter.CONTENT_TYPE;

import java.io.IOException;
import java.util.Locale;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,11 @@
package org.opensearch.sql.sql;

import static org.opensearch.sql.legacy.TestsConstants.TEST_INDEX_BANK_RAW_SANITIZE;
import static org.opensearch.sql.protocol.response.format.FlatResponseFormatter.CONTENT_TYPE;
import static org.opensearch.sql.protocol.response.format.RawResponseFormatter.CONTENT_TYPE;

import java.io.IOException;
import java.util.Locale;
import java.util.Map;
import org.junit.Test;
import org.opensearch.client.Request;
import org.opensearch.client.Response;
Expand Down Expand Up @@ -41,6 +42,25 @@ public void rawFormatWithPipeFieldTest() {
result);
}

@Test
public void rawFormatPrettyWithPipeFieldTest() {
String result =
executeQuery(
String.format(
Locale.ROOT, "SELECT firstname, lastname FROM %s", TEST_INDEX_BANK_RAW_SANITIZE),
"raw",
Map.of("pretty", "true"));
assertEquals(
StringUtils.format(
"firstname |lastname %n"
+ "+Amber JOHnny|Duke Willmington+%n"
+ "-Hattie |Bond- %n"
+ "=Nanette |Bates= %n"
+ "@Dale |Adams@ %n"
+ "@Elinor |\"Ratliff|||\" %n"),
result);
}

@Test
public void contentHeaderTest() throws IOException {
String query =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -169,7 +169,7 @@ private ResponseListener<QueryResponse> createQueryResponseListener(
} else if (format.equals(Format.CSV)) {
formatter = new CsvResponseFormatter(request.sanitize());
} else if (format.equals(Format.RAW)) {
formatter = new RawResponseFormatter();
formatter = new RawResponseFormatter(request.pretty());
} else {
formatter = new JdbcResponseFormatter(PRETTY);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,12 +5,45 @@

package org.opensearch.sql.protocol.response.format;

public class CsvResponseFormatter extends FlatResponseFormatter {
import org.opensearch.sql.protocol.response.QueryResult;

/** Response formatter to format response to csv format. */
public class CsvResponseFormatter implements ResponseFormatter<QueryResult> {
public static final String CONTENT_TYPE = "plain/text; charset=UTF-8";
private final String separator;
private final boolean sanitize;

public CsvResponseFormatter() {
super(",", true);
this(",", true);
}

public CsvResponseFormatter(boolean sanitize) {
super(",", sanitize);
this(",", sanitize);
}

public CsvResponseFormatter(String separator, boolean sanitize) {
this.separator = separator;
this.sanitize = sanitize;
}

@Override
public String format(QueryResult response) {
FlatResponseBase flatResponse;
if (sanitize) {
flatResponse = new FlatResponseWithSanitizer(response, separator);
} else {
flatResponse = new FlatResponseBase(response, separator);
}
return flatResponse.format();
}

@Override
public String format(Throwable t) {
return ErrorFormatter.prettyFormat(t);
}

@Override
public String contentType() {
return CONTENT_TYPE;
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.sql.protocol.response.format;

import com.google.common.collect.ImmutableList;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;
import lombok.Getter;
import org.opensearch.sql.protocol.response.QueryResult;

@Getter
public class FlatResponseBase {
protected static final String INTERLINE_SEPARATOR = System.lineSeparator();

private final QueryResult response;
protected final String separator;

private final List<String> headers;
private final List<List<String>> data;

FlatResponseBase(QueryResult response, String separator) {
this.response = response;
this.separator = separator;
this.headers = getOriginalHeaders(response);
this.data = getOriginalData(response);
}

public String format() {
List<String> headersAndData = new ArrayList<>();
headersAndData.add(getHeaderLine());
headersAndData.addAll(getDataLines());
return String.join(INTERLINE_SEPARATOR, headersAndData);
}

protected String getHeaderLine() {
return String.join(separator, headers);
}

private List<String> getOriginalHeaders(QueryResult response) {
ImmutableList.Builder<String> headers = ImmutableList.builder();
response.columnNameTypes().forEach((column, type) -> headers.add(column));
List<String> result = headers.build();
return formatHeaders(result);
}

protected List<String> getDataLines() {
return data.stream().map(v -> String.join(separator, v)).collect(Collectors.toList());
}

private List<List<String>> getOriginalData(QueryResult response) {
ImmutableList.Builder<List<String>> dataLines = new ImmutableList.Builder<>();
response
.iterator()
.forEachRemaining(
row -> {
ImmutableList.Builder<String> line = new ImmutableList.Builder<>();
// replace null values with empty string
Arrays.asList(row).forEach(val -> line.add(val == null ? "" : val.toString()));
dataLines.add(line.build());
});
List<List<String>> result = dataLines.build();
return formatData(result);
}

protected List<String> formatHeaders(List<String> headers) {
return headers.stream()
.map(cell -> quoteIfRequired(separator, cell))
.collect(Collectors.toList());
}

protected List<List<String>> formatData(List<List<String>> lines) {
List<List<String>> result = new ArrayList<>();
for (List<String> line : lines) {
result.add(
line.stream().map(cell -> quoteIfRequired(separator, cell)).collect(Collectors.toList()));
}
return result;
}

protected String quoteIfRequired(String separator, String cell) {
final String quote = "\"";
return cell.contains(separator) ? quote + cell.replaceAll("\"", "\"\"") + quote : cell;
}
}

This file was deleted.

Loading
Loading