Repository: googleapis/google-cloud-go
Pull request #4945 (2021-10-04 17:40:52)

feat(bigquery): add support for AvroOptions in external data config


Patch:
diff --git a/bigquery/external.go b/bigquery/external.go
index 500a2927ad..1083a19904 100644
--- a/bigquery/external.go
+++ b/bigquery/external.go
@@ -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:
@@ -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
diff --git a/bigquery/file.go b/bigquery/file.go
index 2bfe9e44a6..cb404588be 100644
--- a/bigquery/file.go
+++ b/bigquery/file.go
@@ -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) {
@@ -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()
 }
 
@@ -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,


Test patch:
diff --git a/bigquery/external_test.go b/bigquery/external_test.go
index 75aa3ab909..4a8b3ec7d6 100644
--- a/bigquery/external_test.go
+++ b/bigquery/external_test.go
@@ -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)
diff --git a/bigquery/file_test.go b/bigquery/file_test.go
index 0b0c47aebc..bb55d8abea 100644
--- a/bigquery/file_test.go
+++ b/bigquery/file_test.go
@@ -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{}