diff --git a/go/adbc/driver/bigquery/driver.go b/go/adbc/driver/bigquery/driver.go index fdb25358fc..4acdb16e2a 100644 --- a/go/adbc/driver/bigquery/driver.go +++ b/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" diff --git a/go/adbc/driver/bigquery/job_stats.go b/go/adbc/driver/bigquery/job_stats.go new file mode 100644 index 0000000000..541c4cb92f --- /dev/null +++ b/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) +} diff --git a/go/adbc/driver/bigquery/record_reader.go b/go/adbc/driver/bigquery/record_reader.go index 301003c0fd..c7816e25e6 100644 --- a/go/adbc/driver/bigquery/record_reader.go +++ b/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) } diff --git a/go/adbc/driver/bigquery/statement.go b/go/adbc/driver/bigquery/statement.go index ffc9fbfd15..3fc064ec37 100644 --- a/go/adbc/driver/bigquery/statement.go +++ b/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.