feat(transport): add support for setting quota project with envvar

Fixes: googleapis#1891
codyoss · Mar 6, 2023 · d0d85c2 · d0d85c2
1 parent 5da4d6a
commit d0d85c2
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 15 deletions.
diff --git a/internal/creds.go b/internal/creds.go
@@ -10,13 +10,16 @@ import (
 	"errors"
 	"fmt"
 	"io/ioutil"
+	"os"
 
 	"golang.org/x/oauth2"
 	"google.golang.org/api/internal/impersonate"
 
 	"golang.org/x/oauth2/google"
 )
 
+const quotaProjectEnvVar = "GOOGLE_CLOUD_QUOTA_PROJECT"
+
 // Creds returns credential information obtained from DialSettings, or if none, then
 // it returns default credential information.
 func Creds(ctx context.Context, ds *DialSettings) (*google.Credentials, error) {
@@ -131,14 +134,22 @@ func selfSignedJWTTokenSource(data []byte, ds *DialSettings) (oauth2.TokenSource
 	}
 }
 
-// QuotaProjectFromCreds returns the quota project from the JSON blob in the provided credentials.
-//
-// NOTE(cbro): consider promoting this to a field on google.Credentials.
-func QuotaProjectFromCreds(cred *google.Credentials) string {
+// GetQuotaProject retrieves quota project with precedence being: client option,
+// environment variable, creds file.
+func GetQuotaProject(creds *google.Credentials, clientOpt string) string {
+	if clientOpt != "" {
+		return clientOpt
+	}
+	if env := os.Getenv(quotaProjectEnvVar); env != "" {
+		return env
+	}
+	if creds == nil {
+		return ""
+	}
 	var v struct {
 		QuotaProject string `json:"quota_project_id"`
 	}
-	if err := json.Unmarshal(cred.JSON, &v); err != nil {
+	if err := json.Unmarshal(creds.JSON, &v); err != nil {
 		return ""
 	}
 	return v.QuotaProject

diff --git a/transport/grpc/dial.go b/transport/grpc/dial.go
@@ -155,14 +155,10 @@ func dial(ctx context.Context, insecure bool, o *internal.DialSettings) (*grpc.C
 			return nil, err
 		}
 
-		if o.QuotaProject == "" {
-			o.QuotaProject = internal.QuotaProjectFromCreds(creds)
-		}
-
 		grpcOpts = append(grpcOpts,
 			grpc.WithPerRPCCredentials(grpcTokenSource{
 				TokenSource:   oauth.TokenSource{creds.TokenSource},
-				quotaProject:  o.QuotaProject,
+				quotaProject:  internal.GetQuotaProject(creds, o.QuotaProject),
 				requestReason: o.RequestReason,
 			}),
 		)

diff --git a/transport/http/dial.go b/transport/http/dial.go
@@ -66,7 +66,6 @@ func newTransport(ctx context.Context, base http.RoundTripper, settings *interna
 	paramTransport := &parameterTransport{
 		base:          base,
 		userAgent:     settings.UserAgent,
-		quotaProject:  settings.QuotaProject,
 		requestReason: settings.RequestReason,
 	}
 	var trans http.RoundTripper = paramTransport
@@ -75,6 +74,7 @@ func newTransport(ctx context.Context, base http.RoundTripper, settings *interna
 	case settings.NoAuth:
 		// Do nothing.
 	case settings.APIKey != "":
+		paramTransport.quotaProject = internal.GetQuotaProject(nil, settings.QuotaProject)
 		trans = &transport.APIKey{
 			Transport: trans,
 			Key:       settings.APIKey,
@@ -84,10 +84,7 @@ func newTransport(ctx context.Context, base http.RoundTripper, settings *interna
 		if err != nil {
 			return nil, err
 		}
-		if paramTransport.quotaProject == "" {
-			paramTransport.quotaProject = internal.QuotaProjectFromCreds(creds)
-		}
-
+		paramTransport.quotaProject = internal.GetQuotaProject(creds, settings.QuotaProject)
 		ts := creds.TokenSource
 		if settings.ImpersonationConfig == nil && settings.TokenSource != nil {
 			ts = settings.TokenSource