Skip to content

Commit

Permalink
fix: add support for repeated record query parameters (#2698)
Browse files Browse the repository at this point in the history 
Add support for array of struct query parameters which are used to query repeated record fields

Fixes #2485 ☕️
  • Loading branch information
@obada-ab
obada-ab committed May 25, 2023
1 parent 499af3f commit 51aff50
Show file tree
Hide file tree
Showing 7 changed files with 580 additions and 5 deletions.
1 change: 1 addition & 0 deletions README.md
View file
Expand Up @@ -127,6 +127,7 @@ Samples are in the [`samples/`](https://github.com/googleapis/java-bigquery/tree
| Copy Multiple Tables | [source code](https://github.com/googleapis/java-bigquery/blob/main/samples/snippets/src/main/java/com/example/bigquery/CopyMultipleTables.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-bigquery&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/bigquery/CopyMultipleTables.java) |
| Copy Table | [source code](https://github.com/googleapis/java-bigquery/blob/main/samples/snippets/src/main/java/com/example/bigquery/CopyTable.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-bigquery&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/bigquery/CopyTable.java) |
| Copy Table Cmek | [source code](https://github.com/googleapis/java-bigquery/blob/main/samples/snippets/src/main/java/com/example/bigquery/CopyTableCmek.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-bigquery&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/bigquery/CopyTableCmek.java) |
| Create And Query Repeated Record Field | [source code](https://github.com/googleapis/java-bigquery/blob/main/samples/snippets/src/main/java/com/example/bigquery/CreateAndQueryRepeatedRecordField.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-bigquery&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/bigquery/CreateAndQueryRepeatedRecordField.java) |
| Create Clustered Table | [source code](https://github.com/googleapis/java-bigquery/blob/main/samples/snippets/src/main/java/com/example/bigquery/CreateClusteredTable.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-bigquery&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/bigquery/CreateClusteredTable.java) |
| Create Dataset | [source code](https://github.com/googleapis/java-bigquery/blob/main/samples/snippets/src/main/java/com/example/bigquery/CreateDataset.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-bigquery&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/bigquery/CreateDataset.java) |
| Create Dataset Aws | [source code](https://github.com/googleapis/java-bigquery/blob/main/samples/snippets/src/main/java/com/example/bigquery/CreateDatasetAws.java) | [![Open in Cloud Shell][shell_img]](https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/googleapis/java-bigquery&page=editor&open_in_editor=samples/snippets/src/main/java/com/example/bigquery/CreateDatasetAws.java) |
Expand Down
18 changes: 14 additions & 4 deletions google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryParameterValue.java
View file
Expand Up @@ -349,7 +349,11 @@ public static <T> QueryParameterValue array(T[] array, Class<T> clazz) {
public static <T> QueryParameterValue array(T[] array, StandardSQLTypeName type) {
List<QueryParameterValue> listValues = new ArrayList<>();
for (T obj : array) {
listValues.add(QueryParameterValue.of(obj, type));
if (type == StandardSQLTypeName.STRUCT) {
listValues.add((QueryParameterValue) obj);
} else {
listValues.add(QueryParameterValue.of(obj, type));
}
}
return QueryParameterValue.newBuilder()
.setArrayValues(listValues)
Expand Down Expand Up @@ -522,9 +526,15 @@ QueryParameterType toTypePb() {
QueryParameterType typePb = new QueryParameterType();
typePb.setType(getType().toString());
if (getArrayType() != null) {
QueryParameterType arrayTypePb = new QueryParameterType();
arrayTypePb.setType(getArrayType().toString());
typePb.setArrayType(arrayTypePb);
List<QueryParameterValue> values = getArrayValues();
if (getArrayType() == StandardSQLTypeName.STRUCT && values != null && values.size() != 0) {
QueryParameterType structType = values.get(0).toTypePb();
typePb.setArrayType(structType);
} else {
QueryParameterType arrayTypePb = new QueryParameterType();
arrayTypePb.setType(getArrayType().toString());
typePb.setArrayType(arrayTypePb);
}
}
if (getStructTypes() != null) {
List<QueryParameterType.StructTypes> structTypes = new ArrayList<>();
Expand Down
43 changes: 43 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/QueryParameterValueTest.java
View file
Expand Up @@ -28,6 +28,7 @@
import java.math.BigDecimal;
import java.text.ParseException;
import java.time.Period;
import java.util.ArrayList;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
Expand Down Expand Up @@ -563,6 +564,48 @@ public void testNestedStruct() {
assertThat(nestedRecordField.getStructValues().size()).isEqualTo(structValue.size());
}

@Test
public void testStructArray() {
Boolean[] boolValues = new Boolean[] {true, false};
Integer[] intValues = new Integer[] {15, 20};
String[] stringValues = new String[] {"test-string", "test-string2"};
List<ImmutableMap<String, QueryParameterValue>> fieldMaps = new ArrayList<>();
List<QueryParameterValue> tuples = new ArrayList<>();
for (int i = 0; i < 2; i++) {
QueryParameterValue booleanField = QueryParameterValue.bool(boolValues[i]);
QueryParameterValue integerField = QueryParameterValue.int64(intValues[i]);
QueryParameterValue stringField = QueryParameterValue.string(stringValues[i]);
ImmutableMap<String, QueryParameterValue> fieldMap =
ImmutableMap.of(
"booleanField",
booleanField,
"integerField",
integerField,
"stringField",
stringField);
fieldMaps.add(fieldMap);
QueryParameterValue recordField = QueryParameterValue.struct(fieldMap);
tuples.add(recordField);
}
QueryParameterValue repeatedRecordField =
QueryParameterValue.array(tuples.toArray(), StandardSQLTypeName.STRUCT);
com.google.api.services.bigquery.model.QueryParameterValue parameterValue =
repeatedRecordField.toValuePb();
QueryParameterType parameterType = repeatedRecordField.toTypePb();
QueryParameterValue queryParameterValue =
QueryParameterValue.fromPb(parameterValue, parameterType);
assertThat(queryParameterValue.getValue()).isNull();
assertThat(queryParameterValue.getType()).isEqualTo(StandardSQLTypeName.ARRAY);
assertThat(queryParameterValue.getArrayType()).isEqualTo(StandardSQLTypeName.STRUCT);
assertThat(queryParameterValue.getArrayValues().size()).isEqualTo(2);
for (int i = 0; i < 2; i++) {
QueryParameterValue record = queryParameterValue.getArrayValues().get(i);
assertThat(record.getType()).isEqualTo(StandardSQLTypeName.STRUCT);
assertThat(record.getStructTypes()).isNotNull();
assertThat(record.getStructValues()).isEqualTo(fieldMaps.get(i));
}
}

private static void assertArrayDataEquals(
String[] expectedValues,
StandardSQLTypeName expectedType,
Expand Down
246 changes: 246 additions & 0 deletions google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
View file
Expand Up @@ -156,6 +156,7 @@
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
Expand Down Expand Up @@ -231,6 +232,31 @@ public class ITBigQueryTest {
.setMode(Field.Mode.REQUIRED)
.setDescription("RecordDescription")
.build();

private static final Field REPEATED_RECORD_FIELD_SCHEMA =
Field.newBuilder(
"Addresses",
LegacySQLTypeName.RECORD,
Field.newBuilder("Status", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build(),
Field.newBuilder("Address", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build(),
Field.newBuilder("City", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build(),
Field.newBuilder("State", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build(),
Field.newBuilder("Zip", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build(),
Field.newBuilder("NumberOfYears", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build())
.setMode(Field.Mode.REPEATED)
.build();
private static final Field INTEGER_FIELD_SCHEMA =
Field.newBuilder("IntegerField", LegacySQLTypeName.INTEGER)
.setMode(Field.Mode.NULLABLE)
Expand Down Expand Up @@ -422,6 +448,18 @@ public class ITBigQueryTest {
.setMode(Field.Mode.NULLABLE)
.build());

private static final Schema REPEATED_RECORD_TABLE_SCHEMA =
Schema.of(
Field.newBuilder("ID", LegacySQLTypeName.STRING).setMode(Field.Mode.NULLABLE).build(),
Field.newBuilder("FirstName", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build(),
Field.newBuilder("LastName", LegacySQLTypeName.STRING)
.setMode(Field.Mode.NULLABLE)
.build(),
Field.newBuilder("DOB", LegacySQLTypeName.DATE).setMode(Field.Mode.NULLABLE).build(),
REPEATED_RECORD_FIELD_SCHEMA);

private static final Schema SIMPLE_SCHEMA = Schema.of(STRING_FIELD_SCHEMA);
private static final Schema QUERY_RESULT_SCHEMA =
Schema.of(
Expand Down Expand Up @@ -4062,6 +4100,214 @@ public void testStructNamedQueryParameters() throws InterruptedException {
}
}

@Test
public void testRepeatedRecordNamedQueryParameters() throws InterruptedException {
String[] stringValues = new String[] {"test-stringField", "test-stringField2"};
List<QueryParameterValue> tuples = new ArrayList<>();
for (int i = 0; i < 2; i++) {
QueryParameterValue stringValue = QueryParameterValue.string(stringValues[i]);
Map<String, QueryParameterValue> struct = new HashMap<>();
struct.put("stringField", stringValue);
QueryParameterValue recordValue = QueryParameterValue.struct(struct);
tuples.add(recordValue);
}

QueryParameterValue repeatedRecord =
QueryParameterValue.array(tuples.toArray(), StandardSQLTypeName.STRUCT);
String query = "SELECT @repeatedRecordField AS repeatedRecord";
QueryJobConfiguration config =
QueryJobConfiguration.newBuilder(query)
.setDefaultDataset(DATASET)
.setUseLegacySql(false)
.addNamedParameter("repeatedRecordField", repeatedRecord)
.build();
TableResult result = bigquery.query(config);
assertEquals(1, Iterables.size(result.getValues()));

FieldList subSchema = result.getSchema().getFields().get("repeatedRecord").getSubFields();
for (FieldValueList values : result.iterateAll()) {
for (FieldValue value : values) {
assertEquals(FieldValue.Attribute.REPEATED, value.getAttribute());
assertEquals(2, value.getRepeatedValue().size());
for (int i = 0; i < 2; i++) {
FieldValue record = value.getRepeatedValue().get(i);
assertEquals(FieldValue.Attribute.RECORD, record.getAttribute());
FieldValueList recordValue = record.getRecordValue();
assertEquals(
stringValues[i],
FieldValueList.of(recordValue, subSchema).get("stringField").getValue());
}
}
}
}

@Test
public void testUnnestRepeatedRecordNamedQueryParameter() throws InterruptedException {
Boolean[] boolValues = new Boolean[] {true, false};
List<QueryParameterValue> tuples = new ArrayList<>();
for (int i = 0; i < 2; i++) {
QueryParameterValue boolValue = QueryParameterValue.bool(boolValues[i]);
Map<String, QueryParameterValue> struct = new HashMap<>();
struct.put("boolField", boolValue);
QueryParameterValue recordValue = QueryParameterValue.struct(struct);
tuples.add(recordValue);
}

QueryParameterValue repeatedRecord =
QueryParameterValue.array(tuples.toArray(), StandardSQLTypeName.STRUCT);
String query =
"SELECT * FROM (SELECT STRUCT("
+ boolValues[0]
+ " AS boolField) AS repeatedRecord) WHERE repeatedRecord IN UNNEST(@repeatedRecordField)";
QueryJobConfiguration config =
QueryJobConfiguration.newBuilder(query)
.setDefaultDataset(DATASET)
.setUseLegacySql(false)
.addNamedParameter("repeatedRecordField", repeatedRecord)
.build();
TableResult result = bigquery.query(config);
assertEquals(1, Iterables.size(result.getValues()));

FieldList subSchema = result.getSchema().getFields().get("repeatedRecord").getSubFields();
for (FieldValueList values : result.iterateAll()) {
for (FieldValue value : values) {
assertEquals(FieldValue.Attribute.RECORD, value.getAttribute());
FieldValueList recordValue = value.getRecordValue();
assertEquals(
boolValues[0],
FieldValueList.of(recordValue, subSchema).get("boolField").getBooleanValue());
}
}
}

@Test
public void testUnnestRepeatedRecordNamedQueryParameterFromDataset() throws InterruptedException {
TableId tableId = TableId.of(DATASET, "test_repeated_record_table");
setUpRepeatedRecordTable(tableId);

List<QueryParameterValue> tuples = new ArrayList<>();
QueryParameterValue statusValue = QueryParameterValue.string("single");
QueryParameterValue addressValue = QueryParameterValue.string("123 this lane");
QueryParameterValue cityValue = QueryParameterValue.string("Toronto");
QueryParameterValue stateValue = QueryParameterValue.string("ON");
QueryParameterValue zipValue = QueryParameterValue.string("1h2j34");
QueryParameterValue numberOfYearsValue = QueryParameterValue.string("3");

Map<String, QueryParameterValue> struct = new LinkedHashMap<>();
struct.put("statusValue", statusValue);
struct.put("addressValue", addressValue);
struct.put("cityValue", cityValue);
struct.put("stateValue", stateValue);
struct.put("zipValue", zipValue);
struct.put("numberOfYearsValue", numberOfYearsValue);
QueryParameterValue recordValue = QueryParameterValue.struct(struct);
tuples.add(recordValue);

QueryParameterValue repeatedRecord =
QueryParameterValue.array(tuples.toArray(), StandardSQLTypeName.STRUCT);

String query =
"SELECT * FROM "
+ tableId.getTable()
+ ", UNNEST(@repeatedRecord) AS TEMP where TEMP IN UNNEST(addresses);";
QueryJobConfiguration queryConfig =
QueryJobConfiguration.newBuilder(query)
.setDefaultDataset(DATASET)
.setUseLegacySql(false)
.addNamedParameter("repeatedRecord", repeatedRecord)
.build();
TableResult results = bigquery.query(queryConfig);

assertEquals(1, Iterables.size(results.getValues()));
for (FieldValueList values : results.iterateAll()) {
assertEquals("1", values.get("ID").getStringValue());
assertEquals("first_name1", values.get("FirstName").getStringValue());
assertEquals(2, values.get("Addresses").getRecordValue().size());
}
}

private void setUpRepeatedRecordTable(TableId tableId) {
StandardTableDefinition tableDefinition =
StandardTableDefinition.of(REPEATED_RECORD_TABLE_SCHEMA);
TableInfo tableInfo = TableInfo.of(tableId, tableDefinition);
bigquery.create(tableInfo);

ImmutableMap.Builder<String, Object> builder1 = ImmutableMap.builder();
builder1.put("ID", "1");
builder1.put("FirstName", "first_name1");
builder1.put("LastName", "last_name1");
builder1.put("DOB", "1995-08-09");
builder1.put(
"Addresses",
ImmutableList.of(
ImmutableMap.of(
"Status", "single",
"Address", "123 this lane",
"City", "Toronto",
"State", "ON",
"Zip", "1h2j34",
"NumberOfYears", "3"),
ImmutableMap.of(
"Status", "couple",
"Address", "345 that lane",
"City", "Maple",
"State", "ON",
"Zip", "1h2j34",
"NumberOfYears", "5")));

ImmutableMap.Builder<String, Object> builder2 = ImmutableMap.builder();
builder2.put("ID", "2");
builder2.put("FirstName", "first_name2");
builder2.put("LastName", "last_name2");
builder2.put("DOB", "1992-03-19");
builder2.put(
"Addresses",
ImmutableList.of(
ImmutableMap.of(
"Status", "single",
"Address", "97 Kota lane",
"City", "Ottawa",
"State", "ON",
"Zip", "1h2j34",
"NumberOfYears", "3"),
ImmutableMap.of(
"Status", "couple",
"Address", "75 Malta lane",
"City", "Victoria",
"State", "AL",
"Zip", "1h2j34",
"NumberOfYears", "5")));

InsertAllRequest request =
InsertAllRequest.newBuilder(tableInfo.getTableId())
.addRow(builder1.build())
.addRow(builder2.build())
.build();
bigquery.insertAll(request);
}

@Test
public void testEmptyRepeatedRecordNamedQueryParameters() throws InterruptedException {
QueryParameterValue[] tuples = {};

QueryParameterValue repeatedRecord =
QueryParameterValue.array(tuples, StandardSQLTypeName.STRUCT);
String query =
"SELECT * FROM (SELECT STRUCT(false AS boolField) AS repeatedRecord) WHERE repeatedRecord IN UNNEST(@repeatedRecordField)";
QueryJobConfiguration config =
QueryJobConfiguration.newBuilder(query)
.setDefaultDataset(DATASET)
.setUseLegacySql(false)
.addNamedParameter("repeatedRecordField", repeatedRecord)
.build();
try {
bigquery.query(config);
fail("an empty array of struct query parameter shouldn't work with 'IN UNNEST'");
} catch (BigQueryException e) {
// Nothing to do
}
}

@Test
public void testStructQuery() throws InterruptedException {
// query into a table
Expand Down

0 comments on commit 51aff50

Please sign in to comment.