some adjustments based on review

snowflakedb · Jan 13, 2025 · 989c8f8 · 989c8f8
1 parent 697705d
commit 989c8f8
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 27 deletions.
diff --git a/azure_storage_client.go b/azure_storage_client.go
@@ -45,15 +45,14 @@ 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{
 				MaxRetries: 60,
 				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
@@ -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
@@ -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 {