Skip to content

Commit 8844e40

Browse files
authored
feat(bigquery): add support for AvroOptions in external data config (#4945)
1 parent 744a753 commit 8844e40

File tree

4 files changed

+55
-0
lines changed

4 files changed

+55
-0
lines changed

Diff for: bigquery/external.go

+25
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,8 @@ func bqToExternalDataConfig(q *bq.ExternalDataConfiguration) (*ExternalDataConfi
143143
e.DecimalTargetTypes = append(e.DecimalTargetTypes, DecimalTargetType(v))
144144
}
145145
switch {
146+
case q.AvroOptions != nil:
147+
e.Options = bqToAvroOptions(q.AvroOptions)
146148
case q.CsvOptions != nil:
147149
e.Options = bqToCSVOptions(q.CsvOptions)
148150
case q.GoogleSheetsOptions != nil:
@@ -165,6 +167,29 @@ type ExternalDataConfigOptions interface {
165167
populateExternalDataConfig(*bq.ExternalDataConfiguration)
166168
}
167169

170+
// AvroOptions are additional options for Avro external data data sources.
171+
type AvroOptions struct {
172+
// UseAvroLogicalTypes indicates whether to interpret logical types as the
173+
// corresponding BigQuery data type (for example, TIMESTAMP), instead of using
174+
// the raw type (for example, INTEGER).
175+
UseAvroLogicalTypes bool
176+
}
177+
178+
func (o *AvroOptions) populateExternalDataConfig(c *bq.ExternalDataConfiguration) {
179+
c.AvroOptions = &bq.AvroOptions{
180+
UseAvroLogicalTypes: o.UseAvroLogicalTypes,
181+
}
182+
}
183+
184+
func bqToAvroOptions(q *bq.AvroOptions) *AvroOptions {
185+
if q == nil {
186+
return nil
187+
}
188+
return &AvroOptions{
189+
UseAvroLogicalTypes: q.UseAvroLogicalTypes,
190+
}
191+
}
192+
168193
// CSVOptions are additional options for CSV external data sources.
169194
type CSVOptions struct {
170195
// AllowJaggedRows causes missing trailing optional columns to be tolerated

Diff for: bigquery/external_test.go

+6
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,12 @@ func TestExternalDataConfig(t *testing.T) {
9191
SourceFormat: Parquet,
9292
DecimalTargetTypes: []DecimalTargetType{BigNumericTargetType, NumericTargetType, StringTargetType},
9393
},
94+
{
95+
SourceFormat: Avro,
96+
Options: &AvroOptions{
97+
UseAvroLogicalTypes: true,
98+
},
99+
},
94100
} {
95101
q := want.toBQ()
96102
got, err := bqToExternalDataConfig(&q)

Diff for: bigquery/file.go

+11
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,9 @@ type FileConfig struct {
7777

7878
// Additional options for Parquet files.
7979
ParquetOptions *ParquetOptions
80+
81+
// Additional options for Avro files.
82+
AvroOptions *AvroOptions
8083
}
8184

8285
func (fc *FileConfig) populateLoadConfig(conf *bq.JobConfigurationLoad) {
@@ -98,6 +101,9 @@ func (fc *FileConfig) populateLoadConfig(conf *bq.JobConfigurationLoad) {
98101
EnableListInference: fc.ParquetOptions.EnableListInference,
99102
}
100103
}
104+
if fc.AvroOptions != nil {
105+
conf.UseAvroLogicalTypes = fc.AvroOptions.UseAvroLogicalTypes
106+
}
101107
conf.Quote = fc.quote()
102108
}
103109

@@ -131,6 +137,11 @@ func (fc *FileConfig) populateExternalDataConfig(conf *bq.ExternalDataConfigurat
131137
if format == CSV {
132138
fc.CSVOptions.populateExternalDataConfig(conf)
133139
}
140+
if fc.AvroOptions != nil {
141+
conf.AvroOptions = &bq.AvroOptions{
142+
UseAvroLogicalTypes: fc.AvroOptions.UseAvroLogicalTypes,
143+
}
144+
}
134145
if fc.ParquetOptions != nil {
135146
conf.ParquetOptions = &bq.ParquetOptions{
136147
EnumAsString: fc.ParquetOptions.EnumAsString,

Diff for: bigquery/file_test.go

+13
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,19 @@ func TestFileConfigPopulateLoadConfig(t *testing.T) {
9696
},
9797
},
9898
},
99+
{
100+
description: "avro",
101+
fileConfig: &FileConfig{
102+
SourceFormat: Avro,
103+
AvroOptions: &AvroOptions{
104+
UseAvroLogicalTypes: true,
105+
},
106+
},
107+
want: &bq.JobConfigurationLoad{
108+
SourceFormat: "AVRO",
109+
UseAvroLogicalTypes: true,
110+
},
111+
},
99112
}
100113
for _, tc := range testcases {
101114
got := &bq.JobConfigurationLoad{}

0 commit comments

Comments
 (0)