From ca2fc7baf1bd6fe995914aa987239c93f4397be6 Mon Sep 17 00:00:00 2001
From: Anshul Khandelwal <12948312+k-anshul@users.noreply.github.com>
Date: Fri, 8 Sep 2023 16:43:15 +0530
Subject: [PATCH] BigQuery storage api credentials fix (#3043)

* bq storage api fix

* credentials should be present in all cases

* small fix
---
 runtime/drivers/bigquery/api.go       | 14 ++++++++++++--
 runtime/drivers/bigquery/bigquery.go  | 10 +++-------
 runtime/drivers/bigquery/sql_store.go | 11 ++++++++---
 3 files changed, 23 insertions(+), 12 deletions(-)

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)
 		}