Skip to content

Commit

Permalink
Merge branch 'main' into issue-9111
Browse files Browse the repository at this point in the history
  • Loading branch information
rahul2393 authored Jan 5, 2024
2 parents 9fa72ab + bbff8ac commit 7358306
Show file tree
Hide file tree
Showing 4 changed files with 47 additions and 3 deletions.
13 changes: 11 additions & 2 deletions bigquery/iterator.go
Original file line number Diff line number Diff line change
Expand Up @@ -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).
Expand Down Expand Up @@ -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
Expand Down
33 changes: 33 additions & 0 deletions bigquery/iterator_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
}
3 changes: 2 additions & 1 deletion bigquery/query.go
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
1 change: 1 addition & 0 deletions cloudquotas/CHANGES.md
Original file line number Diff line number Diff line change
@@ -1 +1,2 @@
# Changes

0 comments on commit 7358306

Please sign in to comment.