Repository: snowflakedb/gosnowflake
Pull request #1420 (2025-05-21 18:35:36)

SNOW-1029631: [gosnowflake][GH#536] optimize PUT memory usage / PUT operations read full input data into memory
### Description

SNOW-XXX Please explain the changes you made here.

### Checklist
- [ ] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary



Patch:
diff --git a/connection_util.go b/connection_util.go
index 1f34882..da06a46 100644
--- a/connection_util.go
+++ b/connection_util.go
@@ -132,7 +132,7 @@ func (sc *snowflakeConn) processFileTransfer(
 	return data, nil
 }
 
-func getFileStream(ctx context.Context) (*bytes.Buffer, error) {
+func getFileStream(ctx context.Context) (io.Reader, error) {
 	s := ctx.Value(fileStreamFile)
 	if s == nil {
 		return nil, nil
@@ -141,9 +141,7 @@ func getFileStream(ctx context.Context) (*bytes.Buffer, error) {
 	if !ok {
 		return nil, errors.New("incorrect io.Reader")
 	}
-	buf := new(bytes.Buffer)
-	_, err := buf.ReadFrom(r)
-	return buf, err
+	return r, nil
 }
 
 func getFileTransferOptions(ctx context.Context) *SnowflakeFileTransferOptions {
diff --git a/errors.go b/errors.go
index a0b9d9a..d3c3eb6 100644
--- a/errors.go
+++ b/errors.go
@@ -313,6 +313,7 @@ const (
 	errMsgCommandNotRecognized               = "%v command not recognized"
 	errMsgLocalPathNotDirectory              = "the local path is not a directory: %v"
 	errMsgFileNotExists                      = "file does not exist: %v"
+	errMsgFailToReadDataFromBuffer           = "failed to read data from buffer. err: %v"
 	errMsgInvalidStageFs                     = "destination location type is not valid: %v"
 	errMsgInternalNotMatchEncryptMaterial    = "number of downloading files doesn't match the encryption materials. files=%v, encmat=%v"
 	errMsgFailedToConvertToS3Client          = "failed to convert interface to s3 client"
diff --git a/file_transfer_agent.go b/file_transfer_agent.go
index a4a426c..3a8d67b 100644
--- a/file_transfer_agent.go
+++ b/file_transfer_agent.go
@@ -113,7 +113,7 @@ type snowflakeFileTransferAgent struct {
 	encryptionMaterial          []*snowflakeFileEncryption
 	stageInfo                   *execResponseStageInfo
 	results                     []*fileMetadata
-	sourceStream                *bytes.Buffer
+	sourceStream                io.Reader
 	srcLocations                []string
 	autoCompress                bool
 	srcCompression              string
@@ -133,6 +133,7 @@ func (sfa *snowflakeFileTransferAgent) execute() error {
 	if err = sfa.parseCommand(); err != nil {
 		return err
 	}
+
 	if err = sfa.initFileMetadata(); err != nil {
 		return err
 	}
@@ -352,17 +353,46 @@ func (sfa *snowflakeFileTransferAgent) initFileMetadata() error {
 				MessageArgs: []interface{}{fileName},
 			}).exceptionTelemetry(sfa.sc)
 		}
+
+		// Handles bulk inserts by checking if sourceStream exists.
+		// - If the file exists locally (PUT command), it saves the stream without loading it into memory.
+		// - If not, treats it as an INSERT converted to PUT for bulk upload.
 		if sfa.sourceStream != nil {
+			//Bulk insert case
 			fileName := sfa.srcFiles[0]
-			srcFileSize := int64(sfa.sourceStream.Len())
-			sfa.fileMetadata = append(sfa.fileMetadata, &fileMetadata{
-				name:              baseName(fileName),
-				srcFileName:       fileName,
-				srcStream:         sfa.sourceStream,
-				srcFileSize:       srcFileSize,
-				stageLocationType: sfa.stageLocationType,
-				stageInfo:         sfa.stageInfo,
-			})
+			fileInfo, err := os.Stat(fileName)
+			if err != nil {
+				buf := new(bytes.Buffer)
+				_, err := buf.ReadFrom(sfa.sourceStream)
+				if err != nil {
+					return (&SnowflakeError{
+						Number:      ErrFileNotExists,
+						SQLState:    sfa.data.SQLState,
+						QueryID:     sfa.data.QueryID,
+						Message:     errMsgFailToReadDataFromBuffer,
+						MessageArgs: []interface{}{fileName},
+					}).exceptionTelemetry(sfa.sc)
+				}
+				sfa.fileMetadata = append(sfa.fileMetadata, &fileMetadata{
+					name:              baseName(fileName),
+					srcFileName:       fileName,
+					srcStream:         buf,
+					fileStream:        sfa.sourceStream,
+					srcFileSize:       int64(buf.Len()),
+					stageLocationType: sfa.stageLocationType,
+					stageInfo:         sfa.stageInfo,
+				})
+			} else {
+				//PUT command with existing file
+				sfa.fileMetadata = append(sfa.fileMetadata, &fileMetadata{
+					name:              baseName(fileName),
+					srcFileName:       fileName,
+					fileStream:        sfa.sourceStream,
+					srcFileSize:       fileInfo.Size(),
+					stageLocationType: sfa.stageLocationType,
+					stageInfo:         sfa.stageInfo,
+				})
+			}
 		} else {
 			for i, fileName := range sfa.srcFiles {
 				fi, err := os.Stat(fileName)
@@ -395,7 +425,6 @@ func (sfa *snowflakeFileTransferAgent) initFileMetadata() error {
 				}
 			}
 		}
-
 		if len(sfa.encryptionMaterial) > 0 {
 			for _, meta := range sfa.fileMetadata {
 				meta.encryptionMaterial = sfa.encryptionMaterial[0]
@@ -866,6 +895,7 @@ func (sfa *snowflakeFileTransferAgent) uploadOneFile(meta *fileMetadata) (*fileM
 	defer os.RemoveAll(tmpDir) // cleanup
 
 	fileUtil := new(snowflakeFileUtil)
+
 	err = compressDataIfRequired(meta, fileUtil, tmpDir)
 	if err != nil {
 		return meta, err
@@ -1233,12 +1263,8 @@ func compressDataIfRequired(meta *fileMetadata, fileUtil *snowflakeFileUtil, tmp
 
 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)
-		}
+	if meta.fileStream != nil {
+		meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForStream(meta.fileStream)
 	} else {
 		meta.sha256Digest, meta.uploadSize, err = fileUtil.getDigestAndSizeForFile(meta.realSrcFileName)
 	}
@@ -1265,6 +1291,5 @@ func encryptDataIfRequired(meta *fileMetadata, ct cloudType) error {
 			meta.realSrcFileName = dataFile
 		}
 	}
-
 	return nil
 }
diff --git a/file_util.go b/file_util.go
index 497cab1..3d8707d 100644
--- a/file_util.go
+++ b/file_util.go
@@ -73,21 +73,22 @@ func (util *snowflakeFileUtil) compressFileWithGzip(fileName string, tmpDir stri
 	return gzipFileName, stat.Size(), err
 }
 
-func (util *snowflakeFileUtil) getDigestAndSizeForStream(stream **bytes.Buffer) (string, int64, error) {
+func (util *snowflakeFileUtil) getDigestAndSizeForStream(stream io.Reader) (string, int64, error) {
 	m := sha256.New()
-	r := getReaderFromBuffer(stream)
 	chunk := make([]byte, fileChunkSize)
+	var total int64
 
 	for {
-		n, err := r.Read(chunk)
+		n, err := stream.Read(chunk)
 		if err == io.EOF {
 			break
 		} else if err != nil {
 			return "", 0, err
 		}
+		total += int64(n)
 		m.Write(chunk[:n])
 	}
-	return base64.StdEncoding.EncodeToString(m.Sum(nil)), int64((*stream).Len()), nil
+	return base64.StdEncoding.EncodeToString(m.Sum(nil)), total, nil
 }
 
 func (util *snowflakeFileUtil) getDigestAndSizeForFile(fileName string) (digest string, size int64, err error) {
@@ -154,6 +155,7 @@ type fileMetadata struct {
 	options            *SnowflakeFileTransferOptions
 
 	/* streaming PUT */
+	fileStream    io.Reader
 	srcStream     *bytes.Buffer
 	realSrcStream *bytes.Buffer
 


Test patch:
diff --git a/file_transfer_agent_test.go b/file_transfer_agent_test.go
index c234eff..c9cbe8b 100644
--- a/file_transfer_agent_test.go
+++ b/file_transfer_agent_test.go
@@ -977,7 +977,7 @@ func testUploadDownloadOneFile(t *testing.T, isStream bool) {
 	if isStream {
 		fileStream, _ := os.Open(uploadFile)
 		ctx := WithFileStream(context.Background(), fileStream)
-		uploadMeta.srcStream, err = getFileStream(ctx)
+		uploadMeta.fileStream, err = getFileStream(ctx)
 		assertNilF(t, err)
 	}
 
diff --git a/local_storage_client_test.go b/local_storage_client_test.go
index acb5b6c..b94b0ad 100644
--- a/local_storage_client_test.go
+++ b/local_storage_client_test.go
@@ -74,7 +74,7 @@ func TestLocalUpload(t *testing.T) {
 	}
 	fileStream, _ := os.Open(fname)
 	ctx := WithFileStream(context.Background(), fileStream)
-	uploadMeta.srcStream, err = getFileStream(ctx)
+	uploadMeta.fileStream, err = getFileStream(ctx)
 	assertNilF(t, err)
 
 	err = localUtil.uploadOneFileWithRetry(&uploadMeta)
diff --git a/put_get_test.go b/put_get_test.go
index 3d454c5..245f198 100644
--- a/put_get_test.go
+++ b/put_get_test.go
@@ -1,6 +1,7 @@
 package gosnowflake
 
 import (
+	"bufio"
 	"bytes"
 	"compress/gzip"
 	"context"
@@ -12,6 +13,7 @@ import (
 	"os/user"
 	"path"
 	"path/filepath"
+	"runtime"
 	"strconv"
 	"strings"
 	"testing"
@@ -920,3 +922,161 @@ func TestPutCancel(t *testing.T) {
 		close(c)
 	})
 }
+
+func TestPutGetLargeFileNonStream(t *testing.T) {
+	testPutGetLargeFile(t, false, true)
+}
+
+func TestPutGetLargeFileNonStreamAutoCompressFalse(t *testing.T) {
+	testPutGetLargeFile(t, false, false)
+}
+
+func TestPutGetLargeFileStream(t *testing.T) {
+	testPutGetLargeFile(t, true, true)
+}
+
+func TestPutGetLargeFileStreamAutoCompressFalse(t *testing.T) {
+	testPutGetLargeFile(t, true, false)
+}
+
+func testPutGetLargeFile(t *testing.T, isStream bool, autoCompress bool) {
+	sourceDir, err := os.Getwd()
+	assertNilF(t, err)
+	fname := filepath.Join(sourceDir, "/test_data/largefile.txt")
+	fnameGet := "largefile.txt.gz"
+	if !autoCompress {
+		fnameGet = "largefile.txt"
+	}
+
+	runDBTest(t, func(dbt *DBTest) {
+		stageDir := "test_put_largefile_" + randomString(10)
+		dbt.mustExec("rm @~/" + stageDir)
+
+		ctx := context.Background()
+		if isStream {
+			f, err := os.Open(fname)
+			assertNilF(t, err)
+			defer func() {
+				assertNilF(t, f.Close())
+			}()
+			ctx = WithFileStream(ctx, f)
+		}
+
+		// PUT test
+		putQuery := fmt.Sprintf("put file://%v @~/%v auto_compress=true overwrite=true", fname, stageDir)
+		if !autoCompress {
+			putQuery = fmt.Sprintf("put file://%v @~/%v auto_compress=false overwrite=true", fname, stageDir)
+		}
+		sqlText := strings.ReplaceAll(putQuery, "\\", "\\\\")
+
+		// Record initial memory stats before PUT
+		var startMemStats, endMemStats runtime.MemStats
+		runtime.ReadMemStats(&startMemStats)
+
+		// Execute PUT command
+		_ = dbt.mustExecContext(ctx, sqlText)
+
+		// Record memory stats after PUT
+		runtime.ReadMemStats(&endMemStats)
+		fmt.Printf("Memory used for PUT command: %d MB\n", (endMemStats.Alloc-startMemStats.Alloc)/1024/1024)
+
+		defer dbt.mustExec("rm @~/" + stageDir)
+		rows := dbt.mustQuery("ls @~/" + stageDir)
+		defer func() {
+			assertNilF(t, rows.Close())
+		}()
+		var file, s1, s2, s3 sql.NullString
+		if rows.Next() {
+			err = rows.Scan(&file, &s1, &s2, &s3)
+			assertNilF(t, err)
+		}
+
+		if !strings.Contains(file.String, fnameGet) {
+			t.Fatalf("should contain file. got: %v", file.String)
+		}
+
+		// GET test
+		var streamBuf bytes.Buffer
+		ctx = context.Background()
+		if isStream {
+			ctx = WithFileTransferOptions(ctx, &SnowflakeFileTransferOptions{GetFileToStream: true})
+			ctx = WithFileGetStream(ctx, &streamBuf)
+		}
+
+		tmpDir := t.TempDir()
+		sql := fmt.Sprintf("get @~/%v/%v 'file://%v'", stageDir, fnameGet, tmpDir)
+		sqlText = strings.ReplaceAll(sql, "\\", "\\\\")
+		rows2 := dbt.mustQueryContext(ctx, sqlText)
+		defer func() {
+			assertNilF(t, rows2.Close())
+		}()
+		for rows2.Next() {
+			err = rows2.Scan(&file, &s1, &s2, &s3)
+			assertNilE(t, err)
+			assertTrueE(t, strings.HasPrefix(file.String, fnameGet), "a file was not downloaded by GET")
+			v, err := strconv.Atoi(s1.String)
+			assertNilE(t, err)
+			if autoCompress {
+				assertEqualE(t, v, 424821, "did not return the right file size")
+			} else {
+				assertEqualE(t, v, 70248250, "did not return the right file size")
+			}
+			assertEqualE(t, s2.String, "DOWNLOADED", "did not return DOWNLOADED status")
+			assertEqualE(t, s3.String, "")
+		}
+
+		var contents string
+		var r io.Reader
+		if autoCompress {
+			// convert the compressed contents to string
+			if isStream {
+				r, err = gzip.NewReader(&streamBuf)
+				assertNilE(t, err)
+			} else {
+				downloadedFile := filepath.Join(tmpDir, fnameGet)
+				f, err := os.Open(downloadedFile)
+				assertNilE(t, err)
+				defer func() {
+					assertNilF(t, f.Close())
+				}()
+				r, err = gzip.NewReader(f)
+				assertNilE(t, err)
+			}
+		} else {
+			if isStream {
+				r = bytes.NewReader(streamBuf.Bytes())
+			} else {
+				downloadedFile := filepath.Join(tmpDir, fnameGet)
+				f, err := os.Open(downloadedFile)
+				assertNilE(t, err)
+				defer func() {
+					assertNilF(t, f.Close())
+				}()
+				r = bufio.NewReader(f)
+			}
+		}
+		for {
+			c := make([]byte, defaultChunkBufferSize)
+			if n, err := r.Read(c); err != nil {
+				if err == io.EOF {
+					contents = contents + string(c[:n])
+					break
+				}
+				t.Error(err)
+			} else {
+				contents = contents + string(c[:n])
+			}
+		}
+
+		// verify the downloaded contents with the original file
+		originalFile, err := os.Open(fname)
+		assertNilF(t, err)
+		defer func() {
+			assertNilF(t, originalFile.Close())
+		}()
+
+		originalContents, err := io.ReadAll(originalFile)
+		assertNilE(t, err)
+		assertEqualF(t, contents, string(originalContents), "data did not match content")
+	})
+}