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

feat: add struct query parameters #223

Merged
merged 6 commits into from
Mar 20, 2020
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 @@ -26,13 +26,16 @@
import com.google.cloud.Timestamp;
import com.google.common.base.Function;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.google.common.io.BaseEncoding;
import java.io.Serializable;
import java.math.BigDecimal;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import javax.annotation.Nullable;
import org.threeten.bp.Instant;
import org.threeten.bp.ZoneOffset;
Expand Down Expand Up @@ -127,12 +130,27 @@ public Builder setArrayValues(List<QueryParameterValue> arrayValues) {

abstract Builder setArrayValuesInner(ImmutableList<QueryParameterValue> arrayValues);

/** Sets struct values. The type must set to STRUCT. */
public Builder setStructValues(Map<String, QueryParameterValue> structValues) {
setStructTypes(ImmutableMap.copyOf(structValues));
return setStructValuesInner(ImmutableMap.copyOf(structValues));
}

abstract Builder setStructValuesInner(Map<String, QueryParameterValue> structValues);

/** Sets the parameter data type. */
public abstract Builder setType(StandardSQLTypeName type);

/** Sets the data type of the array elements. The type must set to ARRAY. */
public abstract Builder setArrayType(StandardSQLTypeName arrayType);

/** Sets the data type of the struct elements. The type must set to STRUCT. */
public Builder setStructTypes(Map<String, QueryParameterValue> structTypes) {
return setStructTypesInner(structTypes);
}

abstract Builder setStructTypesInner(Map<String, QueryParameterValue> structTypes);

/** Creates a {@code QueryParameterValue} object. */
public abstract QueryParameterValue build();
}
Expand All @@ -154,13 +172,31 @@ public List<QueryParameterValue> getArrayValues() {
@Nullable
abstract ImmutableList<QueryParameterValue> getArrayValuesInner();

/** Returns the struct values of this parameter. The returned map, if not null, is immutable. */
@Nullable
public Map<String, QueryParameterValue> getStructValues() {
return getStructValuesInner();
}

@Nullable
abstract Map<String, QueryParameterValue> getStructValuesInner();

/** Returns the data type of this parameter. */
public abstract StandardSQLTypeName getType();

/** Returns the data type of the array elements. */
@Nullable
public abstract StandardSQLTypeName getArrayType();

/** Returns the data type of the struct elements. */
@Nullable
public Map<String, QueryParameterValue> getStructTypes() {
return getStructTypesInner();
}

@Nullable
abstract Map<String, QueryParameterValue> getStructTypesInner();

/** Creates a {@code QueryParameterValue} object with the given value and type. */
public static <T> QueryParameterValue of(T value, Class<T> type) {
return of(value, classToType(type));
Expand Down Expand Up @@ -274,6 +310,17 @@ public static <T> QueryParameterValue array(T[] array, StandardSQLTypeName type)
.build();
}

/**
* Creates a map with {@code QueryParameterValue} object and a type of STRUCT the given struct
* element type.
*/
public static QueryParameterValue struct(Map<String, QueryParameterValue> struct) {
return QueryParameterValue.newBuilder()
.setStructValues(struct)
.setType(StandardSQLTypeName.STRUCT)
.build();
}

private static <T> StandardSQLTypeName classToType(Class<T> type) {
if (Boolean.class.isAssignableFrom(type)) {
return StandardSQLTypeName.BOOL;
Expand Down Expand Up @@ -398,6 +445,14 @@ com.google.api.services.bigquery.model.QueryParameterValue toValuePb() {
valuePb.setArrayValues(
Lists.transform(getArrayValues(), QueryParameterValue.TO_VALUE_PB_FUNCTION));
}
if (getStructValues() != null) {
Map<String, com.google.api.services.bigquery.model.QueryParameterValue> structValues =
new HashMap<>();
for (Map.Entry<String, QueryParameterValue> structValue : getStructValues().entrySet()) {
structValues.put(structValue.getKey(), structValue.getValue().toValuePb());
}
valuePb.setStructValues(structValues);
}
return valuePb;
}

Expand All @@ -409,14 +464,24 @@ QueryParameterType toTypePb() {
arrayTypePb.setType(getArrayType().toString());
typePb.setArrayType(arrayTypePb);
}
if (getStructTypes() != null) {
List<QueryParameterType.StructTypes> structTypes = new ArrayList<>();
for (Map.Entry<String, QueryParameterValue> entry : getStructTypes().entrySet()) {
QueryParameterType.StructTypes structType = new QueryParameterType.StructTypes();
structType.setName(entry.getKey());
structType.setType(entry.getValue().toTypePb());
structTypes.add(structType);
}
typePb.setStructTypes(structTypes);
}
return typePb;
}

static QueryParameterValue fromPb(
com.google.api.services.bigquery.model.QueryParameterValue valuePb,
QueryParameterType typePb) {
Builder valueBuilder = newBuilder();

Map<String, QueryParameterType> parameterTypes = new HashMap<>();
StandardSQLTypeName type = StandardSQLTypeName.valueOf(typePb.getType());
valueBuilder.setType(type);
if (type == StandardSQLTypeName.ARRAY) {
Expand All @@ -431,10 +496,35 @@ static QueryParameterValue fromPb(
}
valueBuilder.setArrayValues(arrayValues.build());
}
} else if (type == StandardSQLTypeName.STRUCT) {
Map<String, QueryParameterValue> structTypes = new HashMap<>();
for (QueryParameterType.StructTypes types : typePb.getStructTypes()) {
structTypes.put(
types.getName(),
QueryParameterValue.newBuilder()
.setType(StandardSQLTypeName.valueOf(types.getType().getType()))
.build());
}
valueBuilder.setStructTypes(structTypes);
if (valuePb == null || valuePb.getStructValues() == null) {
valueBuilder.setStructValues(ImmutableMap.<String, QueryParameterValue>of());
} else {
Map<String, QueryParameterValue> structValues = new HashMap<>();
for (QueryParameterType.StructTypes structType : typePb.getStructTypes()) {
parameterTypes.put(structType.getName(), structType.getType());
}
for (Map.Entry<String, com.google.api.services.bigquery.model.QueryParameterValue>
structValue : valuePb.getStructValues().entrySet()) {
structValues.put(
structValue.getKey(),
QueryParameterValue.fromPb(
structValue.getValue(), parameterTypes.get(structValue.getKey())));
}
valueBuilder.setStructValues(structValues);
}
} else {
valueBuilder.setValue(valuePb == null ? "" : valuePb.getValue());
}

return valueBuilder.build();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,10 +23,13 @@
import static org.threeten.bp.temporal.ChronoField.SECOND_OF_MINUTE;

import com.google.api.services.bigquery.model.QueryParameterType;
import com.google.common.collect.ImmutableMap;
import java.math.BigDecimal;
import java.text.ParseException;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.junit.Test;
import org.threeten.bp.Instant;
import org.threeten.bp.ZoneOffset;
Expand Down Expand Up @@ -378,6 +381,68 @@ public void testFromEmptyArray() {
assertThat(value.getArrayValues()).isEmpty();
}

@Test
public void testStruct() {
QueryParameterValue booleanField = QueryParameterValue.bool(true);
QueryParameterValue integerField = QueryParameterValue.int64(15);
QueryParameterValue stringField = QueryParameterValue.string("test-string");
QueryParameterValue recordField =
QueryParameterValue.struct(
ImmutableMap.of(
"booleanField",
booleanField,
"integerField",
integerField,
"stringField",
stringField));
com.google.api.services.bigquery.model.QueryParameterValue parameterValue =
recordField.toValuePb();
QueryParameterType parameterType = recordField.toTypePb();
QueryParameterValue queryParameterValue =
QueryParameterValue.fromPb(parameterValue, parameterType);
assertThat(queryParameterValue).isEqualTo(recordField);
assertThat(recordField.getValue()).isNull();
assertThat(recordField.getType()).isEqualTo(StandardSQLTypeName.STRUCT);
assertThat(recordField.getStructTypes()).isNotNull();
stephaniewang526 marked this conversation as resolved.
Show resolved Hide resolved
assertThat(recordField.getStructValues()).isNotNull();
}

@Test
public void testNestedStruct() {
QueryParameterValue booleanField = QueryParameterValue.bool(true);
QueryParameterValue integerField = QueryParameterValue.int64(15);
QueryParameterValue stringField = QueryParameterValue.string("test-string");
QueryParameterValue recordField =
QueryParameterValue.struct(
ImmutableMap.of(
"booleanField",
booleanField,
"integerField",
integerField,
"stringField",
stringField));
Map<String, QueryParameterValue> structValue = new HashMap<>();
structValue.put("bool", booleanField);
structValue.put("int", integerField);
structValue.put("string", stringField);
structValue.put("struct", recordField);
QueryParameterValue nestedRecordField = QueryParameterValue.struct(structValue);
com.google.api.services.bigquery.model.QueryParameterValue parameterValue =
nestedRecordField.toValuePb();
QueryParameterType parameterType = nestedRecordField.toTypePb();
QueryParameterValue queryParameterValue =
QueryParameterValue.fromPb(parameterValue, parameterType);
assertThat(queryParameterValue).isEqualTo(nestedRecordField);
assertThat(nestedRecordField.getValue()).isNull();
assertThat(nestedRecordField.getType()).isEqualTo(StandardSQLTypeName.STRUCT);
assertThat(nestedRecordField.getStructTypes().get("struct").getType())
.isEqualTo(StandardSQLTypeName.STRUCT);
assertThat(nestedRecordField.getStructValues().get("struct").getStructValues())
.containsAtLeastEntriesIn(recordField.getStructValues());
assertThat(nestedRecordField.getStructTypes().size()).isEqualTo(structValue.size());
assertThat(nestedRecordField.getStructValues().size()).isEqualTo(structValue.size());
}

private static void assertArrayDataEquals(
String[] expectedValues,
StandardSQLTypeName expectedType,
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -1428,6 +1428,80 @@ public void testNamedQueryParameters() throws InterruptedException {
assertEquals(2, Iterables.size(result.getValues()));
}

@Test
public void testStructNamedQueryParameters() throws InterruptedException {
QueryParameterValue booleanValue = QueryParameterValue.bool(true);
QueryParameterValue stringValue = QueryParameterValue.string("test-stringField");
QueryParameterValue integerValue = QueryParameterValue.int64(10);
Map<String, QueryParameterValue> struct = new HashMap<>();
struct.put("booleanField", booleanValue);
struct.put("integerField", integerValue);
struct.put("stringField", stringValue);
QueryParameterValue recordValue = QueryParameterValue.struct(struct);
String query = "SELECT STRUCT(@recordField) AS record";
QueryJobConfiguration config =
QueryJobConfiguration.newBuilder(query)
.setDefaultDataset(DATASET)
.setUseLegacySql(false)
.addNamedParameter("recordField", recordValue)
.build();
TableResult result = bigquery.query(config);
assertEquals(1, Iterables.size(result.getValues()));
stephaniewang526 marked this conversation as resolved.
Show resolved Hide resolved
for (FieldValueList values : result.iterateAll()) {
for (FieldValue value : values) {
for (FieldValue record : value.getRecordValue()) {
assertEquals(FieldValue.Attribute.RECORD, record.getAttribute());
assertEquals(true, record.getRecordValue().get(0).getBooleanValue());
assertEquals(10, record.getRecordValue().get(1).getLongValue());
assertEquals("test-stringField", record.getRecordValue().get(2).getStringValue());
}
}
}
}

@Test
public void testNestedStructNamedQueryParameters() throws InterruptedException {
QueryParameterValue booleanValue = QueryParameterValue.bool(true);
QueryParameterValue stringValue = QueryParameterValue.string("test-stringField");
QueryParameterValue integerValue = QueryParameterValue.int64(10);
Map<String, QueryParameterValue> struct = new HashMap<>();
struct.put("booleanField", booleanValue);
struct.put("integerField", integerValue);
struct.put("stringField", stringValue);
QueryParameterValue recordValue = QueryParameterValue.struct(struct);
Map<String, QueryParameterValue> structValue = new HashMap<>();
structValue.put("bool", booleanValue);
structValue.put("int", integerValue);
structValue.put("string", stringValue);
structValue.put("struct", recordValue);
QueryParameterValue nestedRecordField = QueryParameterValue.struct(structValue);
String query = "SELECT STRUCT(@nestedRecordField) AS record";
QueryJobConfiguration config =
QueryJobConfiguration.newBuilder(query)
.setDefaultDataset(DATASET)
.setUseLegacySql(false)
.addNamedParameter("nestedRecordField", nestedRecordField)
.build();
TableResult result = bigquery.query(config);
assertEquals(1, Iterables.size(result.getValues()));
for (FieldValueList values : result.iterateAll()) {
for (FieldValue value : values) {
assertEquals(FieldValue.Attribute.RECORD, value.getAttribute());
for (FieldValue record : value.getRecordValue()) {
assertEquals(
true, record.getRecordValue().get(0).getRecordValue().get(0).getBooleanValue());
assertEquals(10, record.getRecordValue().get(0).getRecordValue().get(1).getLongValue());
assertEquals(
"test-stringField",
record.getRecordValue().get(0).getRecordValue().get(2).getStringValue());
assertEquals(true, record.getRecordValue().get(1).getBooleanValue());
assertEquals("test-stringField", record.getRecordValue().get(2).getStringValue());
assertEquals(10, record.getRecordValue().get(3).getLongValue());
}
}
}
}

@Test
public void testBytesParameter() throws Exception {
String query = "SELECT BYTE_LENGTH(@p) AS length";
Expand Down