fix: add timeout and disable option for BigQuery Storage Read API

go/bigquery_database.go

@@ -61,6 +61,16 @@ type databaseImpl struct {
 
 	bulkIngestMethod      string
 	bulkIngestCompression string
+
+	// queryBackendAPI selects which BigQuery API is used to read query results.
+	// Valid values: OptionValueQueryBackendAPIStorageRead (default), OptionValueQueryBackendAPIJobs.
+	// See issue #66. The REST fallback is not yet implemented; selecting "jobs"
+	// currently returns an error when results are read.
+	queryBackendAPI string
+
+	// storageReadAPIEndpoint overrides the Storage Read API gRPC endpoint.
+	// Empty string means use the default Google endpoint.
+	storageReadAPIEndpoint string
 }
 
 func (d *databaseImpl) Open(ctx context.Context) (adbc.Connection, error) {
@@ -86,6 +96,8 @@ func (d *databaseImpl) Open(ctx context.Context) (adbc.Connection, error) {
 		quotaProject:               d.quotaProject,
 		bulkIngestMethod:           d.bulkIngestMethod,
 		bulkIngestCompression:      d.bulkIngestCompression,
+		queryBackendAPI:            d.queryBackendAPI,
+		storageReadAPIEndpoint:     d.storageReadAPIEndpoint,
 	}
 
 	err := conn.newClient(ctx)
@@ -148,6 +160,13 @@ func (d *databaseImpl) GetOption(key string) (string, error) {
 			return OptionValueCompressionNone, nil
 		}
 		return d.bulkIngestCompression, nil
+	case OptionStringQueryBackendAPI:
+		if d.queryBackendAPI == "" {
+			return OptionValueQueryBackendAPIStorageRead, nil
+		}
+		return d.queryBackendAPI, nil
+	case OptionStringStorageReadAPIEndpoint:
+		return d.storageReadAPIEndpoint, nil
 	default:
 		return d.DatabaseImplBase.GetOption(key)
 	}
@@ -279,6 +298,17 @@ func (d *databaseImpl) SetOption(key string, value string) error {
 			}
 		}
 		d.bulkIngestCompression = value
+	case OptionStringQueryBackendAPI:
+		if value != OptionValueQueryBackendAPIStorageRead &&
+			value != OptionValueQueryBackendAPIJobs {
+			return adbc.Error{
+				Code: adbc.StatusInvalidArgument,
+				Msg:  fmt.Sprintf("[bq] invalid value for %s: %q (expected %q or %q)", key, value, OptionValueQueryBackendAPIStorageRead, OptionValueQueryBackendAPIJobs),
+			}
+		}
+		d.queryBackendAPI = value
+	case OptionStringStorageReadAPIEndpoint:
+		d.storageReadAPIEndpoint = value
 	default:
 		return d.DatabaseImplBase.SetOption(key, value)
 	}

go/connection.go

@@ -82,6 +82,17 @@ type connectionImpl struct {
 	bulkIngestMethod      string
 	bulkIngestCompression string
 
+	// queryBackendAPI selects which BigQuery API is used to read query results.
+	// Valid values: OptionValueQueryBackendAPIStorageRead (default), OptionValueQueryBackendAPIJobs.
+	// See issue #66. The REST fallback is not yet implemented; selecting "jobs"
+	// currently returns an error when results are read.
+	queryBackendAPI string
+
+	// storageReadAPIEndpoint overrides the Storage Read API gRPC endpoint.
+	// Empty string means use the default Google endpoint.
+	// Useful for testing with a local fake server or BigQuery emulator.
+	storageReadAPIEndpoint string
+
 	client *bigquery.Client
 }
 
@@ -673,6 +684,17 @@ func (c *connectionImpl) SetOption(key string, value string) error {
 			}
 		}
 		c.bulkIngestCompression = value
+	case OptionStringQueryBackendAPI:
+		if value != OptionValueQueryBackendAPIStorageRead &&
+			value != OptionValueQueryBackendAPIJobs {
+			return adbc.Error{
+				Code: adbc.StatusInvalidArgument,
+				Msg:  fmt.Sprintf("[bq] invalid value for %s: %q (expected %q or %q)", key, value, OptionValueQueryBackendAPIStorageRead, OptionValueQueryBackendAPIJobs),
+			}
+		}
+		c.queryBackendAPI = value
+	case OptionStringStorageReadAPIEndpoint:
+		c.storageReadAPIEndpoint = value
 	default:
 		return c.ConnectionImplBase.SetOption(key, value)
 	}
