Repository: googleapis/google-cloud-go
Pull request #2484 (2020-06-18 17:12:33)

bigquery: refactor RowIterator to support fast query path
This CL significantly revamps the rowiterator.

Previously, row iteration necessitated a reference to a Table object.
For jobs, this meant we would consult job statistics to find the
reference to the destination table.

With this change, we switch to using the jobs.getQueryResults RPC to get
rows from a job, while still allowing table reads to use the existing
tabledata.list mechanism.

This change also includes support for a caching abstraction in the row
iterator, which is added but not yet leveraged.  This will be wired up
in subsequent PRs that change the query execution path.

Additionally, this PR address some other reported issues against the library as the use of jobs.getQueryResults allows the service to make decisions about the source of results (e.g. in cases where a scripting/multistatement query returns resultsets from a child job).

Fixes: #1960, #1974

shollyman: still issues with integration tests, tracking those down

Patch:
diff --git a/bigquery/iterator.go b/bigquery/iterator.go
index 298143cbb4..6d28ea02f4 100644
--- a/bigquery/iterator.go
+++ b/bigquery/iterator.go
@@ -16,37 +16,37 @@ package bigquery
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"reflect"
 
 	bq "google.golang.org/api/bigquery/v2"
+	"google.golang.org/api/googleapi"
 	"google.golang.org/api/iterator"
 )
 
 // Construct a RowIterator.
