diff --git a/bigquery/iterator.go b/bigquery/iterator.go index 6339bd52b541..098a94a8e80c 100644 --- a/bigquery/iterator.go +++ b/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 diff --git a/bigquery/iterator_test.go b/bigquery/iterator_test.go index 99d299da3201..faf23b816a82 100644 --- a/bigquery/iterator_test.go +++ b/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) + } + } +} diff --git a/bigquery/query.go b/bigquery/query.go index a44e93227d72..1641ee84429f 100644 --- a/bigquery/query.go +++ b/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,