@@ -832,10 +854,30 @@ func (c *connectionImpl) newClient(ctx context.Context) error {
 		client.Location = c.location
 	}
 
-	// Use original authOptions without custom endpoint for Storage Read API
-	err = client.EnableStorageReadClient(ctx, authOptions...)
-	if err != nil {
-		return errToAdbcErr(adbc.StatusIO, err, "enable storage read client")
+	// EnableStorageReadClient opens a gRPC connection to the BigQuery Storage
+	// Read API. Guard with a 30s timeout so the driver fails fast instead of
+	// hanging indefinitely if the call does not return.
+	//
+	// When OptionStringQueryBackendAPI is set to "jobs", we skip initialising
+	// the Storage Read client entirely. Note: the actual REST-based fallback
+	// for reading query results is not yet implemented (see issue #66); reads
+	// will return an explicit error in that case.
+	if c.queryBackendAPI != OptionValueQueryBackendAPIJobs {
+		storageCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+		defer cancel()
+		storageAuthOptions := authOptions
+		if c.storageReadAPIEndpoint != "" {
+			storageAuthOptions = append(storageAuthOptions, option.WithEndpoint(c.storageReadAPIEndpoint))
+		}
+		if err := client.EnableStorageReadClient(storageCtx, storageAuthOptions...); err != nil {
+			if storageCtx.Err() != nil {
+				c.Logger.Warn("BigQuery Storage Read API timed out after 30s. See " + OptionStringQueryBackendAPI + " (issue #66) for the future REST fallback.")
+			} else {
+				c.Logger.Warn("BigQuery Storage Read API unavailable", "err", err)
+			}
+			// non-fatal: continue without Storage Read API; record_reader will
+			// surface an appropriate error when a caller tries to stream rows.
+		}
 	}
 
 	c.client = client

go/driver.go

@@ -70,6 +70,27 @@ const (
 	OptionStringQueryCreateDisposition = "adbc.bigquery.sql.query.create_disposition"
 	OptionStringQueryWriteDisposition  = "adbc.bigquery.sql.query.write_disposition"
 	OptionBoolQueryDisableQueryCache   = "adbc.bigquery.sql.query.disable_query_cache"
+
+	// OptionStringQueryBackendAPI selects which BigQuery API is used to read
+	// query results. See issue #66 for design discussion.
+	//
+	// Valid values:
+	//   - OptionValueQueryBackendAPIStorageRead (default): use the Storage Read
+	//     API (gRPC/HTTP2). Required for Arrow-format streaming reads.
+	//   - OptionValueQueryBackendAPIJobs: select the REST Jobs API. NOTE: the
+	//     actual REST fallback is not yet implemented in this driver; setting
+	//     this value currently returns an error when reading results. The option
+	//     is added as infrastructure so callers can opt in once the fallback is
+	//     implemented in a follow-up PR.
+	OptionStringQueryBackendAPI             = "adbc.bigquery.query.backend_api"
+	OptionValueQueryBackendAPIStorageRead   = "storage_read"
+	OptionValueQueryBackendAPIJobs          = "jobs"
+
+	// OptionStringStorageReadAPIEndpoint overrides the endpoint used for the
+	// BigQuery Storage Read API (gRPC). Defaults to the public Google endpoint.
+	// Useful for testing with a local fake server or BigQuery emulator.
+	// Format: "host:port" (e.g. "localhost:9443").
+	OptionStringStorageReadAPIEndpoint = "adbc.bigquery.sql.storage_read_api_endpoint"
 	OptionBoolDisableFlattenedResults  = "adbc.bigquery.sql.query.disable_flattened_results"
 	OptionBoolQueryAllowLargeResults   = "adbc.bigquery.sql.query.allow_large_results"
 	OptionStringQueryPriority          = "adbc.bigquery.sql.query.priority"

go/driver_test.go

@@ -1659,6 +1659,59 @@ func TestAuthTypeConsolidation(t *testing.T) {
 	}
 }
 
+func TestQueryBackendAPIOption(t *testing.T) {
+	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer mem.AssertSize(t, 0)
+
+	drv := driver.NewDriver(mem)
+
+	// Default value should be "storage_read".
+	db, err := drv.NewDatabase(nil)
+	require.NoError(t, err, "create database")
+	defer validation.CheckedClose(t, db)
+
+	getSetDB, ok := db.(adbc.GetSetOptions)
+	require.True(t, ok, "database should implement adbc.GetSetOptions")
+
+	val, err := getSetDB.GetOption(driver.OptionStringQueryBackendAPI)
+	require.NoError(t, err, "get default option")
+	assert.Equal(t, driver.OptionValueQueryBackendAPIStorageRead, val, "default backend")
+
+	// Setting "jobs" via SetOptions should work.
+	require.NoError(t,
+		db.SetOptions(map[string]string{driver.OptionStringQueryBackendAPI: driver.OptionValueQueryBackendAPIJobs}),
+		"set jobs backend")
+	val, err = getSetDB.GetOption(driver.OptionStringQueryBackendAPI)
+	require.NoError(t, err)
+	assert.Equal(t, driver.OptionValueQueryBackendAPIJobs, val)
+
+	// Setting "storage_read" explicitly should also work.
+	require.NoError(t,
+		db.SetOptions(map[string]string{driver.OptionStringQueryBackendAPI: driver.OptionValueQueryBackendAPIStorageRead}),
+		"set storage_read backend")
+	val, err = getSetDB.GetOption(driver.OptionStringQueryBackendAPI)
+	require.NoError(t, err)
+	assert.Equal(t, driver.OptionValueQueryBackendAPIStorageRead, val)
+
+	// Setting via NewDatabase should work.
+	db2, err := drv.NewDatabase(map[string]string{
+		driver.OptionStringQueryBackendAPI: driver.OptionValueQueryBackendAPIJobs,
+	})
+	require.NoError(t, err, "create database with option")
+	defer validation.CheckedClose(t, db2)
+
+	getSetDB2, ok := db2.(adbc.GetSetOptions)
+	require.True(t, ok)
+	val, err = getSetDB2.GetOption(driver.OptionStringQueryBackendAPI)
+	require.NoError(t, err)
+	assert.Equal(t, driver.OptionValueQueryBackendAPIJobs, val)
+
+	// Invalid value should return an error.
+	err = db.SetOptions(map[string]string{driver.OptionStringQueryBackendAPI: "rest"})
+	require.Error(t, err, "invalid value should error")
+	assert.Contains(t, err.Error(), "invalid value")
+}
+
 type BigQueryTestSuite struct {
 	suite.Suite
 	project string

go/record_reader.go

@@ -26,9 +26,11 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
 	"log"
 	"log/slog"
 	"sync/atomic"
+	"time"
 
 	"cloud.google.com/go/bigquery"
 	"github.com/apache/arrow-adbc/go/adbc"
@@ -61,7 +63,19 @@ func checkContext(ctx context.Context, maybeErr error) error {
 	return ctx.Err()
 }
 
-func runQuery(ctx context.Context, logger *slog.Logger, query *bigquery.Query, executeUpdate bool) (bigquery.ArrowIterator, int64, error) {
+func runQuery(ctx context.Context, logger *slog.Logger, queryBackendAPI string, query *bigquery.Query, executeUpdate bool) (bigquery.ArrowIterator, int64, error) {
+	// Jobs API backend is not yet implemented (see issue #66). Fail fast with
+	// a clear error rather than hanging or returning a misleading permission
+	// error downstream.
+	if !executeUpdate && queryBackendAPI == OptionValueQueryBackendAPIJobs {
+		return nil, -1, adbc.Error{
+			Code: adbc.StatusNotImplemented,
+			Msg: "[bq] " + OptionStringQueryBackendAPI + "=" + OptionValueQueryBackendAPIJobs +
+				" is not yet implemented. The REST-based fallback for reading query results has not been added to this driver (see issue #66). " +
+				"Use " + OptionStringQueryBackendAPI + "=" + OptionValueQueryBackendAPIStorageRead + " (the default) until the fallback is available.",
+		}
+	}
+
 	job, err := query.Run(ctx)
 	if err != nil {
 		return nil, -1, errToAdbcErr(adbc.StatusInternal, err, "run query")
@@ -122,8 +136,47 @@ func runQuery(ctx context.Context, logger *slog.Logger, query *bigquery.Query, e
 				Msg:  "[bq] Arrow reader requires roles/bigquery.readSessionUser, see https://github.com/apache/arrow-adbc/issues/3282",
 			}
 		}
-		if arrowIterator, err = iter.ArrowIterator(); err != nil {
-			return nil, -1, errToAdbcErr(adbc.StatusInternal, err, "read Arrow query results")
+
+		// ArrowIterator() opens a gRPC read session via BigQuery Storage API
+		// (CreateReadSession RPC). The call does not accept a context, and
+		// has been observed to hang indefinitely in some environments. Apply
+		// a 30s timeout so the caller gets a clear error instead of blocking
+		// forever. The goroutine may outlive the timeout since the underlying
+		// gRPC call cannot be cancelled from here, but it will be cleaned up
+		// when the connection is closed.
+		type arrowIterResult struct {
+			ai  bigquery.ArrowIterator
+			err error
+		}
+		ch := make(chan arrowIterResult, 1)
+		go func() {
+			ai, err := iter.ArrowIterator()
+			ch <- arrowIterResult{ai, err}
+		}()
+
+		storageCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+		defer cancel()
+
+		select {
+		case r := <-ch:
+			if r.err != nil {
+				return nil, -1, errToAdbcErr(adbc.StatusInternal, r.err, "read Arrow query results")
+			}
+			arrowIterator = r.ai
+		case <-storageCtx.Done():
+			if ctx.Err() != nil {
+				// Parent context was cancelled by the caller — propagate it.
+				return nil, -1, adbc.Error{
+					Code: adbc.StatusCancelled,
+					Msg:  fmt.Sprintf("[bq] context cancelled while opening BigQuery Storage read session: %s", ctx.Err()),
+				}
+			}
+			// Our own 30s timeout fired — the Storage Read API is unreachable.
+			return nil, -1, adbc.Error{
+				Code: adbc.StatusTimeout,
+				Msg: "[bq] BigQuery Storage Read API timed out after 30s. " +
+					"See " + OptionStringQueryBackendAPI + " (issue #66) for the future REST fallback.",
+			}
 		}
 	} else {
 		arrowIterator = emptyArrowIterator{iter.Schema}
@@ -167,8 +220,8 @@ func getQueryParameter(values arrow.RecordBatch, row int, parameterMode string)
 	return parameters, nil
 }
 
-func runPlainQuery(ctx context.Context, logger *slog.Logger, query *bigquery.Query, alloc memory.Allocator, resultRecordBufferSize int) (bigqueryRdr *reader, totalRows int64, err error) {
-	arrowIterator, totalRows, err := runQuery(ctx, logger, query, false)
+func runPlainQuery(ctx context.Context, logger *slog.Logger, queryBackendAPI string, query *bigquery.Query, alloc memory.Allocator, resultRecordBufferSize int) (bigqueryRdr *reader, totalRows int64, err error) {
+	arrowIterator, totalRows, err := runQuery(ctx, logger, queryBackendAPI, query, false)
 	if err != nil {
 		return nil, -1, err
 	}
@@ -213,7 +266,7 @@ func runPlainQuery(ctx context.Context, logger *slog.Logger, query *bigquery.Que
 	return bigqueryRdr, totalRows, nil
 }
 
-func queryRecordWithSchemaCallback(ctx context.Context, logger *slog.Logger, group *errgroup.Group, query *bigquery.Query, rec arrow.RecordBatch, ch chan arrow.RecordBatch, parameterMode string, alloc memory.Allocator, rdrSchema func(schema *arrow.Schema)) (int64, error) {
+func queryRecordWithSchemaCallback(ctx context.Context, logger *slog.Logger, queryBackendAPI string, group *errgroup.Group, query *bigquery.Query, rec arrow.RecordBatch, ch chan arrow.RecordBatch, parameterMode string, alloc memory.Allocator, rdrSchema func(schema *arrow.Schema)) (int64, error) {
 	totalRows := int64(-1)
 	for i := range int(rec.NumRows()) {
 		parameters, err := getQueryParameter(rec, i, parameterMode)
@@ -224,7 +277,7 @@ func queryRecordWithSchemaCallback(ctx context.Context, logger *slog.Logger, gro
 			query.Parameters = parameters
 		}
 
-		arrowIterator, rows, err := runQuery(ctx, logger, query, false)
+		arrowIterator, rows, err := runQuery(ctx, logger, queryBackendAPI, query, false)
 		if err != nil {
 			return -1, err
 		}
@@ -249,9 +302,9 @@ func queryRecordWithSchemaCallback(ctx context.Context, logger *slog.Logger, gro
 
 // 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, logger *slog.Logger, query *bigquery.Query, boundParameters array.RecordReader, parameterMode string, alloc memory.Allocator, resultRecordBufferSize, prefetchConcurrency int) (bigqueryRdr *reader, totalRows int64, err error) {
+func newRecordReader(ctx context.Context, logger *slog.Logger, queryBackendAPI string, query *bigquery.Query, boundParameters array.RecordReader, parameterMode string, alloc memory.Allocator, resultRecordBufferSize, prefetchConcurrency int) (bigqueryRdr *reader, totalRows int64, err error) {
 	if boundParameters == nil {
-		return runPlainQuery(ctx, logger, query, alloc, resultRecordBufferSize)
+		return runPlainQuery(ctx, logger, queryBackendAPI, query, alloc, resultRecordBufferSize)
 	}
 	defer boundParameters.Release()
 
@@ -287,7 +340,7 @@ func newRecordReader(ctx context.Context, logger *slog.Logger, query *bigquery.Q
 		// Each call to Record() on the record reader is allowed to release the previous record
 		// and since we're doing this sequentially
 		// we don't need to call rec.Retain() here and call call rec.Release() in queryRecordWithSchemaCallback
-		batchRows, err := queryRecordWithSchemaCallback(ctx, logger, group, query, rec, ch, parameterMode, alloc, func(schema *arrow.Schema) {
+		batchRows, err := queryRecordWithSchemaCallback(ctx, logger, queryBackendAPI, group, query, rec, ch, parameterMode, alloc, func(schema *arrow.Schema) {
 			bigqueryRdr.schema = schema
 		})
 		if err != nil {

go/statement.go

@@ -382,7 +382,7 @@ func (st *statement) ExecuteQuery(ctx context.Context) (array.RecordReader, int6
 		}
 	}
 
-	rr, totalRows, err := newRecordReader(ctx, st.cnxn.Logger, st.query(), st.params, st.parameterMode, st.cnxn.Alloc, st.resultRecordBufferSize, st.prefetchConcurrency)
+	rr, totalRows, err := newRecordReader(ctx, st.cnxn.Logger, st.cnxn.queryBackendAPI, st.query(), st.params, st.parameterMode, st.cnxn.Alloc, st.resultRecordBufferSize, st.prefetchConcurrency)
 	st.params = nil
 	return rr, totalRows, err
 }
@@ -396,7 +396,7 @@ func (st *statement) ExecuteUpdate(ctx context.Context) (int64, error) {
 	}
 
 	if st.params == nil {
-		_, totalRows, err := runQuery(ctx, st.cnxn.Logger, st.query(), true)
+		_, totalRows, err := runQuery(ctx, st.cnxn.Logger, st.cnxn.queryBackendAPI, st.query(), true)
 		if err != nil {
 			return -1, err
 		}
@@ -418,7 +418,7 @@ func (st *statement) ExecuteUpdate(ctx context.Context) (int64, error) {
 					st.queryConfig.Parameters = parameters
 				}
 
-				_, currentRows, err := runQuery(ctx, st.cnxn.Logger, st.query(), true)
+				_, currentRows, err := runQuery(ctx, st.cnxn.Logger, st.cnxn.queryBackendAPI, st.query(), true)
 				if err != nil {
 					return -1, err
 				}