diff --git a/runtime/drivers/bigquery/api.go b/runtime/drivers/bigquery/api.go index 59fdce53ede..7f7c0879847 100644 --- a/runtime/drivers/bigquery/api.go +++ b/runtime/drivers/bigquery/api.go @@ -11,7 +11,12 @@ import ( const defaultPageSize = 20 func (c *Connection) ListDatasets(ctx context.Context, req *runtimev1.BigQueryListDatasetsRequest) ([]string, string, error) { - client, err := c.createClient(ctx, &sourceProperties{ProjectID: bigquery.DetectProjectID}) + opts, err := c.clientOption(ctx) + if err != nil { + return nil, "", err + } + + client, err := bigquery.NewClient(ctx, bigquery.DetectProjectID, opts...) if err != nil { return nil, "", err } @@ -36,7 +41,12 @@ func (c *Connection) ListDatasets(ctx context.Context, req *runtimev1.BigQueryLi } func (c *Connection) ListTables(ctx context.Context, req *runtimev1.BigQueryListTablesRequest) ([]string, string, error) { - client, err := c.createClient(ctx, &sourceProperties{ProjectID: bigquery.DetectProjectID}) + opts, err := c.clientOption(ctx) + if err != nil { + return nil, "", err + } + + client, err := bigquery.NewClient(ctx, bigquery.DetectProjectID, opts...) if err != nil { return nil, "", err } diff --git a/runtime/drivers/bigquery/bigquery.go b/runtime/drivers/bigquery/bigquery.go index b2a274249b5..9f7a89f1fbc 100644 --- a/runtime/drivers/bigquery/bigquery.go +++ b/runtime/drivers/bigquery/bigquery.go @@ -2,7 +2,6 @@ package bigquery import ( "context" - "errors" "fmt" "os" "strings" @@ -219,13 +218,10 @@ func parseSourceProperties(props map[string]any) (*sourceProperties, error) { return conf, err } -func (c *Connection) createClient(ctx context.Context, props *sourceProperties) (*bigquery.Client, error) { +func (c *Connection) clientOption(ctx context.Context) ([]option.ClientOption, error) { creds, err := gcputil.Credentials(ctx, c.config.SecretJSON, c.config.AllowHostAccess) if err != nil { - if !errors.Is(err, gcputil.ErrNoCredentials) { - return nil, err - } - return bigquery.NewClient(ctx, props.ProjectID) + return nil, err } - return bigquery.NewClient(ctx, props.ProjectID, option.WithCredentials(creds)) + return []option.ClientOption{option.WithCredentials(creds)}, nil } diff --git a/runtime/drivers/bigquery/sql_store.go b/runtime/drivers/bigquery/sql_store.go index a972dd3d2ce..93922a5a291 100644 --- a/runtime/drivers/bigquery/sql_store.go +++ b/runtime/drivers/bigquery/sql_store.go @@ -37,7 +37,12 @@ func (c *Connection) QueryAsFiles(ctx context.Context, props map[string]any, sql return nil, err } - client, err := c.createClient(ctx, srcProps) + opts, err := c.clientOption(ctx) + if err != nil { + return nil, err + } + + client, err := bigquery.NewClient(ctx, srcProps.ProjectID, opts...) if err != nil { if strings.Contains(err.Error(), "unable to detect projectID") { return nil, fmt.Errorf("projectID not detected in credentials. Please set `project_id` in source yaml") @@ -45,7 +50,7 @@ func (c *Connection) QueryAsFiles(ctx context.Context, props map[string]any, sql return nil, fmt.Errorf("failed to create bigquery client: %w", err) } - if err := client.EnableStorageReadClient(ctx); err != nil { + if err := client.EnableStorageReadClient(ctx, opts...); err != nil { client.Close() return nil, err } @@ -60,7 +65,7 @@ func (c *Connection) QueryAsFiles(ctx context.Context, props map[string]any, sql // the query results are always cached in a temporary table that storage api can use // there are some exceptions when results aren't cached // so we also try without storage api - client, err = c.createClient(ctx, srcProps) + client, err = bigquery.NewClient(ctx, srcProps.ProjectID, opts...) if err != nil { return nil, fmt.Errorf("failed to create bigquery client: %w", err) }