feat(bigquery): support per-job reservation assignment (#12078)

shollyman · web-flow · commit c9cebcceebc5 · 2025-04-29T12:20:00.000-07:00
* feat(bigquery): support per-job reservation assignment

This PR adds support for explicitly selecting a reservation to use when executing a job. It also includes support for optimized query execution (running via jobs.query vs jobs.insert).
diff --git a/bigquery/copy.go b/bigquery/copy.go
@@ -73,6 +73,13 @@ type CopyConfig struct {
 	// Experimental: this option is experimental and may be modified or removed in future versions,
 	// regardless of any other documented package stability guarantees.
 	JobTimeout time.Duration
+
+	// The reservation that job would use. User can specify a reservation to
+	// execute the job. If reservation is not set, reservation is determined
+	// based on the rules defined by the reservation assignments. The expected
+	// format is
+	// `projects/{project}/locations/{location}/reservations/{reservation}`.
+	Reservation string
 }
 
 func (c *CopyConfig) toBQ() *bq.JobConfiguration {
@@ -91,6 +98,7 @@ func (c *CopyConfig) toBQ() *bq.JobConfiguration {
 			OperationType:                      string(c.OperationType),
 		},
 		JobTimeoutMs: c.JobTimeout.Milliseconds(),
+		Reservation:  c.Reservation,
 	}
 }
 
@@ -103,6 +111,7 @@ func bqToCopyConfig(q *bq.JobConfiguration, c *Client) *CopyConfig {
 		DestinationEncryptionConfig: bqToEncryptionConfig(q.Copy.DestinationEncryptionConfiguration),
 		OperationType:               TableCopyOperationType(q.Copy.OperationType),
 		JobTimeout:                  time.Duration(q.JobTimeoutMs) * time.Millisecond,
+		Reservation:                 q.Reservation,
 	}
 	for _, t := range q.Copy.SourceTables {
 		cc.Srcs = append(cc.Srcs, bqToTable(t, c))
diff --git a/bigquery/copy_test.go b/bigquery/copy_test.go
@@ -88,13 +88,15 @@ func TestCopy(t *testing.T) {
 				WriteDisposition:            WriteTruncate,
 				DestinationEncryptionConfig: &EncryptionConfig{KMSKeyName: "keyName"},
 				Labels:                      map[string]string{"a": "b"},
+				Reservation:                 "reservation/1",
 			},
 			want: func() *bq.Job {
 				j := defaultCopyJob()
 				j.Configuration.Labels = map[string]string{"a": "b"}
 				j.Configuration.Copy.CreateDisposition = "CREATE_NEVER"
 				j.Configuration.Copy.WriteDisposition = "WRITE_TRUNCATE"
 				j.Configuration.Copy.DestinationEncryptionConfiguration = &bq.EncryptionConfiguration{KmsKeyName: "keyName"}
+				j.Configuration.Reservation = "reservation/1"
 				return j
 			}(),
 		},
diff --git a/bigquery/extract.go b/bigquery/extract.go
@@ -57,6 +57,13 @@ type ExtractConfig struct {
 	// Experimental: this option is experimental and may be modified or removed in future versions,
 	// regardless of any other documented package stability guarantees.
 	JobTimeout time.Duration
+
+	// The reservation that job would use. User can specify a reservation to
+	// execute the job. If reservation is not set, reservation is determined
+	// based on the rules defined by the reservation assignments. The expected
+	// format is
+	// `projects/{project}/locations/{location}/reservations/{reservation}`.
+	Reservation string
 }
 
 func (e *ExtractConfig) toBQ() *bq.JobConfiguration {
@@ -77,6 +84,7 @@ func (e *ExtractConfig) toBQ() *bq.JobConfiguration {
 			UseAvroLogicalTypes: e.UseAvroLogicalTypes,
 		},
 		JobTimeoutMs: e.JobTimeout.Milliseconds(),
+		Reservation:  e.Reservation,
 	}
 	if e.Src != nil {
 		cfg.Extract.SourceTable = e.Src.toBQ()
@@ -106,6 +114,7 @@ func bqToExtractConfig(q *bq.JobConfiguration, c *Client) *ExtractConfig {
 		SrcModel:            bqToModel(qe.SourceModel, c),
 		UseAvroLogicalTypes: qe.UseAvroLogicalTypes,
 		JobTimeout:          time.Duration(q.JobTimeoutMs) * time.Millisecond,
+		Reservation:         q.Reservation,
 	}
 }
 
diff --git a/bigquery/extract_test.go b/bigquery/extract_test.go
@@ -85,13 +85,15 @@ func TestExtract(t *testing.T) {
 				DisableHeader: true,
 				Labels:        map[string]string{"a": "b"},
 				JobTimeout:    8 * time.Second,
+				Reservation:   "reservation/1",
 			},
 			want: func() *bq.Job {
 				j := defaultExtractJob()
 				j.Configuration.Labels = map[string]string{"a": "b"}
 				j.Configuration.JobTimeoutMs = 8000
 				f := false
 				j.Configuration.Extract.PrintHeader = &f
+				j.Configuration.Reservation = "reservation/1"
 				return j
 			}(),
 		},
