fix(bigquery): respect context during query execution (#7693)

This PR updates query polling and insertion to better respect context status.  Key change is to push context to the discovery library, but there are some small refactors to make use of context passing more consistent at various callsites.

Fixes: https://togithub.com/googleapis/google-cloud-go/issues/7604

Author: shollyman

bigquery/bigquery.go

@@ -173,7 +173,7 @@ func (c *Client) insertJob(ctx context.Context, job *bq.Job, media io.Reader) (*
 // Due to differences in options it supports, it cannot be used for all existing
 // jobs.insert requests that are query jobs.
 func (c *Client) runQuery(ctx context.Context, queryRequest *bq.QueryRequest) (*bq.QueryResponse, error) {
-	call := c.bqs.Jobs.Query(c.projectID, queryRequest)
+	call := c.bqs.Jobs.Query(c.projectID, queryRequest).Context(ctx)
 	setClientHeader(call.Header())
 
 	var res *bq.QueryResponse

bigquery/iam.go

@@ -58,12 +58,12 @@ func (c *bqIAMClient) GetWithVersion(ctx context.Context, resource string, reque
 			RequestedPolicyVersion: int64(requestedPolicyVersion),
 		},
 	}
-	call := c.bqs.Tables.GetIamPolicy(resource, iamReq)
+	call := c.bqs.Tables.GetIamPolicy(resource, iamReq).Context(ctx)
 	setClientHeader(call.Header())
 
 	var bqp *bq.Policy
 	err = runWithRetry(ctx, func() error {
-		bqp, err = call.Context(ctx).Do()
+		bqp, err = call.Do()
 		return err
 	})
 	if err != nil {
@@ -77,10 +77,10 @@ func (c *bqIAMClient) Set(ctx context.Context, resource string, p *iampb.Policy)
 	defer func() { trace.EndSpan(ctx, err) }()
 
 	bqp := iamToBigQueryPolicy(p)
-	call := c.bqs.Tables.SetIamPolicy(resource, &bq.SetIamPolicyRequest{Policy: bqp})
+	call := c.bqs.Tables.SetIamPolicy(resource, &bq.SetIamPolicyRequest{Policy: bqp}).Context(ctx)
 	setClientHeader(call.Header())
 	return runWithRetry(ctx, func() error {
-		_, err := call.Context(ctx).Do()
+		_, err := call.Do()
 		return err
 	})
 }
@@ -89,12 +89,12 @@ func (c *bqIAMClient) Test(ctx context.Context, resource string, perms []string)
 	ctx = trace.StartSpan(ctx, "cloud.google.com/go/bigquery.IAM.Test")
 	defer func() { trace.EndSpan(ctx, err) }()
 
-	call := c.bqs.Tables.TestIamPermissions(resource, &bq.TestIamPermissionsRequest{Permissions: perms})
+	call := c.bqs.Tables.TestIamPermissions(resource, &bq.TestIamPermissionsRequest{Permissions: perms}).Context(ctx)
 	setClientHeader(call.Header())
 
 	var res *bq.TestIamPermissionsResponse
 	err = runWithRetry(ctx, func() error {
-		res, err = call.Context(ctx).Do()
+		res, err = call.Do()
 		return err
 	})
 	if err != nil {

bigquery/inserter.go

@@ -177,8 +177,7 @@ func (u *Inserter) putMulti(ctx context.Context, src []ValueSaver) error {
 	if req == nil {
 		return nil
 	}
-	call := u.t.c.bqs.Tabledata.InsertAll(u.t.ProjectID, u.t.DatasetID, u.t.TableID, req)
-	call = call.Context(ctx)
+	call := u.t.c.bqs.Tabledata.InsertAll(u.t.ProjectID, u.t.DatasetID, u.t.TableID, req).Context(ctx)
 	setClientHeader(call.Header())
 	var res *bq.TableDataInsertAllResponse
 	err = runWithRetry(ctx, func() (err error) {

bigquery/integration_test.go

@@ -357,6 +357,26 @@ func TestIntegration_JobFrom(t *testing.T) {
 
 }
 
+func TestIntegration_QueryContextTimeout(t *testing.T) {
+	if client == nil {
+		t.Skip("Integration tests skipped")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
+	defer cancel()
+
+	q := client.Query("select count(*) from unnest(generate_array(1,1000000)), unnest(generate_array(1, 1000)) as foo")
+	q.DisableQueryCache = true
+	before := time.Now()
+	_, err := q.Read(ctx)
+	if err != context.DeadlineExceeded {
+		t.Errorf("Read() error, wanted %v, got %v", context.DeadlineExceeded, err)
+	}
+	wantMaxDur := 500 * time.Millisecond
+	if d := time.Since(before); d > wantMaxDur {
+		t.Errorf("return duration too long, wanted max %v got %v", wantMaxDur, d)
+	}
+}
+
 func TestIntegration_SnapshotRestoreClone(t *testing.T) {
 
 	if client == nil {
@@ -581,6 +601,7 @@ func TestIntegration_RangePartitioning(t *testing.T) {
 		t.Errorf("Range.Interval: got %v, wanted %v", gotInt64, wantedRange.Interval)
 	}
 }
+
 func TestIntegration_RemoveTimePartitioning(t *testing.T) {
 	if client == nil {
 		t.Skip("Integration tests skipped")

bigquery/iterator.go

@@ -292,7 +292,7 @@ func fetchTableResultPage(ctx context.Context, src *rowSource, schema Schema, st
 func fetchJobResultPage(ctx context.Context, src *rowSource, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
 	// reduce data transfered by leveraging api projections
 	projectedFields := []googleapi.Field{"rows", "pageToken", "totalRows"}
-	call := src.j.c.bqs.Jobs.GetQueryResults(src.j.projectID, src.j.jobID).Location(src.j.location)
+	call := src.j.c.bqs.Jobs.GetQueryResults(src.j.projectID, src.j.jobID).Location(src.j.location).Context(ctx)
 	if schema == nil {
 		// only project schema if we weren't supplied one.
 		projectedFields = append(projectedFields, "schema")
@@ -309,7 +309,7 @@ func fetchJobResultPage(ctx context.Context, src *rowSource, schema Schema, star
 	}
 	var res *bq.GetQueryResultsResponse
 	err := runWithRetry(ctx, func() (err error) {
-		res, err = call.Context(ctx).Do()
+		res, err = call.Do()
 		return err
 	})
 	if err != nil {