chore(bigquery): downgrade bigquery from 1.61.0 to 1.60.0

go.mod

@@ -6,12 +6,13 @@ toolchain go1.22.2
 
 require (
  cloud.google.com/go v0.112.2
- cloud.google.com/go/bigquery v1.61.0
+ cloud.google.com/go/bigquery v1.60.0
  github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5
  github.com/aws/aws-sdk-go-v2 v1.26.1
  github.com/aws/aws-sdk-go-v2/config v1.27.11
  github.com/aws/aws-sdk-go-v2/credentials v1.17.11
  github.com/aws/aws-sdk-go-v2/service/redshiftdata v1.25.4
+ github.com/cenkalti/backoff/v4 v4.3.0
  github.com/databricks/databricks-sql-go v1.5.5
  github.com/dlclark/regexp2 v1.11.0
  github.com/gliderlabs/ssh v0.3.7
@@ -49,6 +50,7 @@ require (
  github.com/Microsoft/go-winio v0.6.1 // indirect
  github.com/Nvveen/Gotty v0.0.0-20120604004816-cd527374f1e5 // indirect
  github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be // indirect
+ github.com/apache/arrow/go/v14 v14.0.2 // indirect
  github.com/apache/arrow/go/v15 v15.0.2 // indirect
  github.com/apache/arrow/go/v16 v16.0.0 // indirect
  github.com/apache/thrift v0.19.0 // indirect
@@ -68,7 +70,6 @@ require (
  github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.4 // indirect
  github.com/aws/aws-sdk-go-v2/service/sts v1.28.6 // indirect
  github.com/aws/smithy-go v1.20.2 // indirect
- github.com/cenkalti/backoff/v4 v4.3.0 // indirect
  github.com/containerd/continuity v0.4.3 // indirect
  github.com/coreos/go-oidc/v3 v3.5.0 // indirect
  github.com/danieljoos/wincred v1.1.2 // indirect
@@ -142,7 +143,7 @@ require (
  golang.org/x/time v0.5.0 // indirect
  golang.org/x/tools v0.20.0 // indirect
  golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
- google.golang.org/genproto v0.0.0-20240401170217-c3f982113cda // indirect
+ google.golang.org/genproto v0.0.0-20240227224415-6ceb2ff114de // indirect
  google.golang.org/genproto/googleapis/api v0.0.0-20240415180920-8c6c420018be // indirect
  google.golang.org/genproto/googleapis/rpc v0.0.0-20240429193739-8cf5692501f6 // indirect
  google.golang.org/grpc v1.63.2 // indirect

go.sum

@@ -5,8 +5,8 @@ cloud.google.com/go/auth v0.4.1 h1:Z7YNIhlWRtrnKlZke7z3GMqzvuYzdc2z98F9D1NV5Hg=
 cloud.google.com/go/auth v0.4.1/go.mod h1:QVBuVEKpCn4Zp58hzRGvL0tjRGU0YqdRTdCHM1IHnro=
 cloud.google.com/go/auth/oauth2adapt v0.2.2 h1:+TTV8aXpjeChS9M+aTtN/TjdQnzJvmzKFt//oWu7HX4=
 cloud.google.com/go/auth/oauth2adapt v0.2.2/go.mod h1:wcYjgpZI9+Yu7LyYBg4pqSiaRkfEK3GQcpb7C/uyF1Q=
-cloud.google.com/go/bigquery v1.61.0 h1:w2Goy9n6gh91LVi6B2Sc+HpBl8WbWhIyzdvVvrAuEIw=
-cloud.google.com/go/bigquery v1.61.0/go.mod h1:PjZUje0IocbuTOdq4DBOJLNYB0WF3pAKBHzAYyxCwFo=
+cloud.google.com/go/bigquery v1.60.0 h1:kA96WfgvCbkqfLnr7xI5uEfJ4h4FrnkdEb0yty0KSZo=
+cloud.google.com/go/bigquery v1.60.0/go.mod h1:Clwk2OeC0ZU5G5LDg7mo+h8U7KlAa5v06z5rptKdM3g=
 cloud.google.com/go/compute/metadata v0.2.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k=
 cloud.google.com/go/compute/metadata v0.3.0 h1:Tz+eQXMEqDIKRsmY3cHTL6FVaynIjX2QxYC4trgAKZc=
 cloud.google.com/go/compute/metadata v0.3.0/go.mod h1:zFmK7XCadkQkj6TtorcaGlCW1hT1fIilQDwofLpJ20k=
@@ -49,6 +49,8 @@ github.com/ahmetb/dlog v0.0.0-20170105205344-4fb5f8204f26 h1:3YVZUqkoev4mL+aCwVO
 github.com/ahmetb/dlog v0.0.0-20170105205344-4fb5f8204f26/go.mod h1:ymXt5bw5uSNu4jveerFxE0vNYxF8ncqbptntMaFMg3k=
 github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be h1:9AeTilPcZAjCFIImctFaOjnTIavg87rW78vTPkQqLI8=
 github.com/anmitsu/go-shlex v0.0.0-20200514113438-38f4b401e2be/go.mod h1:ySMOLuWl6zY27l47sB3qLNK6tF2fkHG55UZxx8oIVo4=
+github.com/apache/arrow/go/v14 v14.0.2 h1:N8OkaJEOfI3mEZt07BIkvo4sC6XDbL+48MBPWO5IONw=
+github.com/apache/arrow/go/v14 v14.0.2/go.mod h1:u3fgh3EdgN/YQ8cVQRguVW3R+seMybFg8QBQ5LU+eBY=
 github.com/apache/arrow/go/v15 v15.0.2 h1:60IliRbiyTWCWjERBCkO1W4Qun9svcYoZrSLcyOsMLE=
 github.com/apache/arrow/go/v15 v15.0.2/go.mod h1:DGXsR3ajT524njufqf95822i+KTh+yea1jass9YXgjA=
 github.com/apache/arrow/go/v16 v16.0.0 h1:qRLbJRPj4zaseZrjbDHa7mUoZDDIU+4pu+mE2Lucs5g=
@@ -202,8 +204,8 @@ github.com/google/go-cmp v0.5.8/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeN
 github.com/google/go-cmp v0.5.9/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
-github.com/google/martian/v3 v3.3.3 h1:DIhPTQrbPkgs2yJYdXU/eNACCG5DVQjySNRNlflZ9Fc=
-github.com/google/martian/v3 v3.3.3/go.mod h1:iEPrYcgCF7jA9OtScMFQyAlZZ4YXTKEtJ1E6RWzmBA0=
+github.com/google/martian/v3 v3.3.2 h1:IqNFLAmvJOgVlpdEBiQbDc2EwKW77amAycfTuWKdfvw=
+github.com/google/martian/v3 v3.3.2/go.mod h1:oBOf6HBosgwRXnUGWUB05QECsc6uvmMiJ3+6W4l/CUk=
 github.com/google/s2a-go v0.1.7 h1:60BLSyTrOV4/haCDW4zb1guZItoSq8foHCXrAnjBo/o=
 github.com/google/s2a-go v0.1.7/go.mod h1:50CgR4k1jNlWBu4UfS4AcfhVe1r6pdZPygJ3R8F0Qdw=
 github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 h1:El6M4kTTCOh6aBiKaUGG7oYTSPP8MxqL4YI3kZKwcP4=
@@ -493,8 +495,8 @@ google.golang.org/appengine v1.6.7/go.mod h1:8WjMMxjGQR8xUklV/ARdw2HLXBOI7O7uCID
 google.golang.org/genproto v0.0.0-20180817151627-c66870c02cf8/go.mod h1:JiN7NxoALGmiZfu7CAH4rXhgtRTLTxftemlI0sWmxmc=
 google.golang.org/genproto v0.0.0-20190819201941-24fa4b261c55/go.mod h1:DMBHOl98Agz4BDEuKkezgsaosCRResVns1a3J2ZsMNc=
 google.golang.org/genproto v0.0.0-20200526211855-cb27e3aa2013/go.mod h1:NbSheEEYHJ7i3ixzK3sjbqSGDJWnxyFXZblF3eUsNvo=
-google.golang.org/genproto v0.0.0-20240401170217-c3f982113cda h1:wu/KJm9KJwpfHWhkkZGohVC6KRrc1oJNr4jwtQMOQXw=
-google.golang.org/genproto v0.0.0-20240401170217-c3f982113cda/go.mod h1:g2LLCvCeCSir/JJSWosk19BR4NVxGqHUC6rxIRsd7Aw=
+google.golang.org/genproto v0.0.0-20240227224415-6ceb2ff114de h1:F6qOa9AZTYJXOUEr4jDysRDLrm4PHePlge4v4TGAlxY=
+google.golang.org/genproto v0.0.0-20240227224415-6ceb2ff114de/go.mod h1:VUhTRKeHn9wwcdrk73nvdC9gF178Tzhmt/qyaFcPLSo=
 google.golang.org/genproto/googleapis/api v0.0.0-20240415180920-8c6c420018be h1:Zz7rLWqp0ApfsR/l7+zSHhY3PMiH2xqgxlfYfAfNpoU=
 google.golang.org/genproto/googleapis/api v0.0.0-20240415180920-8c6c420018be/go.mod h1:dvdCTIoAGbkWbcIKBniID56/7XHTt6WfxXNMxuziJ+w=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20240429193739-8cf5692501f6 h1:DujSIu+2tC9Ht0aPNA7jgj23Iq8Ewi5sgkQ++wdvonE=

sqlconnect/internal/bigquery/db.go

@@ -30,6 +30,9 @@ func NewDB(configJSON json.RawMessage) (*DB, error) {
 
  db := sql.OpenDB(driver.NewConnector(
  config.ProjectID,
+ driver.Config{
+ JobRateLimitExceededRetryEnabled: true,
+ },
  option.WithCredentialsJSON([]byte(config.CredentialsJSON))),
  )
 

sqlconnect/internal/bigquery/driver/connection.go

@@ -5,12 +5,16 @@ import (
  "database/sql/driver"
  "errors"
  "fmt"
+ "strings"
+ "time"
 
  "cloud.google.com/go/bigquery"
+ "github.com/cenkalti/backoff/v4"
  "google.golang.org/api/iterator"
 )
 
 type bigQueryConnection struct {
+ config Config
  ctx context.Context
  client *bigquery.Client
  closed bool
@@ -76,3 +80,28 @@ func (bigQueryConnection) CheckNamedValue(*driver.NamedValue) error {
 func (connection *bigQueryConnection) BigqueryClient() *bigquery.Client {
  return connection.client
 }
+
+// readWithBackoff will retry the read operation if the error is [jobRateLimitExceeded] and the config is set to retry rate limit errors
+//
+// TODO: this should no longer be needed once this fix is released:
+// https://github.com/googleapis/google-cloud-go/pull/9726
+func (connection *bigQueryConnection) readWithBackoff(ctx context.Context, query *bigquery.Query) (it *bigquery.RowIterator, err error) {
+ if !connection.config.JobRateLimitExceededRetryEnabled {
+ return query.Read(ctx)
+ }
+ // mimicking google's own retry backoff settings
+ // https://github.com/googleapis/google-cloud-go/blob/b2e704d9d287445304d2b6030b6e35a4eb8be80a/bigquery/bigquery.go#L236
+ retry := backoff.WithContext(backoff.NewExponentialBackOff(
+ backoff.WithInitialInterval(1*time.Second),
+ backoff.WithMaxInterval(32*time.Second),
+ backoff.WithMultiplier(2),
+ ), ctx)
+ _ = backoff.Retry(func() error {
+ it, err = query.Read(ctx)
+ if err != nil && (strings.Contains(err.Error(), "jobRateLimitExceeded")) {
+ return err
+ }
+ return nil
+ }, retry)
+ return it, err
+}

sqlconnect/internal/bigquery/driver/connector.go

@@ -8,15 +8,21 @@ import (
  "google.golang.org/api/option"
 )
 
-func NewConnector(projectID string, opts ...option.ClientOption) driver.Connector {
+type Config struct {
+ JobRateLimitExceededRetryEnabled bool // Enable jobRateLimitExceeded retries: default false
+}
+
+func NewConnector(projectID string, config Config, opts ...option.ClientOption) driver.Connector {
  return &bigQueryConnector{
  projectID: projectID,
+ config: config,
  opts: opts,
  }
 }
 
 type bigQueryConnector struct {
  projectID string
+ config Config
  opts []option.ClientOption
 }
 
@@ -27,6 +33,7 @@ func (c *bigQueryConnector) Connect(ctx context.Context) (driver.Conn, error) {
  }
 
  return &bigQueryConnection{
+ config: c.config,
  ctx: ctx,
  client: client,
  }, nil

sqlconnect/internal/bigquery/driver/driver_test.go

@@ -34,7 +34,7 @@ func TestBigqueryDriver(t *testing.T) {
  require.NoError(t, json.Unmarshal([]byte(configJSON), &c))
 
  t.Run("OpenDB", func(t *testing.T) {
- db := sql.OpenDB(driver.NewConnector(c.ProjectID, option.WithCredentialsJSON([]byte(c.CredentialsJSON))))
+ db := sql.OpenDB(driver.NewConnector(c.ProjectID, driver.Config{}, option.WithCredentialsJSON([]byte(c.CredentialsJSON))))
  t.Cleanup(func() {
  require.NoError(t, db.Close(), "it should be able to close the database connection")
  })
@@ -225,9 +225,32 @@ func TestBigqueryDriver(t *testing.T) {
  })
  })
 
+ // TODO: this test should start failing once this fix is released:
+ // https://github.com/googleapis/google-cloud-go/pull/9726
  t.Run("jobRateLimitExceeded", func(t *testing.T) {
+ t.Run("selecting from information schema from multiple go-routines should lead to a jobRateLimitExceeded error", func(t *testing.T) {
+ db := sql.OpenDB(driver.NewConnector(c.ProjectID, driver.Config{}, option.WithCredentialsJSON([]byte(c.CredentialsJSON))))
+ t.Cleanup(func() {
+ require.NoError(t, db.Close(), "it should be able to close the database connection")
+ })
+ ctx, cancel := context.WithCancel(context.Background())
+ t.Cleanup(cancel)
+
+ g, ctx := errgroup.WithContext(ctx)
+ g.SetLimit(10)
+ for i := 0; i < 200; i++ {
+ g.Go(func() error {
+ _, err := db.ExecContext(ctx, "SELECT * FROM INFORMATION_SCHEMA.SCHEMATA WHERE SCHEMA_NAME = 'invalid'")
+ return err
+ })
+ }
+ err := g.Wait()
+ require.Error(t, err, "it should return an error")
+ require.Contains(t, err.Error(), "jobRateLimitExceeded", "it should return a jobRateLimitExceeded error")
+ })
+
  t.Run("selecting from information schema from multiple go-routines and having rate limit retries enabled, shoudln't lead to a jobRateLimitExceeded error", func(t *testing.T) {
- db := sql.OpenDB(driver.NewConnector(c.ProjectID, option.WithCredentialsJSON([]byte(c.CredentialsJSON))))
+ db := sql.OpenDB(driver.NewConnector(c.ProjectID, driver.Config{JobRateLimitExceededRetryEnabled: true}, option.WithCredentialsJSON([]byte(c.CredentialsJSON))))
  t.Cleanup(func() {
  require.NoError(t, db.Close(), "it should be able to close the database connection")
  })

sqlconnect/internal/bigquery/driver/statement.go

@@ -41,7 +41,7 @@ func (statement *bigQueryStatement) ExecContext(ctx context.Context, args []driv
  return nil, err
  }
 
- rowIterator, err := query.Read(ctx)
+ rowIterator, err := statement.connection.readWithBackoff(ctx, query)
  if err != nil {
  return nil, err
  }
@@ -55,7 +55,7 @@ func (statement *bigQueryStatement) QueryContext(ctx context.Context, args []dri
  return nil, err
  }
 
- rowIterator, err := query.Read(ctx)
+ rowIterator, err := statement.connection.readWithBackoff(ctx, query)
  if err != nil {
  return nil, err
  }