From 18ff8ffce2581e208eb6bf7d86ee1e4764d6d199 Mon Sep 17 00:00:00 2001 From: djmantri Date: Wed, 21 Jan 2026 13:40:07 -0800 Subject: [PATCH] Add DatasetProvider interface for BigQuery Table API access This adds an exported DatasetProvider interface that allows callers to access the underlying *bigquery.Dataset from a database/sql connection. This enables operations not supported by the standard SQL interface, such as updating table clustering via the BigQuery Table API. --- driver/connection.go | 31 ++++++++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/driver/connection.go b/driver/connection.go index 2d8303a..dbbe7b3 100644 --- a/driver/connection.go +++ b/driver/connection.go @@ -1,12 +1,38 @@ package driver import ( - "cloud.google.com/go/bigquery" "context" "database/sql/driver" "fmt" + + "cloud.google.com/go/bigquery" ) +// DatasetProvider is an interface that provides access to the underlying +// BigQuery dataset. This interface can be used to access the BigQuery client +// for operations that are not supported by the standard database/sql interface, +// such as updating table clustering. +// +// Usage with database/sql: +// +// conn, err := db.Conn(ctx) +// if err != nil { +// return err +// } +// defer conn.Close() +// +// err = conn.Raw(func(driverConn interface{}) error { +// if provider, ok := driverConn.(driver.DatasetProvider); ok { +// dataset := provider.GetDataset() +// // do operations on the dataset +// } +// return errors.New("connection does not implement DatasetProvider") +// }) +type DatasetProvider interface { + // GetDataset returns the BigQuery dataset for the current connection. + GetDataset() *bigquery.Dataset +} + type bigQueryConnection struct { ctx context.Context client *bigquery.Client @@ -16,6 +42,9 @@ type bigQueryConnection struct { dataset *bigquery.Dataset } +// Ensure bigQueryConnection implements DatasetProvider +var _ DatasetProvider = (*bigQueryConnection)(nil) + func (connection *bigQueryConnection) GetDataset() *bigquery.Dataset { if connection.dataset != nil { return connection.dataset