Skip to content

Commit 62baf56

Browse files
authored
feat(bigquery): add media options to LoadConfig (#8640)
1 parent 22ab199 commit 62baf56

File tree

3 files changed

+13
-4
lines changed

3 files changed

+13
-4
lines changed

bigquery/bigquery.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -140,11 +140,11 @@ func (c *Client) Close() error {
140140
}
141141

142142
// Calls the Jobs.Insert RPC and returns a Job.
143-
func (c *Client) insertJob(ctx context.Context, job *bq.Job, media io.Reader) (*Job, error) {
143+
func (c *Client) insertJob(ctx context.Context, job *bq.Job, media io.Reader, mediaOpts ...googleapi.MediaOption) (*Job, error) {
144144
call := c.bqs.Jobs.Insert(c.projectID, job).Context(ctx)
145145
setClientHeader(call.Header())
146146
if media != nil {
147-
call.Media(media)
147+
call.Media(media, mediaOpts...)
148148
}
149149
var res *bq.Job
150150
var err error

bigquery/integration_test.go

+6-1
Original file line numberDiff line numberDiff line change
@@ -1461,6 +1461,10 @@ func TestIntegration_Load(t *testing.T) {
14611461
loader := table.LoaderFrom(rs)
14621462
loader.WriteDisposition = WriteTruncate
14631463
loader.Labels = map[string]string{"test": "go"}
1464+
loader.MediaOptions = []googleapi.MediaOption{
1465+
googleapi.ContentType("text/csv"),
1466+
googleapi.ChunkSize(googleapi.MinUploadChunkSize),
1467+
}
14641468
job, err := loader.Run(ctx)
14651469
if err != nil {
14661470
t.Fatal(err)
@@ -1480,7 +1484,8 @@ func TestIntegration_Load(t *testing.T) {
14801484
cmp.AllowUnexported(Table{}),
14811485
cmpopts.IgnoreUnexported(Client{}, ReaderSource{}),
14821486
// returned schema is at top level, not in the config
1483-
cmpopts.IgnoreFields(FileConfig{}, "Schema"))
1487+
cmpopts.IgnoreFields(FileConfig{}, "Schema"),
1488+
cmpopts.IgnoreFields(LoadConfig{}, "MediaOptions"))
14841489
if diff != "" {
14851490
t.Errorf("got=-, want=+:\n%s", diff)
14861491
}

bigquery/load.go

+5-1
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121

2222
"cloud.google.com/go/internal/trace"
2323
bq "google.golang.org/api/bigquery/v2"
24+
"google.golang.org/api/googleapi"
2425
)
2526

2627
// LoadConfig holds the configuration for a load job.
@@ -101,6 +102,9 @@ type LoadConfig struct {
101102

102103
// ConnectionProperties are optional key-values settings.
103104
ConnectionProperties []*ConnectionProperty
105+
106+
// MediaOptions stores options for customizing media upload.
107+
MediaOptions []googleapi.MediaOption
104108
}
105109

106110
func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) {
@@ -210,7 +214,7 @@ func (l *Loader) Run(ctx context.Context) (j *Job, err error) {
210214
defer func() { trace.EndSpan(ctx, err) }()
211215

212216
job, media := l.newJob()
213-
return l.c.insertJob(ctx, job, media)
217+
return l.c.insertJob(ctx, job, media, l.LoadConfig.MediaOptions...)
214218
}
215219

216220
func (l *Loader) newJob() (*bq.Job, io.Reader) {

0 commit comments

Comments
 (0)