(feat drivers/bigquery) Introduces jobStats for executeQuery

go/adbc/driver/bigquery/driver.go

@@ -174,6 +174,12 @@ const (
 	// query job if it fails to run.
 	OptionBoolQueryLinkFailedJob = "adbc.bigquery.sql.query.link_failed_job"
 
+	// OptionBoolStatementFetchJobStats instructs the driver to fetch full
+	// BigQuery job statistics
+	//
+	// Since these stats costs extra API calls, it is disabled by default.
+	OptionBoolStatementFetchJobStats = "adbc.bigquery.statement.fetch_job_stats"
+
 	OptionStringCopyTableSource           = "adbc.bigquery.copy_table.source"
 	OptionStringCopyTableDestination      = "adbc.bigquery.copy_table.destination"
 	OptionStringCopyTableWriteDisposition = "adbc.bigquery.copy_table.write_disposition"

go/adbc/driver/bigquery/job_stats.go

@@ -0,0 +1,137 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package bigquery
+
+import (
+	"context"
+	"strconv"
+
+	"cloud.google.com/go/bigquery"
+
+	"github.com/apache/arrow-go/v18/arrow"
+)
+
+// BIGQUERY:* schema-metadata keys — kept in sync with fs's
+// crates/dbt-adapter/src/record_batch.rs.
+const (
+	MetadataKeyBigqueryQueryID             = "BIGQUERY:query_id"
+	MetadataKeyBigqueryProjectID           = "BIGQUERY:project_id"
+	MetadataKeyBigqueryLocation            = "BIGQUERY:location"
+	MetadataKeyBigqueryStatementType       = "BIGQUERY:statement_type"
+	MetadataKeyBigqueryNumDMLAffectedRows  = "BIGQUERY:num_dml_affected_rows"
+	MetadataKeyBigqueryTotalBytesProcessed = "BIGQUERY:total_bytes_processed"
+	MetadataKeyBigqueryTotalBytesBilled    = "BIGQUERY:total_bytes_billed"
+	MetadataKeyBigquerySlotMillis          = "BIGQUERY:slot_ms"
+)
+
+// jobStats holds post-execution statistics from a BigQuery job. Identity
+// fields (JobID / ProjectID / Location) come straight from the *Job; the
+// rest come from QueryStatistics and are zero/empty when the job did not
+// produce query statistics (e.g. non-query jobs, or when the caller opted
+// out of stats fetching).
+type jobStats struct {
+	JobID              string
+	ProjectID          string
+	Location           string
+	StatementType      string
+	NumDMLAffectedRows int64
+	BytesProcessed     int64
+	BytesBilled        int64
+	SlotMillis         int64
+}
+
+func newJobStats(job *bigquery.Job) jobStats {
+	return jobStats{
+		JobID:     job.ID(),
+		ProjectID: job.ProjectID(),
+		Location:  job.Location(),
+	}
+}
+
+// fetch pulls the latest server-side JobStatus (one extra API call) and
+// populates the stats. Errors are swallowed: stats are best-effort so they
+// never fail query execution.
+func (s *jobStats) fetch(ctx context.Context, job *bigquery.Job) {
+	status, err := job.Status(ctx)
+	if err != nil || status == nil {
+		return
+	}
+	s.fromStatus(ctx, status)
+}
+
+// fromStatus populates the stats from an already-fetched JobStatus.
+func (s *jobStats) fromStatus(ctx context.Context, status *bigquery.JobStatus) {
+	if status.Statistics == nil {
+		return
+	}
+	s.BytesProcessed = status.Statistics.TotalBytesProcessed
+	qs, ok := status.Statistics.Details.(*bigquery.QueryStatistics)
+	if !ok {
+		return
+	}
+	s.BytesBilled = qs.TotalBytesBilled
+	s.NumDMLAffectedRows = qs.NumDMLAffectedRows
+	s.StatementType = qs.StatementType
+	s.SlotMillis = qs.SlotMillis
+	if qs.TotalBytesProcessed != 0 {
+		s.BytesProcessed = qs.TotalBytesProcessed
+	}
+
+	// For CREATE_TABLE_AS_SELECT it additionally issues a get_table call on
+	// the DDL target so NumDMLAffectedRows carries the destination row count,
+	// BigQuery does not populate NumDMLAffectedRows for DDL.
+	// reference: https://github.com/dbt-labs/dbt-adapters/blob/9fce78f44db248ba33832c0f65c884a5139c0169/dbt-bigquery/src/dbt/adapters/bigquery/connections.py#L345-L346
+	if qs.StatementType == "CREATE_TABLE_AS_SELECT" && qs.DDLTargetTable != nil {
+		if md, err := qs.DDLTargetTable.Metadata(ctx); err == nil {
+			s.NumDMLAffectedRows = int64(md.NumRows)
+		}
+	}
+}
+
+// attachToSchema returns a new schema with BIGQUERY:* metadata keys added
+// for every populated stat. When called with a nil receiver the schema is
+// returned unchanged (aside from a defensive metadata copy) — this lets
+// callers pass through opt-out queries without a branch.
+//
+// Numeric fields are emitted even when zero so consumers can distinguish
+// "stat available, value is 0" from "stat not provided".
+func (s *jobStats) attachToSchema(schema *arrow.Schema) *arrow.Schema {
+	meta := schema.Metadata().ToMap()
+	if s == nil {
+		finalMeta := arrow.MetadataFrom(meta)
+		return arrow.NewSchema(schema.Fields(), &finalMeta)
+	}
+	if s.JobID != "" {
+		meta[MetadataKeyBigqueryQueryID] = s.JobID
+	}
+	if s.ProjectID != "" {
+		meta[MetadataKeyBigqueryProjectID] = s.ProjectID
+	}
+	if s.Location != "" {
+		meta[MetadataKeyBigqueryLocation] = s.Location
+	}
+	if s.StatementType != "" {
+		meta[MetadataKeyBigqueryStatementType] = s.StatementType
+	}
+	meta[MetadataKeyBigqueryNumDMLAffectedRows] = strconv.FormatInt(s.NumDMLAffectedRows, 10)
+	meta[MetadataKeyBigqueryTotalBytesProcessed] = strconv.FormatInt(s.BytesProcessed, 10)
+	meta[MetadataKeyBigqueryTotalBytesBilled] = strconv.FormatInt(s.BytesBilled, 10)
+	meta[MetadataKeyBigquerySlotMillis] = strconv.FormatInt(s.SlotMillis, 10)
+	finalMeta := arrow.MetadataFrom(meta)
+	return arrow.NewSchema(schema.Fields(), &finalMeta)
+}

go/adbc/driver/bigquery/record_reader.go

@@ -38,10 +38,6 @@ import (
 	"google.golang.org/api/iterator"
 )
 
-const (
-	MetadataKeyBigqueryQueryID = "BIGQUERY:query_id"
-)
-
 type reader struct {
 	refCount   int64
 	schema     *arrow.Schema
@@ -64,14 +60,18 @@ func checkContext(ctx context.Context, maybeErr error) error {
 	return ctx.Err()
 }
 
-func runQuery(ctx context.Context, query *bigquery.Query, executeUpdate bool, linkFailedJob bool, alloc memory.Allocator) (bigquery.ArrowIterator, int64, error) {
+func runUpdate(ctx context.Context, query *bigquery.Query) (int64, error) {
+	if _, err := query.Run(ctx); err != nil {
+		return -1, err
+	}
+	return 0, nil
+}
+
+func runQuery(ctx context.Context, query *bigquery.Query, linkFailedJob bool, alloc memory.Allocator, outStats *jobStats) (bigquery.ArrowIterator, int64, error) {
 	job, err := query.Run(ctx)
 	if err != nil {
 		return nil, -1, err
 	}
-	if executeUpdate {
-		return nil, 0, nil
-	}
 
 	// The project id, location, and job id are all URL-safe:
 	// Project id and job id can only contain URL safe characters: https://cloud.google.com/bigquery/docs/reference/rest/v2/JobReference
@@ -109,6 +109,10 @@ func runQuery(ctx context.Context, query *bigquery.Query, executeUpdate bool, li
 
 	// Store job ID in query for adding to metadata
 	query.JobID = job.ID()
+	if outStats != nil {
+		*outStats = newJobStats(job)
+		outStats.fetch(ctx, job)
+	}
 	return arrowIterator, totalRows, nil
 }
 
@@ -134,8 +138,15 @@ func getQueryParameter(values arrow.RecordBatch, row int, parameterMode string)
 	return parameters, nil
 }
 
-func runPlainQuery(ctx context.Context, query *bigquery.Query, alloc memory.Allocator, resultRecordBufferSize int, linkFailedJob bool) (bigqueryRdr *reader, totalRows int64, err error) {
-	arrowIterator, totalRows, err := runQuery(ctx, query, false, linkFailedJob, alloc)
+func runPlainQuery(ctx context.Context, query *bigquery.Query, alloc memory.Allocator, resultRecordBufferSize int, linkFailedJob bool, fetchJobStats bool) (bigqueryRdr *reader, totalRows int64, err error) {
+	// Only allocate a stats struct (which triggers job.Status in runQuery) if
+	// the caller opted in. Otherwise runQuery skips the extra API call and the
+	// schema is returned without BIGQUERY:* stats metadata.
+	var statsPtr *jobStats
+	if fetchJobStats {
+		statsPtr = &jobStats{}
+	}
+	arrowIterator, totalRows, err := runQuery(ctx, query, linkFailedJob, alloc, statsPtr)
 	if err != nil {
 		return nil, -1, err
 	}
@@ -156,7 +167,7 @@ func runPlainQuery(ctx context.Context, query *bigquery.Query, alloc memory.Allo
 		}
 	}()
 
-	schema := schemaWithQueryId(rdr.Schema(), query)
+	schema := statsPtr.attachToSchema(rdr.Schema())
 
 	bigqueryRdr = &reader{
 		refCount:   1,
@@ -181,7 +192,7 @@ func runPlainQuery(ctx context.Context, query *bigquery.Query, alloc memory.Allo
 	return bigqueryRdr, totalRows, nil
 }
 
-func queryRecordWithSchemaCallback(ctx context.Context, group *errgroup.Group, query *bigquery.Query, rec arrow.RecordBatch, ch chan arrow.RecordBatch, parameterMode string, alloc memory.Allocator, rdrSchema func(schema *arrow.Schema), linkFailedJob bool) (int64, error) {
+func queryRecordWithSchemaCallback(ctx context.Context, group *errgroup.Group, query *bigquery.Query, rec arrow.RecordBatch, ch chan arrow.RecordBatch, parameterMode string, alloc memory.Allocator, rdrSchema func(schema *arrow.Schema), linkFailedJob bool, fetchJobStats bool) (int64, error) {
 	totalRows := int64(-1)
 	for i := 0; i < int(rec.NumRows()); i++ {
 		parameters, err := getQueryParameter(rec, i, parameterMode)
@@ -192,7 +203,11 @@ func queryRecordWithSchemaCallback(ctx context.Context, group *errgroup.Group, q
 			query.Parameters = parameters
 		}
 
-		arrowIterator, rows, err := runQuery(ctx, query, false, linkFailedJob, alloc)
+		var statsPtr *jobStats
+		if fetchJobStats {
+			statsPtr = &jobStats{}
+		}
+		arrowIterator, rows, err := runQuery(ctx, query, linkFailedJob, alloc, statsPtr)
 		if err != nil {
 			return -1, err
 		}
@@ -201,7 +216,7 @@ func queryRecordWithSchemaCallback(ctx context.Context, group *errgroup.Group, q
 		if err != nil {
 			return -1, err
 		}
-		schema := schemaWithQueryId(rdr.Schema(), query)
+		schema := statsPtr.attachToSchema(rdr.Schema())
 		rdrSchema(schema)
 		group.Go(func() error {
 			defer rdr.Release()
@@ -218,9 +233,9 @@ func queryRecordWithSchemaCallback(ctx context.Context, group *errgroup.Group, q
 
 // kicks off a goroutine for each endpoint and returns a reader which
 // gathers all of the records as they come in.
-func newRecordReader(ctx context.Context, query *bigquery.Query, boundParameters array.RecordReader, parameterMode string, alloc memory.Allocator, resultRecordBufferSize, prefetchConcurrency int, linkFailedJob bool) (bigqueryRdr *reader, totalRows int64, err error) {
+func newRecordReader(ctx context.Context, query *bigquery.Query, boundParameters array.RecordReader, parameterMode string, alloc memory.Allocator, resultRecordBufferSize, prefetchConcurrency int, linkFailedJob bool, fetchJobStats bool) (bigqueryRdr *reader, totalRows int64, err error) {
 	if boundParameters == nil {
-		return runPlainQuery(ctx, query, alloc, resultRecordBufferSize, linkFailedJob)
+		return runPlainQuery(ctx, query, alloc, resultRecordBufferSize, linkFailedJob, fetchJobStats)
 	}
 	defer boundParameters.Release()
 
@@ -258,7 +273,7 @@ func newRecordReader(ctx context.Context, query *bigquery.Query, boundParameters
 		// we don't need to call rec.Retain() here and call call rec.Release() in queryRecordWithSchemaCallback
 		batchRows, err := queryRecordWithSchemaCallback(ctx, group, query, rec, ch, parameterMode, alloc, func(schema *arrow.Schema) {
 			bigqueryRdr.schema = schema
-		}, linkFailedJob)
+		}, linkFailedJob, fetchJobStats)
 		if err != nil {
 			return nil, -1, err
 		}
@@ -269,14 +284,6 @@ func newRecordReader(ctx context.Context, query *bigquery.Query, boundParameters
 	return bigqueryRdr, totalRows, nil
 }
 
-func schemaWithQueryId(schema *arrow.Schema, query *bigquery.Query) *arrow.Schema {
-	meta := schema.Metadata().ToMap()
-	meta[MetadataKeyBigqueryQueryID] = query.JobID
-	finalMeta := arrow.MetadataFrom(meta)
-
-	return arrow.NewSchema(schema.Fields(), &finalMeta)
-}
-
 func (r *reader) Retain() {
 	atomic.AddInt64(&r.refCount, 1)
 }

go/adbc/driver/bigquery/statement.go

@@ -125,6 +125,9 @@ type statement struct {
 
 	// Wrap errors with a link to failed job
 	linkFailedJob bool
+
+	// Fetch BigQuery job statistics. Off by default.
+	fetchJobStats bool
 }
 
 func (st *statement) GetOptionBytes(key string) ([]byte, error) {
@@ -217,6 +220,8 @@ func (st *statement) GetOption(key string) (string, error) {
 		return strconv.FormatBool(st.queryConfig.CreateSession), nil
 	case OptionBoolQueryLinkFailedJob:
 		return strconv.FormatBool(st.linkFailedJob), nil
+	case OptionBoolStatementFetchJobStats:
+		return strconv.FormatBool(st.fetchJobStats), nil
 	case OptionBoolUseStorageApiDisabledClient:
 		return strconv.FormatBool(st.useStorageApiDisabledClient), nil
 	case OptionStringIngestFileDelimiter:
@@ -458,6 +463,13 @@ func (st *statement) SetOption(key string, v string) error {
 		} else {
 			return err
 		}
+	case OptionBoolStatementFetchJobStats:
+		val, err := strconv.ParseBool(v)
+		if err == nil {
+			st.fetchJobStats = val
+		} else {
+			return err
+		}
 	case OptionStringNotebookExecuteJobGscPath:
 		st.createNotebookExecuteJobGscPath = v
 	case OptionStringNotebookExecuteJobModelFileName:
@@ -582,7 +594,7 @@ func (st *statement) ExecuteQuery(ctx context.Context) (array.RecordReader, int6
 	}
 
 	ctx = context.WithValue(ctx, ContextKeyUseStorageApiDisabledClient, st.useStorageApiDisabledClient)
-	return newRecordReader(ctx, st.query(), rdr, st.parameterMode, st.cnxn.Alloc, st.resultRecordBufferSize, st.prefetchConcurrency, st.linkFailedJob)
+	return newRecordReader(ctx, st.query(), rdr, st.parameterMode, st.cnxn.Alloc, st.resultRecordBufferSize, st.prefetchConcurrency, st.linkFailedJob, st.fetchJobStats)
 }
 
 // ExecuteUpdate executes a statement that does not generate a result
@@ -594,33 +606,27 @@ func (st *statement) ExecuteUpdate(ctx context.Context) (int64, error) {
 	}
 
 	if boundParameters == nil {
-		_, totalRows, err := runQuery(ctx, st.query(), true, st.linkFailedJob, st.alloc)
-		if err != nil {
-			return -1, err
-		}
-		return totalRows, nil
-	} else {
-		totalRows := int64(0)
-		for boundParameters.Next() {
-			values := boundParameters.RecordBatch()
-			for i := 0; i < int(values.NumRows()); i++ {
-				parameters, err := getQueryParameter(values, i, st.parameterMode)
-				if err != nil {
-					return -1, err
-				}
-				if parameters != nil {
-					st.queryConfig.Parameters = parameters
-				}
-
-				_, currentRows, err := runQuery(ctx, st.query(), true, st.linkFailedJob, st.alloc)
-				if err != nil {
-					return -1, err
-				}
-				totalRows += currentRows
+		return runUpdate(ctx, st.query())
+	}
+	totalRows := int64(0)
+	for boundParameters.Next() {
+		values := boundParameters.RecordBatch()
+		for i := 0; i < int(values.NumRows()); i++ {
+			parameters, err := getQueryParameter(values, i, st.parameterMode)
+			if err != nil {
+				return -1, err
+			}
+			if parameters != nil {
+				st.queryConfig.Parameters = parameters
+			}
+			currentRows, err := runUpdate(ctx, st.query())
+			if err != nil {
+				return -1, err
 			}
+			totalRows += currentRows
 		}
-		return totalRows, nil
 	}
+	return totalRows, nil
 }
 
 // ExecuteSchema gets the schema of the result set of a query without executing it.