Skip to content

Commit a3c0aca

Browse files
shollymanshollyman
authored
feat(bigquery): add support for MaxSlots (#12958)
1 parent 33e1474 commit a3c0aca

11 files changed

Lines changed: 461 additions & 70 deletions

File tree

‎bigquery/copy.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,15 @@ type CopyConfig struct {
8080
// format is
8181
// `projects/{project}/locations/{location}/reservations/{reservation}`.
8282
Reservation string
83+
84+
// A target limit on the rate of slot consumption by this query. If set to a
85+
// value > 0, BigQuery will attempt to limit the rate of slot consumption by
86+
// this query to keep it below the configured limit, even if the query is
87+
// eligible for more slots based on fair scheduling. The unused slots will be
88+
// available for other jobs and queries to use.
89+
//
90+
// Note: This feature is not yet generally available.
91+
MaxSlots int32
8392
}
8493

8594
func (c *CopyConfig) toBQ() *bq.JobConfiguration {
@@ -99,6 +108,7 @@ func (c *CopyConfig) toBQ() *bq.JobConfiguration {
99108
},
100109
JobTimeoutMs: c.JobTimeout.Milliseconds(),
101110
Reservation: c.Reservation,
111+
MaxSlots: int64(c.MaxSlots),
102112
}
103113
}
104114

