feat(go/driver/bigquery): support specifying credentials json filepath

apache · Apr 16, 2024 · 525d7c2 · 525d7c2
1 parent a3adeb0
commit 525d7c2
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 8 deletions.
diff --git a/go/adbc/driver/bigquery/bigquery_database.go b/go/adbc/driver/bigquery/bigquery_database.go
@@ -28,8 +28,9 @@ import (
 type databaseImpl struct {
 	driverbase.DatabaseImplBase
 
-	authType  string
-	projectID *string
+	authType    string
+	credentials *string
+	projectID   *string
 
 	alloc memory.Allocator
 }
@@ -56,6 +57,12 @@ func (d *databaseImpl) GetOption(key string) (string, error) {
 	switch key {
 	case OptionStringAuthType:
 		return d.authType, nil
+	case OptionStringCredentials:
+		if d.credentials == nil {
+			return "", nil
+		} else {
+			return *d.credentials, nil
+		}
 	case OptionStringProjectID:
 		if d.projectID == nil {
 			return "", adbc.Error{
@@ -98,7 +105,19 @@ func (d *databaseImpl) SetOptions(options map[string]string) error {
 		v := v // copy into loop scope
 		switch k {
 		case OptionStringAuthType:
-			d.authType = v
+			switch v {
+			case OptionValueAuthTypeDefault:
+				d.authType = v
+			case OptionValueAuthTypeCredentialsFile:
+				d.authType = v
+			default:
+				return adbc.Error{
+					Code: adbc.StatusInvalidArgument,
+					Msg:  fmt.Sprintf("unknown database auth type value `%s`", v),
+				}
+			}
+		case OptionStringCredentials:
+			d.credentials = &v
 		case OptionStringProjectID:
 			d.projectID = &v
 		default:
@@ -114,7 +133,19 @@ func (d *databaseImpl) SetOptions(options map[string]string) error {
 func (d *databaseImpl) SetOption(key string, value string) error {
 	switch key {
 	case OptionStringAuthType:
-		d.authType = value
+		switch value {
+		case OptionValueAuthTypeDefault:
+			d.authType = value
+		case OptionValueAuthTypeCredentialsFile:
+			d.authType = value
+		default:
+			return adbc.Error{
+				Code: adbc.StatusInvalidArgument,
+				Msg:  fmt.Sprintf("unknown database auth type value `%s`", value),
+			}
+		}
+	case OptionStringCredentials:
+		d.credentials = &value
 	case OptionStringProjectID:
 		d.projectID = &value
 	default:

diff --git a/go/adbc/driver/bigquery/driver.go b/go/adbc/driver/bigquery/driver.go
@@ -29,8 +29,11 @@ import (
 // under the License.
 
 const (
-	OptionStringAuthType  = "adbc.bigquery.sql.auth_type"
-	OptionStringProjectID = "adbc.bigquery.sql.project_id"
+	OptionStringAuthType               = "adbc.bigquery.sql.auth_type"
+	OptionStringCredentials            = "adbc.bigquery.sql.credentials"
+	OptionStringProjectID              = "adbc.bigquery.sql.project_id"
+	OptionValueAuthTypeDefault         = "default"
+	OptionValueAuthTypeCredentialsFile = "credentials_file"
 
 	OptionStringQueryDestinationTable  = "adbc.bigquery.sql.query.destination_table"
 	OptionStringQueryDefaultProjectID  = "adbc.bigquery.sql.query.default_project_id"
@@ -88,6 +91,7 @@ func (d *driverImpl) NewDatabase(opts map[string]string) (adbc.Database, error)
 	db := &databaseImpl{
 		DatabaseImplBase: driverbase.NewDatabaseImplBase(&d.DriverImplBase),
 		alloc:            d.alloc,
+		authType:         OptionValueAuthTypeDefault,
 	}
 	if err := db.SetOptions(opts); err != nil {
 		return nil, err

diff --git a/go/adbc/driver/bigquery/statement.go b/go/adbc/driver/bigquery/statement.go
@@ -25,6 +25,7 @@ import (
 	"github.com/apache/arrow/go/v16/arrow"
 	"github.com/apache/arrow/go/v16/arrow/array"
 	"github.com/apache/arrow/go/v16/arrow/memory"
+	"google.golang.org/api/option"
 	"time"
 )
 
@@ -339,11 +340,15 @@ func (st *statement) ExecuteQuery(ctx context.Context) (array.RecordReader, int6
 	if err != nil {
 		return nil, -1, err
 	}
-	client, err := bigquery.NewClient(ctx, projectID)
+	authType, err := st.GetOption(OptionStringAuthType)
 	if err != nil {
 		return nil, -1, err
 	}
-	err = client.EnableStorageReadClient(ctx)
+	credentials, err := st.GetOption(OptionStringCredentials)
+	if err != nil {
+		return nil, -1, err
+	}
+	client, err := newClient(ctx, projectID, authType, credentials)
 	if err != nil {
 		return nil, -1, err
 	}
@@ -355,6 +360,31 @@ func (st *statement) ExecuteQuery(ctx context.Context) (array.RecordReader, int6
 	return reader, -1, nil
 }
 
+func newClient(ctx context.Context, projectID, authType, credentials string) (*bigquery.Client, error) {
+	var client *bigquery.Client
+	switch authType {
+	case OptionValueAuthTypeCredentialsFile:
+		client, err := bigquery.NewClient(ctx, projectID, option.WithCredentialsFile(credentials))
+		if err != nil {
+			return nil, err
+		}
+		err = client.EnableStorageReadClient(ctx, option.WithCredentialsFile(credentials))
+		if err != nil {
+			return nil, err
+		}
+	default:
+		client, err := bigquery.NewClient(ctx, projectID)
+		if err != nil {
+			return nil, err
+		}
+		err = client.EnableStorageReadClient(ctx)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return client, nil
+}
+
 // ExecuteUpdate executes a statement that does not generate a result
 // set. It returns the number of rows affected if known, otherwise -1.
 func (st *statement) ExecuteUpdate(ctx context.Context) (int64, error) {