diff --git a/mmv1/third_party/terraform/utils/bigtable_client_factory.go b/mmv1/third_party/terraform/utils/bigtable_client_factory.go index 7fe1a9ac205a..46ce6ece2b44 100644 --- a/mmv1/third_party/terraform/utils/bigtable_client_factory.go +++ b/mmv1/third_party/terraform/utils/bigtable_client_factory.go @@ -2,6 +2,7 @@ package google import ( "context" + "os" "cloud.google.com/go/bigtable" "golang.org/x/oauth2" @@ -9,14 +10,36 @@ import ( ) type BigtableClientFactory struct { - UserAgent string - TokenSource oauth2.TokenSource + UserAgent string + TokenSource oauth2.TokenSource + BillingProject string + UserProjectOverride bool } func (s BigtableClientFactory) NewInstanceAdminClient(project string) (*bigtable.InstanceAdminClient, error) { - return bigtable.NewInstanceAdminClient(context.Background(), project, option.WithTokenSource(s.TokenSource), option.WithUserAgent(s.UserAgent)) + var opts []option.ClientOption + if requestReason := os.Getenv("CLOUDSDK_CORE_REQUEST_REASON"); requestReason != "" { + opts = append(opts, option.WithRequestReason(requestReason)) + } + + if s.UserProjectOverride && s.BillingProject != "" { + opts = append(opts, option.WithQuotaProject(s.BillingProject)) + } + + opts = append(opts, option.WithTokenSource(s.TokenSource), option.WithUserAgent(s.UserAgent)) + return bigtable.NewInstanceAdminClient(context.Background(), project, opts...) } func (s BigtableClientFactory) NewAdminClient(project, instance string) (*bigtable.AdminClient, error) { - return bigtable.NewAdminClient(context.Background(), project, instance, option.WithTokenSource(s.TokenSource), option.WithUserAgent(s.UserAgent)) + var opts []option.ClientOption + if requestReason := os.Getenv("CLOUDSDK_CORE_REQUEST_REASON"); requestReason != "" { + opts = append(opts, option.WithRequestReason(requestReason)) + } + + if s.UserProjectOverride && s.BillingProject != "" { + opts = append(opts, option.WithQuotaProject(s.BillingProject)) + } + + opts = append(opts, option.WithTokenSource(s.TokenSource), option.WithUserAgent(s.UserAgent)) + return bigtable.NewAdminClient(context.Background(), project, instance, opts...) } diff --git a/mmv1/third_party/terraform/utils/config.go.erb b/mmv1/third_party/terraform/utils/config.go.erb index 06f8de9da1d9..a0dcc7bc80d9 100644 --- a/mmv1/third_party/terraform/utils/config.go.erb +++ b/mmv1/third_party/terraform/utils/config.go.erb @@ -809,8 +809,10 @@ func (c *Config) NewCloudIdentityClient(userAgent string) *cloudidentity.Service func (c *Config) BigTableClientFactory(userAgent string) *BigtableClientFactory { bigtableClientFactory := &BigtableClientFactory{ - UserAgent: userAgent, - TokenSource: c.tokenSource, + UserAgent: userAgent, + TokenSource: c.tokenSource, + BillingProject: c.BillingProject, + UserProjectOverride: c.UserProjectOverride, } return bigtableClientFactory