From 868d07dcdd70a9011b5d14f9abe18147c412a456 Mon Sep 17 00:00:00 2001
From: shollyman <shollyman@google.com>
Date: Wed, 15 Sep 2021 13:54:11 -0700
Subject: [PATCH] feat: expose the query source of a rowiterator via
 SourceJob() (#4748)

RowIterator can currently come from calling Read() on a Table,
Read() on a Job, and Read() directly from a Query.  The third invocation
is used for fast path queries and doesn't return a query identifier
suitable for gathering statistics.

This PR exposes a SourceJob() method on RowIterator to address this
issue.  Users still get the benefits of optimized query execution, but
can get the reference to a Job for looking up additional statistics
should they need it.

Fixes: https://github.com/googleapis/google-cloud-go/issues/4745
---
 bigquery/integration_test.go | 23 +++++++++++++++++++++
 bigquery/iterator.go         | 17 ++++++++++++++++
 bigquery/iterator_test.go    | 39 ++++++++++++++++++++++++++++++++++++
 3 files changed, 79 insertions(+)

diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
index 0df0587f1ed..2526392b37c 100644
--- a/bigquery/integration_test.go
+++ b/bigquery/integration_test.go
@@ -2197,6 +2197,29 @@ func TestIntegration_LegacyQuery(t *testing.T) {
 	}
 }
 
+func TestIntegration_IteratorSource(t *testing.T) {
+	if client == nil {
+		t.Skip("Integration tests skipped")
+	}
+	ctx := context.Background()
+	q := client.Query("SELECT 17 as foo")
+	it, err := q.Read(ctx)
+	if err != nil {
+		t.Errorf("Read: %v", err)
+	}
+	src := it.SourceJob()
+	if src == nil {
+		t.Errorf("wanted source job, got nil")
+	}
+	status, err := src.Status(ctx)
+	if err != nil {
+		t.Errorf("Status: %v", err)
+	}
+	if status == nil {
+		t.Errorf("got nil status")
+	}
+}
+
 func TestIntegration_QueryExternalHivePartitioning(t *testing.T) {
 	if client == nil {
 		t.Skip("Integration tests skipped")
diff --git a/bigquery/iterator.go b/bigquery/iterator.go
index fce76a3c283..4bfa3405b96 100644
--- a/bigquery/iterator.go
+++ b/bigquery/iterator.go
@@ -63,6 +63,23 @@ type RowIterator struct {
 	structLoader structLoader // used to populate a pointer to a struct
 }
 
+// SourceJob returns an instance of a Job if the RowIterator is backed by a query,
+// or a nil.
+func (ri *RowIterator) SourceJob() *Job {
+	if ri.src == nil {
+		return nil
+	}
+	if ri.src.j == nil {
+		return nil
+	}
+	return &Job{
+		c:         ri.src.j.c,
+		projectID: ri.src.j.projectID,
+		location:  ri.src.j.location,
+		jobID:     ri.src.j.jobID,
+	}
+}
+
 // 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).
diff --git a/bigquery/iterator_test.go b/bigquery/iterator_test.go
index 14a21112f6b..99d299da320 100644
--- a/bigquery/iterator_test.go
+++ b/bigquery/iterator_test.go
@@ -471,3 +471,42 @@ func TestIteratorNextTypes(t *testing.T) {
 		}
 	}
 }
+
+func TestIteratorSourceJob(t *testing.T) {
+	testcases := []struct {
+		description string
+		src         *rowSource
+		wantJob     *Job
+	}{
+		{
+			description: "nil source",
+			src:         nil,
+			wantJob:     nil,
+		},
+		{
+			description: "empty source",
+			src:         &rowSource{},
+			wantJob:     nil,
+		},
+		{
+			description: "table source",
+			src:         &rowSource{t: &Table{ProjectID: "p", DatasetID: "d", TableID: "t"}},
+			wantJob:     nil,
+		},
+		{
+			description: "job source",
+			src:         &rowSource{j: &Job{projectID: "p", location: "l", jobID: "j"}},
+			wantJob:     &Job{projectID: "p", location: "l", jobID: "j"},
+		},
+	}
+
+	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.SourceJob()
+		// AllowUnexported because of the embedded client reference, which we're ignoring.
+		if !cmp.Equal(got, tc.wantJob, cmp.AllowUnexported(Job{})) {
+			t.Errorf("%s: mismatch on SourceJob, got %v want %v", tc.description, got, tc.wantJob)
+		}
+	}
+}