ruby_snowflake_client 1.0.0

data/ext/vendor/github.com/snowflakedb/gosnowflake/file_transfer_agent.go

@@ -0,0 +1,1228 @@
+// Copyright (c) 2021-2022 Snowflake Computing Inc. All rights reserved.
+
+package gosnowflake
+
+//lint:file-ignore U1000 Ignore all unused code
+
+import (
+	"bytes"
+	"context"
+	"database/sql/driver"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"math"
+	"net/url"
+	"os"
+	"path/filepath"
+	"regexp"
+	"runtime"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/aws/aws-sdk-go-v2/service/s3"
+	"github.com/aws/smithy-go"
+	"github.com/gabriel-vasile/mimetype"
+)
+
+type (
+	cloudType   string
+	commandType string
+)
+
+const (
+	fileProtocol              = "file://"
+	dataSizeThreshold int64   = 64 * 1024 * 1024
+	injectWaitPut             = 0
+	isWindows                 = runtime.GOOS == "windows"
+	mb                float64 = 1024.0 * 1024.0
+)
+
+const (
+	uploadCommand   commandType = "UPLOAD"
+	downloadCommand commandType = "DOWNLOAD"
+	unknownCommand  commandType = "UNKNOWN"
+
+	putRegexp string = `(?i)^(?:/\*.*\*/\s*)*put\s+`
+	getRegexp string = `(?i)^(?:/\*.*\*/\s*)*get\s+`
+)
+
+const (
+	s3Client    cloudType = "S3"
+	azureClient cloudType = "AZURE"
+	gcsClient   cloudType = "GCS"
+	local       cloudType = "LOCAL_FS"
+)
+
+type resultStatus int
+
+const (
+	errStatus resultStatus = iota
+	uploaded
+	downloaded
+	skipped
+	renewToken
+	renewPresignedURL
+	notFoundFile
+	needRetry
+	needRetryWithLowerConcurrency
+)
+
+func (rs resultStatus) String() string {
+	return [...]string{"ERROR", "UPLOADED", "DOWNLOADED", "SKIPPED",
+		"RENEW_TOKEN", "RENEW_PRESIGNED_URL", "NOT_FOUND_FILE", "NEED_RETRY",
+		"NEED_RETRY_WITH_LOWER_CONCURRENCY"}[rs]
+}
+
+func (rs resultStatus) isSet() bool {
+	return uploaded <= rs && rs <= needRetryWithLowerConcurrency
+}
+
+// SnowflakeFileTransferOptions enables users to specify options regarding
+// files transfers such as PUT/GET
+type SnowflakeFileTransferOptions struct {
+	showProgressBar    bool
+	RaisePutGetError   bool
+	MultiPartThreshold int64
+
+	/* streaming PUT */
+	compressSourceFromStream bool
+
+	/* PUT */
+	DisablePutOverwrite     bool
+	putCallback             *snowflakeProgressPercentage
+	putAzureCallback        *snowflakeProgressPercentage
+	putCallbackOutputStream *io.Writer
+
+	/* GET */
+	getCallback             *snowflakeProgressPercentage
+	getAzureCallback        *snowflakeProgressPercentage
+	getCallbackOutputStream *io.Writer
+}
+
+type snowflakeFileTransferAgent struct {
+	sc                          *snowflakeConn
+	data                        *execResponseData
+	command                     string
+	commandType                 commandType
+	stageLocationType           cloudType
+	fileMetadata                []*fileMetadata
+	encryptionMaterial          []*snowflakeFileEncryption
+	stageInfo                   *execResponseStageInfo
+	results                     []*fileMetadata
+	sourceStream                *bytes.Buffer
+	srcLocations                []string
+	autoCompress                bool
+	srcCompression              string
+	parallel                    int64
+	overwrite                   bool
+	srcFiles                    []string
+	localLocation               string
+	srcFileToEncryptionMaterial map[string]*snowflakeFileEncryption
+	useAccelerateEndpoint       bool
+	presignedURLs               []string
+	options                     *SnowflakeFileTransferOptions
+}
+
+func (sfa *snowflakeFileTransferAgent) execute() error {
+	var err error
+	if err = sfa.parseCommand(); err != nil {
+		return err
+	}
+	if err = sfa.initFileMetadata(); err != nil {
+		return err
+	}
+
+	if sfa.commandType == uploadCommand {
+		if err = sfa.processFileCompressionType(); err != nil {
+			return err
+		}
+	}
+
+	if err = sfa.transferAccelerateConfig(); err != nil {
+		return err
+	}
+
+	if sfa.commandType == downloadCommand {
+		if _, err = os.Stat(sfa.localLocation); os.IsNotExist(err) {
+			if err = os.MkdirAll(sfa.localLocation, os.ModePerm); err != nil {
+				return err
+			}
+		}
+	}
+
+	if sfa.stageLocationType == local {
+		if _, err = os.Stat(sfa.stageInfo.Location); os.IsNotExist(err) {
+			if err = os.MkdirAll(sfa.stageInfo.Location, os.ModePerm); err != nil {
+				return err
+			}
+		}
+	}
+
+	if err = sfa.updateFileMetadataWithPresignedURL(); err != nil {
+		return err
+	}
+
+	smallFileMetas := make([]*fileMetadata, 0)
+	largeFileMetas := make([]*fileMetadata, 0)
+
+	for _, meta := range sfa.fileMetadata {
+		meta.overwrite = sfa.overwrite
+		meta.sfa = sfa
+		meta.options = sfa.options
+		if sfa.stageLocationType != local {
+			sizeThreshold := sfa.options.MultiPartThreshold
+			meta.options.MultiPartThreshold = sizeThreshold
+			if meta.srcFileSize > sizeThreshold {
+				meta.parallel = sfa.parallel
+				largeFileMetas = append(largeFileMetas, meta)
+			} else {
+				meta.parallel = 1
+				smallFileMetas = append(smallFileMetas, meta)
+			}
+		} else {
+			meta.parallel = 1
+			smallFileMetas = append(smallFileMetas, meta)
+		}
+	}
+
+	if sfa.commandType == uploadCommand {
+		if err = sfa.upload(largeFileMetas, smallFileMetas); err != nil {
+			return err
+		}
+	} else {
+		if err = sfa.download(largeFileMetas, smallFileMetas); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) parseCommand() error {
+	var err error
+	if sfa.data.Command != "" {
+		sfa.commandType = commandType(sfa.data.Command)
+	} else {
+		sfa.commandType = unknownCommand
+	}
+
+	sfa.initEncryptionMaterial()
+	if len(sfa.data.SrcLocations) == 0 {
+		return (&SnowflakeError{
+			Number:   ErrInvalidStageLocation,
+			SQLState: sfa.data.SQLState,
+			QueryID:  sfa.data.QueryID,
+			Message:  "failed to parse location",
+		}).exceptionTelemetry(sfa.sc)
+	}
+	sfa.srcLocations = sfa.data.SrcLocations
+
+	if sfa.commandType == uploadCommand {
+		if sfa.sourceStream != nil {
+			sfa.srcFiles = sfa.srcLocations // streaming PUT
+		} else {
+			sfa.srcFiles, err = sfa.expandFilenames(sfa.srcLocations)
+			if err != nil {
+				return err
+			}
+		}
+		sfa.autoCompress = sfa.data.AutoCompress
+		sfa.srcCompression = strings.ToLower(sfa.data.SourceCompression)
+	} else {
+		sfa.srcFiles = sfa.srcLocations
+		sfa.srcFileToEncryptionMaterial = make(map[string]*snowflakeFileEncryption)
+		if len(sfa.data.SrcLocations) == len(sfa.encryptionMaterial) {
+			for i, srcFile := range sfa.srcFiles {
+				sfa.srcFileToEncryptionMaterial[srcFile] = sfa.encryptionMaterial[i]
+			}
+		} else if len(sfa.encryptionMaterial) != 0 {
+			return (&SnowflakeError{
+				Number:      ErrInternalNotMatchEncryptMaterial,
+				SQLState:    sfa.data.SQLState,
+				QueryID:     sfa.data.QueryID,
+				Message:     errMsgInternalNotMatchEncryptMaterial,
+				MessageArgs: []interface{}{len(sfa.data.SrcLocations), len(sfa.encryptionMaterial)},
+			}).exceptionTelemetry(sfa.sc)
+		}
+
+		sfa.localLocation, err = expandUser(sfa.data.LocalLocation)
+		if err != nil {
+			return err
+		}
+		if fi, err := os.Stat(sfa.localLocation); err != nil || !fi.IsDir() {
+			return (&SnowflakeError{
+				Number:      ErrLocalPathNotDirectory,
+				SQLState:    sfa.data.SQLState,
+				QueryID:     sfa.data.QueryID,
+				Message:     errMsgLocalPathNotDirectory,
+				MessageArgs: []interface{}{sfa.localLocation},
+			}).exceptionTelemetry(sfa.sc)
+		}
+	}
+
+	sfa.parallel = 1
+	if sfa.data.Parallel != 0 {
+		sfa.parallel = sfa.data.Parallel
+	}
+	sfa.overwrite = !sfa.options.DisablePutOverwrite
+	sfa.stageLocationType = cloudType(strings.ToUpper(sfa.data.StageInfo.LocationType))
+	sfa.stageInfo = &sfa.data.StageInfo
+	sfa.presignedURLs = make([]string, 0)
+	if len(sfa.data.PresignedURLs) != 0 {
+		sfa.presignedURLs = sfa.data.PresignedURLs
+	}
+
+	if sfa.getStorageClient(sfa.stageLocationType) == nil {
+		return (&SnowflakeError{
+			Number:      ErrInvalidStageFs,
+			SQLState:    sfa.data.SQLState,
+			QueryID:     sfa.data.QueryID,
+			Message:     errMsgInvalidStageFs,
+			MessageArgs: []interface{}{sfa.stageLocationType},
+		}).exceptionTelemetry(sfa.sc)
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) initEncryptionMaterial() {
+	sfa.encryptionMaterial = make([]*snowflakeFileEncryption, 0)
+	wrapper := sfa.data.EncryptionMaterial
+
+	if sfa.commandType == uploadCommand {
+		if wrapper.QueryID != "" {
+			sfa.encryptionMaterial = append(sfa.encryptionMaterial, &wrapper.snowflakeFileEncryption)
+		}
+	} else {
+		for _, encmat := range wrapper.EncryptionMaterials {
+			if encmat.QueryID != "" {
+				sfa.encryptionMaterial = append(sfa.encryptionMaterial, &encmat)
+			}
+		}
+	}
+}
+
+func (sfa *snowflakeFileTransferAgent) expandFilenames(locations []string) ([]string, error) {
+	canonicalLocations := make([]string, 0)
+	for _, fileName := range locations {
+		if sfa.commandType == uploadCommand {
+			var err error
+			fileName, err = expandUser(fileName)
+			if err != nil {
+				return []string{}, err
+			}
+			if !filepath.IsAbs(fileName) {
+				cwd, err := getDirectory()
+				if err != nil {
+					return []string{}, err
+				}
+				fileName = filepath.Join(cwd, fileName)
+			}
+			if isWindows && len(fileName) > 2 && fileName[0] == '/' && fileName[2] == ':' {
+				// Windows path: /C:/data/file1.txt where it starts with slash
+				// followed by a drive letter and colon.
+				fileName = fileName[1:]
+			}
+			files, err := filepath.Glob(fileName)
+			if err != nil {
+				return []string{}, err
+			}
+			canonicalLocations = append(canonicalLocations, files...)
+		} else {
+			canonicalLocations = append(canonicalLocations, fileName)
+		}
+	}
+	return canonicalLocations, nil
+}
+
+func (sfa *snowflakeFileTransferAgent) initFileMetadata() error {
+	sfa.fileMetadata = []*fileMetadata{}
+	if sfa.commandType == uploadCommand {
+		if len(sfa.srcFiles) == 0 {
+			fileName := sfa.data.SrcLocations
+			return (&SnowflakeError{
+				Number:      ErrFileNotExists,
+				SQLState:    sfa.data.SQLState,
+				QueryID:     sfa.data.QueryID,
+				Message:     errMsgFileNotExists,
+				MessageArgs: []interface{}{fileName},
+			}).exceptionTelemetry(sfa.sc)
+		}
+		if sfa.sourceStream != nil {
+			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,
+			})
+		} else {
+			for i, fileName := range sfa.srcFiles {
+				fi, err := os.Stat(fileName)
+				if os.IsNotExist(err) {
+					return (&SnowflakeError{
+						Number:      ErrFileNotExists,
+						SQLState:    sfa.data.SQLState,
+						QueryID:     sfa.data.QueryID,
+						Message:     errMsgFileNotExists,
+						MessageArgs: []interface{}{fileName},
+					}).exceptionTelemetry(sfa.sc)
+				} else if fi.IsDir() {
+					return (&SnowflakeError{
+						Number:      ErrFileNotExists,
+						SQLState:    sfa.data.SQLState,
+						QueryID:     sfa.data.QueryID,
+						Message:     errMsgFileNotExists,
+						MessageArgs: []interface{}{fileName},
+					}).exceptionTelemetry(sfa.sc)
+				}
+				sfa.fileMetadata = append(sfa.fileMetadata, &fileMetadata{
+					name:              baseName(fileName),
+					srcFileName:       fileName,
+					srcFileSize:       fi.Size(),
+					stageLocationType: sfa.stageLocationType,
+					stageInfo:         sfa.stageInfo,
+				})
+				if len(sfa.encryptionMaterial) > 0 {
+					sfa.fileMetadata[i].encryptionMaterial = sfa.encryptionMaterial[0]
+				}
+			}
+		}
+
+		if len(sfa.encryptionMaterial) > 0 {
+			for _, meta := range sfa.fileMetadata {
+				meta.encryptionMaterial = sfa.encryptionMaterial[0]
+			}
+		}
+	} else if sfa.commandType == downloadCommand {
+		for _, fileName := range sfa.srcFiles {
+			if len(fileName) > 0 {
+				firstPathSep := strings.Index(fileName, "/")
+				dstFileName := fileName
+				if firstPathSep >= 0 {
+					dstFileName = fileName[firstPathSep+1:]
+				}
+				sfa.fileMetadata = append(sfa.fileMetadata, &fileMetadata{
+					name:              baseName(fileName),
+					srcFileName:       fileName,
+					dstFileName:       dstFileName,
+					stageLocationType: sfa.stageLocationType,
+					stageInfo:         sfa.stageInfo,
+					localLocation:     sfa.localLocation,
+				})
+			}
+		}
+		// TODO is this necessary?
+		for _, meta := range sfa.fileMetadata {
+			fileName := meta.srcFileName
+			if val, ok := sfa.srcFileToEncryptionMaterial[fileName]; ok {
+				meta.encryptionMaterial = val
+			}
+		}
+	}
+
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) processFileCompressionType() error {
+	var userSpecifiedSourceCompression *compressionType
+	var autoDetect bool
+	if sfa.srcCompression == "auto_detect" {
+		autoDetect = true
+	} else if sfa.srcCompression == "none" {
+		autoDetect = false
+	} else {
+		userSpecifiedSourceCompression = lookupByMimeSubType(sfa.srcCompression)
+		if userSpecifiedSourceCompression == nil || !userSpecifiedSourceCompression.isSupported {
+			return (&SnowflakeError{
+				Number:      ErrCompressionNotSupported,
+				SQLState:    sfa.data.SQLState,
+				QueryID:     sfa.data.QueryID,
+				Message:     errMsgFeatureNotSupported,
+				MessageArgs: []interface{}{userSpecifiedSourceCompression},
+			}).exceptionTelemetry(sfa.sc)
+		}
+		autoDetect = false
+	}
+
+	gzipCompression := compressionTypes["GZIP"]
+	for _, meta := range sfa.fileMetadata {
+		fileName := meta.srcFileName
+		var currentFileCompressionType *compressionType
+		if autoDetect {
+			currentFileCompressionType = lookupByExtension(filepath.Ext(fileName))
+			if currentFileCompressionType == nil {
+				var mtype *mimetype.MIME
+				var err error
+				if meta.srcStream != nil {
+					r := getReaderFromBuffer(&meta.srcStream)
+					mtype, err = mimetype.DetectReader(r)
+					if err != nil {
+						return err
+					}
+					io.ReadAll(r) // flush out tee buffer
+				} else {
+					mtype, err = mimetype.DetectFile(fileName)
+					if err != nil {
+						return err
+					}
+				}
+				currentFileCompressionType = lookupByExtension(mtype.Extension())
+			}
+
+			if currentFileCompressionType != nil && !currentFileCompressionType.isSupported {
+				return (&SnowflakeError{
+					Number:      ErrCompressionNotSupported,
+					SQLState:    sfa.data.SQLState,
+					QueryID:     sfa.data.QueryID,
+					Message:     errMsgFeatureNotSupported,
+					MessageArgs: []interface{}{userSpecifiedSourceCompression},
+				}).exceptionTelemetry(sfa.sc)
+			}
+		} else {
+			currentFileCompressionType = userSpecifiedSourceCompression
+		}
+
+		if currentFileCompressionType != nil {
+			meta.srcCompressionType = currentFileCompressionType
+			if currentFileCompressionType.isSupported {
+				meta.dstCompressionType = currentFileCompressionType
+				meta.requireCompress = false
+				meta.dstFileName = meta.name
+			} else {
+				return (&SnowflakeError{
+					Number:      ErrCompressionNotSupported,
+					SQLState:    sfa.data.SQLState,
+					QueryID:     sfa.data.QueryID,
+					Message:     errMsgFeatureNotSupported,
+					MessageArgs: []interface{}{userSpecifiedSourceCompression},
+				}).exceptionTelemetry(sfa.sc)
+			}
+		} else {
+			meta.requireCompress = sfa.autoCompress
+			meta.srcCompressionType = nil
+			if sfa.autoCompress {
+				dstFileName := meta.name + compressionTypes["GZIP"].fileExtension
+				meta.dstFileName = dstFileName
+				meta.dstCompressionType = gzipCompression
+			} else {
+				meta.dstFileName = meta.name
+				meta.dstCompressionType = nil
+			}
+		}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) updateFileMetadataWithPresignedURL() error {
+	// presigned URL only applies to GCS
+	if sfa.stageLocationType == gcsClient {
+		if sfa.commandType == uploadCommand {
+			filePathToBeReplaced := sfa.getLocalFilePathFromCommand(sfa.command)
+			for _, meta := range sfa.fileMetadata {
+				filePathToBeReplacedWith := strings.TrimRight(filePathToBeReplaced, meta.dstFileName) + meta.dstFileName
+				commandWithSingleFile := strings.ReplaceAll(sfa.command, filePathToBeReplaced, filePathToBeReplacedWith)
+				req := execRequest{
+					SQLText: commandWithSingleFile,
+				}
+				headers := getHeaders()
+				headers[httpHeaderAccept] = headerContentTypeApplicationJSON
+				jsonBody, err := json.Marshal(req)
+				if err != nil {
+					return err
+				}
+				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
+				}
+
+				if data.Data.StageInfo != (execResponseStageInfo{}) {
+					meta.stageInfo = &data.Data.StageInfo
+					meta.presignedURL = nil
+					if meta.stageInfo.PresignedURL != "" {
+						meta.presignedURL, err = url.Parse(meta.stageInfo.PresignedURL)
+						if err != nil {
+							return err
+						}
+					}
+				}
+			}
+		} else if sfa.commandType == downloadCommand {
+			for i, meta := range sfa.fileMetadata {
+				if len(sfa.presignedURLs) > 0 {
+					var err error
+					meta.presignedURL, err = url.Parse(sfa.presignedURLs[i])
+					if err != nil {
+						return err
+					}
+				} else {
+					meta.presignedURL = nil
+				}
+			}
+		} else {
+			return (&SnowflakeError{
+				Number:      ErrCommandNotRecognized,
+				SQLState:    sfa.data.SQLState,
+				QueryID:     sfa.data.QueryID,
+				Message:     errMsgCommandNotRecognized,
+				MessageArgs: []interface{}{sfa.commandType},
+			}).exceptionTelemetry(sfa.sc)
+		}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) transferAccelerateConfig() error {
+	if sfa.stageLocationType == s3Client {
+		s3Util := new(snowflakeS3Client)
+		s3Loc, err := s3Util.extractBucketNameAndPath(sfa.stageInfo.Location)
+		if err != nil {
+			return err
+		}
+		s3Cli, err := s3Util.createClient(sfa.stageInfo, false)
+		if err != nil {
+			return err
+		}
+		client, ok := s3Cli.(*s3.Client)
+		if !ok {
+			return (&SnowflakeError{
+				Number:   ErrFailedToConvertToS3Client,
+				SQLState: sfa.data.SQLState,
+				QueryID:  sfa.data.QueryID,
+				Message:  errMsgFailedToConvertToS3Client,
+			}).exceptionTelemetry(sfa.sc)
+		}
+		ret, err := client.GetBucketAccelerateConfiguration(context.Background(), &s3.GetBucketAccelerateConfigurationInput{
+			Bucket: &s3Loc.bucketName,
+		})
+		sfa.useAccelerateEndpoint = ret != nil && ret.Status == "Enabled"
+		if err != nil {
+			var ae smithy.APIError
+			if errors.As(err, &ae) {
+				if ae.ErrorCode() == "AccessDenied" {
+					return nil
+				} else if ae.ErrorCode() == "MethodNotAllowed" {
+					return nil
+				}
+			}
+			return err
+		}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) getLocalFilePathFromCommand(command string) string {
+	if len(command) == 0 || !strings.Contains(command, fileProtocol) {
+		return ""
+	}
+	if !regexp.MustCompile(putRegexp).Match([]byte(command)) {
+		return ""
+	}
+
+	filePathBeginIdx := strings.Index(command, fileProtocol)
+	isFilePathQuoted := command[filePathBeginIdx-1] == '\''
+	filePathBeginIdx += len(fileProtocol)
+	var filePathEndIdx int
+	filePath := ""
+
+	if isFilePathQuoted {
+		filePathEndIdx = filePathBeginIdx + strings.Index(command[filePathBeginIdx:], "'")
+		if filePathEndIdx > filePathBeginIdx {
+			filePath = command[filePathBeginIdx:filePathEndIdx]
+		}
+	} else {
+		indexList := make([]int, 0)
+		delims := []rune{' ', '\n', ';'}
+		for _, delim := range delims {
+			index := strings.Index(command[filePathBeginIdx:], string(delim))
+			if index != -1 {
+				indexList = append(indexList, index)
+			}
+		}
+		filePathEndIdx = -1
+		if getMin(indexList) != -1 {
+			filePathEndIdx = getMin(indexList)
+		}
+
+		if filePathEndIdx > filePathBeginIdx {
+			filePath = command[filePathBeginIdx:filePathEndIdx]
+		} else {
+			filePath = command[filePathBeginIdx:]
+		}
+	}
+	return filePath
+}
+
+func (sfa *snowflakeFileTransferAgent) upload(
+	largeFileMetadata []*fileMetadata,
+	smallFileMetadata []*fileMetadata) error {
+	client, err := sfa.getStorageClient(sfa.stageLocationType).
+		createClient(sfa.stageInfo, sfa.useAccelerateEndpoint)
+	if err != nil {
+		return err
+	}
+	for _, meta := range smallFileMetadata {
+		meta.client = client
+	}
+	for _, meta := range largeFileMetadata {
+		meta.client = client
+	}
+
+	if len(smallFileMetadata) > 0 {
+		logger.Infof("uploading %v small files", len(smallFileMetadata))
+		if err = sfa.uploadFilesParallel(smallFileMetadata); err != nil {
+			return err
+		}
+	}
+	if len(largeFileMetadata) > 0 {
+		logger.Infof("uploading %v large files", len(largeFileMetadata))
+		if err = sfa.uploadFilesSequential(largeFileMetadata); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) download(
+	largeFileMetadata []*fileMetadata,
+	smallFileMetadata []*fileMetadata) error {
+	client, err := sfa.getStorageClient(sfa.stageLocationType).
+		createClient(sfa.stageInfo, sfa.useAccelerateEndpoint)
+	if err != nil {
+		return err
+	}
+	for _, meta := range smallFileMetadata {
+		meta.client = client
+	}
+	for _, meta := range largeFileMetadata {
+		meta.client = client
+	}
+
+	if len(smallFileMetadata) > 0 {
+		logger.WithContext(sfa.sc.ctx).Infof("downloading %v small files", len(smallFileMetadata))
+		if err = sfa.downloadFilesParallel(smallFileMetadata); err != nil {
+			return err
+		}
+	}
+	if len(largeFileMetadata) > 0 {
+		logger.WithContext(sfa.sc.ctx).Infof("downloading %v large files", len(largeFileMetadata))
+		if err = sfa.downloadFilesSequential(largeFileMetadata); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) uploadFilesParallel(fileMetas []*fileMetadata) error {
+	idx := 0
+	fileMetaLen := len(fileMetas)
+	var err error
+	for idx < fileMetaLen {
+		endOfIdx := intMin(fileMetaLen, idx+int(sfa.parallel))
+		targetMeta := fileMetas[idx:endOfIdx]
+		for {
+			var wg sync.WaitGroup
+			results := make([]*fileMetadata, len(targetMeta))
+			errors := make([]error, len(targetMeta))
+			for i, meta := range targetMeta {
+				wg.Add(1)
+				go func(k int, m *fileMetadata) {
+					defer wg.Done()
+					results[k], errors[k] = sfa.uploadOneFile(m)
+				}(i, meta)
+			}
+			wg.Wait()
+
+			retryMeta := make([]*fileMetadata, 0)
+			for i, result := range results {
+				result.errorDetails = errors[i]
+				if result.resStatus == renewToken || result.resStatus == renewPresignedURL {
+					retryMeta = append(retryMeta, result)
+				} else {
+					sfa.results = append(sfa.results, result)
+				}
+			}
+			if len(retryMeta) == 0 {
+				break
+			}
+
+			needRenewToken := false
+			for _, result := range retryMeta {
+				if result.resStatus == renewToken {
+					needRenewToken = true
+				}
+			}
+
+			if needRenewToken {
+				client, err := sfa.renewExpiredClient()
+				if err != nil {
+					return err
+				}
+				for _, result := range retryMeta {
+					result.client = client
+				}
+				if endOfIdx < fileMetaLen {
+					for i := idx + int(sfa.parallel); i < fileMetaLen; i++ {
+						fileMetas[i].client = client
+					}
+				}
+			}
+
+			for _, result := range retryMeta {
+				if result.resStatus == renewPresignedURL {
+					sfa.updateFileMetadataWithPresignedURL()
+					break
+				}
+			}
+			targetMeta = retryMeta
+		}
+		if endOfIdx == fileMetaLen {
+			break
+		}
+		idx += int(sfa.parallel)
+	}
+	return err
+}
+
+func (sfa *snowflakeFileTransferAgent) uploadFilesSequential(fileMetas []*fileMetadata) error {
+	idx := 0
+	fileMetaLen := len(fileMetas)
+	for idx < fileMetaLen {
+		res, err := sfa.uploadOneFile(fileMetas[idx])
+		if err != nil {
+			return err
+		}
+
+		if res.resStatus == renewToken {
+			client, err := sfa.renewExpiredClient()
+			if err != nil {
+				return err
+			}
+			for i := idx; i < fileMetaLen; i++ {
+				fileMetas[i].client = client
+			}
+			continue
+		} else if res.resStatus == renewPresignedURL {
+			sfa.updateFileMetadataWithPresignedURL()
+			continue
+		}
+
+		sfa.results = append(sfa.results, res)
+		idx++
+		if injectWaitPut > 0 {
+			time.Sleep(injectWaitPut)
+		}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) uploadOneFile(meta *fileMetadata) (*fileMetadata, error) {
+	meta.realSrcFileName = meta.srcFileName
+	tmpDir, err := os.MkdirTemp("", "")
+	if err != nil {
+		return nil, err
+	}
+	meta.tmpDir = tmpDir
+	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
+		}
+	}
+
+	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)
+	}
+	if err != nil {
+		return meta, err
+	}
+
+	client := sfa.getStorageClient(sfa.stageLocationType)
+	if err = client.uploadOneFileWithRetry(meta); err != nil {
+		return meta, err
+	}
+	return meta, nil
+}
+
+func (sfa *snowflakeFileTransferAgent) downloadFilesParallel(fileMetas []*fileMetadata) error {
+	idx := 0
+	fileMetaLen := len(fileMetas)
+	var err error
+	for idx < fileMetaLen {
+		endOfIdx := intMin(fileMetaLen, idx+int(sfa.parallel))
+		targetMeta := fileMetas[idx:endOfIdx]
+		for {
+			var wg sync.WaitGroup
+			results := make([]*fileMetadata, len(targetMeta))
+			errors := make([]error, len(targetMeta))
+			for i, meta := range targetMeta {
+				wg.Add(1)
+				go func(k int, m *fileMetadata) {
+					defer wg.Done()
+					results[k], errors[k] = sfa.downloadOneFile(m)
+				}(i, meta)
+			}
+			wg.Wait()
+
+			retryMeta := make([]*fileMetadata, 0)
+			for i, result := range results {
+				result.errorDetails = errors[i]
+				if result.resStatus == renewToken || result.resStatus == renewPresignedURL {
+					retryMeta = append(retryMeta, result)
+				} else {
+					sfa.results = append(sfa.results, result)
+				}
+			}
+			if len(retryMeta) == 0 {
+				break
+			}
+			logger.WithContext(sfa.sc.ctx).Infof("%v retries found", len(retryMeta))
+
+			needRenewToken := false
+			for _, result := range retryMeta {
+				if result.resStatus == renewToken {
+					needRenewToken = true
+				}
+				logger.WithContext(sfa.sc.ctx).Infof(
+					"retying download file %v with status %v",
+					result.name, result.resStatus)
+			}
+
+			if needRenewToken {
+				client, err := sfa.renewExpiredClient()
+				if err != nil {
+					return err
+				}
+				for _, result := range retryMeta {
+					result.client = client
+				}
+				if endOfIdx < fileMetaLen {
+					for i := idx + int(sfa.parallel); i < fileMetaLen; i++ {
+						fileMetas[i].client = client
+					}
+				}
+			}
+
+			for _, result := range retryMeta {
+				if result.resStatus == renewPresignedURL {
+					sfa.updateFileMetadataWithPresignedURL()
+					break
+				}
+			}
+			targetMeta = retryMeta
+		}
+		if endOfIdx == fileMetaLen {
+			break
+		}
+		idx += int(sfa.parallel)
+	}
+	return err
+}
+
+func (sfa *snowflakeFileTransferAgent) downloadFilesSequential(fileMetas []*fileMetadata) error {
+	idx := 0
+	fileMetaLen := len(fileMetas)
+	for idx < fileMetaLen {
+		res, err := sfa.downloadOneFile(fileMetas[idx])
+		if err != nil {
+			return err
+		}
+
+		if res.resStatus == renewToken {
+			client, err := sfa.renewExpiredClient()
+			if err != nil {
+				return err
+			}
+			for i := idx; i < fileMetaLen; i++ {
+				fileMetas[i].client = client
+			}
+			continue
+		} else if res.resStatus == renewPresignedURL {
+			sfa.updateFileMetadataWithPresignedURL()
+			continue
+		}
+
+		sfa.results = append(sfa.results, res)
+		idx++
+		if injectWaitPut > 0 {
+			time.Sleep(injectWaitPut)
+		}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) downloadOneFile(meta *fileMetadata) (*fileMetadata, error) {
+	tmpDir, err := os.MkdirTemp("", "")
+	if err != nil {
+		return nil, err
+	}
+	meta.tmpDir = tmpDir
+	defer os.RemoveAll(tmpDir) // cleanup
+	client := sfa.getStorageClient(sfa.stageLocationType)
+	if err = client.downloadOneFile(meta); err != nil {
+		meta.dstFileSize = -1
+		if !meta.resStatus.isSet() {
+			meta.resStatus = errStatus
+		}
+		meta.errorDetails = fmt.Errorf(err.Error() + ", file=" + meta.dstFileName)
+		return meta, err
+	}
+	return meta, nil
+}
+
+func (sfa *snowflakeFileTransferAgent) getStorageClient(stageLocationType cloudType) storageUtil {
+	if stageLocationType == local {
+		return &localUtil{}
+	} else if stageLocationType == s3Client || stageLocationType == azureClient || stageLocationType == gcsClient {
+		return &remoteStorageUtil{}
+	}
+	return nil
+}
+
+func (sfa *snowflakeFileTransferAgent) renewExpiredClient() (cloudClient, error) {
+	data, err := sfa.sc.exec(
+		sfa.sc.ctx,
+		sfa.command,
+		false,
+		false,
+		false,
+		[]driver.NamedValue{})
+	if err != nil {
+		return nil, err
+	}
+	storageClient := sfa.getStorageClient(sfa.stageLocationType)
+	return storageClient.createClient(&data.Data.StageInfo, sfa.useAccelerateEndpoint)
+}
+
+func (sfa *snowflakeFileTransferAgent) result() (*execResponse, error) {
+	// inherit old response data
+	data := sfa.data
+	rowset := make([]fileTransferResultType, 0)
+	if sfa.commandType == uploadCommand {
+		if len(sfa.results) > 0 {
+			for _, meta := range sfa.results {
+				var srcCompressionType, dstCompressionType *compressionType
+				if meta.srcCompressionType != nil {
+					srcCompressionType = meta.srcCompressionType
+				} else {
+					srcCompressionType = &compressionType{
+						name: "NONE",
+					}
+				}
+				if meta.dstCompressionType != nil {
+					dstCompressionType = meta.dstCompressionType
+				} else {
+					dstCompressionType = &compressionType{
+						name: "NONE",
+					}
+				}
+				errorDetails := meta.errorDetails
+				srcFileSize := meta.srcFileSize
+				dstFileSize := meta.dstFileSize
+				if sfa.options.RaisePutGetError && errorDetails != nil {
+					return nil, (&SnowflakeError{
+						Number:   ErrFailedToUploadToStage,
+						SQLState: sfa.data.SQLState,
+						QueryID:  sfa.data.QueryID,
+						Message:  errorDetails.Error(),
+					}).exceptionTelemetry(sfa.sc)
+				}
+				rowset = append(rowset, fileTransferResultType{
+					meta.name,
+					meta.srcFileName,
+					meta.dstFileName,
+					srcFileSize,
+					dstFileSize,
+					srcCompressionType,
+					dstCompressionType,
+					meta.resStatus,
+					meta.errorDetails,
+				})
+			}
+			sort.Slice(rowset, func(i, j int) bool {
+				return rowset[i].srcFileName < rowset[j].srcFileName
+			})
+			ccrs := make([][]*string, 0, len(rowset))
+			for _, rs := range rowset {
+				srcFileSize := fmt.Sprintf("%v", rs.srcFileSize)
+				dstFileSize := fmt.Sprintf("%v", rs.dstFileSize)
+				resStatus := rs.resStatus.String()
+				errorStr := ""
+				if rs.errorDetails != nil {
+					errorStr = rs.errorDetails.Error()
+				}
+				ccrs = append(ccrs, []*string{
+					&rs.srcFileName,
+					&rs.dstFileName,
+					&srcFileSize,
+					&dstFileSize,
+					&rs.srcCompressionType.name,
+					&rs.dstCompressionType.name,
+					&resStatus,
+					&errorStr,
+				})
+			}
+			data.RowSet = ccrs
+			cc := make([]chunkRowType, len(ccrs))
+			populateJSONRowSet(cc, ccrs)
+			data.QueryResultFormat = "json"
+			rt := []execResponseRowType{
+				{Name: "source", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+				{Name: "target", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+				{Name: "source_size", ByteLength: 64, Length: 64, Type: "FIXED", Scale: 0, Nullable: false},
+				{Name: "target_size", ByteLength: 64, Length: 64, Type: "FIXED", Scale: 0, Nullable: false},
+				{Name: "source_compression", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+				{Name: "target_compression", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+				{Name: "status", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+				{Name: "message", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+			}
+			data.RowType = rt
+			return &execResponse{Data: *data, Success: true}, nil
+		}
+	} else { // DOWNLOAD
+		if len(sfa.results) > 0 {
+			for _, meta := range sfa.results {
+				dstFileSize := meta.dstFileSize
+				errorDetails := meta.errorDetails
+				if sfa.options.RaisePutGetError && errorDetails != nil {
+					return nil, (&SnowflakeError{
+						Number:   ErrFailedToDownloadFromStage,
+						SQLState: sfa.data.SQLState,
+						QueryID:  sfa.data.QueryID,
+						Message:  errorDetails.Error(),
+					}).exceptionTelemetry(sfa.sc)
+				}
+
+				rowset = append(rowset, fileTransferResultType{
+					"", "", meta.dstFileName, 0, dstFileSize,
+					nil, nil, meta.resStatus, meta.errorDetails,
+				})
+			}
+			sort.Slice(rowset, func(i, j int) bool {
+				return rowset[i].srcFileName < rowset[j].srcFileName
+			})
+			ccrs := make([][]*string, 0, len(rowset))
+			for _, rs := range rowset {
+				dstFileSize := fmt.Sprintf("%v", rs.dstFileSize)
+				resStatus := rs.resStatus.String()
+				errorStr := ""
+				if rs.errorDetails != nil {
+					errorStr = rs.errorDetails.Error()
+				}
+				ccrs = append(ccrs, []*string{
+					&rs.dstFileName,
+					&dstFileSize,
+					&resStatus,
+					&errorStr,
+				})
+			}
+			data.RowSet = ccrs
+			cc := make([]chunkRowType, len(ccrs))
+			populateJSONRowSet(cc, ccrs)
+			data.QueryResultFormat = "json"
+			rt := []execResponseRowType{
+				{Name: "file", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+				{Name: "size", ByteLength: 64, Length: 64, Type: "FIXED", Scale: 0, Nullable: false},
+				{Name: "status", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+				{Name: "message", ByteLength: 10000, Length: 10000, Type: "TEXT", Scale: 0, Nullable: false},
+			}
+			data.RowType = rt
+			return &execResponse{Data: *data, Success: true}, nil
+		}
+	}
+	return nil, (&SnowflakeError{
+		Number:   ErrNotImplemented,
+		SQLState: sfa.data.SQLState,
+		QueryID:  sfa.data.QueryID,
+		Message:  errMsgNotImplemented,
+	}).exceptionTelemetry(sfa.sc)
+}
+
+func isFileTransfer(query string) bool {
+	putRe := regexp.MustCompile(putRegexp)
+	getRe := regexp.MustCompile(getRegexp)
+	return putRe.Match([]byte(query)) || getRe.Match([]byte(query))
+}
+
+type snowflakeProgressPercentage struct {
+	filename        string
+	fileSize        float64
+	outputStream    *io.Writer
+	showProgressBar bool
+	seenSoFar       int64
+	done            bool
+	startTime       time.Time
+}
+
+func (spp *snowflakeProgressPercentage) call(bytesAmount int64) {
+	if spp.outputStream != nil {
+		spp.seenSoFar += bytesAmount
+		percentage := percent(spp.seenSoFar, spp.fileSize)
+		if !spp.done {
+			spp.done = updateProgress(spp.filename, spp.startTime, spp.fileSize, percentage, spp.outputStream, spp.showProgressBar)
+		}
+	}
+}
+
+func percent(seenSoFar int64, size float64) float64 {
+	if float64(seenSoFar) >= size || size <= 0 {
+		return 1.0
+	}
+	return float64(seenSoFar) / size
+}
+
+func updateProgress(filename string, startTime time.Time, totalSize float64, progress float64, outputStream *io.Writer, showProgressBar bool) bool {
+	barLength := 10
+	totalSize /= mb
+	status := ""
+	elapsedTime := time.Since(startTime)
+
+	var throughput float64
+	if elapsedTime != 0.0 {
+		throughput = totalSize / elapsedTime.Seconds()
+	}
+
+	if progress < 0 {
+		progress = 0
+		status = "Halt...\r\n"
+	}
+	if progress >= 1 {
+		status = fmt.Sprintf("Done (%.3fs, %.2fMB/s)", elapsedTime.Seconds(), throughput)
+	}
+	if status == "" && showProgressBar {
+		status = fmt.Sprintf("(%.3fsm %.2fMB/s)", elapsedTime.Seconds(), throughput)
+	}
+	if status != "" {
+		block := int(math.Round(float64(barLength) * progress))
+		text := fmt.Sprintf("\r%v(%.2fMB): [%v] %.2f%% %v ", filename, totalSize, strings.Repeat("#", block)+strings.Repeat("-", barLength-block), progress*100, status)
+		(*outputStream).Write([]byte(text))
+	}
+	return progress == 1.0
+}