Allow Specify AWS Region (#1184)

Author: Tang8330

Diff for: clients/iceberg/staging.go

@@ -37,6 +37,7 @@ func (s Store) uploadToS3(ctx context.Context, fp string) (string, error) {
 		FilePath:                   fp,
 		OverrideAWSAccessKeyID:     s.config.Iceberg.S3Tables.AwsAccessKeyID,
 		OverrideAWSAccessKeySecret: s.config.Iceberg.S3Tables.AwsSecretAccessKey,
+		Region:                     s.config.Iceberg.S3Tables.Region,
 	})
 
 	if err != nil {

Diff for: clients/redshift/redshift.go

@@ -140,6 +140,13 @@ func LoadRedshift(ctx context.Context, cfg config.Config, _store *db.Store) (*St
 		Store:             store,
 	}
 
+	// TODO: Move this and the lower required env into a separate function
+	for _, requiredEnv := range []string{"AWS_REGION"} {
+		if os.Getenv(requiredEnv) == "" {
+			return nil, fmt.Errorf("required environment variable %q is not set", requiredEnv)
+		}
+	}
+
 	if cfg.Redshift.RoleARN != "" {
 		for _, requiredEnv := range []string{"AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"} {
 			if os.Getenv(requiredEnv) == "" {

Diff for: clients/redshift/staging.go

@@ -54,6 +54,7 @@ func (s *Store) PrepareTemporaryTable(ctx context.Context, tableData *optimizati
 		OptionalS3Prefix: s.optionalS3Prefix,
 		Bucket:           s.bucket,
 		FilePath:         fp,
+		Region:           os.Getenv("AWS_REGION"),
 	}
 
 	if s._awsCredentials != nil {

Diff for: clients/s3/s3.go

@@ -135,6 +135,7 @@ func (s *Store) Merge(ctx context.Context, tableData *optimization.TableData) (b
 		FilePath:                   fp,
 		OverrideAWSAccessKeyID:     s.config.S3.AwsAccessKeyID,
 		OverrideAWSAccessKeySecret: s.config.S3.AwsSecretAccessKey,
+		Region:                     s.config.S3.AwsRegion,
 	}); err != nil {
 		return false, fmt.Errorf("failed to upload file to s3: %w", err)
 	}

Diff for: lib/awslib/s3.go

@@ -1,6 +1,7 @@
 package awslib
 
 import (
+	"cmp"
 	"context"
 	"fmt"
 	"os"
@@ -19,6 +20,7 @@ type UploadArgs struct {
 	OverrideAWSAccessKeyID     string
 	OverrideAWSAccessKeySecret string
 	OverrideAWSSessionToken    string
+	Region                     string
 }
 
 // UploadLocalFileToS3 - takes a filepath with the file and bucket and optional expiry
@@ -27,11 +29,12 @@ func UploadLocalFileToS3(ctx context.Context, args UploadArgs) (string, error) {
 	var cfg aws.Config
 	var err error
 
+	awsRegion := cmp.Or(args.Region, os.Getenv("AWS_REGION"))
 	if args.OverrideAWSAccessKeyID != "" && args.OverrideAWSAccessKeySecret != "" {
 		creds := credentials.NewStaticCredentialsProvider(args.OverrideAWSAccessKeyID, args.OverrideAWSAccessKeySecret, args.OverrideAWSSessionToken)
-		cfg, err = config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(creds))
+		cfg, err = config.LoadDefaultConfig(ctx, config.WithCredentialsProvider(creds), config.WithRegion(awsRegion))
 	} else {
-		cfg, err = config.LoadDefaultConfig(ctx)
+		cfg, err = config.LoadDefaultConfig(ctx, config.WithRegion(awsRegion))
 	}
 
 	if err != nil {

Diff for: lib/config/destination_types.go

@@ -50,6 +50,7 @@ type S3Settings struct {
 	Bucket             string                   `yaml:"bucket"`
 	AwsAccessKeyID     string                   `yaml:"awsAccessKeyID"`
 	AwsSecretAccessKey string                   `yaml:"awsSecretAccessKey"`
+	AwsRegion          string                   `yaml:"awsRegion"`
 	OutputFormat       constants.S3OutputFormat `yaml:"outputFormat"`
 }