fix test failure TestUploadDownloadOneFileRequireCompressStream

snowflakedb · Dec 5, 2024 · a841c0e · a841c0e
1 parent b963ef0
commit a841c0e
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 12 deletions.
diff --git a/connection_util.go b/connection_util.go
@@ -151,13 +151,9 @@ func getFileStream(ctx context.Context) (*bytes.Buffer, error) {
 
 	// read a small amount of data to check if file stream will be used
 	buf := make([]byte, defaultStringBufferSize)
-	for {
-		_, err := r.Read(buf)
-		if err != nil {
-			return nil, err
-		} else {
-			break
-		}
+	_, err := r.Read(buf)
+	if err != nil {
+		return nil, err
 	}
 	return bytes.NewBuffer(buf), nil
 }

diff --git a/file_transfer_agent.go b/file_transfer_agent.go
@@ -857,7 +857,7 @@ func (sfa *snowflakeFileTransferAgent) uploadOneFile(meta *fileMetadata) (*fileM
 	}
 
 	if meta.srcStream != nil {
-		meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForStream(&meta.realSrcStream, &meta.srcStream, sfa.ctx)
+		meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForStream(sfa.ctx, &meta.realSrcStream, &meta.srcStream)
 	} else {
 		meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForFile(meta.realSrcFileName)
 	}

diff --git a/file_transfer_agent_test.go b/file_transfer_agent_test.go
@@ -7,7 +7,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/aws/aws-sdk-go-v2/service/s3"
 	"io"
 	"net/url"
 	"os"
@@ -18,6 +17,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+
 	"github.com/aws/smithy-go"
 )
 
@@ -980,8 +981,8 @@ func testUploadDownloadOneFile(t *testing.T, isStream bool) {
 
 	if isStream {
 		fileStream, _ := os.Open(uploadFile)
-		ctx := WithFileStream(context.Background(), fileStream)
-		uploadMeta.srcStream, err = getFileStream(ctx)
+		sfa.ctx = WithFileStream(context.Background(), fileStream)
+		uploadMeta.srcStream, err = getFileStream(sfa.ctx)
 		assertNilF(t, err)
 	}
 

diff --git a/file_util.go b/file_util.go
@@ -88,7 +88,7 @@ func (util *snowflakeFileUtil) compressFileWithGzip(fileName string, tmpDir stri
 	return gzipFileName, stat.Size(), err
 }
 
-func (util *snowflakeFileUtil) getDigestAndSizeForStream(realSrcStream **bytes.Buffer, srcStream **bytes.Buffer, ctx context.Context) (string, int64, error) {
+func (util *snowflakeFileUtil) getDigestAndSizeForStream(ctx context.Context, realSrcStream **bytes.Buffer, srcStream **bytes.Buffer) (string, int64, error) {
 	var r io.Reader
 	var stream **bytes.Buffer
 	if realSrcStream != nil {