Refactor aws sdk client creation

Signed-off-by: Phan Duc <[email protected]>
kedacore · Oct 28, 2022 · cf048ce · cf048ce
1 parent fad174a
commit cf048ce
Show file tree

Hide file tree

Showing 6 changed files with 68 additions and 149 deletions.
diff --git a/pkg/scalers/aws_cloudwatch_scaler.go b/pkg/scalers/aws_cloudwatch_scaler.go
@@ -8,9 +8,6 @@ import (
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"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/cloudwatch"
 	"github.com/aws/aws-sdk-go/service/cloudwatch/cloudwatchiface"
 	"github.com/go-logr/logr"
@@ -113,32 +110,11 @@ func getFloatMetadataValue(metadata map[string]string, key string, required bool
 }
 
 func createCloudwatchClient(metadata *awsCloudwatchMetadata) *cloudwatch.CloudWatch {
-	sess := session.Must(session.NewSession(&aws.Config{
-		Region:   aws.String(metadata.awsRegion),
-		Endpoint: aws.String(metadata.awsEndpoint),
-	}))
+	sess, config := getAwsConfig(metadata.awsRegion,
+		metadata.awsEndpoint,
+		metadata.awsAuthorization)
 
-	var cloudwatchClient *cloudwatch.CloudWatch
-	if metadata.awsAuthorization.podIdentityOwner {
-		creds := credentials.NewStaticCredentials(metadata.awsAuthorization.awsAccessKeyID, metadata.awsAuthorization.awsSecretAccessKey, metadata.awsAuthorization.awsSessionToken)
-
-		if metadata.awsAuthorization.awsRoleArn != "" {
-			creds = stscreds.NewCredentials(sess, metadata.awsAuthorization.awsRoleArn)
-		}
-
-		cloudwatchClient = cloudwatch.New(sess, &aws.Config{
-			Region:      aws.String(metadata.awsRegion),
-			Endpoint:    aws.String(metadata.awsEndpoint),
-			Credentials: creds,
-		})
-	} else {
-		cloudwatchClient = cloudwatch.New(sess, &aws.Config{
-			Region:   aws.String(metadata.awsRegion),
-			Endpoint: aws.String(metadata.awsEndpoint),
-		})
-	}
-
-	return cloudwatchClient
+	return cloudwatch.New(sess, config)
 }
 
 func parseAwsCloudwatchMetadata(config *ScalerConfig) (*awsCloudwatchMetadata, error) {

diff --git a/pkg/scalers/aws_iam_authorization.go → pkg/scalers/aws_common.go b/pkg/scalers/aws_iam_authorization.go → pkg/scalers/aws_common.go
@@ -1,6 +1,13 @@
 package scalers
 
-import "fmt"
+import (
+	"fmt"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"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"
+)
 
 type awsAuthorizationMetadata struct {
 	awsRoleArn string
@@ -12,6 +19,43 @@ type awsAuthorizationMetadata struct {
 	podIdentityOwner bool
 }
 
+type awsConfigMetadata struct {
+	awsRegion        string
+	awsEndpoint      string
+	awsAuthorization awsAuthorizationMetadata
+}
+
+func getAwsConfig(awsRegion string, awsEndpoint string, awsAuthorization awsAuthorizationMetadata) (*session.Session, *aws.Config) {
+	metadata := &awsConfigMetadata{
+		awsRegion:        awsRegion,
+		awsEndpoint:      awsEndpoint,
+		awsAuthorization: awsAuthorization}
+
+	sess := session.Must(session.NewSession(&aws.Config{
+		Region:   aws.String(metadata.awsRegion),
+		Endpoint: aws.String(metadata.awsEndpoint),
+	}))
+
+	if !metadata.awsAuthorization.podIdentityOwner {
+		return sess, &aws.Config{
+			Region:   aws.String(metadata.awsRegion),
+			Endpoint: aws.String(metadata.awsEndpoint),
+		}
+	}
+
+	creds := credentials.NewStaticCredentials(metadata.awsAuthorization.awsAccessKeyID, metadata.awsAuthorization.awsSecretAccessKey, "")
+
+	if metadata.awsAuthorization.awsRoleArn != "" {
+		creds = stscreds.NewCredentials(sess, metadata.awsAuthorization.awsRoleArn)
+	}
+
+	return sess, &aws.Config{
+		Region:      aws.String(metadata.awsRegion),
+		Endpoint:    aws.String(metadata.awsEndpoint),
+		Credentials: creds,
+	}
+}
+
 func getAwsAuthorization(authParams, metadata, resolvedEnv map[string]string) (awsAuthorizationMetadata, error) {
 	meta := awsAuthorizationMetadata{}
 

diff --git a/pkg/scalers/aws_dynamodb_scaler.go b/pkg/scalers/aws_dynamodb_scaler.go
@@ -8,9 +8,6 @@ import (
 	"strconv"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"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/dynamodb"
 	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
 	"github.com/go-logr/logr"
@@ -147,36 +144,12 @@ func parseAwsDynamoDBMetadata(config *ScalerConfig) (*awsDynamoDBMetadata, error
 	return &meta, nil
 }
 
-func createDynamoDBClient(meta *awsDynamoDBMetadata) *dynamodb.DynamoDB {
-	sess := session.Must(session.NewSession(&aws.Config{
-		Region:   aws.String(meta.awsRegion),
-		Endpoint: aws.String(meta.awsEndpoint),
-	}))
+func createDynamoDBClient(metadata *awsDynamoDBMetadata) *dynamodb.DynamoDB {
+	sess, config := getAwsConfig(metadata.awsRegion,
+		metadata.awsEndpoint,
+		metadata.awsAuthorization)
 
-	var dbClient *dynamodb.DynamoDB
-
-	if !meta.awsAuthorization.podIdentityOwner {
-		dbClient = dynamodb.New(sess, &aws.Config{
-			Region:   aws.String(meta.awsRegion),
-			Endpoint: aws.String(meta.awsEndpoint),
-		})
-
-		return dbClient
-	}
-
-	creds := credentials.NewStaticCredentials(meta.awsAuthorization.awsAccessKeyID, meta.awsAuthorization.awsSecretAccessKey, "")
-
-	if meta.awsAuthorization.awsRoleArn != "" {
-		creds = stscreds.NewCredentials(sess, meta.awsAuthorization.awsRoleArn)
-	}
-
-	dbClient = dynamodb.New(sess, &aws.Config{
-		Region:      aws.String(meta.awsRegion),
-		Endpoint:    aws.String(meta.awsEndpoint),
-		Credentials: creds,
-	})
-
-	return dbClient
+	return dynamodb.New(sess, config)
 }
 
 func (s *awsDynamoDBScaler) GetMetrics(ctx context.Context, metricName string, metricSelector labels.Selector) ([]external_metrics.ExternalMetricValue, error) {

diff --git a/pkg/scalers/aws_dynamodb_streams_scaler.go b/pkg/scalers/aws_dynamodb_streams_scaler.go
@@ -5,10 +5,6 @@ import (
 	"fmt"
 	"strconv"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"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/dynamodb"
 	"github.com/aws/aws-sdk-go/service/dynamodb/dynamodbiface"
 	"github.com/aws/aws-sdk-go/service/dynamodbstreams"
@@ -127,39 +123,16 @@ func parseAwsDynamoDBStreamsMetadata(config *ScalerConfig, logger logr.Logger) (
 }
 
 func createClientsForDynamoDBStreamsScaler(metadata *awsDynamoDBStreamsMetadata) (*dynamodb.DynamoDB, *dynamodbstreams.DynamoDBStreams) {
-	sess := session.Must(session.NewSession(&aws.Config{
-		Region:   aws.String(metadata.awsRegion),
-		Endpoint: aws.String(metadata.awsEndpoint),
-	}))
+	sess, config := getAwsConfig(metadata.awsRegion,
+		metadata.awsEndpoint,
+		metadata.awsAuthorization)
 
 	var dbClient *dynamodb.DynamoDB
 	var dbStreamClient *dynamodbstreams.DynamoDBStreams
 
-	if metadata.awsAuthorization.podIdentityOwner {
-		creds := credentials.NewStaticCredentials(metadata.awsAuthorization.awsAccessKeyID, metadata.awsAuthorization.awsSecretAccessKey, metadata.awsAuthorization.awsSessionToken)
-		if metadata.awsAuthorization.awsRoleArn != "" {
-			creds = stscreds.NewCredentials(sess, metadata.awsAuthorization.awsRoleArn)
-		}
-		dbClient = dynamodb.New(sess, &aws.Config{
-			Region:      aws.String(metadata.awsRegion),
-			Endpoint:    aws.String(metadata.awsEndpoint),
-			Credentials: creds,
-		})
-		dbStreamClient = dynamodbstreams.New(sess, &aws.Config{
-			Region:      aws.String(metadata.awsRegion),
-			Endpoint:    aws.String(metadata.awsEndpoint),
-			Credentials: creds,
-		})
-	} else {
-		dbClient = dynamodb.New(sess, &aws.Config{
-			Region:   aws.String(metadata.awsRegion),
-			Endpoint: aws.String(metadata.awsEndpoint),
-		})
-		dbStreamClient = dynamodbstreams.New(sess, &aws.Config{
-			Region:   aws.String(metadata.awsRegion),
-			Endpoint: aws.String(metadata.awsEndpoint),
-		})
-	}
+	dbClient = dynamodb.New(sess, config)
+	dbStreamClient = dynamodbstreams.New(sess, config)
+
 	return dbClient, dbStreamClient
 }
 

diff --git a/pkg/scalers/aws_kinesis_stream_scaler.go b/pkg/scalers/aws_kinesis_stream_scaler.go
@@ -5,10 +5,6 @@ import (
 	"fmt"
 	"strconv"
 
-	"github.com/aws/aws-sdk-go/aws"
-	"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/kinesis"
 	"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
 	"github.com/go-logr/logr"
@@ -116,31 +112,11 @@ func parseAwsKinesisStreamMetadata(config *ScalerConfig, logger logr.Logger) (*a
 }
 
 func createKinesisClient(metadata *awsKinesisStreamMetadata) *kinesis.Kinesis {
-	sess := session.Must(session.NewSession(&aws.Config{
-		Region:   aws.String(metadata.awsRegion),
-		Endpoint: aws.String(metadata.awsEndpoint),
-	}))
+	sess, config := getAwsConfig(metadata.awsRegion,
+		metadata.awsEndpoint,
+		metadata.awsAuthorization)
 
-	var kinesisClinent *kinesis.Kinesis
-	if metadata.awsAuthorization.podIdentityOwner {
-		creds := credentials.NewStaticCredentials(metadata.awsAuthorization.awsAccessKeyID, metadata.awsAuthorization.awsSecretAccessKey, metadata.awsAuthorization.awsSessionToken)
-
-		if metadata.awsAuthorization.awsRoleArn != "" {
-			creds = stscreds.NewCredentials(sess, metadata.awsAuthorization.awsRoleArn)
-		}
-
-		kinesisClinent = kinesis.New(sess, &aws.Config{
-			Region:      aws.String(metadata.awsRegion),
-			Endpoint:    aws.String(metadata.awsEndpoint),
-			Credentials: creds,
-		})
-	} else {
-		kinesisClinent = kinesis.New(sess, &aws.Config{
-			Region:   aws.String(metadata.awsRegion),
-			Endpoint: aws.String(metadata.awsEndpoint),
-		})
-	}
-	return kinesisClinent
+	return kinesis.New(sess, config)
 }
 
 // IsActive determines if we need to scale from zero

diff --git a/pkg/scalers/aws_sqs_queue_scaler.go b/pkg/scalers/aws_sqs_queue_scaler.go
@@ -8,9 +8,6 @@ import (
 	"strings"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"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/sqs"
 	"github.com/aws/aws-sdk-go/service/sqs/sqsiface"
 	"github.com/go-logr/logr"
@@ -163,31 +160,11 @@ func parseAwsSqsQueueMetadata(config *ScalerConfig, logger logr.Logger) (*awsSqs
 }
 
 func createSqsClient(metadata *awsSqsQueueMetadata) *sqs.SQS {
-	sess := session.Must(session.NewSession(&aws.Config{
-		Region:   aws.String(metadata.awsRegion),
-		Endpoint: aws.String(metadata.awsEndpoint),
-	}))
+	sess, config := getAwsConfig(metadata.awsRegion,
+		metadata.awsEndpoint,
+		metadata.awsAuthorization)
 
-	var sqsClient *sqs.SQS
-	if metadata.awsAuthorization.podIdentityOwner {
-		creds := credentials.NewStaticCredentials(metadata.awsAuthorization.awsAccessKeyID, metadata.awsAuthorization.awsSecretAccessKey, metadata.awsAuthorization.awsSessionToken)
-
-		if metadata.awsAuthorization.awsRoleArn != "" {
-			creds = stscreds.NewCredentials(sess, metadata.awsAuthorization.awsRoleArn)
-		}
-
-		sqsClient = sqs.New(sess, &aws.Config{
-			Region:      aws.String(metadata.awsRegion),
-			Endpoint:    aws.String(metadata.awsEndpoint),
-			Credentials: creds,
-		})
-	} else {
-		sqsClient = sqs.New(sess, &aws.Config{
-			Region:   aws.String(metadata.awsRegion),
-			Endpoint: aws.String(metadata.awsEndpoint),
-		})
-	}
-	return sqsClient
+	return sqs.New(sess, config)
 }
 
 // IsActive determines if we need to scale from zero