From bbff8ac817fb95af219c588bdadc0ca1784c5c0c Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Fri, 5 Jan 2024 10:44:15 -0800
Subject: [PATCH] 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      | 13 +++++++++++--
 bigquery/iterator_test.go | 33 +++++++++++++++++++++++++++++++++
 bigquery/query.go         |  3 ++-
 3 files changed, 46 insertions(+), 3 deletions(-)

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,