support masterpayer ARN

pkg/cloud/awsprovider.go

@@ -26,6 +26,7 @@ import (
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/aws/credentials"
+	"github.com/aws/aws-sdk-go/aws/credentials/stscreds"
 	"github.com/aws/aws-sdk-go/aws/session"
 	"github.com/aws/aws-sdk-go/service/athena"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -282,6 +283,7 @@ type AwsAthenaInfo struct {
 	ServiceKeyName   string `json:"serviceKeyName"`
 	ServiceKeySecret string `json:"serviceKeySecret"`
 	AccountID        string `json:"projectID"`
+	MasterPayerARN   string `json:"masterPayerARN"`
 }
 
 func (aws *AWS) GetManagementPlatform() (string, error) {
@@ -352,6 +354,9 @@ func (aws *AWS) UpdateConfig(r io.Reader, updateType string) (*CustomPricing, er
 			if a.ServiceKeySecret != "" {
 				c.ServiceKeySecret = a.ServiceKeySecret
 			}
+			if a.MasterPayerARN != "" {
+				c.MasterPayerARN = a.MasterPayerARN
+			}
 			c.AthenaProjectID = a.AccountID
 		} else {
 			a := make(map[string]interface{})
@@ -1429,6 +1434,13 @@ func (a *AWS) QueryAthenaBillingData(query string) (*athena.GetQueryResultsOutpu
 	}
 	s := session.Must(session.NewSession(c))
 	svc := athena.New(s)
+	if customPricing.MasterPayerARN != "" {
+		creds := stscreds.NewCredentials(s, customPricing.MasterPayerARN)
+		svc = athena.New(s, &aws.Config{
+			Region:      region,
+			Credentials: creds,
+		})
+	}
 
 	var e athena.StartQueryExecutionInput
 

pkg/cloud/provider.go

@@ -140,6 +140,7 @@ type CustomPricing struct {
 	AthenaRegion          string            `json:"athenaRegion"`
 	AthenaDatabase        string            `json:"athenaDatabase"`
 	AthenaTable           string            `json:"athenaTable"`
+	MasterPayerARN        string            `json:"masterPayerARN"`
 	BillingDataDataset    string            `json:"billingDataDataset,omitempty"`
 	CustomPricesEnabled   string            `json:"customPricesEnabled"`
 	DefaultIdle           string            `json:"defaultIdle"`