Repository: snowflakedb/gosnowflake
Pull request #1192 (2024-08-01 14:26:49)

SNOW-1029646: Add WithFileGetStream that supports downloading a file into stream
### Description

SNOW-1029646: [gosnowflake][GH#654] WithFileStream support for GET queries
https://github.com/snowflakedb/snowflake-sdks-drivers-issues-teamwork/issues/855

Added the context `WithFileGetStream` to support downloading a file into in-memory stream.

### Checklist
- [x] Created tests which fail without the change (if possible)
- [x] Extended the README / documentation, if necessary


codecov[bot]: ## [Codecov](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?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 `73.42657%` with `38 lines` in your changes missing coverage. Please review.
> Project coverage is 82.51%. Comparing base [(`52bb7bc`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/52bb7bc8b5315b4add85a2712a28fca5ea289072?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) to head [(`cb97f86`)](https://app.codecov.io/gh/snowflakedb/gosnowflake/commit/cb97f86b7749ebcc70dd6d6dba07530483a97291?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb).

| [Files](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) | Patch % | Lines |
|---|---|---|
| [encrypt\_util.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&filepath=encrypt_util.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-ZW5jcnlwdF91dGlsLmdv) | 64.44% | [16 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |
| [gcs\_storage\_client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&filepath=gcs_storage_client.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Z2NzX3N0b3JhZ2VfY2xpZW50Lmdv) | 63.63% | [8 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |
| [azure\_storage\_client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&filepath=azure_storage_client.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-YXp1cmVfc3RvcmFnZV9jbGllbnQuZ28=) | 77.77% | [6 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |
| [storage\_client.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?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=) | 71.42% | [6 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb) |
| [connection\_util.go](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?src=pr&el=tree&filepath=connection_util.go&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=snowflakedb#diff-Y29ubmVjdGlvbl91dGlsLmdv) | 80.00% | [2 Missing :warning: ](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?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    #1192      +/-   ##
==========================================
- Coverage   82.55%   82.51%   -0.04%     
==========================================
  Files          54       54              
  Lines       12823    12920      +97     
==========================================
+ Hits        10586    10661      +75     
- Misses       2237     2259      +22     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/snowflakedb/gosnowflake/pull/1192?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).


Patch:
diff --git a/azure_storage_client.go b/azure_storage_client.go
index dce7f8d..307c7b2 100644
--- a/azure_storage_client.go
+++ b/azure_storage_client.go
@@ -34,6 +34,7 @@ type azureAPI interface {
 	UploadStream(ctx context.Context, body io.Reader, o *azblob.UploadStreamOptions) (azblob.UploadStreamResponse, error)
 	UploadFile(ctx context.Context, file *os.File, o *azblob.UploadFileOptions) (azblob.UploadFileResponse, error)
 	DownloadFile(ctx context.Context, file *os.File, o *blob.DownloadFileOptions) (int64, error)
+	DownloadStream(ctx context.Context, o *blob.DownloadStreamOptions) (azblob.DownloadStreamResponse, error)
 	GetProperties(ctx context.Context, o *blob.GetPropertiesOptions) (blob.GetPropertiesResponse, error)
 }
 
@@ -276,16 +277,29 @@ func (util *snowflakeAzureClient) nativeDownloadFile(
 	if meta.mockAzureClient != nil {
 		blobClient = meta.mockAzureClient
 	}
-	f, err := os.OpenFile(fullDstFileName, os.O_CREATE|os.O_WRONLY, readWriteFileMode)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-	_, err = blobClient.DownloadFile(
-		context.Background(), f, &azblob.DownloadFileOptions{
-			Concurrency: uint16(maxConcurrency)})
-	if err != nil {
-		return err
+	if meta.options.getFileToStream {
+		blobDownloadResponse, err := blobClient.DownloadStream(context.Background(), &azblob.DownloadStreamOptions{})
+		if err != nil {
+			return err
+		}
+		retryReader := blobDownloadResponse.NewRetryReader(context.Background(), &azblob.RetryReaderOptions{})
+		defer retryReader.Close()
+		_, err = meta.dstStream.ReadFrom(retryReader)
+		if err != nil {
+			return err
+		}
+	} else {
+		f, err := os.OpenFile(fullDstFileName, os.O_CREATE|os.O_WRONLY, readWriteFileMode)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		_, err = blobClient.DownloadFile(
+			context.Background(), f, &azblob.DownloadFileOptions{
+				Concurrency: uint16(maxConcurrency)})
+		if err != nil {
+			return err
+		}
 	}
 	meta.resStatus = downloaded
 	return nil
diff --git a/connection_util.go b/connection_util.go
index 1672235..31b97c9 100644
--- a/connection_util.go
+++ b/connection_util.go
@@ -5,6 +5,7 @@ package gosnowflake
 import (
 	"bytes"
 	"context"
+	"errors"
 	"fmt"
 	"io"
 	"os"
@@ -88,10 +89,11 @@ func (sc *snowflakeConn) processFileTransfer(
 	isInternal bool) (
 	*execResponse, error) {
 	sfa := snowflakeFileTransferAgent{
-		sc:      sc,
-		data:    &data.Data,
-		command: query,
-		options: new(SnowflakeFileTransferOptions),
+		sc:           sc,
+		data:         &data.Data,
+		command:      query,
+		options:      new(SnowflakeFileTransferOptions),
+		streamBuffer: new(bytes.Buffer),
 	}
 	if fs := getFileStream(ctx); fs != nil {
 		sfa.sourceStream = fs
@@ -112,6 +114,11 @@ func (sc *snowflakeConn) processFileTransfer(
 	if err != nil {
 		return nil, err
 	}
+	if sfa.options.getFileToStream {
+		if err := writeFileStream(ctx, sfa.streamBuffer); err != nil {
+			return nil, err
+		}
+	}
 	return data, nil
 }
 
@@ -138,6 +145,19 @@ func getFileTransferOptions(ctx context.Context) *SnowflakeFileTransferOptions {
 	return o
 }
 
+func writeFileStream(ctx context.Context, streamBuf *bytes.Buffer) error {
+	s := ctx.Value(fileGetStream)
+	w, ok := s.(io.Writer)
+	if !ok {
+		return errors.New("expected an io.Writer")
+	}
+	_, err := streamBuf.WriteTo(w)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
 func (sc *snowflakeConn) populateSessionParameters(parameters []nameValueParameter) {
 	// other session parameters (not all)
 	logger.WithContext(sc.ctx).Infof("params: %#v", parameters)
diff --git a/doc.go b/doc.go
index 2ad5fc9..ec18651 100644
--- a/doc.go
+++ b/doc.go
@@ -1254,6 +1254,19 @@ an absolute path rather than a relative path. For example:
 
 	db.Query("GET @~ file:///tmp/my_data_file auto_compress=false overwrite=false")
 
+To download a file into an in-memory stream (rather than a file) use code similar to the code below.
+
+	var streamBuf bytes.Buffer
+	ctx := WithFileTransferOptions(context.Background(), &SnowflakeFileTransferOptions{getFileToStream: true})
+	ctx = WithFileGetStream(ctx, &streamBuf)
+
+	sql := "get @~/data1.txt.gz file:///tmp/testData"
+	dbt.mustExecContext(ctx, sql)
+	// streamBuf is now filled with the stream. Use bytes.NewReader(streamBuf.Bytes()) to read uncompressed stream or
+	// use gzip.NewReader(&streamBuf) for to read compressed stream.
+
+Note: GET statements are not supported for multi-statement queries.
+
 Specifying temporary directory for encryption and compression:
 
 Putting and getting requires compression and/or encryption, which is done in the OS temporary directory.
diff --git a/encrypt_util.go b/encrypt_util.go
index 0817989..481ebc7 100644
--- a/encrypt_util.go
+++ b/encrypt_util.go
@@ -190,46 +190,51 @@ func encryptFile(
 	return meta, tmpOutputFile.Name(), nil
 }
 
-func decryptFile(
+func decryptFileKey(
 	metadata *encryptMetadata,
-	sfe *snowflakeFileEncryption,
-	filename string,
-	chunkSize int,
-	tmpDir string) (
-	string, error) {
-	if chunkSize == 0 {
-		chunkSize = aes.BlockSize * 4 * 1024
-	}
+	sfe *snowflakeFileEncryption) ([]byte, []byte, error) {
 	decodedKey, err := base64.StdEncoding.DecodeString(sfe.QueryStageMasterKey)
 	if err != nil {
-		return "", err
+		return nil, nil, err
 	}
 	keyBytes, err := base64.StdEncoding.DecodeString(metadata.key) // encrypted file key
 	if err != nil {
-		return "", err
+		return nil, nil, err
 	}
 	ivBytes, err := base64.StdEncoding.DecodeString(metadata.iv)
 	if err != nil {
-		return "", err
+		return nil, nil, err
 	}
 
 	// decrypt file key
 	decryptedKey := make([]byte, len(keyBytes))
 	if err = decryptECB(decryptedKey, keyBytes, decodedKey); err != nil {
-		return "", err
+		return nil, nil, err
 	}
 	decryptedKey, err = paddingTrim(decryptedKey)
 	if err != nil {
-		return "", err
+		return nil, nil, err
 	}
 
-	// decrypt file
+	return decryptedKey, ivBytes, err
+}
+
+func initCBC(decryptedKey []byte, ivBytes []byte) (cipher.BlockMode, error) {
 	block, err := aes.NewCipher(decryptedKey)
 	if err != nil {
-		return "", err
+		return nil, err
 	}
 	mode := cipher.NewCBCDecrypter(block, ivBytes)
 
+	return mode, err
+}
+
+func decryptFile(
+	metadata *encryptMetadata,
+	sfe *snowflakeFileEncryption,
+	filename string,
+	chunkSize int,
+	tmpDir string) (string, error) {
 	tmpOutputFile, err := os.CreateTemp(tmpDir, baseName(filename)+"#")
 	if err != nil {
 		return "", err
@@ -240,11 +245,37 @@ func decryptFile(
 		return "", err
 	}
 	defer infile.Close()
+	totalFileSize, err := decryptStream(metadata, sfe, chunkSize, infile, tmpOutputFile)
+	if err != nil {
+		return "", err
+	}
+	tmpOutputFile.Truncate(int64(totalFileSize))
+	return tmpOutputFile.Name(), nil
+}
+
+func decryptStream(
+	metadata *encryptMetadata,
+	sfe *snowflakeFileEncryption,
+	chunkSize int,
+	src io.Reader,
+	out io.Writer) (int, error) {
+	if chunkSize == 0 {
+		chunkSize = aes.BlockSize * 4 * 1024
+	}
+	decryptedKey, ivBytes, err := decryptFileKey(metadata, sfe)
+	if err != nil {
+		return 0, err
+	}
+	mode, err := initCBC(decryptedKey, ivBytes)
+	if err != nil {
+		return 0, err
+	}
+
 	var totalFileSize int
 	var prevChunk []byte
 	for {
 		chunk := make([]byte, chunkSize)
-		n, err := infile.Read(chunk)
+		n, err := src.Read(chunk)
 		if n == 0 || err != nil {
 			break
 		} else if n%aes.BlockSize != 0 {
@@ -255,17 +286,16 @@ func decryptFile(
 		totalFileSize += n
 		chunk = chunk[:n]
 		mode.CryptBlocks(chunk, chunk)
-		tmpOutputFile.Write(chunk)
+		out.Write(chunk)
 		prevChunk = chunk
 	}
 	if err != nil {
-		return "", err
+		return 0, err
 	}
 	if prevChunk != nil {
 		totalFileSize -= paddingOffset(prevChunk)
 	}
-	tmpOutputFile.Truncate(int64(totalFileSize))
-	return tmpOutputFile.Name(), nil
+	return totalFileSize, err
 }
 
 type materialDescriptor struct {
diff --git a/file_transfer_agent.go b/file_transfer_agent.go
index 7d894b3..03cd1dc 100644
--- a/file_transfer_agent.go
+++ b/file_transfer_agent.go
@@ -91,6 +91,9 @@ type SnowflakeFileTransferOptions struct {
 	/* streaming PUT */
 	compressSourceFromStream bool
 
+	/* streaming GET */
+	getFileToStream bool
+
 	/* PUT */
 	putCallback             *snowflakeProgressPercentage
 	putAzureCallback        *snowflakeProgressPercentage
@@ -124,6 +127,7 @@ type snowflakeFileTransferAgent struct {
 	useAccelerateEndpoint       bool
 	presignedURLs               []string
 	options                     *SnowflakeFileTransferOptions
+	streamBuffer                *bytes.Buffer
 }
 
 func (sfa *snowflakeFileTransferAgent) execute() error {
@@ -411,6 +415,7 @@ func (sfa *snowflakeFileTransferAgent) initFileMetadata() error {
 					name:              baseName(fileName),
 					srcFileName:       fileName,
 					dstFileName:       dstFileName,
+					dstStream:         new(bytes.Buffer),
 					stageLocationType: sfa.stageLocationType,
 					stageInfo:         sfa.stageInfo,
 					localLocation:     sfa.localLocation,
diff --git a/file_util.go b/file_util.go
index 13740bb..3e18cae 100644
--- a/file_util.go
+++ b/file_util.go
@@ -137,6 +137,9 @@ type fileMetadata struct {
 	srcStream     *bytes.Buffer
 	realSrcStream *bytes.Buffer
 
+	/* streaming GET */
+	dstStream *bytes.Buffer
+
 	/* GCS */
 	presignedURL                *url.URL
 	gcsFileHeaderDigest         string
diff --git a/gcs_storage_client.go b/gcs_storage_client.go
index b52f0e3..b59708c 100644
--- a/gcs_storage_client.go
+++ b/gcs_storage_client.go
@@ -322,13 +322,24 @@ func (util *snowflakeGcsClient) nativeDownloadFile(
 		return meta.lastError
 	}
 
-	f, err := os.OpenFile(fullDstFileName, os.O_CREATE|os.O_WRONLY, readWriteFileMode)
-	if err != nil {
-		return err
-	}
-	defer f.Close()
-	if _, err = io.Copy(f, resp.Body); err != nil {
-		return err
+	if meta.options.getFileToStream {
+		if _, err := io.Copy(meta.dstStream, resp.Body); err != nil {
+			return err
+		}
+	} else {
+		f, err := os.OpenFile(fullDstFileName, os.O_CREATE|os.O_WRONLY, readWriteFileMode)
+		if err != nil {
+			return err
+		}
+		defer f.Close()
+		if _, err = io.Copy(f, resp.Body); err != nil {
+			return err
+		}
+		fi, err := os.Stat(fullDstFileName)
+		if err != nil {
+			return err
+		}
+		meta.srcFileSize = fi.Size()
 	}
 
 	var encryptMeta encryptMetadata
@@ -348,12 +359,6 @@ func (util *snowflakeGcsClient) nativeDownloadFile(
 			}
 		}
 	}
-
-	fi, err := os.Stat(fullDstFileName)
-	if err != nil {
-		return err
-	}
-	meta.srcFileSize = fi.Size()
 	meta.resStatus = downloaded
 	meta.gcsFileHeaderDigest = resp.Header.Get(gcsMetadataSfcDigest)
 	meta.gcsFileHeaderContentLength = resp.ContentLength
diff --git a/s3_storage_client.go b/s3_storage_client.go
index 37799c8..342d0c9 100644
--- a/s3_storage_client.go
+++ b/s3_storage_client.go
@@ -247,10 +247,22 @@ func (util *snowflakeS3Client) nativeDownloadFile(
 	if meta.mockDownloader != nil {
 		downloader = meta.mockDownloader
 	}
-	if _, err = downloader.Download(context.Background(), f, &s3.GetObjectInput{
-		Bucket: s3Obj.Bucket,
-		Key:    s3Obj.Key,
-	}); err != nil {
+
+	if meta.options.getFileToStream {
+		buf := manager.NewWriteAtBuffer([]byte{})
+		_, err = downloader.Download(context.Background(), buf, &s3.GetObjectInput{
+			Bucket: s3Obj.Bucket,
+			Key:    s3Obj.Key,
+		})
+		meta.dstStream = bytes.NewBuffer(buf.Bytes())
+	} else {
+		_, err = downloader.Download(context.Background(), f, &s3.GetObjectInput{
+			Bucket: s3Obj.Bucket,
+			Key:    s3Obj.Key,
+		})
+	}
+
+	if err != nil {
 		var ae smithy.APIError
 		if errors.As(err, &ae) {
 			if ae.ErrorCode() == expiredToken {
diff --git a/storage_client.go b/storage_client.go
index ee746a6..550be0d 100644
--- a/storage_client.go
+++ b/storage_client.go
@@ -218,17 +218,30 @@ func (rsu *remoteStorageUtil) downloadOneFile(meta *fileMetadata) error {
 						return err
 					}
 				}
-				tmpDstFileName, err := decryptFile(header.encryptionMetadata,
-					meta.encryptionMaterial, fullDstFileName, 0, meta.tmpDir)
-				if err != nil {
-					return err
-				}
-				if err = os.Rename(tmpDstFileName, fullDstFileName); err != nil {
-					return err
+				if meta.options.getFileToStream {
+					totalFileSize, err := decryptStream(header.encryptionMetadata,
+						meta.encryptionMaterial, 0, meta.dstStream, meta.sfa.streamBuffer)
+					if err != nil {
+						return err
+					}
+					meta.sfa.streamBuffer.Truncate(totalFileSize)
+					meta.dstFileSize = int64(totalFileSize)
+				} else {
+					tmpDstFileName, err := decryptFile(header.encryptionMetadata,
+						meta.encryptionMaterial, fullDstFileName, 0, meta.tmpDir)
+					if err != nil {
+						return err
+					}
+					if err = os.Rename(tmpDstFileName, fullDstFileName); err != nil {
+						return err
+					}
 				}
+
 			}
-			if fi, err := os.Stat(fullDstFileName); err == nil {
-				meta.dstFileSize = fi.Size()
+			if !meta.options.getFileToStream {
+				if fi, err := os.Stat(fullDstFileName); err == nil {
+					meta.dstFileSize = fi.Size()
+				}
 			}
 			return nil
 		}
diff --git a/util.go b/util.go
index 4bc8a5a..53fef8e 100644
--- a/util.go
+++ b/util.go
@@ -25,6 +25,7 @@ const (
 	snowflakeRequestIDKey            contextKey = "SNOWFLAKE_REQUEST_ID"
 	fetchResultByID                  contextKey = "SF_FETCH_RESULT_BY_ID"
 	fileStreamFile                   contextKey = "STREAMING_PUT_FILE"
+	fileGetStream                    contextKey = "STREAMING_GET_FILE"
 	fileTransferOptions              contextKey = "FILE_TRANSFER_OPTIONS"
 	enableHigherPrecision            contextKey = "ENABLE_HIGHER_PRECISION"
 	enableArrowBatchesUtf8Validation contextKey = "ENABLE_ARROW_BATCHES_UTF8_VALIDATION"
@@ -82,6 +83,11 @@ func WithFileStream(ctx context.Context, reader io.Reader) context.Context {
 	return context.WithValue(ctx, fileStreamFile, reader)
 }
 
+// WithFileGetStream returns a context that contains the address of the file stream to be GET
+func WithFileGetStream(ctx context.Context, writer io.Writer) context.Context {
+	return context.WithValue(ctx, fileGetStream, writer)
+}
+
 // WithFileTransferOptions returns a context that contains the address of file transfer options
 func WithFileTransferOptions(ctx context.Context, options *SnowflakeFileTransferOptions) context.Context {
 	return context.WithValue(ctx, fileTransferOptions, options)


Test patch:
diff --git a/azure_storage_client_test.go b/azure_storage_client_test.go
index 601b989..ddf45b5 100644
--- a/azure_storage_client_test.go
+++ b/azure_storage_client_test.go
@@ -109,10 +109,11 @@ func TestUnitDetectAzureTokenExpireError(t *testing.T) {
 }
 
 type azureObjectAPIMock struct {
-	UploadStreamFunc  func(ctx context.Context, body io.Reader, o *azblob.UploadStreamOptions) (azblob.UploadStreamResponse, error)
-	UploadFileFunc    func(ctx context.Context, file *os.File, o *azblob.UploadFileOptions) (azblob.UploadFileResponse, error)
-	DownloadFileFunc  func(ctx context.Context, file *os.File, o *blob.DownloadFileOptions) (int64, error)
-	GetPropertiesFunc func(ctx context.Context, o *blob.GetPropertiesOptions) (blob.GetPropertiesResponse, error)
+	UploadStreamFunc   func(ctx context.Context, body io.Reader, o *azblob.UploadStreamOptions) (azblob.UploadStreamResponse, error)
+	UploadFileFunc     func(ctx context.Context, file *os.File, o *azblob.UploadFileOptions) (azblob.UploadFileResponse, error)
+	DownloadFileFunc   func(ctx context.Context, file *os.File, o *blob.DownloadFileOptions) (int64, error)
+	DownloadStreamFunc func(ctx context.Context, o *blob.DownloadStreamOptions) (azblob.DownloadStreamResponse, error)
+	GetPropertiesFunc  func(ctx context.Context, o *blob.GetPropertiesOptions) (blob.GetPropertiesResponse, error)
 }
 
 func (c *azureObjectAPIMock) UploadStream(ctx context.Context, body io.Reader, o *azblob.UploadStreamOptions) (azblob.UploadStreamResponse, error) {
@@ -131,6 +132,10 @@ func (c *azureObjectAPIMock) DownloadFile(ctx context.Context, file *os.File, o
 	return c.DownloadFileFunc(ctx, file, o)
 }
 
+func (c *azureObjectAPIMock) DownloadStream(ctx context.Context, o *blob.DownloadStreamOptions) (azblob.DownloadStreamResponse, error) {
+	return c.DownloadStreamFunc(ctx, o)
+}
+
 func TestUploadFileWithAzureUploadFailedError(t *testing.T) {
 	info := execResponseStageInfo{
 		Location:     "azblob/storage/users/456/",
diff --git a/put_get_test.go b/put_get_test.go
index b94555d..2c03822 100644
--- a/put_get_test.go
+++ b/put_get_test.go
@@ -236,26 +236,23 @@ func TestPutLocalFile(t *testing.T) {
 	})
 }
 
-func TestPutWithAutoCompressFalse(t *testing.T) {
-	if runningOnGithubAction() && !runningOnAWS() {
-		t.Skip("skipping non aws environment")
-	}
+func TestPutGetWithAutoCompressFalse(t *testing.T) {
 	tmpDir := t.TempDir()
 	testData := filepath.Join(tmpDir, "data.txt")
 	f, err := os.Create(testData)
 	if err != nil {
 		t.Error(err)
 	}
-	f.WriteString("test1,test2\ntest3,test4")
+	originalContents := "test1,test2\ntest3,test4"
+	f.WriteString(originalContents)
 	f.Sync()
 	defer f.Close()
 
 	runDBTest(t, func(dbt *DBTest) {
-		if _, err = dbt.exec("use role sysadmin"); err != nil {
-			t.Skip("snowflake admin account not accessible")
-		}
 		dbt.mustExec("rm @~/test_put_uncompress_file")
-		sqlText := fmt.Sprintf("put file://%v @~/test_put_uncompress_file auto_compress=FALSE", testData)
+
+		// PUT test
+		sqlText := fmt.Sprintf("put 'file://%v' @~/test_put_uncompress_file auto_compress=FALSE", testData)
 		sqlText = strings.ReplaceAll(sqlText, "\\", "\\\\")
 		dbt.mustExec(sqlText)
 		defer dbt.mustExec("rm @~/test_put_uncompress_file")
@@ -263,16 +260,45 @@ func TestPutWithAutoCompressFalse(t *testing.T) {
 		defer rows.Close()
 		var file, s1, s2, s3 string
 		if rows.Next() {
-			if err := rows.Scan(&file, &s1, &s2, &s3); err != nil {
-				t.Fatal(err)
-			}
+			err = rows.Scan(&file, &s1, &s2, &s3)
+			assertNilE(t, err)
 		}
-		if !strings.Contains(file, "test_put_uncompress_file/data.txt") {
-			t.Fatalf("should contain file. got: %v", file)
+		assertTrueF(t, strings.Contains(file, "test_put_uncompress_file/data.txt"), fmt.Sprintf("should contain file. got: %v", file))
+		assertFalseF(t, strings.Contains(file, "data.txt.gz"), fmt.Sprintf("should not contain file. got: %v", file))
+
+		// GET test
+		var streamBuf bytes.Buffer
+		ctx := WithFileTransferOptions(context.Background(), &SnowflakeFileTransferOptions{getFileToStream: true})
+		ctx = WithFileGetStream(ctx, &streamBuf)
+		sql := fmt.Sprintf("get @~/test_put_uncompress_file/data.txt 'file://%v'", tmpDir)
+		sqlText = strings.ReplaceAll(sql, "\\", "\\\\")
+		rows2 := dbt.mustQueryContext(ctx, sqlText)
+		defer rows2.Close()
+		for rows2.Next() {
+			err = rows2.Scan(&file, &s1, &s2, &s3)
+			assertNilE(t, err)
+			assertTrueE(t, strings.HasPrefix(file, "data.txt"), "a file was not downloaded by GET")
+			v, err := strconv.Atoi(s1)
+			assertNilE(t, err)
+			assertEqualE(t, v, 23, "did not return the right file size")
+			assertEqualE(t, s2, "DOWNLOADED", "did not return DOWNLOADED status")
+			assertEqualE(t, s3, "")
 		}
-		if strings.Contains(file, "data.txt.gz") {
-			t.Fatalf("should not contain file. got: %v", file)
+		var contents string
+		r := bytes.NewReader(streamBuf.Bytes())
+		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])
+			}
 		}
+		assertEqualE(t, contents, originalContents)
 	})
 }
 
@@ -436,9 +462,14 @@ func testPutGet(t *testing.T, isStream bool) {
 		dbt.mustExec(fmt.Sprintf(`copy into @%%%v from %v file_format=(type=csv
 			compression='gzip')`, tableName, tableName))
 
+		var streamBuf bytes.Buffer
+		if isStream {
+			ctx = WithFileTransferOptions(ctx, &SnowflakeFileTransferOptions{getFileToStream: true})
+			ctx = WithFileGetStream(ctx, &streamBuf)
+		}
 		sql = fmt.Sprintf("get @%%%v 'file://%v'", tableName, tmpDir)
 		sqlText = strings.ReplaceAll(sql, "\\", "\\\\")
-		rows2 := dbt.mustQuery(sqlText)
+		rows2 := dbt.mustQueryContext(ctx, sqlText)
 		defer rows2.Close()
 		for rows2.Next() {
 			if err = rows2.Scan(&s0, &s1, &s2, &s3); err != nil {
@@ -458,38 +489,51 @@ func testPutGet(t *testing.T, isStream bool) {
 			}
 		}
 
-		files, err := filepath.Glob(filepath.Join(tmpDir, "data_*"))
-		if err != nil {
-			t.Fatal(err)
-		}
-		fileName := files[0]
-		f, err := os.Open(fileName)
-		if err != nil {
-			t.Error(err)
-		}
-		defer f.Close()
-		gz, err := gzip.NewReader(f)
-		if err != nil {
-			t.Error(err)
-		}
-		defer gz.Close()
 		var contents string
-		for {
-			c := make([]byte, defaultChunkBufferSize)
-			if n, err := gz.Read(c); err != nil {
-				if err == io.EOF {
+		if isStream {
+			gz, err := gzip.NewReader(&streamBuf)
+			assertNilE(t, err)
+			defer gz.Close()
+			for {
+				c := make([]byte, defaultChunkBufferSize)
+				if n, err := gz.Read(c); err != nil {
+					if err == io.EOF {
+						contents = contents + string(c[:n])
+						break
+					}
+					t.Error(err)
+				} else {
 					contents = contents + string(c[:n])
-					break
 				}
-				t.Error(err)
-			} else {
-				contents = contents + string(c[:n])
 			}
-		}
+		} else {
+			files, err := filepath.Glob(filepath.Join(tmpDir, "data_*"))
+			if err != nil {
+				t.Fatal(err)
+			}
+			fileName := files[0]
+			f, err := os.Open(fileName)
+			assertNilE(t, err)
+			defer f.Close()
 
-		if contents != originalContents {
-			t.Error("output is different from the original file")
+			gz, err := gzip.NewReader(f)
+			assertNilE(t, err)
+			defer gz.Close()
+
+			for {
+				c := make([]byte, defaultChunkBufferSize)
+				if n, err := gz.Read(c); err != nil {
+					if err == io.EOF {
+						contents = contents + string(c[:n])
+						break
+					}
+					t.Error(err)
+				} else {
+					contents = contents + string(c[:n])
+				}
+			}
 		}
+		assertEqualE(t, contents, originalContents, "output is different from the original contents")
 	})
 }
 func TestPutGetGcsDownscopedCredential(t *testing.T) {
@@ -612,14 +656,14 @@ func TestPutGetGcsDownscopedCredential(t *testing.T) {
 	})
 }
 
-func TestPutLargeFile(t *testing.T) {
+func TestPutGetLargeFile(t *testing.T) {
 	sourceDir, err := os.Getwd()
-	if err != nil {
-		t.Fatal(err)
-	}
+	assertNilF(t, err)
 
 	runDBTest(t, func(dbt *DBTest) {
 		dbt.mustExec("rm @~/test_put_largefile")
+
+		// PUT test
 		putQuery := fmt.Sprintf("put file://%v/test_data/largefile.txt @%v", sourceDir, "~/test_put_largefile")
 		sqlText := strings.ReplaceAll(putQuery, "\\", "\\\\")
 		dbt.mustExec(sqlText)
@@ -628,15 +672,59 @@ func TestPutLargeFile(t *testing.T) {
 		defer rows.Close()
 		var file, s1, s2, s3 string
 		if rows.Next() {
-			if err := rows.Scan(&file, &s1, &s2, &s3); err != nil {
-				t.Fatal(err)
-			}
+			rows.Scan(&file, &s1, &s2, &s3)
+			assertNilF(t, err)
 		}
 
 		if !strings.Contains(file, "largefile.txt.gz") {
 			t.Fatalf("should contain file. got: %v", file)
 		}
 
+		// GET test with stream
+		var streamBuf bytes.Buffer
+		ctx := WithFileTransferOptions(context.Background(), &SnowflakeFileTransferOptions{getFileToStream: true})
+		ctx = WithFileGetStream(ctx, &streamBuf)
+		sql := fmt.Sprintf("get @%v 'file://%v'", "~/test_put_largefile/largefile.txt.gz", t.TempDir())
+		sqlText = strings.ReplaceAll(sql, "\\", "\\\\")
+		rows2 := dbt.mustQueryContext(ctx, sqlText)
+		defer rows2.Close()
+		for rows2.Next() {
+			err = rows2.Scan(&file, &s1, &s2, &s3)
+			assertNilE(t, err)
+			assertTrueE(t, strings.HasPrefix(file, "largefile.txt.gz"), "a file was not downloaded by GET")
+			v, err := strconv.Atoi(s1)
+			assertNilE(t, err)
+			assertEqualE(t, v, 424821, "did not return the right file size")
+			assertEqualE(t, s2, "DOWNLOADED", "did not return DOWNLOADED status")
+			assertEqualE(t, s3, "")
+		}
+
+		// convert the compressed stream to string
+		var contents string
+		gz, err := gzip.NewReader(&streamBuf)
+		assertNilE(t, err)
+		defer gz.Close()
+		for {
+			c := make([]byte, defaultChunkBufferSize)
+			if n, err := gz.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 stream with the original file
+		fname := filepath.Join(sourceDir, "/test_data/largefile.txt")
+		f, err := os.Open(fname)
+		assertNilE(t, err)
+		defer f.Close()
+		originalContents, err := io.ReadAll(f)
+		assertNilE(t, err)
+		assertEqualF(t, contents, string(originalContents), "data did not match content")
 	})
 }