feat: Add user agent

Author: daixba

dth/client.go

@@ -28,7 +28,6 @@ import (
 	"sync"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
-	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/credentials"
 	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/aws/aws-sdk-go-v2/service/s3/types"
@@ -85,13 +84,9 @@ func getEndpointURL(region, sourceType string) (url string) {
 // NewS3Client creates a S3Client instance
 func NewS3Client(ctx context.Context, bucket, prefix, endpoint, region, sourceType string, cred *S3Credentials) *S3Client {
 
-	// config, err := config.LoadDefaultConfig(ctx, config.WithRegion(region))
-	config, err := config.LoadDefaultConfig(ctx)
-	if err != nil {
-		log.Fatalf("Failed to load default SDK config to create S3 client - %s\n", err.Error())
-	}
+	cfg := loadDefaultConfig(ctx)
 
-	client := s3.NewFromConfig(config, func(o *s3.Options) {
+	client := s3.NewFromConfig(cfg, func(o *s3.Options) {
 		// retryer := retry.AddWithMaxBackoffDelay(retry.NewStandard(), time.Second*5)
 		// o.Retryer = retryer
 		if region != "" {

dth/config.go

@@ -16,6 +16,16 @@ limitations under the License.
 
 package dth
 
+import (
+	"context"
+	"log"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	mw "github.com/aws/aws-sdk-go-v2/aws/middleware"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/smithy-go/middleware"
+)
+
 const (
 	// MaxRetries when failed used globally
 	// No need an option of this.
@@ -62,3 +72,18 @@ type JobConfig struct {
 	SrcInCurrentAccount, DestInCurrentAccount                                     bool
 	*JobOptions
 }
+
+// loadDefaultConfig reads the SDK's default external configurations with custom user agent for API tracking purpose
+func loadDefaultConfig(ctx context.Context) (cfg aws.Config) {
+	cfg, err := config.LoadDefaultConfig(
+		ctx,
+		// config.WithClientLogMode(aws.LogRequest|aws.LogRequestWithBody),
+		config.WithAPIOptions([]func(*middleware.Stack) error{
+			mw.AddUserAgentKey("AWS/SO8001/v2.0.0"),
+		}),
+	)
+	if err != nil {
+		log.Fatalf("Failed to load default SDK config to create client - %s\n", err.Error())
+	}
+	return cfg
+}

dth/service.go

@@ -23,7 +23,6 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/aws/aws-sdk-go-v2/config"
 	"github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue"
 	"github.com/aws/aws-sdk-go-v2/service/dynamodb"
 	dtype "github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
@@ -60,12 +59,8 @@ type SecretService struct {
 
 // NewSecretService is a helper func to create a SecretService instance
 func NewSecretService(ctx context.Context) (*SecretService, error) {
-	cfg, err := config.LoadDefaultConfig(ctx)
-	if err != nil {
-		// Don't need to quit the process
-		log.Printf("Failed to load default SDK config to create SSM client - %s\n", err.Error())
-		return nil, err
-	}
+
+	cfg := loadDefaultConfig(ctx)
 
 	// Create an Amazon DynamoDB client.
 	client := sm.NewFromConfig(cfg)
@@ -77,10 +72,8 @@ func NewSecretService(ctx context.Context) (*SecretService, error) {
 
 // NewDBService is a helper func to create a DBService instance
 func NewDBService(ctx context.Context, tableName string) (*DBService, error) {
-	cfg, err := config.LoadDefaultConfig(ctx)
-	if err != nil {
-		log.Fatalf("Failed to load default SDK config to create DynamoDB client - %s\n", err.Error())
-	}
+
+	cfg := loadDefaultConfig(ctx)
 
 	// Create an Amazon DynamoDB client.
 	client := dynamodb.NewFromConfig(cfg)
@@ -94,11 +87,7 @@ func NewDBService(ctx context.Context, tableName string) (*DBService, error) {
 // NewSqsService is a helper func to create a SqsService instance
 func NewSqsService(ctx context.Context, queueName string) (*SqsService, error) {
 
-	// Only to use default config here
-	cfg, err := config.LoadDefaultConfig(ctx)
-	if err != nil {
-		log.Fatalf("Failed to load default SDK config to create SQS client - %s\n", err.Error())
-	}
+	cfg := loadDefaultConfig(ctx)
 
 	client := sqs.NewFromConfig(cfg)
 	input := &sqs.GetQueueUrlInput{

go.mod

@@ -8,9 +8,10 @@ require (
 	github.com/aws/aws-sdk-go-v2/credentials v1.2.1
 	github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue v1.1.1
 	github.com/aws/aws-sdk-go-v2/service/dynamodb v1.3.1
-	github.com/aws/aws-sdk-go-v2/service/s3 v1.8.0
+	github.com/aws/aws-sdk-go-v2/service/s3 v1.9.0
 	github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.3.1
 	github.com/aws/aws-sdk-go-v2/service/sqs v1.4.1
+	github.com/aws/smithy-go v1.4.0
 	github.com/spf13/cobra v1.1.3
 	github.com/spf13/viper v1.7.1
 )

go.sum

@@ -42,8 +42,8 @@ github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.1.1 h1:l7pDLsmOGr
 github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.1.1/go.mod h1:2+ehJPkdIdl46VCj67Emz/EH2hpebHZtaLdzqg+sWOI=
 github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.3.1 h1:VH1Y4k+IZ5kcRVqSNw7eAkXyfS7k2/ibKjrNtbhYhV4=
 github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.3.1/go.mod h1:IpjxfORBAFfkMM0VEx5gPPnEy6WV4Hk0F/+zb/SUWyw=
-github.com/aws/aws-sdk-go-v2/service/s3 v1.8.0 h1:rljno3viFN46b59CbjkIqYwxEAzk4naLe+djOb/exLs=
-github.com/aws/aws-sdk-go-v2/service/s3 v1.8.0/go.mod h1:zHCjYoODbYRLz/iFicYswq1gRoxBnHvpY5h2Vg3/tJ4=
+github.com/aws/aws-sdk-go-v2/service/s3 v1.9.0 h1:FZ5UL5aiybSJKiJglPT7YMMwc431IgOX5gvlFAzSjzs=
+github.com/aws/aws-sdk-go-v2/service/s3 v1.9.0/go.mod h1:zHCjYoODbYRLz/iFicYswq1gRoxBnHvpY5h2Vg3/tJ4=
 github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.3.1 h1:atHdsCczZyM/y9QIoCQnxudoKk8+ya2EPIplDOofkjw=
 github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.3.1/go.mod h1:ayQUSrG5QyIl2jRSB0YnoJ1e9swNxsBWaCK3hNI2caI=
 github.com/aws/aws-sdk-go-v2/service/sqs v1.4.1 h1:BBW0EAZrJh3pesbs9m7ghgC8kylLjTPmWN3yG7POjnA=