Repository: snowflakedb/gosnowflake
Pull request #415 (2021-04-22 04:56:18)

SNOW-294155 Azure/GCP PUT Support
### Description
PUT & streaming PUT support for Azure/GCP

### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [x] Extended the README / documentation, if necessary



Patch:
diff --git a/azure_util.go b/azure_util.go
index 0f52ff4..d04696f 100644
--- a/azure_util.go
+++ b/azure_util.go
@@ -4,8 +4,15 @@ package gosnowflake
 
 import (
 	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io/ioutil"
+	"net/http"
 	"net/url"
+	"os"
 	"strings"
+	"time"
 
 	"github.com/Azure/azure-storage-blob-go/azblob"
 )
@@ -13,44 +20,147 @@ import (
 type snowflakeAzureUtil struct {
 }
 
-// storageUtil implementation
-func (util *snowflakeAzureUtil) createClient(info *execResponseStageInfo, useAccelerateEndpoint bool) cloudClient {
-	sasToken := info.Creds.AzureSasToken
-	if sasToken != "" && strings.HasPrefix(sasToken, "?") {
-		sasToken = sasToken[1:]
-	}
-	endPoint := info.EndPoint
-	endPoint = strings.TrimPrefix(endPoint, "blob.")
-
-	key, _ := azblob.NewSharedKeyCredential(info.StorageAccount, sasToken)
-	p := azblob.NewPipeline(key, azblob.PipelineOptions{})
-	u, _ := url.Parse("https://" + info.StorageAccount + ".blob." + endPoint)
-	serviceURL := azblob.NewServiceURL(*u, p)
-	containerURL := serviceURL.NewContainerURL("mycontainer")
-	containerURL.Create(context.Background(), azblob.Metadata{}, azblob.PublicAccessNone)
-	return containerURL
+type azureLocation struct {
+	containerName string
+	path          string
 }
 
-// storageUtil implementation
-func (util *snowflakeAzureUtil) uploadOneFileWithRetry(meta *fileMetadata) error {
-	// TODO SNOW-294155
-	return nil
-}
+func (util *snowflakeAzureUtil) createClient(info *execResponseStageInfo, _ bool) cloudClient {
+	sasToken := info.Creds.AzureSasToken
+	p := azblob.NewPipeline(azblob.NewAnonymousCredential(), azblob.PipelineOptions{
+		Retry: azblob.RetryOptions{
+			Policy:     azblob.RetryPolicyExponential,
+			MaxTries:   60,
+			RetryDelay: 2 * time.Second,
+		},
+	})
 
-// storageUtil implementation
-func (util *snowflakeAzureUtil) downloadOneFile() {
-	// TODO SNOW-294151
+	u, _ := url.Parse(fmt.Sprintf("https://%s.%s/%s%s", info.StorageAccount, info.EndPoint, info.Path, sasToken))
+	containerURL := azblob.NewContainerURL(*u, p)
+	return &containerURL
 }
 
 // cloudUtil implementation
 func (util *snowflakeAzureUtil) getFileHeader(meta *fileMetadata, filename string) *fileHeader {
-	// TODO SNOW-294155
-	return &fileHeader{}
+	container, ok := meta.client.(*azblob.ContainerURL)
+	if !ok {
+		return nil
+	}
+
+	b := container.NewBlockBlobURL(filename)
+	resp, err := b.GetProperties(context.Background(), azblob.BlobAccessConditions{}, azblob.ClientProvidedKeyOptions{})
+	if err != nil {
+		var se azblob.StorageError
+		if errors.As(err, &se) {
+			if se.ServiceCode() == azblob.ServiceCodeBlobNotFound {
+				meta.resStatus = notFoundFile
+				return &fileHeader{}
+			} else if se.Response().StatusCode == 403 {
+				meta.resStatus = renewToken
+				return nil
+			}
+		}
+		meta.resStatus = errStatus
+		return nil
+	}
+
+	meta.resStatus = uploaded
+	metadata := resp.NewMetadata()
+	var encData encryptionData
+	if err = json.Unmarshal([]byte(metadata["encryptiondata"]), &encData); err != nil {
+		return nil
+	}
+	encryptionMetadata := encryptMetadata{
+		encData.WrappedContentKey.EncryptionKey,
+		encData.ContentEncryptionIV,
+		metadata["matdesc"],
+	}
+
+	return &fileHeader{
+		metadata["sfcdigest"],
+		int64(len(metadata)),
+		&encryptionMetadata,
+	}
 }
 
 // cloudUtil implementation
-func (util *snowflakeAzureUtil) uploadFile(dataFile string, meta *fileMetadata, encryptMeta *encryptMetadata, maxConcurrency int, multiPartThreshold int64) error {
-	// TODO SNOW-294155
+func (util *snowflakeAzureUtil) uploadFile(
+	dataFile string,
+	meta *fileMetadata,
+	encryptMeta *encryptMetadata,
+	maxConcurrency int,
+	multiPartThreshold int64) error {
+	azureMeta := map[string]string{
+		"sfcdigest": meta.sha256Digest,
+	}
+	if encryptMeta != nil {
+		ed := &encryptionData{
+			EncryptionMode: "FullBlob",
+			WrappedContentKey: contentKey{
+				"symmKey1",
+				encryptMeta.key,
+				"AES_CBC_256",
+			},
+			EncryptionAgent: encryptionAgent{
+				"1.0",
+				"AES_CBC_128",
+			},
+			ContentEncryptionIV: encryptMeta.iv,
+			KeyWrappingMetadata: keyMetadata{
+				"Java 5.3.0",
+			},
+		}
+		metadata, _ := json.Marshal(ed)
+		azureMeta["encryptiondata"] = string(metadata)
+		azureMeta["matdesc"] = encryptMeta.matdesc
+	}
+
+	azureLoc := util.extractContainerNameAndPath(meta.stageInfo.Location)
+	path := azureLoc.path + strings.TrimLeft(meta.dstFileName, "/")
+	azContainerURL, ok := meta.client.(*azblob.ContainerURL)
+	if !ok {
+		return &SnowflakeError{
+			Message: "failed to cast to azure client",
+		}
+	}
+
+	var err error
+	blobURL := azContainerURL.NewBlockBlobURL(path)
+	if meta.srcStream != nil {
+		uploadSrc := meta.srcStream
+		if meta.realSrcStream != nil {
+			uploadSrc = meta.realSrcStream
+		}
+		_, err = azblob.UploadStreamToBlockBlob(context.Background(), uploadSrc, blobURL, azblob.UploadStreamToBlockBlobOptions{
+			BufferSize: uploadSrc.Len(),
+			Metadata:   azureMeta,
+		})
+	} else {
+		f, _ := os.OpenFile(dataFile, os.O_RDONLY, os.ModePerm)
+		fi, _ := f.Stat()
+		_, err = azblob.UploadFileToBlockBlob(context.Background(), f, blobURL, azblob.UploadToBlockBlobOptions{
+			BlockSize:   fi.Size(),
+			Parallelism: uint16(maxConcurrency),
+			Metadata:    azureMeta,
+		})
+	}
+	if err != nil {
+		var se azblob.StorageError
+		if errors.As(err, &se) {
+			if se.Response().StatusCode == 403 && util.detectAzureTokenExpireError(se.Response()) {
+				meta.resStatus = renewToken
+			} else {
+				meta.resStatus = needRetry
+				meta.lastError = err
+			}
+			return err
+		}
+		meta.resStatus = errStatus
+		return err
+	}
+
+	meta.dstFileSize = meta.uploadSize
+	meta.resStatus = uploaded
 	return nil
 }
 
@@ -58,3 +168,31 @@ func (util *snowflakeAzureUtil) uploadFile(dataFile string, meta *fileMetadata,
 func (util *snowflakeAzureUtil) nativeDownloadFile() {
 	// TODO SNOW-294151
 }
+
+func (util *snowflakeAzureUtil) extractContainerNameAndPath(location string) *azureLocation {
+	stageLocation := expandUser(location)
+	containerName := stageLocation
+	path := ""
+
+	if strings.Contains(stageLocation, "/") {
+		containerName = stageLocation[:strings.Index(stageLocation, "/")]
+		path = stageLocation[strings.Index(stageLocation, "/")+1:]
+		if path != "" && !strings.HasSuffix(path, "/") {
+			path += "/"
+		}
+	}
+	return &azureLocation{containerName, path}
+}
+
+func (util *snowflakeAzureUtil) detectAzureTokenExpireError(resp *http.Response) bool {
+	if resp.StatusCode != 403 {
+		return false
+	}
+	azureErr, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return false
+	}
+	errStr := string(azureErr)
+	return strings.Contains(errStr, "Signature not valid in the specified time frame") ||
+		strings.Contains(errStr, "Server failed to authenticate the request")
+}
diff --git a/bind_uploader.go b/bind_uploader.go
index 08ad6a3..7abc393 100644
--- a/bind_uploader.go
+++ b/bind_uploader.go
@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"context"
 	"database/sql/driver"
+	"fmt"
 	"reflect"
 	"strconv"
 	"strings"
@@ -72,20 +73,15 @@ func (bu *bindUploader) uploadStreamInternal(inputStream *bytes.Buffer, compress
 		}
 	}
 
-	var putCommand strings.Builder
 	// use a placeholder for source file
-	putCommand.WriteString("put file:///tmp/placeholder ")
-	// add stage name surrounded by quotations in case special chars are used in directory name
-	putCommand.WriteString("'")
-	putCommand.WriteString(stageName)
-	putCommand.WriteString("'")
-	putCommand.WriteString(" overwrite=true")
+	putCommand := fmt.Sprintf("put 'file:///tmp/placeholder' '%v' overwrite=true", stageName)
+	// for Windows queries
+	putCommand = strings.ReplaceAll(putCommand, "\\", "\\\\")
 	// prepare context for PUT command
 	ctx := WithFileStream(bu.ctx, inputStream)
 	ctx = WithFileTransferOptions(ctx, &SnowflakeFileTransferOptions{
 		compressSourceFromStream: compressData})
-	return bu.sc.exec(
-		ctx, putCommand.String(), false, true, false, []driver.NamedValue{})
+	return bu.sc.exec(ctx, putCommand, false, true, false, []driver.NamedValue{})
 }
 
 func (bu *bindUploader) createStageIfNeeded() error {
diff --git a/encrypt_util.go b/encrypt_util.go
index 132fb33..c5d43c6 100644
--- a/encrypt_util.go
+++ b/encrypt_util.go
@@ -242,3 +242,26 @@ func paddingOffset(src []byte) int {
 	length := len(src)
 	return int(src[length-1])
 }
+
+type contentKey struct {
+	KeyID         string `json:"KeyId,omitempty"`
+	EncryptionKey string `json:"EncryptedKey,omitempty"`
+	Algorithm     string `json:"Algorithm,omitempty"`
+}
+
+type encryptionAgent struct {
+	Protocol            string `json:"Protocol,omitempty"`
+	EncryptionAlgorithm string `json:"EncryptionAlgorithm,omitempty"`
+}
+
+type keyMetadata struct {
+	EncryptionLibrary string `json:"EncryptionLibrary,omitempty"`
+}
+
+type encryptionData struct {
+	EncryptionMode      string          `json:"EncryptionMode,omitempty"`
+	WrappedContentKey   contentKey      `json:"WrappedContentKey,omitempty"`
+	EncryptionAgent     encryptionAgent `json:"EncryptionAgent,omitempty"`
+	ContentEncryptionIV string          `json:"ContentEncryptionIV,omitempty"`
+	KeyWrappingMetadata keyMetadata     `json:"KeyWrappingMetadata,omitempty"`
+}
diff --git a/file_transfer_agent.go b/file_transfer_agent.go
index ad2780d..6fec110 100644
--- a/file_transfer_agent.go
+++ b/file_transfer_agent.go
@@ -8,6 +8,7 @@ import (
 	"bytes"
 	"context"
 	"database/sql/driver"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -146,6 +147,8 @@ func (sfa *snowflakeFileTransferAgent) execute() error {
 		}
 	}
 
+	sfa.updateFileMetadataWithPresignedURL()
+
 	smallFileMetas := make([]*fileMetadata, 0)
 	largeFileMetas := make([]*fileMetadata, 0)
 
@@ -470,22 +473,28 @@ func (sfa *snowflakeFileTransferAgent) processFileCompressionType() error {
 }
 
 func (sfa *snowflakeFileTransferAgent) updateFileMetadataWithPresignedURL() error {
-	storageClient := sfa.getStorageClient(sfa.stageLocationType)
-
-	// presigned URL only applies to remote storage
-	var client interface{} = storageClient
-	if _, ok := client.(remoteStorageUtil); !ok {
-		return &SnowflakeError{}
-	}
-
 	// presigned URL only applies to GCS
-	if _, ok := client.(snowflakeGcsUtil); ok {
+	if sfa.stageLocationType == gcsClient {
 		if sfa.commandType == uploadCommand {
 			filePathToBeReplaced := sfa.getLocalFilePathFromCommand(sfa.command)
 			for _, meta := range sfa.fileMetadata {
-				filePathToBeReplacedWith := meta.dstFileName
+				filePathToBeReplacedWith := strings.TrimRight(filePathToBeReplaced, meta.dstFileName) + meta.dstFileName
 				commandWithSingleFile := strings.ReplaceAll(sfa.command, filePathToBeReplaced, filePathToBeReplacedWith)
-				data, err := sfa.sc.exec(sfa.sc.ctx, commandWithSingleFile, false, false, false, []driver.NamedValue{})
+				req := execRequest{
+					SQLText: commandWithSingleFile,
+				}
+				headers := getHeaders()
+				headers[httpHeaderAccept] = headerContentTypeApplicationJSON
+				jsonBody, _ := json.Marshal(req)
+				data, err := sfa.sc.rest.FuncPostQuery(
+					sfa.sc.ctx,
+					sfa.sc.rest,
+					&url.Values{},
+					headers,
+					jsonBody,
+					sfa.sc.rest.RequestTimeout,
+					getOrGenerateRequestIDFromContext(sfa.sc.ctx),
+					sfa.sc.cfg)
 				if err != nil {
 					return err
 				}
@@ -543,7 +552,7 @@ func (sfa *snowflakeFileTransferAgent) getLocalFilePathFromCommand(command strin
 	if len(command) == 0 || !strings.Contains(command, fileProtocol) {
 		return ""
 	}
-	if regexp.MustCompile(putRegexp).Match([]byte(command)) {
+	if !regexp.MustCompile(putRegexp).Match([]byte(command)) {
 		return ""
 	}
 
@@ -555,7 +564,7 @@ func (sfa *snowflakeFileTransferAgent) getLocalFilePathFromCommand(command strin
 	var filePathEndIdx = 0
 
 	if isFilePathQuoted {
-		filePathEndIdx = strings.Index(command[filePathBeginIdx:], "'")
+		filePathEndIdx = filePathBeginIdx + strings.Index(command[filePathBeginIdx:], "'")
 		if filePathEndIdx > filePathBeginIdx {
 			filePath = command[filePathBeginIdx:filePathEndIdx]
 		}
@@ -881,12 +890,12 @@ type snowflakeProgressPercentage struct {
 	fileSize        float64
 	outputStream    *io.Writer
 	showProgressBar bool
-	seenSoFar       int
+	seenSoFar       int64
 	done            bool
 	startTime       time.Time
 }
 
-func (spp *snowflakeProgressPercentage) call(bytesAmount int) {
+func (spp *snowflakeProgressPercentage) call(bytesAmount int64) {
 	if spp.outputStream != nil {
 		spp.seenSoFar += bytesAmount
 		percentage := percent(spp.seenSoFar, spp.fileSize)
@@ -896,7 +905,7 @@ func (spp *snowflakeProgressPercentage) call(bytesAmount int) {
 	}
 }
 
-func percent(seenSoFar int, size float64) float64 {
+func percent(seenSoFar int64, size float64) float64 {
 	if float64(seenSoFar) >= size || size <= 0 {
 		return 1.0
 	}
diff --git a/gcs_util.go b/gcs_util.go
index 3eb9971..d3ec54f 100644
--- a/gcs_util.go
+++ b/gcs_util.go
@@ -1,40 +1,215 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
 package gosnowflake
 
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"os"
+	"strconv"
+	"strings"
+)
+
+const (
+	gcsMetadataPrefix             = "x-goog-meta-"
+	gcsMetadataSfcDigest          = gcsMetadataPrefix + sfcDigest
+	gcsMetadataMatdescKey         = gcsMetadataPrefix + "matdesc"
+	gcsMetadataEncryptionDataProp = gcsMetadataPrefix + "encryptiondata"
+	gcsFileHeaderDigest           = "gcs-file-header-digest"
+)
+
 type snowflakeGcsUtil struct {
 }
 
-// storageUtil implementation
-func (util *snowflakeGcsUtil) createClient(info *execResponseStageInfo, useAccelerateEndpoint bool) cloudClient {
+type gcsLocation struct {
+	bucketName string
+	path       string
+}
+
+func (util *snowflakeGcsUtil) createClient(info *execResponseStageInfo, _ bool) cloudClient {
 	securityToken := info.Creds.GcsAccessToken
 	var client cloudClient
 	if securityToken != "" {
 		client = securityToken
 	} else {
-		client = nil
+		client = ""
 	}
 	return client
 }
 
-// storageUtil implementation
-func (util *snowflakeGcsUtil) uploadOneFileWithRetry(meta *fileMetadata) error {
-	// TODO SNOW-294155
-	return nil
-}
-
-// storageUtil implementation
-func (util *snowflakeGcsUtil) downloadOneFile() {
-	// TODO SNOW-294151
-}
-
 // cloudUtil implementation
 func (util *snowflakeGcsUtil) getFileHeader(meta *fileMetadata, filename string) *fileHeader {
-	// TODO SNOW-294155
-	return &fileHeader{}
+	if meta.resStatus == uploaded || meta.resStatus == downloaded {
+		return &fileHeader{
+			digest:             meta.gcsFileHeaderDigest,
+			contentLength:      meta.gcsFileHeaderContentLength,
+			encryptionMetadata: meta.gcsFileHeaderEncryptionMeta,
+		}
+	}
+	if meta.presignedURL != nil {
+		meta.resStatus = notFoundFile
+	} else {
+		URL := util.generateFileURL(meta.stageInfo.Location, strings.TrimLeft(filename, "/"))
+		accessToken := meta.client.(string)
+		gcsHeaders := map[string]string{
+			"Authorization": "Bearer " + accessToken,
+		}
+
+		req, err := http.NewRequest("HEAD", URL.String(), nil)
+		if err != nil {
+			return nil
+		}
+		for k, v := range gcsHeaders {
+			req.Header.Add(k, v)
+		}
+		client := &http.Client{}
+		resp, _ := client.Do(req)
+		if resp.StatusCode != http.StatusOK {
+			if resp.StatusCode == 403 || resp.StatusCode == 408 || resp.StatusCode == 429 || resp.StatusCode == 500 || resp.StatusCode == 503 {
+				meta.lastError = fmt.Errorf(resp.Status)
+				meta.resStatus = needRetry
+				return nil
+			}
+			if resp.StatusCode == 404 {
+				meta.resStatus = notFoundFile
+			} else if util.isTokenExpired(resp) {
+				meta.lastError = fmt.Errorf(resp.Status)
+				meta.resStatus = renewToken
+			}
+			meta.lastError = fmt.Errorf(resp.Status)
+			meta.resStatus = errStatus
+			return nil
+		}
+
+		digest := resp.Header.Get(gcsMetadataSfcDigest)
+		contentLength, _ := strconv.Atoi(resp.Header.Get("content-length"))
+		var encryptionMeta *encryptMetadata
+		if resp.Header.Get(gcsMetadataEncryptionDataProp) != "" {
+			var encryptData *encryptionData
+			json.Unmarshal([]byte(resp.Header.Get(gcsMetadataEncryptionDataProp)), encryptData)
+			if encryptData != nil {
+				encryptionMeta = &encryptMetadata{
+					key: encryptData.WrappedContentKey.EncryptionKey,
+					iv:  encryptData.ContentEncryptionIV,
+				}
+				if resp.Header.Get(gcsMetadataMatdescKey) != "" {
+					encryptionMeta.matdesc = resp.Header.Get(gcsMetadataMatdescKey)
+				}
+			}
+		}
+		meta.resStatus = uploaded
+		return &fileHeader{
+			digest:             digest,
+			contentLength:      int64(contentLength),
+			encryptionMetadata: encryptionMeta,
+		}
+	}
+	return nil
 }
 
 // cloudUtil implementation
 func (util *snowflakeGcsUtil) uploadFile(dataFile string, meta *fileMetadata, encryptMeta *encryptMetadata, maxConcurrency int, multiPartThreshold int64) error {
-	// TODO SNOW-294155
+	uploadURL := meta.presignedURL
+	var accessToken string
+
+	if uploadURL == nil {
+		uploadURL = util.generateFileURL(meta.stageInfo.Location, strings.TrimLeft(meta.dstFileName, "/"))
+		accessToken = meta.client.(string)
+	}
+
+	var contentEncoding string
+	if meta.dstCompressionType != nil {
+		contentEncoding = strings.ToLower(meta.dstCompressionType.name)
+	}
+
+	if contentEncoding == "gzip" {
+		contentEncoding = ""
+	}
+
+	gcsHeaders := make(map[string]string)
+	gcsHeaders[httpHeaderContentEncoding] = contentEncoding
+	gcsHeaders[gcsMetadataSfcDigest] = meta.sha256Digest
+	if accessToken != "" {
+		gcsHeaders["Authorization"] = "Bearer " + accessToken
+	}
+
+	if encryptMeta != nil {
+		encryptData := encryptionData{
+			"FullBlob",
+			contentKey{
+				"symmKey1",
+				encryptMeta.key,
+				"AES_CBC_256",
+			},
+			encryptionAgent{
+				"1.0",
+				"AES_CBC_256",
+			},
+			encryptMeta.iv,
+			keyMetadata{
+				"Java 5.3.0",
+			},
+		}
+		b, _ := json.Marshal(&encryptData)
+		gcsHeaders[gcsMetadataEncryptionDataProp] = string(b)
+		gcsHeaders[gcsMetadataMatdescKey] = encryptMeta.matdesc
+	}
+
+	var uploadSrc io.Reader
+	if meta.srcStream != nil {
+		uploadSrc = meta.srcStream
+		if meta.realSrcStream != nil {
+			uploadSrc = meta.realSrcStream
+		}
+	} else {
+		uploadSrc, _ = os.OpenFile(dataFile, os.O_RDONLY, os.ModePerm)
+	}
+
+	req, err := http.NewRequest("PUT", uploadURL.String(), uploadSrc)
+	if err != nil {
+		return err
+	}
+	for k, v := range gcsHeaders {
+		req.Header.Add(k, v)
+	}
+	client := &http.Client{}
+	resp, _ := client.Do(req)
+	if resp.StatusCode != http.StatusOK {
+		meta.lastError = fmt.Errorf(resp.Status)
+		if resp.StatusCode == 403 || resp.StatusCode == 408 || resp.StatusCode == 429 || resp.StatusCode == 500 || resp.StatusCode == 503 {
+			meta.lastError = fmt.Errorf(resp.Status)
+			meta.resStatus = needRetry
+		} else if accessToken == "" && resp.StatusCode == 400 && meta.lastError == nil {
+			meta.lastError = fmt.Errorf(resp.Status)
+			meta.resStatus = renewPresignedURL
+		} else if accessToken != "" && util.isTokenExpired(resp) {
+			meta.lastError = fmt.Errorf(resp.Status)
+			meta.resStatus = renewToken
+		}
+		return meta.lastError
+	}
+
+	if meta.options.putCallback != nil {
+		meta.options.putCallback = &snowflakeProgressPercentage{
+			filename:        dataFile,
+			fileSize:        float64(meta.srcFileSize),
+			outputStream:    meta.options.putCallbackOutputStream,
+			showProgressBar: meta.options.showProgressBar,
+		}
+	}
+
+	meta.dstFileSize = meta.uploadSize
+	meta.resStatus = uploaded
+
+	meta.gcsFileHeaderDigest = gcsHeaders[gcsFileHeaderDigest]
+	meta.gcsFileHeaderContentLength = meta.uploadSize
+	if err = json.Unmarshal([]byte(gcsHeaders[gcsMetadataEncryptionDataProp]), &encryptMeta); err != nil {
+		return err
+	}
+	meta.gcsFileHeaderEncryptionMeta = encryptMeta
 	return nil
 }
 
@@ -42,3 +217,27 @@ func (util *snowflakeGcsUtil) uploadFile(dataFile string, meta *fileMetadata, en
 func (util *snowflakeGcsUtil) nativeDownloadFile() {
 	// TODO SNOW-294151
 }
+
+func (util *snowflakeGcsUtil) extractBucketNameAndPath(location string) *gcsLocation {
+	containerName := location
+	var path string
+	if strings.Contains(location, "/") {
+		containerName = location[:strings.Index(location, "/")]
+		path = location[strings.Index(location, "/")+1:]
+		if path != "" && strings.HasSuffix(location, "/") {
+			path += "/"
+		}
+	}
+	return &gcsLocation{containerName, path}
+}
+
+func (util *snowflakeGcsUtil) generateFileURL(stageLocation string, filename string) *url.URL {
+	gcsLoc := util.extractBucketNameAndPath(stageLocation)
+	fullFilePath := gcsLoc.path + filename
+	URL, _ := url.Parse("https://storage.googleapis.com/" + gcsLoc.bucketName + "/" + url.QueryEscape(fullFilePath))
+	return URL
+}
+
+func (util *snowflakeGcsUtil) isTokenExpired(resp *http.Response) bool {
+	return resp.StatusCode == 401
+}
diff --git a/s3_util.go b/s3_util.go
index 78bdd11..6c80d53 100644
--- a/s3_util.go
+++ b/s3_util.go
@@ -2,8 +2,6 @@
 
 package gosnowflake
 
-//lint:file-ignore U1000 SNOW-294151
-
 import (
 	"bytes"
 	"context"
@@ -30,7 +28,6 @@ const (
 )
 
 type snowflakeS3Util struct {
-	meta *fileMetadata
 }
 
 type s3Location struct {
@@ -38,7 +35,6 @@ type s3Location struct {
 	s3Path     string
 }
 
-// storageUtil implementation
 func (util *snowflakeS3Util) createClient(info *execResponseStageInfo, useAccelerateEndpoint bool) cloudClient {
 	stageCredentials := info.Creds
 	var resolver s3.EndpointResolver
@@ -57,11 +53,6 @@ func (util *snowflakeS3Util) createClient(info *execResponseStageInfo, useAccele
 	})
 }
 
-// storageUtil implementation
-func (util *snowflakeS3Util) downloadOneFile() {
-	// TODO SNOW-294151
-}
-
 type s3HeaderAPI interface {
 	HeadObject(ctx context.Context, params *s3.HeadObjectInput, optFns ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
 }
@@ -147,7 +138,7 @@ func (util *snowflakeS3Util) uploadFile(
 	var uploader s3UploadAPI
 	uploader = manager.NewUploader(client, func(u *manager.Uploader) {
 		u.Concurrency = maxConcurrency
-		u.PartSize = multiPartThreshold
+		u.PartSize = int64Max(multiPartThreshold, manager.DefaultUploadPartSize)
 	})
 	if meta.mockUploader != nil {
 		uploader = meta.mockUploader
diff --git a/util.go b/util.go
index 98e4d39..3e0244b 100644
--- a/util.go
+++ b/util.go
@@ -107,6 +107,13 @@ func intMax(a, b int) int {
 	return b
 }
 
+func int64Max(a, b int64) int64 {
+	if a > b {
+		return a
+	}
+	return b
+}
+
 func getMin(arr []int) int {
 	if len(arr) == 0 {
 		return -1


Test patch:
diff --git a/driver_test.go b/driver_test.go
index 414ba31..2f1f617 100644
--- a/driver_test.go
+++ b/driver_test.go
@@ -376,6 +376,14 @@ func runningOnAWS() bool {
 	return os.Getenv("CLOUD_PROVIDER") == "AWS"
 }
 
+func runningOnAzure() bool {
+	return os.Getenv("CLOUD_PROVIDER") == "AZURE"
+}
+
+func runningOnGCP() bool {
+	return os.Getenv("CLOUD_PROVIDER") == "GCP"
+}
+
 func TestBogusUserPasswordParameters(t *testing.T) {
 	invalidDNS := fmt.Sprintf("%s:%s@%s", "bogus", pass, host)
 	invalidUserPassErrorTests(invalidDNS, t)
diff --git a/put_get_test.go b/put_get_test.go
index 5641f66..f3fc698 100644
--- a/put_get_test.go
+++ b/put_get_test.go
@@ -66,7 +66,7 @@ func TestPutError(t *testing.T) {
 
 func TestPercentage(t *testing.T) {
 	testcases := []struct {
-		seen     int
+		seen     int64
 		size     float64
 		expected float64
 	}{
@@ -395,7 +395,7 @@ func TestPutOverwrite(t *testing.T) {
 			}
 		}
 		if s6 != skipped.String() {
-			t.Errorf("expected SKIPPED, got %v", s6)
+			t.Fatalf("expected SKIPPED, got %v", s6)
 		}
 
 		f, _ = os.Open(testData)
diff --git a/put_get_with_azure_test.go b/put_get_with_azure_test.go
new file mode 100644
index 0000000..af55668
--- /dev/null
+++ b/put_get_with_azure_test.go
@@ -0,0 +1,80 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"bytes"
+	"compress/gzip"
+	"context"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func TestPutFileWithAzure(t *testing.T) {
+	if runningOnGithubAction() && !runningOnAzure() {
+		t.Skip("skipping non azure environment")
+	}
+	testPutWithAzure(t, false)
+}
+
+func TestPutStreamWithAzure(t *testing.T) {
+	if runningOnGithubAction() && !runningOnAzure() {
+		t.Skip("skipping non azure environment")
+	}
+	testPutWithAzure(t, true)
+}
+
+func testPutWithAzure(t *testing.T, isStream bool) {
+	tmpDir, _ := ioutil.TempDir("", "azure_put")
+	defer os.RemoveAll(tmpDir)
+	fname := filepath.Join(tmpDir, "test_put_get_with_azure.txt.gz")
+	originalContents := "123,test1\n456,test2\n"
+	tableName := randomString(5)
+
+	var b bytes.Buffer
+	gzw := gzip.NewWriter(&b)
+	gzw.Write([]byte(originalContents))
+	gzw.Close()
+	if err := ioutil.WriteFile(fname, b.Bytes(), os.ModePerm); err != nil {
+		t.Fatal("could not write to gzip file")
+	}
+
+	runTests(t, dsn, func(dbt *DBTest) {
+		dbt.mustExec("create or replace table " + tableName + " (a int, b string)")
+		fileStream, _ := os.OpenFile(fname, os.O_RDONLY, os.ModePerm)
+		defer func() {
+			dbt.mustExec("drop table " + tableName)
+			if fileStream != nil {
+				fileStream.Close()
+			}
+		}()
+
+		var rows *RowsExtended
+		sql := "put 'file://%v' @%%%v auto_compress=true parallel=30"
+		if isStream {
+			sqlText := fmt.Sprintf(sql, strings.ReplaceAll(fname, "\\", "\\\\"), tableName)
+			rows = dbt.mustQueryContext(WithFileStream(context.Background(), fileStream), sqlText)
+		} else {
+			sqlText := fmt.Sprintf(sql, strings.ReplaceAll(fname, "\\", "\\\\"), tableName)
+			rows = dbt.mustQuery(sqlText)
+		}
+
+		var s0, s1, s2, s3, s4, s5, s6, s7 string
+		if rows.Next() {
+			if err := rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
+				t.Fatal(err)
+			}
+		}
+		if s6 != uploaded.String() {
+			t.Fatalf("expected %v, got: %v", uploaded, s6)
+		}
+
+		dbt.mustExec("copy into " + tableName)
+		dbt.mustExec("rm @%" + tableName)
+		dbt.mustQueryAssertCount("ls @%"+tableName, 0)
+	})
+}
diff --git a/put_get_with_gcp_test.go b/put_get_with_gcp_test.go
new file mode 100644
index 0000000..6ae3478
--- /dev/null
+++ b/put_get_with_gcp_test.go
@@ -0,0 +1,56 @@
+// Copyright (c) 2021 Snowflake Computing Inc. All right reserved.
+
+package gosnowflake
+
+import (
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io/ioutil"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+)
+
+func TestPutWithGCP(t *testing.T) {
+	if runningOnGithubAction() && !runningOnGCP() {
+		t.Skip("skipping non gcp environment")
+	}
+	tmpDir, _ := ioutil.TempDir("", "gcp_put")
+	defer os.RemoveAll(tmpDir)
+	fname := filepath.Join(tmpDir, "test_put_get_with_gcp.txt.gz")
+	originalContents := "123,test1\n456,test2\n"
+
+	var b bytes.Buffer
+	gzw := gzip.NewWriter(&b)
+	gzw.Write([]byte(originalContents))
+	gzw.Close()
+	err := ioutil.WriteFile(fname, b.Bytes(), os.ModePerm)
+	if err != nil {
+		t.Fatal("could not write to gzip file")
+	}
+	tableName := randomString(5)
+
+	runTests(t, dsn, func(dbt *DBTest) {
+		dbt.mustExec("create or replace table " + tableName + " (a int, b string)")
+		defer dbt.mustExec("drop table " + tableName)
+		sql := "put 'file://%v' @%%%v auto_compress=true parallel=30"
+		sqlText := fmt.Sprintf(sql, strings.ReplaceAll(fname, "\\", "\\\\"),
+			tableName)
+		rows := dbt.mustQuery(sqlText)
+		var s0, s1, s2, s3, s4, s5, s6, s7 string
+		if rows.Next() {
+			if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
+				t.Fatal(err)
+			}
+		}
+		if s6 != uploaded.String() {
+			t.Fatalf("expected %v, got: %v", uploaded, s6)
+		}
+
+		dbt.mustExec("copy into " + tableName)
+		dbt.mustExec("rm @%" + tableName)
+		dbt.mustQueryAssertCount("ls @%"+tableName, 0)
+	})
+}