Skip to content

Commit 10c1961

Browse files
stephaniewang526stephaniewang526
authored
feat: add support for AvroOptions (#1630)
New AvroOptions class and wired into ExternalTableDefinition
1 parent 6bdbc2d commit 10c1961

File tree

7 files changed

+198
-4
lines changed

7 files changed

+198
-4
lines changed

‎google-cloud-bigquery/clirr-ignored-differences.xml

+10
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<!-- see http://www.mojohaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
3+
<differences>
4+
<difference>
5+
<differenceType>7006</differenceType>
6+
<className>com/google/cloud/bigquery/FormatOptions</className>
7+
<method>com.google.cloud.bigquery.FormatOptions avro()</method>
8+
<to>com.google.cloud.bigquery.AvroOptions</to>
9+
</difference>
10+
</differences>

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/AvroOptions.java

+112
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import com.google.common.base.MoreObjects;
20+
import java.util.Objects;
21+
22+
/**
23+
* Google BigQuery options for AVRO format. This class wraps some properties of AVRO files used by
24+
* BigQuery to parse external data.
25+
*/
26+
public final class AvroOptions extends FormatOptions {
27+
28+
private static final long serialVersionUID = 2293570529308612712L;
29+
30+
private final Boolean useAvroLogicalTypes;
31+
32+
public static final class Builder {
33+
34+
private Boolean useAvroLogicalTypes;
35+
36+
private Builder() {}
37+
38+
private Builder(AvroOptions avroOptions) {
39+
this.useAvroLogicalTypes = avroOptions.useAvroLogicalTypes;
40+
}
41+
42+
/**
43+
* [Optional] Sets whether BigQuery should interpret logical types as the corresponding BigQuery
44+
* data type (for example, TIMESTAMP), instead of using the raw type (for example, INTEGER).
45+
*/
46+
public Builder setUseAvroLogicalTypes(boolean useAvroLogicalTypes) {
47+
this.useAvroLogicalTypes = useAvroLogicalTypes;
48+
return this;
49+
}
50+
51+
/** Creates a {@code AvroOptions} object. */
52+
public AvroOptions build() {
53+
return new AvroOptions(this);
54+
}
55+
}
56+
57+
private AvroOptions(Builder builder) {
58+
super(FormatOptions.AVRO);
59+
this.useAvroLogicalTypes = builder.useAvroLogicalTypes;
60+
}
61+
62+
/**
63+
* Returns whether BigQuery should interpret logical types as the corresponding BigQuery data type
64+
* (for example, TIMESTAMP), instead of using the raw type (for example, INTEGER).
65+
*/
66+
public Boolean useAvroLogicalTypes() {
67+
return useAvroLogicalTypes;
68+
}
69+
70+
public Builder toBuilder() {
71+
return new Builder(this);
72+
}
73+
74+
@Override
75+
public String toString() {
76+
return MoreObjects.toStringHelper(this)
77+
.add("type", getType())
78+
.add("useAvroLogicalTypes", useAvroLogicalTypes)
79+
.toString();
80+
}
81+
82+
@Override
83+
public int hashCode() {
84+
return Objects.hash(getType(), useAvroLogicalTypes);
85+
}
86+
87+
@Override
88+
public boolean equals(Object obj) {
89+
return obj == this
90+
|| obj instanceof AvroOptions && Objects.equals(toPb(), ((AvroOptions) obj).toPb());
91+
}
92+
93+
com.google.api.services.bigquery.model.AvroOptions toPb() {
94+
com.google.api.services.bigquery.model.AvroOptions avroOptions =
95+
new com.google.api.services.bigquery.model.AvroOptions();
96+
avroOptions.setUseAvroLogicalTypes(useAvroLogicalTypes);
97+
return avroOptions;
98+
}
99+
100+
/** Returns a builder for a AvroOptions object. */
101+
public static AvroOptions.Builder newBuilder() {
102+
return new AvroOptions.Builder();
103+
}
104+
105+
static AvroOptions fromPb(com.google.api.services.bigquery.model.AvroOptions avroOptions) {
106+
Builder builder = newBuilder();
107+
if (avroOptions.getUseAvroLogicalTypes() != null) {
108+
builder.setUseAvroLogicalTypes(avroOptions.getUseAvroLogicalTypes());
109+
}
110+
return builder.build();
111+
}
112+
}

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ExternalTableDefinition.java

+9
Original file line numberDiff line numberDiff line change
@@ -300,6 +300,9 @@ com.google.api.services.bigquery.model.ExternalDataConfiguration toExternalDataC
300300
if (getDecimalTargetTypes() != null) {
301301
externalConfigurationPb.setDecimalTargetTypes(getDecimalTargetTypes());
302302
}
303+
if (getFormatOptions() != null && FormatOptions.AVRO.equals(getFormatOptions().getType())) {
304+
externalConfigurationPb.setAvroOptions(((AvroOptions) getFormatOptions()).toPb());
305+
}
303306
if (getFormatOptions() != null && FormatOptions.CSV.equals(getFormatOptions().getType())) {
304307
externalConfigurationPb.setCsvOptions(((CsvOptions) getFormatOptions()).toPb());
305308
}
@@ -459,6 +462,9 @@ static ExternalTableDefinition fromPb(Table tablePb) {
459462
builder.setConnectionId(externalDataConfiguration.getConnectionId());
460463
}
461464
builder.setIgnoreUnknownValues(externalDataConfiguration.getIgnoreUnknownValues());
465+
if (externalDataConfiguration.getAvroOptions() != null) {
466+
builder.setFormatOptions(AvroOptions.fromPb(externalDataConfiguration.getAvroOptions()));
467+
}
462468
if (externalDataConfiguration.getCsvOptions() != null) {
463469
builder.setFormatOptions(CsvOptions.fromPb(externalDataConfiguration.getCsvOptions()));
464470
}
@@ -508,6 +514,9 @@ static ExternalTableDefinition fromExternalDataConfiguration(
508514
if (externalDataConfiguration.getIgnoreUnknownValues() != null) {
509515
builder.setIgnoreUnknownValues(externalDataConfiguration.getIgnoreUnknownValues());
510516
}
517+
if (externalDataConfiguration.getAvroOptions() != null) {
518+
builder.setFormatOptions(AvroOptions.fromPb(externalDataConfiguration.getAvroOptions()));
519+
}
511520
if (externalDataConfiguration.getCsvOptions() != null) {
512521
builder.setFormatOptions(CsvOptions.fromPb(externalDataConfiguration.getCsvOptions()));
513522
}

‎google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/FormatOptions.java

+4-2
Original file line numberDiff line numberDiff line change
@@ -91,8 +91,8 @@ public static FormatOptions datastoreBackup() {
9191
}
9292

9393
/** Default options for AVRO format. */
94-
public static FormatOptions avro() {
95-
return new FormatOptions(AVRO);
94+
public static AvroOptions avro() {
95+
return AvroOptions.newBuilder().build();
9696
}
9797

9898
/** Default options for BIGTABLE format. */
@@ -120,6 +120,8 @@ public static FormatOptions of(String format) {
120120
checkArgument(!isNullOrEmpty(format), "Provided format is null or empty");
121121
if (format.equals(CSV)) {
122122
return csv();
123+
} else if (format.equals(AVRO)) {
124+
return avro();
123125
} else if (format.equals(DATASTORE_BACKUP)) {
124126
return datastoreBackup();
125127
} else if (format.equals(GOOGLE_SHEETS)) {

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/AvroOptionsTest.java

+56
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
/*
2+
* Copyright 2021 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package com.google.cloud.bigquery;
18+
19+
import static org.junit.Assert.assertEquals;
20+
21+
import org.junit.Test;
22+
23+
public class AvroOptionsTest {
24+
25+
private static final Boolean USE_AVRO_LOGICAL_TYPES = true;
26+
private static final AvroOptions AVRO_OPTIONS =
27+
AvroOptions.newBuilder().setUseAvroLogicalTypes(USE_AVRO_LOGICAL_TYPES).build();
28+
29+
@Test
30+
public void testToBuilder() {
31+
compareAvroOptions(AVRO_OPTIONS, AVRO_OPTIONS.toBuilder().build());
32+
AvroOptions avroOptions = AVRO_OPTIONS.toBuilder().setUseAvroLogicalTypes(false).build();
33+
assertEquals(false, avroOptions.useAvroLogicalTypes());
34+
avroOptions = avroOptions.toBuilder().setUseAvroLogicalTypes(true).build();
35+
compareAvroOptions(AVRO_OPTIONS, avroOptions);
36+
}
37+
38+
@Test
39+
public void testBuilder() {
40+
assertEquals(FormatOptions.AVRO, AVRO_OPTIONS.getType());
41+
assertEquals(USE_AVRO_LOGICAL_TYPES, AVRO_OPTIONS.useAvroLogicalTypes());
42+
}
43+
44+
@Test
45+
public void testToAndFromPb() {
46+
compareAvroOptions(AVRO_OPTIONS, AvroOptions.fromPb(AVRO_OPTIONS.toPb()));
47+
AvroOptions avroOptions =
48+
AvroOptions.newBuilder().setUseAvroLogicalTypes(USE_AVRO_LOGICAL_TYPES).build();
49+
compareAvroOptions(avroOptions, AvroOptions.fromPb(avroOptions.toPb()));
50+
}
51+
52+
private void compareAvroOptions(AvroOptions expected, AvroOptions value) {
53+
assertEquals(expected, value);
54+
assertEquals(expected.useAvroLogicalTypes(), value.useAvroLogicalTypes());
55+
}
56+
}

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ExternalTableDefinitionTest.java

+5
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,7 @@ public class ExternalTableDefinitionTest {
5050
private static final String COMPRESSION = "GZIP";
5151
private static final String CONNECTION_ID = "123456789";
5252
private static final Boolean AUTODETECT = true;
53+
private static final AvroOptions AVRO_OPTIONS = AvroOptions.newBuilder().build();
5354
private static final CsvOptions CSV_OPTIONS = CsvOptions.newBuilder().build();
5455
private static final HivePartitioningOptions HIVE_PARTITIONING_OPTIONS =
5556
HivePartitioningOptions.newBuilder()
@@ -67,6 +68,9 @@ public class ExternalTableDefinitionTest {
6768
.setHivePartitioningOptions(HIVE_PARTITIONING_OPTIONS)
6869
.build();
6970

71+
private static final ExternalTableDefinition EXTERNAL_TABLE_DEFINITION_AVRO =
72+
ExternalTableDefinition.newBuilder(SOURCE_URIS, TABLE_SCHEMA, AVRO_OPTIONS).build();
73+
7074
@Test
7175
public void testToBuilder() {
7276
compareExternalTableDefinition(
@@ -109,6 +113,7 @@ public void testBuilder() {
109113
assertEquals(TableDefinition.Type.EXTERNAL, EXTERNAL_TABLE_DEFINITION.getType());
110114
assertEquals(COMPRESSION, EXTERNAL_TABLE_DEFINITION.getCompression());
111115
assertEquals(CONNECTION_ID, EXTERNAL_TABLE_DEFINITION.getConnectionId());
116+
assertEquals(AVRO_OPTIONS, EXTERNAL_TABLE_DEFINITION_AVRO.getFormatOptions());
112117
assertEquals(CSV_OPTIONS, EXTERNAL_TABLE_DEFINITION.getFormatOptions());
113118
assertEquals(IGNORE_UNKNOWN_VALUES, EXTERNAL_TABLE_DEFINITION.ignoreUnknownValues());
114119
assertEquals(MAX_BAD_RECORDS, EXTERNAL_TABLE_DEFINITION.getMaxBadRecords());

‎google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/LoadJobConfigurationTest.java

+2-2
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ public class LoadJobConfigurationTest {
5656
ImmutableList.of(SchemaUpdateOption.ALLOW_FIELD_ADDITION);
5757
private static final Schema TABLE_SCHEMA = Schema.of(FIELD_SCHEMA);
5858
private static final Boolean AUTODETECT = true;
59-
private static final Boolean USERAVROLOGICALTYPES = true;
59+
private static final Boolean USE_AVRO_LOGICAL_TYPES = true;
6060
private static final EncryptionConfiguration JOB_ENCRYPTION_CONFIGURATION =
6161
EncryptionConfiguration.newBuilder().setKmsKeyName("KMS_KEY_1").build();
6262
private static final TimePartitioning TIME_PARTITIONING = TimePartitioning.of(Type.DAY);
@@ -128,7 +128,7 @@ public class LoadJobConfigurationTest {
128128
.setDestinationEncryptionConfiguration(JOB_ENCRYPTION_CONFIGURATION)
129129
.setTimePartitioning(TIME_PARTITIONING)
130130
.setClustering(CLUSTERING)
131-
.setUseAvroLogicalTypes(USERAVROLOGICALTYPES)
131+
.setUseAvroLogicalTypes(USE_AVRO_LOGICAL_TYPES)
132132
.setLabels(LABELS)
133133
.setJobTimeoutMs(TIMEOUT)
134134
.setRangePartitioning(RANGE_PARTITIONING)

0 commit comments

Comments
 (0)