github.com/snowflakedb/gosnowflake@v1.9.0/local_storage_client.go (about)

     1  // Copyright (c) 2021-2022 Snowflake Computing Inc. All rights reserved.
     2  
     3  package gosnowflake
     4  
     5  import (
     6  	"bufio"
     7  	"fmt"
     8  	"io"
     9  	"os"
    10  	"path"
    11  	"path/filepath"
    12  	"strings"
    13  )
    14  
    15  type localUtil struct {
    16  }
    17  
    18  func (util *localUtil) createClient(_ *execResponseStageInfo, _ bool) (cloudClient, error) {
    19  	return nil, nil
    20  }
    21  
    22  func (util *localUtil) uploadOneFileWithRetry(meta *fileMetadata) error {
    23  	var frd *bufio.Reader
    24  	if meta.srcStream != nil {
    25  		b := meta.srcStream
    26  		if meta.realSrcStream != nil {
    27  			b = meta.realSrcStream
    28  		}
    29  		frd = bufio.NewReader(b)
    30  	} else {
    31  		f, err := os.Open(meta.realSrcFileName)
    32  		if err != nil {
    33  			return err
    34  		}
    35  		defer f.Close()
    36  		frd = bufio.NewReader(f)
    37  	}
    38  
    39  	user, err := expandUser(meta.stageInfo.Location)
    40  	if err != nil {
    41  		return err
    42  	}
    43  	if !meta.overwrite {
    44  		if _, err := os.Stat(filepath.Join(user, meta.dstFileName)); err == nil {
    45  			meta.dstFileSize = 0
    46  			meta.resStatus = skipped
    47  			return nil
    48  		}
    49  	}
    50  	output, err := os.OpenFile(filepath.Join(user, meta.dstFileName), os.O_CREATE|os.O_WRONLY, readWriteFileMode)
    51  	if err != nil {
    52  		return err
    53  	}
    54  	defer output.Close()
    55  	data := make([]byte, meta.uploadSize)
    56  	for {
    57  		n, err := frd.Read(data)
    58  		if err != nil && err != io.EOF {
    59  			return err
    60  		}
    61  		if n == 0 {
    62  			break
    63  		}
    64  
    65  		if _, err = output.Write(data); err != nil {
    66  			return err
    67  		}
    68  	}
    69  	meta.dstFileSize = meta.uploadSize
    70  	meta.resStatus = uploaded
    71  	return nil
    72  }
    73  
    74  func (util *localUtil) downloadOneFile(meta *fileMetadata) error {
    75  	srcFileName := meta.srcFileName
    76  	if strings.HasPrefix(meta.srcFileName, fmt.Sprintf("%b", os.PathSeparator)) {
    77  		srcFileName = srcFileName[1:]
    78  	}
    79  	user, err := expandUser(meta.stageInfo.Location)
    80  	if err != nil {
    81  		return err
    82  	}
    83  	fullSrcFileName := path.Join(user, srcFileName)
    84  	user, err = expandUser(meta.localLocation)
    85  	if err != nil {
    86  		return err
    87  	}
    88  	fullDstFileName := path.Join(user, baseName(meta.dstFileName))
    89  	baseDir, err := getDirectory()
    90  	if err != nil {
    91  		return err
    92  	}
    93  	if _, err = os.Stat(baseDir); os.IsNotExist(err) {
    94  		if err = os.MkdirAll(baseDir, os.ModePerm); err != nil {
    95  			return err
    96  		}
    97  	}
    98  
    99  	data, err := os.ReadFile(fullSrcFileName)
   100  	if err != nil {
   101  		return err
   102  	}
   103  	if err = os.WriteFile(fullDstFileName, data, readWriteFileMode); err != nil {
   104  		return err
   105  	}
   106  	fi, err := os.Stat(fullDstFileName)
   107  	if err != nil {
   108  		return err
   109  	}
   110  	meta.dstFileSize = fi.Size()
   111  	meta.resStatus = downloaded
   112  	return nil
   113  }