-// If pf is nil, there are no rows in the result set.
-func newRowIterator(ctx context.Context, t *Table, pf pageFetcher) *RowIterator {
+func newRowIterator(ctx context.Context, src *rowSource, pf pageFetcher) *RowIterator {
 	it := &RowIterator{
-		ctx:   ctx,
-		table: t,
-		pf:    pf,
-	}
-	if pf != nil {
-		it.pageInfo, it.nextFunc = iterator.NewPageInfo(
-			it.fetch,
-			func() int { return len(it.rows) },
-			func() interface{} { r := it.rows; it.rows = nil; return r })
+		ctx: ctx,
+		src: src,
+		pf:  pf,
 	}
+	it.pageInfo, it.nextFunc = iterator.NewPageInfo(
+		it.fetch,
+		func() int { return len(it.rows) },
+		func() interface{} { r := it.rows; it.rows = nil; return r })
 	return it
 }
 
 // A RowIterator provides access to the result of a BigQuery lookup.
 type RowIterator struct {
-	ctx      context.Context
-	table    *Table
-	pf       pageFetcher
+	ctx context.Context
+	src *rowSource
+
 	pageInfo *iterator.PageInfo
 	nextFunc func() error
+	pf       pageFetcher
 
 	// StartIndex can be set before the first call to Next. If PageInfo().Token
 	// is also set, StartIndex is ignored.
@@ -63,6 +63,11 @@ type RowIterator struct {
 	structLoader structLoader // used to populate a pointer to a struct
 }
 
+// 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).
+type pageFetcher func(ctx context.Context, _ *rowSource, _ Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error)
+
 // Next loads the next row into dst. Its return value is iterator.Done if there
 // are no more results. Once Next returns iterator.Done, all subsequent calls
 // will return iterator.Done.
@@ -103,9 +108,6 @@ type RowIterator struct {
 // NullDateTime. You can also use a *[]Value or *map[string]Value to read from a
 // table with NULLs.
 func (it *RowIterator) Next(dst interface{}) error {
-	if it.pf == nil { // There are no rows in the result set.
-		return iterator.Done
-	}
 	var vl ValueLoader
 	switch dst := dst.(type) {
 	case ValueLoader:
@@ -145,19 +147,42 @@ func isStructPtr(x interface{}) bool {
 func (it *RowIterator) PageInfo() *iterator.PageInfo { return it.pageInfo }
 
 func (it *RowIterator) fetch(pageSize int, pageToken string) (string, error) {
-	res, err := it.pf(it.ctx, it.table, it.Schema, it.StartIndex, int64(pageSize), pageToken)
+	res, err := it.pf(it.ctx, it.src, it.Schema, it.StartIndex, int64(pageSize), pageToken)
 	if err != nil {
 		return "", err
 	}
 	it.rows = append(it.rows, res.rows...)
-	it.Schema = res.schema
+	if it.Schema == nil {
+		it.Schema = res.schema
+	}
 	it.TotalRows = res.totalRows
 	return res.pageToken, nil
 }
 
-// A pageFetcher returns a page of rows from a destination table.
-type pageFetcher func(ctx context.Context, _ *Table, _ Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error)
+// rowSource represents one of the multiple sources of data for a row iterator.
+// Rows can be read directly from a BigQuery table or from a job reference.
+// If a job is present, that's treated as the authoritative source.
+//
+// rowSource can also cache results for special situations, primarily for the
+// fast execution query path which can return status, rows, and schema all at
+// once.  Our cache data expectations are as follows:
+//
+// * We can only cache data from the start of a source.
+// * We need to cache schema, rows, and next page token to effective service
+//   a request from cache.
+// * cache references are destroyed as soon as they're interrogated.  We don't
+//   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
+
+	cachedRows      []*bq.TableRow
+	cachedSchema    *bq.TableSchema
+	cachedNextToken string
+}
 
+// fetchPageResult represents a page of rows returned from the backend.
 type fetchPageResult struct {
 	pageToken string
 	rows      [][]Value
@@ -165,8 +190,26 @@ type fetchPageResult struct {
 	schema    Schema
 }
 
-// fetchPage gets a page of rows from t.
-func fetchPage(ctx context.Context, t *Table, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
+// fetchPage is our generalized fetch mechanism.  It interrogates from cache, and
+// then dispatches to either the appropriate job or table-based backend mechanism
+// as needed.
+func fetchPage(ctx context.Context, src *rowSource, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
+	result, err := fetchCachedPage(ctx, src, schema, startIndex, pageSize, pageToken)
+	if err != nil {
+		if err != errNoCacheData {
+			// This likely means something more severe, like a problem with schema.
+			return nil, err
+		}
+		// If we failed to fet data from cache, invoke the appropriate service method.
+		if src.j != nil {
+			return fetchJobResultPage(ctx, src, schema, startIndex, pageSize, pageToken)
+		}
+		return fetchTableResultPage(ctx, src, schema, startIndex, pageSize, pageToken)
+	}
+	return result, nil
+}
+
+func fetchTableResultPage(ctx context.Context, src *rowSource, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
 	// Fetch the table schema in the background, if necessary.
 	errc := make(chan error, 1)
 	if schema != nil {
@@ -175,7 +218,7 @@ func fetchPage(ctx context.Context, t *Table, schema Schema, startIndex uint64,
 		go func() {
 			var bqt *bq.Table
 			err := runWithRetry(ctx, func() (err error) {
-				bqt, err = t.c.bqs.Tables.Get(t.ProjectID, t.DatasetID, t.TableID).
+				bqt, err = src.t.c.bqs.Tables.Get(src.t.ProjectID, src.t.DatasetID, src.t.TableID).
 					Fields("schema").
 					Context(ctx).
 					Do()
@@ -187,7 +230,7 @@ func fetchPage(ctx context.Context, t *Table, schema Schema, startIndex uint64,
 			errc <- err
 		}()
 	}
-	call := t.c.bqs.Tabledata.List(t.ProjectID, t.DatasetID, t.TableID)
+	call := src.t.c.bqs.Tabledata.List(src.t.ProjectID, src.t.DatasetID, src.t.TableID)
 	setClientHeader(call.Header())
 	if pageToken != "" {
 		call.PageToken(pageToken)
@@ -220,3 +263,100 @@ func fetchPage(ctx context.Context, t *Table, schema Schema, startIndex uint64,
 		schema:    schema,
 	}, nil
 }
+
+func fetchJobResultPage(ctx context.Context, src *rowSource, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
+	// reduce data transfered by leveraging api projections
+	projectedFields := []googleapi.Field{"rows", "pageToken", "totalRows"}
+	call := src.j.c.bqs.Jobs.GetQueryResults(src.j.projectID, src.j.jobID).Location(src.j.location)
+	call = call.Fields(projectedFields...)
+	if schema == nil {
+		// only project schema if we weren't supplied one.
+		call = call.Fields("schema")
+	}
+	setClientHeader(call.Header())
+	if pageToken != "" {
+		call.PageToken(pageToken)
+	} else {
+		call.StartIndex(startIndex)
+	}
+	if pageSize > 0 {
+		call.MaxResults(pageSize)
+	}
+	var res *bq.GetQueryResultsResponse
+	err := runWithRetry(ctx, func() (err error) {
+		res, err = call.Context(ctx).Do()
+		return err
+	})
+	if err != nil {
+		return nil, err
+	}
+	// Populate schema in the rowsource if it's missing
+	if schema == nil {
+		schema = bqToSchema(res.Schema)
+	}
+
+	rows, err := convertRows(res.Rows, schema)
+	if err != nil {
+		return nil, err
+	}
+	return &fetchPageResult{
+		pageToken: res.PageToken,
+		rows:      rows,
+		totalRows: uint64(res.TotalRows),
+		schema:    schema,
+	}, nil
+}
+
+var errNoCacheData = errors.New("No rows in rowSource cache")
+
+// fetchCachedPage attempts to service the first page of results.  For the jobs path specifically, we have an
+// opportunity to fetch rows before the iterator is constructed, and thus serve that data as the first request
+// without an unnecessary network round trip.
+func fetchCachedPage(ctx context.Context, src *rowSource, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
+	// we have no cached data
+	if src.cachedRows == nil {
+		return nil, errNoCacheData
+	}
+	// we have no schema for decoding.  convert from the cached representation if available.
+	if schema == nil {
+		if src.cachedSchema == nil {
+			// We can't progress with no schema, destroy references and return a miss.
+			src.cachedRows = nil
+			src.cachedNextToken = ""
+			return nil, errNoCacheData
+		}
+		schema = bqToSchema(src.cachedSchema)
+	}
+	// Only serve from cache where we're confident we know someone's asking for the first page
+	// without having to align data.
+	//
+	// Future consideration: we could service pagesizes smaller than the cache if we're willing to handle generation
+	// of pageTokens for the cache.
+	if pageToken == "" &&
+		startIndex == 0 &&
+		(pageSize == 0 || pageSize == int64(len(src.cachedRows))) {
+		converted, err := convertRows(src.cachedRows, schema)
+		if err != nil {
+			// destroy cache references and return error
+			src.cachedRows = nil
+			src.cachedSchema = nil
+			src.cachedNextToken = ""
+			return nil, err
+		}
+		result := &fetchPageResult{
+			pageToken: src.cachedNextToken,
+			rows:      converted,
+			schema:    schema,
+			totalRows: uint64(len(converted)),
+		}
+		// clear cache references and return response.
+		src.cachedRows = nil
+		src.cachedNextToken = ""
+		return result, nil
+	}
+	// All other cases are invalid.  Destroy any cache references on the way out the door.
+	src.cachedRows = nil
+	src.cachedSchema = nil
+	src.cachedNextToken = ""
+	return nil, errNoCacheData
+}
diff --git a/bigquery/job.go b/bigquery/job.go
index 6a519733ef..edaa136951 100644
--- a/bigquery/job.go
+++ b/bigquery/job.go
@@ -284,21 +284,18 @@ func (j *Job) read(ctx context.Context, waitForQuery func(context.Context, strin
 	if !j.isQuery() {
 		return nil, errors.New("bigquery: cannot read from a non-query job")
 	}
-	destTable := j.config.Query.DestinationTable
-	// The destination table should only be nil if there was a query error.
-	projectID := j.projectID
-	if destTable != nil && projectID != destTable.ProjectId {
-		return nil, fmt.Errorf("bigquery: job project ID is %q, but destination table's is %q", projectID, destTable.ProjectId)
-	}
-	schema, totalRows, err := waitForQuery(ctx, projectID)
+	schema, totalRows, err := waitForQuery(ctx, j.projectID)
 	if err != nil {
 		return nil, err
 	}
-	if destTable == nil {
-		return nil, errors.New("bigquery: query job missing destination table")
+	// Shave off some potential overhead by only retaining the minimal job representation in the iterator.
+	itJob := &Job{
+		c:         j.c,
+		projectID: j.projectID,
+		jobID:     j.jobID,
+		location:  j.location,
 	}
-	dt := bqToTable(destTable, j.c)
-	it := newRowIterator(ctx, dt, pf)
+	it := newRowIterator(ctx, &rowSource{j: itJob}, pf)
 	it.Schema = schema
 	it.TotalRows = totalRows
 	return it, nil
diff --git a/bigquery/table.go b/bigquery/table.go
index b06d2ff10d..05e1270e8f 100644
--- a/bigquery/table.go
+++ b/bigquery/table.go
@@ -614,7 +614,7 @@ func (t *Table) Read(ctx context.Context) *RowIterator {
 }
 
 func (t *Table) read(ctx context.Context, pf pageFetcher) *RowIterator {
-	return newRowIterator(ctx, t, pf)
+	return newRowIterator(ctx, &rowSource{t: t}, pf)
 }
 
 // NeverExpire is a sentinel value used to remove a table'e expiration time.


Test patch:
diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
index 7590f13b4d..77ee8f29e8 100644
--- a/bigquery/integration_test.go
+++ b/bigquery/integration_test.go
@@ -954,12 +954,12 @@ func TestIntegration_SimpleRowResults(t *testing.T) {
 			want:        [][]Value{},
 		},
 		{
-			// Note: currently CTAS returns the rows due to the destination table reference,
-			// but it's not clear that it should.
-			// https://github.com/googleapis/google-cloud-go/issues/1467 for followup.
+			// Previously this would return rows due to the destination reference being present
+			// in the job config, but switching to relying on jobs.getQueryResults allows the
+			// service to decide the behavior.
 			description: "ctas ddl",
 			query:       fmt.Sprintf("CREATE TABLE %s.%s AS SELECT 17 as foo", dataset.DatasetID, tableIDs.New()),
-			want:        [][]Value{{int64(17)}},
+			want:        nil,
 		},
 	}
 	for _, tc := range testCases {
@@ -976,6 +976,48 @@ func TestIntegration_SimpleRowResults(t *testing.T) {
 	}
 }
 
+func TestIntegration_RoutineStoredProcedure(t *testing.T) {
+	// Verifies we're exhibiting documented behavior, where we're expected
+	// to return the last resultset in a script as the response from a script
+	// job.
+	// https://github.com/googleapis/google-cloud-go/issues/1974
+	if client == nil {
+		t.Skip("Integration tests skipped")
+	}
+	ctx := context.Background()
+
+	// Define a simple stored procedure via DDL.
+	routineID := routineIDs.New()
+	routine := dataset.Routine(routineID)
+	sql := fmt.Sprintf(`
+		CREATE OR REPLACE PROCEDURE `+"`%s`"+`(val INT64)
+		BEGIN
+			SELECT CURRENT_TIMESTAMP() as ts;
+			SELECT val * 2 as f2;
+		END`,
+		routine.FullyQualifiedName())
+
+	if err := runQueryJob(ctx, sql); err != nil {
+		t.Fatal(err)
+	}
+	defer routine.Delete(ctx)
+
+	// Invoke the stored procedure.
+	sql = fmt.Sprintf(`
+	CALL `+"`%s`"+`(5)`,
+		routine.FullyQualifiedName())
+
+	q := client.Query(sql)
+	it, err := q.Read(ctx)
+	if err != nil {
+		t.Fatalf("query.Read: %v", err)
+	}
+
+	checkReadAndTotalRows(t,
+		"expect result set from procedure",
+		it, [][]Value{{int64(10)}})
+}
+
 func TestIntegration_InsertAndRead(t *testing.T) {
 	if client == nil {
 		t.Skip("Integration tests skipped")
diff --git a/bigquery/iterator_test.go b/bigquery/iterator_test.go
index 125796d7cc..14a21112f6 100644
--- a/bigquery/iterator_test.go
+++ b/bigquery/iterator_test.go
@@ -21,6 +21,8 @@ import (
 	"testing"
 
 	"cloud.google.com/go/internal/testutil"
+	"github.com/google/go-cmp/cmp"
+	bq "google.golang.org/api/bigquery/v2"
 	"google.golang.org/api/iterator"
 )
 
@@ -35,7 +37,7 @@ type pageFetcherStub struct {
 	err            error
 }
 
-func (pf *pageFetcherStub) fetchPage(ctx context.Context, _ *Table, _ Schema, _ uint64, _ int64, pageToken string) (*fetchPageResult, error) {
+func (pf *pageFetcherStub) fetchPage(ctx context.Context, _ *rowSource, _ Schema, _ uint64, _ int64, pageToken string) (*fetchPageResult, error) {
 	call, ok := pf.fetchResponses[pageToken]
 	if !ok {
 		pf.err = fmt.Errorf("Unexpected page token: %q", pageToken)
@@ -43,6 +45,115 @@ func (pf *pageFetcherStub) fetchPage(ctx context.Context, _ *Table, _ Schema, _
 	return call.result, call.err
 }
 
+func TestRowIteratorCacheBehavior(t *testing.T) {
+
+	testSchema := &bq.TableSchema{
+		Fields: []*bq.TableFieldSchema{
+			{Type: "INTEGER", Name: "field1"},
+			{Type: "STRING", Name: "field2"},
+		},
+	}
+	testRows := []*bq.TableRow{
+		{F: []*bq.TableCell{
+			{V: "1"},
+			{V: "foo"},
+		},
+		},
+	}
+	convertedSchema := bqToSchema(testSchema)
+
+	convertedRows, _ := convertRows(testRows, convertedSchema)
+
+	testCases := []struct {
+		inSource     *rowSource
+		inSchema     Schema
+		inStartIndex uint64
+		inPageSize   int64
+		inPageToken  string
+		wantErr      error
+		wantResult   *fetchPageResult
+	}{
+		{
+			inSource: &rowSource{},
+			wantErr:  errNoCacheData,
+		},
+		{
+			// primary success case: schema in cache
+			inSource: &rowSource{
+				cachedSchema: testSchema,
+				cachedRows:   testRows,
+			},
+			wantResult: &fetchPageResult{
+				totalRows: uint64(len(convertedRows)),
+				schema:    convertedSchema,
+				rows:      convertedRows,
+			},
+		},
+		{
+			// secondary success case: schema provided
+			inSource: &rowSource{
+				cachedRows:      testRows,
+				cachedNextToken: "foo",
+			},
+			inSchema: convertedSchema,
+			wantResult: &fetchPageResult{
+				totalRows: uint64(len(convertedRows)),
+				schema:    convertedSchema,
+				rows:      convertedRows,
+				pageToken: "foo",
+			},
+		},
+		{
+			// misaligned page size.
+			inSource: &rowSource{
+				cachedSchema: testSchema,
+				cachedRows:   testRows,
+			},
+			inPageSize: 99,
+			wantErr:    errNoCacheData,
+		},
+		{
+			// nonzero start.
+			inSource: &rowSource{
+				cachedSchema: testSchema,
+				cachedRows:   testRows,
+			},
+			inStartIndex: 1,
+			wantErr:      errNoCacheData,
+		},
+		{
+			// data without schema
+			inSource: &rowSource{
+				cachedSchema: testSchema,
+				cachedRows:   testRows,
+			},
+			inStartIndex: 1,
+			wantErr:      errNoCacheData,
+		},
+		{
+			// data without schema
+			inSource: &rowSource{
+				cachedSchema: testSchema,
+				cachedRows:   testRows,
+			},
+			inStartIndex: 1,
+			wantErr:      errNoCacheData,
+		},
+	}
+	for _, tc := range testCases {
+		gotResp, gotErr := fetchCachedPage(context.Background(), tc.inSource, tc.inSchema, tc.inStartIndex, tc.inPageSize, tc.inPageToken)
+		if gotErr != tc.wantErr {
+			t.Errorf("err mismatch.  got %v, want %v", gotErr, tc.wantErr)
+		} else {
+			if diff := testutil.Diff(gotResp, tc.wantResult,
+				cmp.AllowUnexported(fetchPageResult{}, rowSource{}, Job{}, Client{}, Table{})); diff != "" {
+				t.Errorf("response diff (got=-, want=+):\n%s", diff)
+			}
+		}
+	}
+
+}
+
 func TestIterator(t *testing.T) {
 	var (
 		iiSchema = Schema{
diff --git a/bigquery/read_test.go b/bigquery/read_test.go
index 883430a00c..66f6136828 100644
--- a/bigquery/read_test.go
+++ b/bigquery/read_test.go
@@ -26,7 +26,7 @@ import (
 )
 
 type pageFetcherArgs struct {
-	table      *Table
+	src        *rowSource
 	schema     Schema
 	startIndex uint64
 	pageSize   int64
@@ -43,9 +43,9 @@ type pageFetcherReadStub struct {
 	calls []pageFetcherArgs
 }
 
-func (s *pageFetcherReadStub) fetchPage(ctx context.Context, t *Table, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
+func (s *pageFetcherReadStub) fetchPage(ctx context.Context, src *rowSource, schema Schema, startIndex uint64, pageSize int64, pageToken string) (*fetchPageResult, error) {
 	s.calls = append(s.calls,
-		pageFetcherArgs{t, schema, startIndex, pageSize, pageToken})
+		pageFetcherArgs{src, schema, startIndex, pageSize, pageToken})
 	result := &fetchPageResult{
 		pageToken: s.pageTokens[pageToken],
 		rows:      s.values[0],
@@ -156,7 +156,7 @@ func TestNoMoreValues(t *testing.T) {
 
 var errBang = errors.New("bang")
 
-func errorFetchPage(context.Context, *Table, Schema, uint64, int64, string) (*fetchPageResult, error) {
+func errorFetchPage(context.Context, *rowSource, Schema, uint64, int64, string) (*fetchPageResult, error) {
 	return nil, errBang
 }
 
@@ -184,11 +184,13 @@ func TestReadTabledataOptions(t *testing.T) {
 		t.Fatal(err)
 	}
 	want := []pageFetcherArgs{{
-		table:     tr,
+		src: &rowSource{
+			t: tr,
+		},
 		pageSize:  5,
 		pageToken: "",
 	}}
-	if diff := testutil.Diff(s.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, pageFetcherReadStub{}, Table{}, Client{})); diff != "" {
+	if diff := testutil.Diff(s.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, pageFetcherReadStub{}, rowSource{}, Table{}, Client{})); diff != "" {
 		t.Errorf("reading (got=-, want=+):\n%s", diff)
 	}
 }
@@ -223,11 +225,18 @@ func TestReadQueryOptions(t *testing.T) {
 	}
 
 	want := []pageFetcherArgs{{
-		table:     bqToTable(tr, c),
+		src: &rowSource{
+			j: &Job{
+				c:         queryJob.c,
+				jobID:     queryJob.jobID,
+				projectID: queryJob.projectID,
+				location:  queryJob.location,
+			},
+		},
 		pageSize:  5,
 		pageToken: "",
 	}}
-	if !testutil.Equal(pf.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, Table{}, Client{})) {
+	if !testutil.Equal(pf.calls, want, cmp.AllowUnexported(pageFetcherArgs{}, rowSource{}, Job{}, Client{})) {
 		t.Errorf("reading: got:\n%v\nwant:\n%v", pf.calls, want)
 	}
 }