Repository: snowflakedb/gosnowflake
Pull request #1364 (2025-04-10 10:45:26)

SNOW-2034185: Remove reencryption when retry PUT upload
### Description

SNOW-2034185: When the retry of the PUT upload to the storage happened, the stream was encrypted again every retry.

### Checklist
- [x] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary


sfc-gh-snowflakedb-snyk-sa: ![](https://res.cloudinary.com/snyk/image/upload/r-d/scm-platform/snyk-pull-requests/pr-banner-default.svg)

### :tada: **Snyk checks have passed. No issues have been found so far.**

:white_check_mark: **security/snyk** check is complete. No issues have been found. [(View Details)](https://app.snyk.io/org/snowflakedb-sca-scanning-public-repo/pr-checks/c53f1771-28c7-447a-84ae-8de89acca169?refs=f12ecb9107183f2b9dd4a959d3b6793d67c22261&source=prChecksComment)

:white_check_mark: **license/snyk** check is complete. No issues have been found. [(View Details)](https://app.snyk.io/org/snowflakedb-sca-scanning-public-repo/pr-checks/c53f1771-28c7-447a-84ae-8de89acca169/license?refs=f12ecb9107183f2b9dd4a959d3b6793d67c22261&source=prChecksComment)



codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1364?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) Report
Attention: Patch coverage is `84.61538%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 82.21%. Comparing base [(`fb3995b`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/fb3995b3f25ebb1491a3522cae82abecb1c58d3e?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) to head [(`4860fe5`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/4860fe5fe545cac432fc345764405a551d6ed52a?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

| [Files with missing lines](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1364?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [storage\_client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1364?src=pr&el=tree&filepath=storage_client.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-c3RvcmFnZV9jbGllbnQuZ28=) | 84.61% | [2 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1364?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #1364      +/-   ##
==========================================
+ Coverage   82.18%   82.21%   +0.02%     
==========================================
  Files          57       57              
  Lines       13916    13918       +2     
==========================================
+ Hits        11437    11442       +5     
+ Misses       2479     2476       -3     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1364?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>

Patch:
diff --git a/azure_storage_client.go b/azure_storage_client.go
index 30a4b40..4772546 100644
--- a/azure_storage_client.go
+++ b/azure_storage_client.go
@@ -1,6 +1,7 @@
 package gosnowflake
 
 import (
+	"cmp"
 	"context"
 	"encoding/json"
 	"errors"
@@ -143,25 +144,24 @@ func (util *snowflakeAzureClient) getFileHeader(meta *fileMetadata, filename str
 func (util *snowflakeAzureClient) uploadFile(
 	dataFile string,
 	meta *fileMetadata,
-	encryptMeta *encryptMetadata,
 	maxConcurrency int,
 	multiPartThreshold int64) error {
 	azureMeta := map[string]*string{
 		"sfcdigest": &meta.sha256Digest,
 	}
-	if encryptMeta != nil {
+	if meta.encryptMeta != nil {
 		ed := &encryptionData{
 			EncryptionMode: "FullBlob",
 			WrappedContentKey: contentKey{
 				"symmKey1",
-				encryptMeta.key,
+				meta.encryptMeta.key,
 				"AES_CBC_256",
 			},
 			EncryptionAgent: encryptionAgent{
 				"1.0",
 				"AES_CBC_128",
 			},
-			ContentEncryptionIV: encryptMeta.iv,
+			ContentEncryptionIV: meta.encryptMeta.iv,
 			KeyWrappingMetadata: keyMetadata{
 				"Java 5.3.0",
 			},
@@ -172,7 +172,7 @@ func (util *snowflakeAzureClient) uploadFile(
 		}
 		encryptionMetadata := string(metadata)
 		azureMeta["encryptiondata"] = &encryptionMetadata
-		azureMeta["matdesc"] = &encryptMeta.matdesc
+		azureMeta["matdesc"] = &meta.encryptMeta.matdesc
 	}
 
 	azureLoc, err := util.extractContainerNameAndPath(meta.stageInfo.Location)
@@ -200,10 +200,7 @@ func (util *snowflakeAzureClient) uploadFile(
 		blobClient = meta.mockAzureClient
 	}
 	if meta.srcStream != nil {
-		uploadSrc := meta.srcStream
-		if meta.realSrcStream != nil {
-			uploadSrc = meta.realSrcStream
-		}
+		uploadSrc := cmp.Or(meta.realSrcStream, meta.srcStream)
 		_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (azblob.UploadStreamResponse, error) {
 			return blobClient.UploadStream(ctx, uploadSrc, &azblob.UploadStreamOptions{
 				BlockSize: int64(uploadSrc.Len()),
diff --git a/file_transfer_agent.go b/file_transfer_agent.go
index 49ba2b0..a4a426c 100644
--- a/file_transfer_agent.go
+++ b/file_transfer_agent.go
@@ -4,6 +4,7 @@ package gosnowflake
 
 import (
 	"bytes"
+	"cmp"
 	"context"
 	"database/sql/driver"
 	"encoding/json"
@@ -865,26 +866,17 @@ func (sfa *snowflakeFileTransferAgent) uploadOneFile(meta *fileMetadata) (*fileM
 	defer os.RemoveAll(tmpDir) // cleanup
 
 	fileUtil := new(snowflakeFileUtil)
-	if meta.requireCompress {
-		if meta.srcStream != nil {
-			meta.realSrcStream, _, err = fileUtil.compressFileWithGzipFromStream(&meta.srcStream)
-		} else {
-			meta.realSrcFileName, _, err = fileUtil.compressFileWithGzip(meta.srcFileName, tmpDir)
-		}
-		if err != nil {
-			return nil, err
-		}
+	err = compressDataIfRequired(meta, fileUtil, tmpDir)
+	if err != nil {
+		return meta, err
 	}
 
-	if meta.srcStream != nil {
-		if meta.realSrcStream != nil {
-			meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForStream(&meta.realSrcStream)
-		} else {
-			meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForStream(&meta.srcStream)
-		}
-	} else {
-		meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForFile(meta.realSrcFileName)
+	err = updateUploadSize(meta, fileUtil)
+	if err != nil {
+		return meta, err
 	}
+
+	err = encryptDataIfRequired(meta, sfa.stageLocationType)
 	if err != nil {
 		return meta, err
 	}
@@ -1226,3 +1218,53 @@ func (spp *snowflakeProgressPercentage) updateProgress(filename string, startTim
 	}
 	return progress == 1.0
 }
+
+func compressDataIfRequired(meta *fileMetadata, fileUtil *snowflakeFileUtil, tmpDir string) error {
+	var err error
+	if meta.requireCompress {
+		if meta.srcStream != nil {
+			meta.realSrcStream, _, err = fileUtil.compressFileWithGzipFromStream(&meta.srcStream)
+		} else {
+			meta.realSrcFileName, _, err = fileUtil.compressFileWithGzip(meta.srcFileName, tmpDir)
+		}
+	}
+	return err
+}
+
+func updateUploadSize(meta *fileMetadata, fileUtil *snowflakeFileUtil) error {
+	var err error
+	if meta.srcStream != nil {
+		if meta.realSrcStream != nil {
+			meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForStream(&meta.realSrcStream)
+		} else {
+			meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForStream(&meta.srcStream)
+		}
+	} else {
+		meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForFile(meta.realSrcFileName)
+	}
+	return err
+}
+
+func encryptDataIfRequired(meta *fileMetadata, ct cloudType) error {
+	if ct != local && meta.encryptionMaterial != nil {
+		var err error
+		if meta.srcStream != nil {
+			var encryptedStream bytes.Buffer
+			srcStream := cmp.Or(meta.realSrcStream, meta.srcStream)
+			meta.encryptMeta, err = encryptStreamCBC(meta.encryptionMaterial, srcStream, &encryptedStream, 0)
+			if err != nil {
+				return err
+			}
+			meta.realSrcStream = &encryptedStream
+		} else {
+			var dataFile string
+			meta.encryptMeta, dataFile, err = encryptFileCBC(meta.encryptionMaterial, meta.realSrcFileName, 0, meta.tmpDir)
+			if err != nil {
+				return err
+			}
+			meta.realSrcFileName = dataFile
+		}
+	}
+
+	return nil
+}
diff --git a/file_util.go b/file_util.go
index 35d260c..497cab1 100644
--- a/file_util.go
+++ b/file_util.go
@@ -129,6 +129,7 @@ type fileMetadata struct {
 	resStatus          resultStatus
 	stageInfo          *execResponseStageInfo
 	encryptionMaterial *snowflakeFileEncryption
+	encryptMeta        *encryptMetadata
 
 	srcFileName        string
 	realSrcFileName    string
diff --git a/gcs_storage_client.go b/gcs_storage_client.go
index 338f7d3..62cf2bd 100644
--- a/gcs_storage_client.go
+++ b/gcs_storage_client.go
@@ -143,7 +143,6 @@ type gcsAPI interface {
 func (util *snowflakeGcsClient) uploadFile(
 	dataFile string,
 	meta *fileMetadata,
-	encryptMeta *encryptMetadata,
 	maxConcurrency int,
 	multiPartThreshold int64) error {
 	uploadURL := meta.presignedURL
@@ -178,19 +177,19 @@ func (util *snowflakeGcsClient) uploadFile(
 		gcsHeaders["Authorization"] = "Bearer " + accessToken
 	}
 
-	if encryptMeta != nil {
+	if meta.encryptMeta != nil {
 		encryptData := encryptionData{
 			"FullBlob",
 			contentKey{
 				"symmKey1",
-				encryptMeta.key,
+				meta.encryptMeta.key,
 				"AES_CBC_256",
 			},
 			encryptionAgent{
 				"1.0",
 				"AES_CBC_256",
 			},
-			encryptMeta.iv,
+			meta.encryptMeta.iv,
 			keyMetadata{
 				"Java 5.3.0",
 			},
@@ -200,7 +199,7 @@ func (util *snowflakeGcsClient) uploadFile(
 			return err
 		}
 		gcsHeaders[gcsMetadataEncryptionDataProp] = string(b)
-		gcsHeaders[gcsMetadataMatdescKey] = encryptMeta.matdesc
+		gcsHeaders[gcsMetadataMatdescKey] = meta.encryptMeta.matdesc
 	}
 
 	var uploadSrc io.Reader
@@ -265,10 +264,10 @@ func (util *snowflakeGcsClient) uploadFile(
 
 	meta.gcsFileHeaderDigest = gcsHeaders[gcsFileHeaderDigest]
 	meta.gcsFileHeaderContentLength = meta.uploadSize
-	if err = json.Unmarshal([]byte(gcsHeaders[gcsMetadataEncryptionDataProp]), &encryptMeta); err != nil {
+	if err = json.Unmarshal([]byte(gcsHeaders[gcsMetadataEncryptionDataProp]), &meta.encryptMeta); err != nil {
 		return err
 	}
-	meta.gcsFileHeaderEncryptionMeta = encryptMeta
+	meta.gcsFileHeaderEncryptionMeta = meta.encryptMeta
 	return nil
 }
 
diff --git a/local_storage_client.go b/local_storage_client.go
index e4bce9b..195f71a 100644
--- a/local_storage_client.go
+++ b/local_storage_client.go
@@ -2,6 +2,7 @@ package gosnowflake
 
 import (
 	"bufio"
+	"cmp"
 	"fmt"
 	"io"
 	"os"
@@ -20,10 +21,7 @@ func (util *localUtil) createClient(_ *execResponseStageInfo, _ bool, _ *Config)
 func (util *localUtil) uploadOneFileWithRetry(meta *fileMetadata) error {
 	var frd *bufio.Reader
 	if meta.srcStream != nil {
-		b := meta.srcStream
-		if meta.realSrcStream != nil {
-			b = meta.realSrcStream
-		}
+		b := cmp.Or(meta.realSrcStream, meta.srcStream)
 		frd = bufio.NewReader(b)
 	} else {
 		f, err := os.Open(meta.realSrcFileName)
diff --git a/s3_storage_client.go b/s3_storage_client.go
index cb1ac6f..1c5d389 100644
--- a/s3_storage_client.go
+++ b/s3_storage_client.go
@@ -2,6 +2,7 @@ package gosnowflake
 
 import (
 	"bytes"
+	"cmp"
 	"context"
 	"errors"
 	"fmt"
@@ -176,17 +177,16 @@ type s3UploadAPI interface {
 func (util *snowflakeS3Client) uploadFile(
 	dataFile string,
 	meta *fileMetadata,
-	encryptMeta *encryptMetadata,
 	maxConcurrency int,
 	multiPartThreshold int64) error {
 	s3Meta := map[string]string{
 		httpHeaderContentType: httpHeaderValueOctetStream,
 		sfcDigest:             meta.sha256Digest,
 	}
-	if encryptMeta != nil {
-		s3Meta[amzIv] = encryptMeta.iv
-		s3Meta[amzKey] = encryptMeta.key
-		s3Meta[amzMatdesc] = encryptMeta.matdesc
+	if meta.encryptMeta != nil {
+		s3Meta[amzIv] = meta.encryptMeta.iv
+		s3Meta[amzKey] = meta.encryptMeta.key
+		s3Meta[amzMatdesc] = meta.encryptMeta.matdesc
 	}
 
 	s3loc, err := util.extractBucketNameAndPath(meta.stageInfo.Location)
@@ -213,10 +213,7 @@ func (util *snowflakeS3Client) uploadFile(
 
 	_, err = withCloudStorageTimeout(util.cfg, func(ctx context.Context) (any, error) {
 		if meta.srcStream != nil {
-			uploadStream := meta.srcStream
-			if meta.realSrcStream != nil {
-				uploadStream = meta.realSrcStream
-			}
+			uploadStream := cmp.Or(meta.realSrcStream, meta.srcStream)
 			return uploader.Upload(ctx, &s3.PutObjectInput{
 				Bucket:   &s3loc.bucketName,
 				Key:      &s3path,
diff --git a/storage_client.go b/storage_client.go
index 5f35ac5..c04edc2 100644
--- a/storage_client.go
+++ b/storage_client.go
@@ -1,7 +1,6 @@
 package gosnowflake
 
 import (
-	"bytes"
 	"fmt"
 	"math"
 	"os"
@@ -26,7 +25,7 @@ type storageUtil interface {
 type cloudUtil interface {
 	createClient(*execResponseStageInfo, bool) (cloudClient, error)
 	getFileHeader(*fileMetadata, string) (*fileHeader, error)
-	uploadFile(string, *fileMetadata, *encryptMetadata, int, int64) error
+	uploadFile(string, *fileMetadata, int, int64) error
 	nativeDownloadFile(*fileMetadata, string, int64) error
 }
 
@@ -60,32 +59,6 @@ func (rsu *remoteStorageUtil) createClient(info *execResponseStageInfo, useAccel
 }
 
 func (rsu *remoteStorageUtil) uploadOneFile(meta *fileMetadata) error {
-	var encryptMeta *encryptMetadata
-	var dataFile string
-	var err error
-	if meta.encryptionMaterial != nil {
-		if meta.srcStream != nil {
-			var encryptedStream bytes.Buffer
-			srcStream := meta.srcStream
-			if meta.realSrcStream != nil {
-				srcStream = meta.realSrcStream
-			}
-			encryptMeta, err = encryptStreamCBC(meta.encryptionMaterial, srcStream, &encryptedStream, 0)
-			if err != nil {
-				return err
-			}
-			meta.realSrcStream = &encryptedStream
-			dataFile = meta.realSrcFileName
-		} else {
-			encryptMeta, dataFile, err = encryptFileCBC(meta.encryptionMaterial, meta.realSrcFileName, 0, meta.tmpDir)
-			if err != nil {
-				return err
-			}
-		}
-	} else {
-		dataFile = meta.realSrcFileName
-	}
-
 	utilClass := rsu.getNativeCloudType(meta.stageInfo.LocationType, meta.sfa.sc.cfg)
 	maxConcurrency := int(meta.parallel)
 	var lastErr error
@@ -94,9 +67,9 @@ func (rsu *remoteStorageUtil) uploadOneFile(meta *fileMetadata) error {
 		if !meta.overwrite {
 			header, err := utilClass.getFileHeader(meta, meta.dstFileName)
 			if meta.resStatus == notFoundFile {
-				err := utilClass.uploadFile(dataFile, meta, encryptMeta, maxConcurrency, meta.options.MultiPartThreshold)
+				err := utilClass.uploadFile(meta.realSrcFileName, meta, maxConcurrency, meta.options.MultiPartThreshold)
 				if err != nil {
-					logger.Warnf("Error uploading %v. err: %v", dataFile, err)
+					logger.Warnf("Error uploading %v. err: %v", meta.realSrcFileName, err)
 				}
 			} else if err != nil {
 				return err
@@ -108,9 +81,9 @@ func (rsu *remoteStorageUtil) uploadOneFile(meta *fileMetadata) error {
 			}
 		}
 		if meta.overwrite || meta.resStatus == notFoundFile {
-			err := utilClass.uploadFile(dataFile, meta, encryptMeta, maxConcurrency, meta.options.MultiPartThreshold)
+			err := utilClass.uploadFile(meta.realSrcFileName, meta, maxConcurrency, meta.options.MultiPartThreshold)
 			if err != nil {
-				logger.Debugf("Error uploading %v. err: %v", dataFile, err)
+				logger.Debugf("Error uploading %v. err: %v", meta.realSrcFileName, err)
 			}
 		}
 		if meta.resStatus == uploaded || meta.resStatus == renewToken || meta.resStatus == renewPresignedURL {
@@ -135,7 +108,7 @@ func (rsu *remoteStorageUtil) uploadOneFile(meta *fileMetadata) error {
 	if lastErr != nil {
 		return lastErr
 	}
-	return fmt.Errorf("unkown error uploading %v", dataFile)
+	return fmt.Errorf("unkown error uploading %v", meta.realSrcFileName)
 }
 
 func (rsu *remoteStorageUtil) uploadOneFileWithRetry(meta *fileMetadata) error {


Test patch:
diff --git a/azure_storage_client_test.go b/azure_storage_client_test.go
index 9d6fb87..c25b5de 100644
--- a/azure_storage_client_test.go
+++ b/azure_storage_client_test.go
@@ -165,6 +165,7 @@ func TestUploadFileWithAzureUploadFailedError(t *testing.T) {
 		dstFileName:        "data1.txt.gz",
 		srcFileName:        path.Join(dir, "/test_data/put_get_1.txt"),
 		encryptionMaterial: &encMat,
+		encryptMeta:        testEncryptionMeta(),
 		overwrite:          true,
 		dstCompressionType: compressionTypes["GZIP"],
 		options: &SnowflakeFileTransferOptions{
@@ -223,6 +224,7 @@ func TestUploadStreamWithAzureUploadFailedError(t *testing.T) {
 		dstFileName:        "data1.txt.gz",
 		srcStream:          bytes.NewBuffer(src),
 		encryptionMaterial: &encMat,
+		encryptMeta:        testEncryptionMeta(),
 		overwrite:          true,
 		dstCompressionType: compressionTypes["GZIP"],
 		options: &SnowflakeFileTransferOptions{
@@ -285,6 +287,7 @@ func TestUploadFileWithAzureUploadTokenExpired(t *testing.T) {
 		stageInfo:          &info,
 		dstFileName:        "data1.txt.gz",
 		srcFileName:        path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:        testEncryptionMeta(),
 		overwrite:          true,
 		dstCompressionType: compressionTypes["GZIP"],
 		options: &SnowflakeFileTransferOptions{
@@ -361,6 +364,7 @@ func TestUploadFileWithAzureUploadNeedsRetry(t *testing.T) {
 		stageInfo:          &info,
 		dstFileName:        "data1.txt.gz",
 		srcFileName:        path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:        testEncryptionMeta(),
 		overwrite:          true,
 		dstCompressionType: compressionTypes["GZIP"],
 		options: &SnowflakeFileTransferOptions{
@@ -569,6 +573,7 @@ func TestUploadFileToAzureClientCastFail(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
diff --git a/file_transfer_agent_test.go b/file_transfer_agent_test.go
index 0d54c2c..6a574e4 100644
--- a/file_transfer_agent_test.go
+++ b/file_transfer_agent_test.go
@@ -1060,3 +1060,239 @@ func TestPutGetRegexShouldIgnoreWhitespaceAtTheBeginning(t *testing.T) {
 		}
 	}
 }
+
+func TestEncryptStream(t *testing.T) {
+	srcBytes := []byte{63, 64, 65}
+	initStr := bytes.NewBuffer(srcBytes)
+
+	for _, tc := range []struct {
+		ct            cloudType
+		encrypt       bool
+		realSrcStream bool
+		encryptMat    bool
+	}{
+		{
+			ct:            s3Client,
+			encrypt:       true,
+			realSrcStream: true,
+			encryptMat:    true,
+		},
+		{
+			ct:            s3Client,
+			encrypt:       true,
+			realSrcStream: false,
+			encryptMat:    true,
+		},
+		{
+			ct:            s3Client,
+			encrypt:       false,
+			realSrcStream: false,
+			encryptMat:    false,
+		},
+		{
+			ct:            azureClient,
+			encrypt:       true,
+			realSrcStream: true,
+			encryptMat:    true,
+		},
+		{
+			ct:            azureClient,
+			encrypt:       true,
+			realSrcStream: false,
+			encryptMat:    true,
+		},
+		{
+			ct:            azureClient,
+			encrypt:       false,
+			realSrcStream: false,
+			encryptMat:    false,
+		},
+		{
+			ct:            gcsClient,
+			encrypt:       true,
+			realSrcStream: true,
+			encryptMat:    true,
+		},
+		{
+			ct:            gcsClient,
+			encrypt:       true,
+			realSrcStream: false,
+			encryptMat:    true,
+		},
+		{
+			ct:            gcsClient,
+			encrypt:       false,
+			realSrcStream: false,
+			encryptMat:    false,
+		},
+		{
+			ct:            local,
+			encrypt:       false,
+			realSrcStream: true,
+			encryptMat:    true,
+		},
+		{
+			ct:            local,
+			encrypt:       false,
+			realSrcStream: true,
+			encryptMat:    false,
+		},
+		{
+			ct:            local,
+			encrypt:       false,
+			realSrcStream: false,
+			encryptMat:    true,
+		},
+		{
+			ct:            local,
+			encrypt:       false,
+			realSrcStream: false,
+			encryptMat:    false,
+		},
+	} {
+		{
+			var encMat *snowflakeFileEncryption = nil
+			if tc.encryptMat {
+				encMat = &snowflakeFileEncryption{
+					QueryStageMasterKey: "abCdEFO0upIT36dAxGsa0w==",
+					QueryID:             "01abc874-0406-1bf0-0000-53b10668e056",
+					SMKID:               92019681909886,
+				}
+			}
+			var realSrcStr *bytes.Buffer = nil
+			if tc.realSrcStream {
+				realSrcStr = initStr
+			}
+			uploadMeta := fileMetadata{
+				name:               "data1.txt.gz",
+				stageLocationType:  tc.ct,
+				noSleepingTime:     true,
+				parallel:           int64(100),
+				client:             nil,
+				sha256Digest:       "123456789abcdef",
+				stageInfo:          nil,
+				dstFileName:        "data1.txt.gz",
+				srcStream:          initStr,
+				realSrcStream:      realSrcStr,
+				overwrite:          true,
+				options:            nil,
+				encryptionMaterial: encMat,
+				mockUploader:       nil,
+				sfa:                nil,
+			}
+
+			t.Run(string(tc.ct)+" encrypt "+strconv.FormatBool(tc.encrypt)+" realSrcStream "+strconv.FormatBool(tc.realSrcStream)+" encryptMat "+strconv.FormatBool(tc.encryptMat), func(t *testing.T) {
+				err := encryptDataIfRequired(&uploadMeta, tc.ct)
+				assertNilF(t, err)
+				if tc.encrypt {
+					assertNotNilF(t, uploadMeta.encryptMeta, "encryption metadata should be present")
+					if tc.realSrcStream {
+						assertNotEqualF(t, uploadMeta.realSrcStream, realSrcStr, "stream should be encrypted")
+					} else {
+						assertNotEqualF(t, uploadMeta.realSrcStream, initStr, "stream should not be encrypted")
+					}
+				} else {
+					assertNilF(t, uploadMeta.encryptMeta, "encryption metadata should be empty")
+					assertEqualF(t, uploadMeta.realSrcStream, realSrcStr, "stream should not be encrypted")
+				}
+			})
+		}
+	}
+}
+
+func TestEncryptFile(t *testing.T) {
+	for _, tc := range []struct {
+		ct         cloudType
+		encrypt    bool
+		encryptMat bool
+	}{
+		{
+			ct:         s3Client,
+			encrypt:    true,
+			encryptMat: true,
+		},
+		{
+			ct:         s3Client,
+			encrypt:    false,
+			encryptMat: false,
+		},
+		{
+			ct:         azureClient,
+			encrypt:    true,
+			encryptMat: true,
+		},
+		{
+			ct:         azureClient,
+			encrypt:    false,
+			encryptMat: false,
+		},
+		{
+			ct:         gcsClient,
+			encrypt:    true,
+			encryptMat: true,
+		},
+		{
+			ct:         gcsClient,
+			encrypt:    false,
+			encryptMat: false,
+		},
+		{
+			ct:         local,
+			encrypt:    false,
+			encryptMat: true,
+		},
+		{
+			ct:         local,
+			encrypt:    false,
+			encryptMat: false,
+		},
+	} {
+		dir, err := os.Getwd()
+		srcF := path.Join(dir, "/test_data/put_get_1.txt")
+		assertNilF(t, err, "error getting current directory")
+
+		var encMat *snowflakeFileEncryption = nil
+		if tc.encryptMat {
+			encMat = &snowflakeFileEncryption{
+				QueryStageMasterKey: "abCdEFO0upIT36dAxGsa0w==",
+				QueryID:             "01abc874-0406-1bf0-0000-53b10668e056",
+				SMKID:               92019681909886,
+			}
+		}
+
+		uploadMeta := fileMetadata{
+			name:               "data1.txt.gz",
+			stageLocationType:  tc.ct,
+			noSleepingTime:     true,
+			parallel:           int64(100),
+			client:             nil,
+			sha256Digest:       "123456789abcdef",
+			stageInfo:          nil,
+			dstFileName:        "data1.txt.gz",
+			srcFileName:        srcF,
+			realSrcFileName:    srcF,
+			overwrite:          true,
+			options:            nil,
+			encryptionMaterial: encMat,
+			mockUploader:       nil,
+			sfa:                nil,
+		}
+
+		t.Run(string(tc.ct)+" encrypt "+strconv.FormatBool(tc.encrypt)+" encryptMat "+strconv.FormatBool(tc.encryptMat), func(t *testing.T) {
+			err := encryptDataIfRequired(&uploadMeta, tc.ct)
+			assertNilF(t, err)
+			if tc.encrypt {
+				assertNotNilF(t, uploadMeta.encryptMeta, "encryption metadata should be present")
+				assertNotEqualF(t, uploadMeta.realSrcFileName, srcF, "file should be encrypted")
+				srcBytes, err := os.ReadFile(srcF)
+				assertNilF(t, err)
+				encBytes, err := os.ReadFile(uploadMeta.realSrcFileName)
+				assertNilF(t, err)
+				assertFalseF(t, bytes.Equal(srcBytes, encBytes), "file contents should differ")
+			} else {
+				assertNilF(t, uploadMeta.encryptMeta, "encryption metadata should be empty")
+				assertEqualF(t, uploadMeta.realSrcFileName, srcF, "file should not be encrypted")
+			}
+		})
+	}
+}
diff --git a/gcs_storage_client_test.go b/gcs_storage_client_test.go
index d8ef7c9..4db54a4 100644
--- a/gcs_storage_client_test.go
+++ b/gcs_storage_client_test.go
@@ -153,6 +153,7 @@ func TestUploadFileWithGcsUploadFailedError(t *testing.T) {
 		stageInfo:          &info,
 		dstFileName:        "data1.txt.gz",
 		srcFileName:        path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:        testEncryptionMeta(),
 		overwrite:          true,
 		dstCompressionType: compressionTypes["GZIP"],
 		options: &SnowflakeFileTransferOptions{
@@ -213,6 +214,7 @@ func TestUploadFileWithGcsUploadFailedWithRetry(t *testing.T) {
 		stageInfo:          &info,
 		dstFileName:        "data1.txt.gz",
 		srcFileName:        path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:        testEncryptionMeta(),
 		overwrite:          true,
 		dstCompressionType: compressionTypes["GZIP"],
 		encryptionMaterial: &encMat,
@@ -280,6 +282,7 @@ func TestUploadFileWithGcsUploadFailedWithTokenExpired(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
@@ -748,6 +751,7 @@ func TestUploadStreamFailed(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcStream:         bytes.NewBuffer(src),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
@@ -797,6 +801,7 @@ func TestUploadFileWithBadRequest(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		lastError:         nil,
 		options: &SnowflakeFileTransferOptions{
@@ -954,6 +959,7 @@ func TestUploadFileToGcsNoStatus(t *testing.T) {
 		stageInfo:          &info,
 		dstFileName:        "data1.txt.gz",
 		srcFileName:        path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:        testEncryptionMeta(),
 		overwrite:          true,
 		dstCompressionType: compressionTypes["GZIP"],
 		encryptionMaterial: &encMat,
@@ -1103,7 +1109,7 @@ func Test_snowflakeGcsClient_uploadFile(t *testing.T) {
 		client:    1,
 		stageInfo: &info,
 	}
-	err := new(snowflakeGcsClient).uploadFile("somedata", &meta, nil, 1, 1)
+	err := new(snowflakeGcsClient).uploadFile("somedata", &meta, 1, 1)
 	if err == nil {
 		t.Error("should have raised an error")
 	}
diff --git a/s3_storage_client_test.go b/s3_storage_client_test.go
index 781dda9..4776b72 100644
--- a/s3_storage_client_test.go
+++ b/s3_storage_client_test.go
@@ -81,6 +81,7 @@ func TestUploadOneFileToS3WSAEConnAborted(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
@@ -157,6 +158,7 @@ func TestUploadOneFileToS3ConnReset(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
@@ -216,6 +218,7 @@ func TestUploadFileWithS3UploadFailedError(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
@@ -598,6 +601,7 @@ func TestUploadFileToS3ClientCastFail(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
@@ -678,6 +682,7 @@ func TestS3UploadRetryWithHeaderNotFound(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcFileName:       path.Join(dir, "/test_data/put_get_1.txt"),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
@@ -739,6 +744,7 @@ func TestS3UploadStreamFailed(t *testing.T) {
 		stageInfo:         &info,
 		dstFileName:       "data1.txt.gz",
 		srcStream:         bytes.NewBuffer(src),
+		encryptMeta:       testEncryptionMeta(),
 		overwrite:         true,
 		options: &SnowflakeFileTransferOptions{
 			MultiPartThreshold: dataSizeThreshold,
diff --git a/storage_file_util_test.go b/storage_file_util_test.go
new file mode 100644
index 0000000..854e98d
--- /dev/null
+++ b/storage_file_util_test.go
@@ -0,0 +1,10 @@
+package gosnowflake
+
+func testEncryptionMeta() *encryptMetadata {
+	const mockMatDesc = "{\"queryid\":\"01abc874-0406-1bf0-0000-53b10668e056\",\"smkid\":\"92019681909886\",\"key\":\"128\"}"
+	return &encryptMetadata{
+		key:     "testencryptedkey12345678910==",
+		iv:      "testIVkey12345678910==",
+		matdesc: mockMatDesc,
+	}
+}