feat(bigquery): expose query id on row iterator if available (#9224)

This feature plumbs through the query ID, if available, as part of the row iterator.  Currently only queries run via the jobs.query RPC have query ID exposed.  This is related to the query preview features exposed in PR 8653.

bigquery/iterator.go

@@ -91,6 +91,14 @@ func (ri *RowIterator) SourceJob() *Job {
 	}
 }
 
+// QueryID returns a query ID if available, or an empty string.
+func (ri *RowIterator) QueryID() string {
+	if ri.src == nil {
+		return ""
+	}
+	return ri.src.queryID
+}
+
 // We declare a function signature for fetching results.  The primary reason
 // for this is to enable us to swap out the fetch function with alternate
 // implementations (e.g. to enable testing).
@@ -210,8 +218,9 @@ func (it *RowIterator) fetch(pageSize int, pageToken string) (string, error) {
 //     want to retain the data unnecessarily, and we expect that the backend
 //     can always provide them if needed.
 type rowSource struct {
-	j *Job
-	t *Table
+	j       *Job
+	t       *Table
+	queryID string
 
 	cachedRows      []*bq.TableRow
 	cachedSchema    *bq.TableSchema

bigquery/iterator_test.go

@@ -510,3 +510,36 @@ func TestIteratorSourceJob(t *testing.T) {
 		}
 	}
 }
+
+func TestIteratorQueryID(t *testing.T) {
+	testcases := []struct {
+		description string
+		src         *rowSource
+		want        string
+	}{
+		{
+			description: "nil source",
+			src:         nil,
+			want:        "",
+		},
+		{
+			description: "empty source",
+			src:         &rowSource{},
+			want:        "",
+		},
+		{
+			description: "populated id",
+			src:         &rowSource{queryID: "foo"},
+			want:        "foo",
+		},
+	}
+
+	for _, tc := range testcases {
+		// Don't pass a page func, we're not reading from the iterator.
+		it := newRowIterator(context.Background(), tc.src, nil)
+		got := it.QueryID()
+		if got != tc.want {
+			t.Errorf("%s: mismatch queryid, got %q want %q", tc.description, got, tc.want)
+		}
+	}
+}

bigquery/query.go

@@ -412,7 +412,8 @@ func (q *Query) Read(ctx context.Context) (it *RowIterator, err error) {
 			}
 		}
 		rowSource := &rowSource{
-			j: minimalJob,
+			j:       minimalJob,
+			queryID: resp.QueryId,
 			// RowIterator can precache results from the iterator to save a lookup.
 			cachedRows:      resp.Rows,
 			cachedSchema:    resp.Schema,