Skip to content

Commit

Permalink
feat(bigquery): add support for AvroOptions in external data config (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
shollyman committed Oct 5, 2021
1 parent 744a753 commit 8844e40
Show file tree
Hide file tree
Showing 4 changed files with 55 additions and 0 deletions.
25 changes: 25 additions & 0 deletions bigquery/external.go
Expand Up @@ -143,6 +143,8 @@ func bqToExternalDataConfig(q *bq.ExternalDataConfiguration) (*ExternalDataConfi
e.DecimalTargetTypes = append(e.DecimalTargetTypes, DecimalTargetType(v))
}
switch {
case q.AvroOptions != nil:
e.Options = bqToAvroOptions(q.AvroOptions)
case q.CsvOptions != nil:
e.Options = bqToCSVOptions(q.CsvOptions)
case q.GoogleSheetsOptions != nil:
Expand All @@ -165,6 +167,29 @@ type ExternalDataConfigOptions interface {
populateExternalDataConfig(*bq.ExternalDataConfiguration)
}

// AvroOptions are additional options for Avro external data data sources.
type AvroOptions struct {
// UseAvroLogicalTypes indicates whether to interpret logical types as the
// corresponding BigQuery data type (for example, TIMESTAMP), instead of using
// the raw type (for example, INTEGER).
UseAvroLogicalTypes bool
}

func (o *AvroOptions) populateExternalDataConfig(c *bq.ExternalDataConfiguration) {
c.AvroOptions = &bq.AvroOptions{
UseAvroLogicalTypes: o.UseAvroLogicalTypes,
}
}

func bqToAvroOptions(q *bq.AvroOptions) *AvroOptions {
if q == nil {
return nil
}
return &AvroOptions{
UseAvroLogicalTypes: q.UseAvroLogicalTypes,
}
}

// CSVOptions are additional options for CSV external data sources.
type CSVOptions struct {
// AllowJaggedRows causes missing trailing optional columns to be tolerated
Expand Down
6 changes: 6 additions & 0 deletions bigquery/external_test.go
Expand Up @@ -91,6 +91,12 @@ func TestExternalDataConfig(t *testing.T) {
SourceFormat: Parquet,
DecimalTargetTypes: []DecimalTargetType{BigNumericTargetType, NumericTargetType, StringTargetType},
},
{
SourceFormat: Avro,
Options: &AvroOptions{
UseAvroLogicalTypes: true,
},
},
} {
q := want.toBQ()
got, err := bqToExternalDataConfig(&q)
Expand Down
11 changes: 11 additions & 0 deletions bigquery/file.go
Expand Up @@ -77,6 +77,9 @@ type FileConfig struct {

// Additional options for Parquet files.
ParquetOptions *ParquetOptions

// Additional options for Avro files.
AvroOptions *AvroOptions
}

func (fc *FileConfig) populateLoadConfig(conf *bq.JobConfigurationLoad) {
Expand All @@ -98,6 +101,9 @@ func (fc *FileConfig) populateLoadConfig(conf *bq.JobConfigurationLoad) {
EnableListInference: fc.ParquetOptions.EnableListInference,
}
}
if fc.AvroOptions != nil {
conf.UseAvroLogicalTypes = fc.AvroOptions.UseAvroLogicalTypes
}
conf.Quote = fc.quote()
}

Expand Down Expand Up @@ -131,6 +137,11 @@ func (fc *FileConfig) populateExternalDataConfig(conf *bq.ExternalDataConfigurat
if format == CSV {
fc.CSVOptions.populateExternalDataConfig(conf)
}
if fc.AvroOptions != nil {
conf.AvroOptions = &bq.AvroOptions{
UseAvroLogicalTypes: fc.AvroOptions.UseAvroLogicalTypes,
}
}
if fc.ParquetOptions != nil {
conf.ParquetOptions = &bq.ParquetOptions{
EnumAsString: fc.ParquetOptions.EnumAsString,
Expand Down
13 changes: 13 additions & 0 deletions bigquery/file_test.go
Expand Up @@ -96,6 +96,19 @@ func TestFileConfigPopulateLoadConfig(t *testing.T) {
},
},
},
{
description: "avro",
fileConfig: &FileConfig{
SourceFormat: Avro,
AvroOptions: &AvroOptions{
UseAvroLogicalTypes: true,
},
},
want: &bq.JobConfigurationLoad{
SourceFormat: "AVRO",
UseAvroLogicalTypes: true,
},
},
}
for _, tc := range testcases {
got := &bq.JobConfigurationLoad{}
Expand Down

0 comments on commit 8844e40

Please sign in to comment.