diff --git a/common/build.gradle b/common/build.gradle index cdabc14f9b..59ba428729 100644 --- a/common/build.gradle +++ b/common/build.gradle @@ -12,4 +12,4 @@ dependencies { compile group: 'com.google.guava', name: 'guava', version:'23.0' testCompile group: 'junit', name: 'junit', version: '4.12' -} \ No newline at end of file +} diff --git a/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResponse.java b/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResponse.java new file mode 100644 index 0000000000..6b1d246f74 --- /dev/null +++ b/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/QueryResponse.java @@ -0,0 +1,83 @@ +/* + * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"). + * You may not use this file except in compliance with the License. + * A copy of the License is located at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * or in the "license" file accompanying this file. This file is distributed + * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either + * express or implied. See the License for the specific language governing + * permissions and limitations under the License. + * + */ + +package com.amazon.opendistroforelasticsearch.sql.protocol.response; + +import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValue; +import com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils; +import lombok.RequiredArgsConstructor; + +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.LinkedHashMap; +import java.util.Map; + +/** + * Query response that encapsulates query results and isolate expression related class from formatter implementation. + */ +@RequiredArgsConstructor +public class QueryResponse implements Iterable { + + /** + * Results which are collection of expression + */ + private final Collection exprValues; + + /** + * Parse column name from results + * @return mapping from column names to its expression type + */ + public Map columnNameTypes() { + if (exprValues.isEmpty()) { + return Collections.emptyMap(); + } + + Map tupleValue = getFirstTupleValue(); + return populateColumnNameAndTypes(tupleValue); + } + + private Map getFirstTupleValue() { + // Assume expression is always tuple on first level and columns of all tuples are exactly same + ExprValue firstValue = exprValues.iterator().next(); + return ExprValueUtils.getTupleValue(firstValue); + } + + private Map populateColumnNameAndTypes(Map tupleValue) { + Map colNameTypes = new LinkedHashMap<>(); // Maintain original order in tuple expression + tupleValue.forEach((name, expr) -> { + colNameTypes.put(name, expr.type().name().toLowerCase()); + }); + return colNameTypes; + } + + @Override + public Iterator iterator() { + // Any chance to avoid copy for json response generation? + return exprValues.stream(). + map(ExprValueUtils::getTupleValue). + map(Map::values). + map(this::convertExprValuesToValues). + iterator(); + } + + private Object[] convertExprValuesToValues(Collection exprValues) { + return exprValues.stream(). + map(ExprValue::value). + toArray(Object[]::new); + } + +} diff --git a/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java b/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java index cf9119deed..e0b62ae567 100644 --- a/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java +++ b/protocol/src/main/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/SimpleJsonResponseFormatter.java @@ -16,6 +16,7 @@ package com.amazon.opendistroforelasticsearch.sql.protocol.response.format; +import com.amazon.opendistroforelasticsearch.sql.protocol.response.QueryResponse; import lombok.Builder; import lombok.Getter; import lombok.RequiredArgsConstructor; @@ -40,18 +41,23 @@ * } * */ -public class SimpleJsonResponseFormatter extends JsonResponseFormatter> { +public class SimpleJsonResponseFormatter extends JsonResponseFormatter { public SimpleJsonResponseFormatter(Style style) { super(style); } @Override - public Object buildJsonObject(List objects) { - return JsonResponse.builder(). - column(new Column("firstname")). - row(new DataRow(new Object[]{"John"})). - build(); + public Object buildJsonObject(QueryResponse response) { + JsonResponse.JsonResponseBuilder json = JsonResponse.builder(); + response.columnNameTypes().forEach((name, type) -> { + json.column(new Column(name, type)); + }); + + for (Object[] values : response) { + json.row(new DataRow(values)); + } + return json.build(); } @Override @@ -76,6 +82,7 @@ public static class JsonResponse { @Getter public static class Column { private final String name; + private final String type; } @RequiredArgsConstructor diff --git a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/JsonResponseFormatterTest.java b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/JsonResponseFormatterTest.java index a47e3dfc53..80ff8fdd9b 100644 --- a/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/JsonResponseFormatterTest.java +++ b/protocol/src/test/java/com/amazon/opendistroforelasticsearch/sql/protocol/response/format/JsonResponseFormatterTest.java @@ -16,10 +16,14 @@ package com.amazon.opendistroforelasticsearch.sql.protocol.response.format; +import com.amazon.opendistroforelasticsearch.sql.protocol.response.QueryResponse; +import com.google.common.collect.ImmutableMap; import org.junit.jupiter.api.Test; +import java.util.Arrays; import java.util.Collections; +import static com.amazon.opendistroforelasticsearch.sql.data.model.ExprValueUtils.tupleValue; import static com.amazon.opendistroforelasticsearch.sql.protocol.response.format.JsonResponseFormatter.Style.COMPACT; import static com.amazon.opendistroforelasticsearch.sql.protocol.response.format.JsonResponseFormatter.Style.PRETTY; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -28,22 +32,49 @@ class JsonResponseFormatterTest { @Test public void formatResponse() { + QueryResponse response = new QueryResponse(Arrays.asList( + tupleValue(ImmutableMap.of("firstname", "John", "age", 20)), + tupleValue(ImmutableMap.of("firstname", "Smith", "age", 30)) + )); SimpleJsonResponseFormatter formatter = new SimpleJsonResponseFormatter(COMPACT); assertEquals( - "{\"schema\":[{\"name\":\"firstname\"}],\"datarows\":[{\"row\":[\"John\"]}]}", - formatter.format(Collections.emptyList()) + "{\"schema\":[{\"name\":\"firstname\",\"type\":\"string\"},{\"name\":\"age\",\"type\":\"integer\"}]," + + "\"datarows\":[{\"row\":[\"John\",20]},{\"row\":[\"Smith\",30]}]}", + formatter.format(response) ); } @Test public void formatResponsePretty() { + QueryResponse response = new QueryResponse(Arrays.asList( + tupleValue(ImmutableMap.of("firstname", "John", "age", 20)), + tupleValue(ImmutableMap.of("firstname", "Smith", "age", 30)) + )); SimpleJsonResponseFormatter formatter = new SimpleJsonResponseFormatter(PRETTY); assertEquals( "{\n" + - " \"schema\": [{\"name\": \"firstname\"}],\n" + - " \"datarows\": [{\"row\": [\"John\"]}]\n" + + " \"schema\": [\n" + + " {\n" + + " \"name\": \"firstname\",\n" + + " \"type\": \"string\"\n" + + " },\n" + + " {\n" + + " \"name\": \"age\",\n" + + " \"type\": \"integer\"\n" + + " }\n" + + " ],\n" + + " \"datarows\": [\n" + + " {\"row\": [\n" + + " \"John\",\n" + + " 20\n" + + " ]},\n" + + " {\"row\": [\n" + + " \"Smith\",\n" + + " 30\n" + + " ]}\n" + + " ]\n" + "}", - formatter.format(Collections.emptyList()) + formatter.format(response) ); }