feat: update S3 client initialization (thomaspoignant#2480)

Co-authored-by: Thomas Poignant <[email protected]>

Author: hoangnv-bkhn

exporter/s3exporterv2/exporter.go

@@ -51,6 +51,13 @@ type Exporter struct {
 	// Default: SNAPPY
 	ParquetCompressionCodec string
 
+	// S3ClientOptions is a list of functional options to configure the S3 client.
+	// Provide additional functional options to further configure the behavior of the client,
+	// such as changing the client's endpoint or adding custom middleware behavior.
+	// For more information about the options, please check:
+	// https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/service/s3#Options
+	S3ClientOptions []func(*s3.Options)
+
 	s3Uploader UploaderAPI
 	init       sync.Once
 	ffLogger   *fflog.FFLogger
@@ -68,7 +75,7 @@ func (f *Exporter) initializeUploader(ctx context.Context) error {
 			f.AwsConfig = &cfg
 		}
 
-		client := s3.NewFromConfig(*f.AwsConfig)
+		client := s3.NewFromConfig(*f.AwsConfig, f.S3ClientOptions...)
 		f.s3Uploader = manager.NewUploader(client)
 	})
 	return initErr

exporter/s3exporterv2/exporter_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/stretchr/testify/assert"
 	"github.com/thomaspoignant/go-feature-flag/exporter"
 	"github.com/thomaspoignant/go-feature-flag/testutils"
