Navigation Menu

Skip to content

Commit

Permalink
feat: add struct query parameters (#223)
Browse files Browse the repository at this point in the history
* feat: struct query parameters

* feat: unit and integration test case

* feat: adjust code

* feat: remove clirr and refactor code

* feat: additional test case

* feat: add asserts for values
  • Loading branch information
Praful Makani committed Mar 20, 2020
1 parent ec6328e commit c1997dd
Show file tree
Hide file tree
Showing 3 changed files with 231 additions and 2 deletions.
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();
}
}
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();
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
Expand Up @@ -1455,6 +1455,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()));
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

0 comments on commit c1997dd

Please sign in to comment.