From b62db7958ca2a985e063428f84010af22b9ae37f Mon Sep 17 00:00:00 2001
From: Alvaro Viebrantz <aviebrantz@google.com>
Date: Tue, 3 Oct 2023 16:25:15 -0400
Subject: [PATCH] feat(bigquery): add media options to LoadConfig

---
 bigquery/bigquery.go         | 4 ++--
 bigquery/integration_test.go | 7 ++++++-
 bigquery/load.go             | 6 +++++-
 3 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/bigquery/bigquery.go b/bigquery/bigquery.go
index a71c4faa3673..94720a912784 100644
--- a/bigquery/bigquery.go
+++ b/bigquery/bigquery.go
@@ -140,11 +140,11 @@ func (c *Client) Close() error {
 }
 
 // Calls the Jobs.Insert RPC and returns a Job.
-func (c *Client) insertJob(ctx context.Context, job *bq.Job, media io.Reader) (*Job, error) {
+func (c *Client) insertJob(ctx context.Context, job *bq.Job, media io.Reader, mediaOpts ...googleapi.MediaOption) (*Job, error) {
 	call := c.bqs.Jobs.Insert(c.projectID, job).Context(ctx)
 	setClientHeader(call.Header())
 	if media != nil {
-		call.Media(media)
+		call.Media(media, mediaOpts...)
 	}
 	var res *bq.Job
 	var err error
diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
index 8c98283ed7ae..64ea202a5923 100644
--- a/bigquery/integration_test.go
+++ b/bigquery/integration_test.go
@@ -1461,6 +1461,10 @@ func TestIntegration_Load(t *testing.T) {
 	loader := table.LoaderFrom(rs)
 	loader.WriteDisposition = WriteTruncate
 	loader.Labels = map[string]string{"test": "go"}
+	loader.MediaOptions = []googleapi.MediaOption{
+		googleapi.ContentType("text/csv"),
+		googleapi.ChunkSize(googleapi.MinUploadChunkSize),
+	}
 	job, err := loader.Run(ctx)
 	if err != nil {
 		t.Fatal(err)
@@ -1480,7 +1484,8 @@ func TestIntegration_Load(t *testing.T) {
 		cmp.AllowUnexported(Table{}),
 		cmpopts.IgnoreUnexported(Client{}, ReaderSource{}),
 		// returned schema is at top level, not in the config
-		cmpopts.IgnoreFields(FileConfig{}, "Schema"))
+		cmpopts.IgnoreFields(FileConfig{}, "Schema"),
+		cmpopts.IgnoreFields(LoadConfig{}, "MediaOptions"))
 	if diff != "" {
 		t.Errorf("got=-, want=+:\n%s", diff)
 	}
diff --git a/bigquery/load.go b/bigquery/load.go
index 8af650d7a869..3693719c4ea4 100644
--- a/bigquery/load.go
+++ b/bigquery/load.go
@@ -21,6 +21,7 @@ import (
 
 	"cloud.google.com/go/internal/trace"
 	bq "google.golang.org/api/bigquery/v2"
+	"google.golang.org/api/googleapi"
 )
 
 // LoadConfig holds the configuration for a load job.
@@ -101,6 +102,9 @@ type LoadConfig struct {
 
 	// ConnectionProperties are optional key-values settings.
 	ConnectionProperties []*ConnectionProperty
+
+	// MediaOptions stores options for customizing media upload.
+	MediaOptions []googleapi.MediaOption
 }
 
 func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) {
@@ -210,7 +214,7 @@ func (l *Loader) Run(ctx context.Context) (j *Job, err error) {
 	defer func() { trace.EndSpan(ctx, err) }()
 
 	job, media := l.newJob()
-	return l.c.insertJob(ctx, job, media)
+	return l.c.insertJob(ctx, job, media, l.LoadConfig.MediaOptions...)
 }
 
 func (l *Loader) newJob() (*bq.Job, io.Reader) {