@@ -112,6 +122,7 @@ func bqToCopyConfig(q *bq.JobConfiguration, c *Client) *CopyConfig {
112122
OperationType: TableCopyOperationType(q.Copy.OperationType),
113123
JobTimeout: time.Duration(q.JobTimeoutMs) * time.Millisecond,
114124
Reservation: q.Reservation,
125+
MaxSlots: int32(q.MaxSlots),
115126
}
116127
for _, t := range q.Copy.SourceTables {
117128
cc.Srcs = append(cc.Srcs, bqToTable(t, c))

‎bigquery/copy_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,7 @@ func TestCopy(t *testing.T) {
8989
DestinationEncryptionConfig: &EncryptionConfig{KMSKeyName: "keyName"},
9090
Labels: map[string]string{"a": "b"},
9191
Reservation: "reservation/1",
92+
MaxSlots: 123,
9293
},
9394
want: func() *bq.Job {
9495
j := defaultCopyJob()
@@ -97,6 +98,7 @@ func TestCopy(t *testing.T) {
9798
j.Configuration.Copy.WriteDisposition = "WRITE_TRUNCATE"
9899
j.Configuration.Copy.DestinationEncryptionConfiguration = &bq.EncryptionConfiguration{KmsKeyName: "keyName"}
99100
j.Configuration.Reservation = "reservation/1"
101+
j.Configuration.MaxSlots = 123
100102
return j
101103
}(),
102104
},

‎bigquery/extract.go‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,15 @@ type ExtractConfig struct {
6464
// format is
6565
// `projects/{project}/locations/{location}/reservations/{reservation}`.
6666
Reservation string
67+
68+
// A target limit on the rate of slot consumption by this query. If set to a
69+
// value > 0, BigQuery will attempt to limit the rate of slot consumption by
70+
// this query to keep it below the configured limit, even if the query is
71+
// eligible for more slots based on fair scheduling. The unused slots will be
72+
// available for other jobs and queries to use.
73+
//
74+
// Note: This feature is not yet generally available.
75+
MaxSlots int32
6776
}
6877

6978
func (e *ExtractConfig) toBQ() *bq.JobConfiguration {
@@ -85,6 +94,7 @@ func (e *ExtractConfig) toBQ() *bq.JobConfiguration {
8594
},
8695
JobTimeoutMs: e.JobTimeout.Milliseconds(),
8796
Reservation: e.Reservation,
97+
MaxSlots: int64(e.MaxSlots),
8898
}
8999
if e.Src != nil {
90100
cfg.Extract.SourceTable = e.Src.toBQ()
@@ -115,6 +125,7 @@ func bqToExtractConfig(q *bq.JobConfiguration, c *Client) *ExtractConfig {
115125
UseAvroLogicalTypes: qe.UseAvroLogicalTypes,
116126
JobTimeout: time.Duration(q.JobTimeoutMs) * time.Millisecond,
117127
Reservation: q.Reservation,
128+
MaxSlots: int32(q.MaxSlots),
118129
}
119130
}
120131

‎bigquery/extract_test.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ func TestExtract(t *testing.T) {
8686
Labels: map[string]string{"a": "b"},
8787
JobTimeout: 8 * time.Second,
8888
Reservation: "reservation/1",
89+
MaxSlots: 234,
8990
},
9091
want: func() *bq.Job {
9192
j := defaultExtractJob()
@@ -94,6 +95,7 @@ func TestExtract(t *testing.T) {
9495
f := false
9596
j.Configuration.Extract.PrintHeader = &f
9697
j.Configuration.Reservation = "reservation/1"
98+
j.Configuration.MaxSlots = 234
9799
return j
98100
}(),
99101
},

‎bigquery/go.mod‎

Lines changed: 22 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -13,33 +13,33 @@ require (
1313
github.com/google/uuid v1.6.0
1414
github.com/googleapis/gax-go/v2 v2.15.0
1515
go.opencensus.io v0.24.0
16-
go.opentelemetry.io/otel v1.36.0
17-
go.opentelemetry.io/otel/sdk v1.36.0
18-
golang.org/x/sync v0.16.0
16+
go.opentelemetry.io/otel v1.37.0
17+
go.opentelemetry.io/otel/sdk v1.37.0
18+
golang.org/x/sync v0.17.0
1919
golang.org/x/xerrors v0.0.0-20240903120638-7835f813f4da
20-
google.golang.org/api v0.247.0
20+
google.golang.org/api v0.250.0
2121
google.golang.org/genproto v0.0.0-20250603155806-513f23925822
2222
google.golang.org/genproto/googleapis/api v0.0.0-20250818200422-3122310a409c
23-
google.golang.org/genproto/googleapis/rpc v0.0.0-20250818200422-3122310a409c
24-
google.golang.org/grpc v1.74.2
25-
google.golang.org/protobuf v1.36.7
23+
google.golang.org/genproto/googleapis/rpc v0.0.0-20250908214217-97024824d090
24+
google.golang.org/grpc v1.75.1
25+
google.golang.org/protobuf v1.36.9
2626
)
2727

2828
require (
2929
cel.dev/expr v0.24.0 // indirect
30-
cloud.google.com/go/auth v0.16.4 // indirect
30+
cloud.google.com/go/auth v0.16.5 // indirect
3131
cloud.google.com/go/auth/oauth2adapt v0.2.8 // indirect
32-
cloud.google.com/go/compute/metadata v0.8.0 // indirect
32+
cloud.google.com/go/compute/metadata v0.8.4 // indirect
3333
cloud.google.com/go/monitoring v1.24.2 // indirect
34-
github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.27.0 // indirect
34+
github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.29.0 // indirect
3535
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.53.0 // indirect
3636
github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.53.0 // indirect
3737
github.com/cespare/xxhash/v2 v2.3.0 // indirect
3838
github.com/cncf/xds/go v0.0.0-20250501225837-2ac532fd4443 // indirect
3939
github.com/envoyproxy/go-control-plane/envoy v1.32.4 // indirect
4040
github.com/envoyproxy/protoc-gen-validate v1.2.1 // indirect
4141
github.com/felixge/httpsnoop v1.0.4 // indirect
42-
github.com/go-jose/go-jose/v4 v4.0.5 // indirect
42+
github.com/go-jose/go-jose/v4 v4.1.1 // indirect
4343
github.com/go-logr/logr v1.4.3 // indirect
4444
github.com/go-logr/stdr v1.2.2 // indirect
4545
github.com/goccy/go-json v0.10.2 // indirect
@@ -58,16 +58,16 @@ require (
5858
go.opentelemetry.io/contrib/detectors/gcp v1.36.0 // indirect
5959
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.61.0 // indirect
6060
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.61.0 // indirect
61-
go.opentelemetry.io/otel/metric v1.36.0 // indirect
62-
go.opentelemetry.io/otel/sdk/metric v1.36.0 // indirect
63-
go.opentelemetry.io/otel/trace v1.36.0 // indirect
64-
golang.org/x/crypto v0.41.0 // indirect
61+
go.opentelemetry.io/otel/metric v1.37.0 // indirect
62+
go.opentelemetry.io/otel/sdk/metric v1.37.0 // indirect
63+
go.opentelemetry.io/otel/trace v1.37.0 // indirect
64+
golang.org/x/crypto v0.42.0 // indirect
6565
golang.org/x/exp v0.0.0-20240719175910-8a7402abbf56 // indirect
66-
golang.org/x/mod v0.26.0 // indirect
67-
golang.org/x/net v0.43.0 // indirect
68-
golang.org/x/oauth2 v0.30.0 // indirect
69-
golang.org/x/sys v0.35.0 // indirect
70-
golang.org/x/text v0.28.0 // indirect
71-
golang.org/x/time v0.12.0 // indirect
72-
golang.org/x/tools v0.35.0 // indirect
66+
golang.org/x/mod v0.27.0 // indirect
67+
golang.org/x/net v0.44.0 // indirect
68+
golang.org/x/oauth2 v0.31.0 // indirect
69+
golang.org/x/sys v0.36.0 // indirect
70+
golang.org/x/text v0.29.0 // indirect
71+
golang.org/x/time v0.13.0 // indirect
72+
golang.org/x/tools v0.36.0 // indirect
7373
)

0 commit comments

Comments
 (0)