fix(bigquery): cache total rows count (#12230)

alvarowolfx · web-flow · commit 202dce02888c · 2025-05-09T17:14:25.000Z
Total rows on fast query path was using the amount of rows returned on the first page, which was incorrect for result sets with multiple pages. Supersedes #11874 Fixes #11873
diff --git a/bigquery/integration_test.go b/bigquery/integration_test.go
@@ -424,6 +424,31 @@ func TestIntegration_QueryContextTimeout(t *testing.T) {
 	}
 }
 
+func TestIntegration_QueryCachedTotalRows(t *testing.T) {
+	if client == nil {
+		t.Skip("Integration tests skipped")
+	}
+	ctx := context.Background()
+
+	q := client.Query("select * from unnest(generate_array(1,1000000))") // force multiple pages
+	it, err := q.Read(ctx)
+	if err != nil {
+		t.Errorf("Read() error: %v", err)
+	}
+	var values []Value
+	err = it.Next(&values)
+	if err == iterator.Done {
+		t.Error("Expected iterator to have values, but found iterator.Done")
+	}
+	if err != nil {
+		t.Errorf("Expected iterator to read a value without errors, but found: %v", err)
+	}
+	expectedTotal := uint64(1000000)
+	if it.TotalRows != expectedTotal {
+		t.Errorf("Expected total rows to be %d, but found: %d", expectedTotal, it.TotalRows)
+	}
+}
+
 func TestIntegration_SnapshotRestoreClone(t *testing.T) {
 
 	if client == nil {
diff --git a/bigquery/iterator.go b/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
 }
diff --git a/bigquery/iterator_test.go b/bigquery/iterator_test.go
@@ -80,8 +80,9 @@ func TestRowIteratorCacheBehavior(t *testing.T) {
 		{
 			// primary success case: schema in cache
 			inSource: &rowSource{
-				cachedSchema: testSchema,
-				cachedRows:   testRows,
+				cachedSchema:    testSchema,
+				cachedRows:      testRows,
+				cachedTotalRows: uint64(len(convertedRows)),
 			},
 			wantResult: &fetchPageResult{
 				totalRows: uint64(len(convertedRows)),
@@ -94,6 +95,7 @@ func TestRowIteratorCacheBehavior(t *testing.T) {
 			inSource: &rowSource{
 				cachedRows:      testRows,
 				cachedNextToken: "foo",
+				cachedTotalRows: uint64(len(convertedRows)),
 			},
 			inSchema: convertedSchema,
 			wantResult: &fetchPageResult{
diff --git a/bigquery/query.go b/bigquery/query.go
@@ -427,6 +427,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
 	}
