fix(bigquery) Fix totalRows returned by fetchCachedPage

The RowIterator.TotalRows is wrong when hitting the cached page on the
fast path (e.g. not changing start index or updating the page info).

fixes: #11873

Author: zecke

bigquery/iterator.go

@@ -229,6 +229,7 @@ type rowSource struct {
 	cachedRows      []*bq.TableRow
 	cachedSchema    *bq.TableSchema
 	cachedNextToken string
+	cachedTotalRows uint64
 }
 
 // fetchPageResult represents a page of rows returned from the backend.
@@ -382,6 +383,7 @@ func fetchCachedPage(src *rowSource, schema Schema, startIndex uint64, pageSize
 			// We can't progress with no schema, destroy references and return a miss.
 			src.cachedRows = nil
 			src.cachedNextToken = ""
+			src.cachedTotalRows = 0
 			return nil, errNoCacheData
 		}
 		schema = bqToSchema(src.cachedSchema)
@@ -400,23 +402,26 @@ func fetchCachedPage(src *rowSource, schema Schema, startIndex uint64, pageSize
 			src.cachedRows = nil
 			src.cachedSchema = nil
 			src.cachedNextToken = ""
+			src.cachedTotalRows = 0
 			return nil, err
 		}
 		result := &fetchPageResult{
 			pageToken: src.cachedNextToken,
 			rows:      converted,
 			schema:    schema,
-			totalRows: uint64(len(converted)),
+			totalRows: src.cachedTotalRows,
 		}
 		// clear cache references and return response.
 		src.cachedRows = nil
 		src.cachedSchema = nil
 		src.cachedNextToken = ""
+		src.cachedTotalRows = 0
 		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 = ""
+	src.cachedTotalRows = 0
 	return nil, errNoCacheData
 }

bigquery/iterator_test.go

@@ -80,11 +80,12 @@ func TestRowIteratorCacheBehavior(t *testing.T) {
 		{
 			// primary success case: schema in cache
 			inSource: &rowSource{
-				cachedSchema: testSchema,
-				cachedRows:   testRows,
+				cachedSchema:    testSchema,
+				cachedRows:      testRows,
+				cachedTotalRows: 10,
 			},
 			wantResult: &fetchPageResult{
-				totalRows: uint64(len(convertedRows)),
+				totalRows: uint64(10),
 				schema:    convertedSchema,
 				rows:      convertedRows,
 			},
@@ -94,10 +95,11 @@ func TestRowIteratorCacheBehavior(t *testing.T) {
 			inSource: &rowSource{
 				cachedRows:      testRows,
 				cachedNextToken: "foo",
+				cachedTotalRows: 20,
 			},
 			inSchema: convertedSchema,
 			wantResult: &fetchPageResult{
-				totalRows: uint64(len(convertedRows)),
+				totalRows: uint64(20),
 				schema:    convertedSchema,
 				rows:      convertedRows,
 				pageToken: "foo",

bigquery/query.go

@@ -418,6 +418,7 @@ func (q *Query) Read(ctx context.Context) (it *RowIterator, err error) {
 			cachedRows:      resp.Rows,
 			cachedSchema:    resp.Schema,
 			cachedNextToken: resp.PageToken,
+			cachedTotalRows: resp.TotalRows,
 		}
 		return newRowIterator(ctx, rowSource, fetchPage), nil
 	}