From 989c8f86ffcd4a127c0b595f6c7f2d9108b8ed0d Mon Sep 17 00:00:00 2001 From: sfc-gh-dszmolka Date: Mon, 13 Jan 2025 17:56:15 +0100 Subject: [PATCH] some adjustments based on review --- azure_storage_client.go | 6 ++---- gcs_storage_client.go | 3 +-- s3_storage_client.go | 25 ++++--------------------- 3 files changed, 7 insertions(+), 27 deletions(-) diff --git a/azure_storage_client.go b/azure_storage_client.go index ef28112d6..b948d35c3 100644 --- a/azure_storage_client.go +++ b/azure_storage_client.go @@ -45,7 +45,6 @@ func (util *snowflakeAzureClient) createClient(info *execResponseStageInfo, _ bo if err != nil { return nil, err } - transport := getTransport(util.cfg) client, err := azblob.NewClientWithNoCredential(u.String(), &azblob.ClientOptions{ ClientOptions: azcore.ClientOptions{ Retry: policy.RetryOptions{ @@ -53,7 +52,7 @@ func (util *snowflakeAzureClient) createClient(info *execResponseStageInfo, _ bo RetryDelay: 2 * time.Second, }, Transport: &http.Client{ - Transport: transport, + Transport: getTransport(util.cfg), }, }, }) @@ -350,10 +349,9 @@ func (util *snowflakeAzureClient) detectAzureTokenExpireError(resp *http.Respons } func createContainerClient(clientURL string, cfg *Config) (*container.Client, error) { - transport := getTransport(cfg) return container.NewClientWithNoCredential(clientURL, &container.ClientOptions{ClientOptions: azcore.ClientOptions{ Transport: &http.Client{ - Transport: transport, + Transport: getTransport(cfg), }, }}) } diff --git a/gcs_storage_client.go b/gcs_storage_client.go index f2f736d4d..207c1ad12 100644 --- a/gcs_storage_client.go +++ b/gcs_storage_client.go @@ -405,9 +405,8 @@ func (util *snowflakeGcsClient) isTokenExpired(resp *http.Response) bool { } func newGcsClient(cfg *Config) gcsAPI { - transport := getTransport(cfg) return &http.Client{ - Transport: transport, + Transport: getTransport(cfg), } } diff --git a/s3_storage_client.go b/s3_storage_client.go index 35c2a4361..61a568bcc 100644 --- a/s3_storage_client.go +++ b/s3_storage_client.go @@ -49,7 +49,6 @@ func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAcce stageCredentials := info.Creds s3Logger := logging.LoggerFunc(s3LoggingFunc) endPoint := getS3CustomEndpoint(info) - transport := getTransport(util.cfg) return s3.New(s3.Options{ Region: info.Region, @@ -60,7 +59,7 @@ func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAcce BaseEndpoint: endPoint, UseAccelerate: useAccelerateEndpoint, HTTPClient: &http.Client{ - Transport: transport, + Transport: getTransport(util.cfg), }, ClientLogMode: S3LoggingMode, Logger: s3Logger, @@ -69,25 +68,9 @@ func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAcce // to be used with S3 transferAccelerateConfigWithUtil func (util *snowflakeS3Client) createClientWithConfig(info *execResponseStageInfo, useAccelerateEndpoint bool, cfg *Config) (cloudClient, error) { - stageCredentials := info.Creds - s3Logger := logging.LoggerFunc(s3LoggingFunc) - endPoint := getS3CustomEndpoint(info) - transport := getTransport(cfg) - - return s3.New(s3.Options{ - Region: info.Region, - Credentials: aws.NewCredentialsCache(credentials.NewStaticCredentialsProvider( - stageCredentials.AwsKeyID, - stageCredentials.AwsSecretKey, - stageCredentials.AwsToken)), - BaseEndpoint: endPoint, - UseAccelerate: useAccelerateEndpoint, - HTTPClient: &http.Client{ - Transport: transport, - }, - ClientLogMode: S3LoggingMode, - Logger: s3Logger, - }), nil + // copy snowflakeFileTransferAgent's config onto the cloud client so we could decide which Transport to use + util.cfg = cfg + return util.createClient(info, useAccelerateEndpoint) } func getS3CustomEndpoint(info *execResponseStageInfo) *string {