@@ -16,12 +17,13 @@ import (
 func TestS3_Export(t *testing.T) {
 	hostname, _ := os.Hostname()
 	type fields struct {
-		Bucket      string
-		AwsConfig   *aws.Config
-		Format      string
-		S3Path      string
-		Filename    string
-		CsvTemplate string
+		Bucket          string
+		AwsConfig       *aws.Config
+		Format          string
+		S3Path          string
+		Filename        string
+		CsvTemplate     string
+		S3ClientOptions []func(*s3.Options)
 	}
 
 	tests := []struct {
@@ -164,19 +166,45 @@ func TestS3_Export(t *testing.T) {
 			},
 			wantErr: true,
 		},
+		{
+			name: "With S3 Client Options",
+			fields: fields{
+				Bucket: "test",
+				S3ClientOptions: []func(*s3.Options){
+					func(o *s3.Options) {
+						o.UseAccelerate = true
+					},
+				},
+			},
+			events: []exporter.FeatureEvent{
+				{
+					Kind: "feature", ContextKind: "anonymousUser", UserKey: "ABCD", CreationDate: 1617970547, Key: "random-key",
+					Variation: "Default", Value: "YO", Default: false, Source: "SERVER",
+				},
+			},
+			expectedFile: "./testdata/all_default.json",
+			expectedName: "^/flag-variation-" + hostname + "-[0-9]*\\.json$",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			s3ManagerMock := testutils.S3ManagerV2Mock{}
 			f := &Exporter{
-				Bucket:      tt.fields.Bucket,
-				AwsConfig:   tt.fields.AwsConfig,
-				Format:      tt.fields.Format,
-				S3Path:      tt.fields.S3Path,
-				Filename:    tt.fields.Filename,
-				CsvTemplate: tt.fields.CsvTemplate,
-				s3Uploader:  &s3ManagerMock,
+				Bucket:          tt.fields.Bucket,
+				AwsConfig:       tt.fields.AwsConfig,
+				Format:          tt.fields.Format,
+				S3Path:          tt.fields.S3Path,
+				Filename:        tt.fields.Filename,
+				CsvTemplate:     tt.fields.CsvTemplate,
+				S3ClientOptions: tt.fields.S3ClientOptions,
+				s3Uploader:      &s3ManagerMock,
+			}
+
+			// Verify that S3ClientOptions are correctly set on the Exporter
+			if tt.fields.S3ClientOptions != nil {
+				assert.Equal(t, tt.fields.S3ClientOptions, f.S3ClientOptions, "S3ClientOptions should be set correctly on the Exporter")
 			}
+
 			err := f.Export(context.Background(), &fflog.FFLogger{LeveledLogger: slog.Default()}, tt.events)
 			if tt.wantErr {
 				assert.Error(t, err, "Export should error")

retriever/s3retrieverv2/retriever.go

@@ -24,6 +24,13 @@ type Retriever struct {
 	// download your feature flag configuration file.
 	AwsConfig *aws.Config
 
+	// S3ClientOptions is a list of functional options to configure the S3 client.
+	// Provide additional functional options to further configure the behavior of the client,
+	// such as changing the client's endpoint or adding custom middleware behavior.
+	// For more information about the options, please check:
+	// https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/service/s3#Options
+	S3ClientOptions []func(*s3.Options)
+
 	// downloader is an internal field, it is the downloader use by the AWS-SDK
 	downloader DownloaderAPI
 	status     retriever.Status
@@ -40,7 +47,7 @@ func (s *Retriever) Init(ctx context.Context, _ *fflog.FFLogger) error {
 			}
 			s.AwsConfig = &cfg
 		}
-		client := s3.NewFromConfig(*s.AwsConfig)
+		client := s3.NewFromConfig(*s.AwsConfig, s.S3ClientOptions...)
 		s.downloader = manager.NewDownloader(client)
 	}
 	s.status = retriever.RetrieverReady

retriever/s3retrieverv2/retriever_test.go

@@ -6,17 +6,19 @@ import (
 	"testing"
 
 	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"github.com/stretchr/testify/assert"
 	"github.com/thomaspoignant/go-feature-flag/retriever"
 	"github.com/thomaspoignant/go-feature-flag/testutils"
 )
 
 func Test_s3Retriever_Retrieve(t *testing.T) {
 	type fields struct {
-		downloader DownloaderAPI
-		bucket     string
-		item       string
-		context    context.Context
+		downloader      DownloaderAPI
+		bucket          string
+		item            string
+		context         context.Context
+		S3ClientOptions []func(*s3.Options)
 	}
 	tests := []struct {
 		name    string
@@ -60,22 +62,46 @@ func Test_s3Retriever_Retrieve(t *testing.T) {
 			want:    "./testdata/flag-config.yaml",
 			wantErr: false,
 		},
+		{
+			name: "With S3 Client Options",
+			fields: fields{
+				downloader: &testutils.S3ManagerV2Mock{
+					TestDataLocation: "./testdata",
+				},
+				bucket: "Bucket",
+				item:   "valid",
+				S3ClientOptions: []func(*s3.Options){
+					func(o *s3.Options) {
+						o.UseAccelerate = true
+					},
+				},
+			},
+			want:    "./testdata/flag-config.yaml",
+			wantErr: false,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			awsConf, _ := config.LoadDefaultConfig(context.TODO())
 			s := Retriever{
-				Bucket:     tt.fields.bucket,
-				Item:       tt.fields.item,
-				AwsConfig:  &awsConf,
-				downloader: tt.fields.downloader,
+				Bucket:          tt.fields.bucket,
+				Item:            tt.fields.item,
+				AwsConfig:       &awsConf,
+				downloader:      tt.fields.downloader,
+				S3ClientOptions: tt.fields.S3ClientOptions,
 			}
 			err := s.Init(context.Background(), nil)
 			assert.NoError(t, err)
 			defer func() {
 				err := s.Shutdown(context.Background())
 				assert.NoError(t, err)
 			}()
+
+			// Verify that S3ClientOptions are correctly set on the Retriever
+			if tt.fields.S3ClientOptions != nil {
+				assert.Equal(t, tt.fields.S3ClientOptions, s.S3ClientOptions, "S3ClientOptions should be set correctly on the Retriever")
+			}
+
 			got, err := s.Retrieve(tt.fields.context)
 			assert.Equal(t, tt.wantErr, err != nil, "Retrieve() error = %v, wantErr %v", err, tt.wantErr)
 			if err == nil {
@@ -118,4 +144,25 @@ func TestRetriever_Init(t *testing.T) {
 		assert.Equal(t, "us-east-1", s.AwsConfig.Region, "Setting the region from the AwsConfig should be used over the environment variable")
 		assert.Equal(t, retriever.RetrieverReady, s.Status())
 	})
+
+	t.Run("With S3 Client Options", func(t *testing.T) {
+		t.Setenv("AWS_REGION", "us-west-2")
+		s := Retriever{
+			Bucket: "TestBucket",
+			Item:   "TestItem",
+			S3ClientOptions: []func(*s3.Options){
+				func(o *s3.Options) {
+					o.UseAccelerate = true
+				},
+			},
+		}
+		err := s.Init(context.Background(), nil)
+		assert.NoError(t, err)
+		assert.NotNil(t, s.AwsConfig)
+		assert.NotNil(t, s.downloader)
+		assert.Equal(t, "us-west-2", s.AwsConfig.Region, "Setting the region from the environment variable should be copied to the aws config")
+		assert.Equal(t, retriever.RetrieverReady, s.Status())
+		assert.NotNil(t, s.S3ClientOptions, "S3ClientOptions should be set")
+		assert.Len(t, s.S3ClientOptions, 1, "S3ClientOptions should have one option")
+	})
 }