diff --git a/bigquery/load.go b/bigquery/load.go
@@ -110,6 +110,13 @@ type LoadConfig struct {
 	// For more information, see:
 	// https://cloud.google.com/bigquery/docs/reference/rest/v2/Job#columnnamecharactermap
 	ColumnNameCharacterMap ColumnNameCharacterMap
+
+	// The reservation that job would use. User can specify a reservation to
+	// execute the job. If reservation is not set, reservation is determined
+	// based on the rules defined by the reservation assignments. The expected
+	// format is
+	// `projects/{project}/locations/{location}/reservations/{reservation}`.
+	Reservation string
 }
 
 func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) {
@@ -140,6 +147,7 @@ func (l *LoadConfig) toBQ() (*bq.JobConfiguration, io.Reader) {
 		config.Load.ConnectionProperties = append(config.Load.ConnectionProperties, v.toBQ())
 	}
 	media := l.Src.populateLoadConfig(config.Load)
+	config.Reservation = l.Reservation
 	return config, media
 }
 
@@ -160,6 +168,7 @@ func bqToLoadConfig(q *bq.JobConfiguration, c *Client) *LoadConfig {
 		ReferenceFileSchemaURI:      q.Load.ReferenceFileSchemaUri,
 		CreateSession:               q.Load.CreateSession,
 		ColumnNameCharacterMap:      ColumnNameCharacterMap(q.Load.ColumnNameCharacterMap),
+		Reservation:                 q.Reservation,
 	}
 	if q.JobTimeoutMs > 0 {
 		lc.JobTimeout = time.Duration(q.JobTimeoutMs) * time.Millisecond
diff --git a/bigquery/load_test.go b/bigquery/load_test.go
@@ -140,7 +140,8 @@ func TestLoad(t *testing.T) {
 				return g
 			}(),
 			config: LoadConfig{
-				JobTimeout: 4 * time.Second,
+				JobTimeout:  4 * time.Second,
+				Reservation: "reservation/1",
 			},
 			want: func() *bq.Job {
 				j := defaultLoadJob()
@@ -149,6 +150,7 @@ func TestLoad(t *testing.T) {
 				j.Configuration.Load.AllowQuotedNewlines = true
 				j.Configuration.Load.IgnoreUnknownValues = true
 				j.Configuration.JobTimeoutMs = 4000
+				j.Configuration.Reservation = "reservation/1"
 				return j
 			}(),
 		},
diff --git a/bigquery/query.go b/bigquery/query.go
@@ -152,6 +152,13 @@ type QueryConfig struct {
 
 	// Force usage of Storage API if client is available. For test scenarios
 	forceStorageAPI bool
+
+	// The reservation that job would use. User can specify a reservation to
+	// execute the job. If reservation is not set, reservation is determined
+	// based on the rules defined by the reservation assignments. The expected
+	// format is
+	// `projects/{project}/locations/{location}/reservations/{reservation}`.
+	Reservation string
 }
 
 func (qc *QueryConfig) toBQ() (*bq.JobConfiguration, error) {
@@ -224,9 +231,10 @@ func (qc *QueryConfig) toBQ() (*bq.JobConfiguration, error) {
 		qconf.ConnectionProperties = bqcp
 	}
 	jc := &bq.JobConfiguration{
-		Labels: qc.Labels,
-		DryRun: qc.DryRun,
-		Query:  qconf,
+		Labels:      qc.Labels,
+		DryRun:      qc.DryRun,
+		Reservation: qc.Reservation,
+		Query:       qconf,
 	}
 	if qc.JobTimeout > 0 {
 		jc.JobTimeoutMs = qc.JobTimeout.Milliseconds()
@@ -255,6 +263,7 @@ func bqToQueryConfig(q *bq.JobConfiguration, c *Client) (*QueryConfig, error) {
 		CreateSession:               qq.CreateSession,
 	}
 	qc.UseStandardSQL = !qc.UseLegacySQL
+	qc.Reservation = q.Reservation
 
 	if len(qq.TableDefinitions) > 0 {
 		qc.TableDefinitions = make(map[string]ExternalData)
@@ -469,6 +478,7 @@ func (q *Query) probeFastPath() (*bq.QueryRequest, error) {
 		UseLegacySql:       &pfalse,
 		MaximumBytesBilled: q.QueryConfig.MaxBytesBilled,
 		RequestId:          uid.NewSpace("request", nil).New(),
+		Reservation:        q.Reservation,
 		Labels:             q.Labels,
 		FormatOptions: &bq.DataFormatOptions{
 			UseInt64Timestamp: true,
diff --git a/bigquery/query_test.go b/bigquery/query_test.go
@@ -357,6 +357,20 @@ func TestQuery(t *testing.T) {
 				return j
 			}(),
 		},
+		{
+			dst: c.Dataset("dataset-id").Table("table-id"),
+			src: &QueryConfig{
+				Q:                "query string",
+				Reservation:      "reservation/1",
+				DefaultProjectID: "def-project-id",
+				DefaultDatasetID: "def-dataset-id",
+			},
+			want: func() *bq.Job {
+				j := defaultQueryJob()
+				j.Configuration.Reservation = "reservation/1"
+				return j
+			}(),
+		},
 	}
 	for i, tc := range testCases {
 		query := c.Query("")
@@ -459,6 +473,7 @@ func TestProbeFastPath(t *testing.T) {
 				Labels: map[string]string{
 					"key": "val",
 				},
+				Reservation: "reservation/1",
 			},
 			wantReq: &bq.QueryRequest{
 				Query:          "foo",
@@ -479,6 +494,7 @@ func TestProbeFastPath(t *testing.T) {
 				FormatOptions: &bq.DataFormatOptions{
 					UseInt64Timestamp: true,
 				},
+				Reservation: "reservation/1",
 			},
 		},
 		{
