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

     1  // Copyright (c) 2023 Snowflake Computing Inc. All rights reserved.
     2  
     3  package gosnowflake
     4  
     5  import (
     6  	"bytes"
     7  	"compress/gzip"
     8  	"context"
     9  	"os"
    10  	"path"
    11  	"path/filepath"
    12  	"testing"
    13  )
    14  
    15  func TestLocalUpload(t *testing.T) {
    16  	tmpDir, err := os.MkdirTemp("", "local_put")
    17  	if err != nil {
    18  		t.Error(err)
    19  	}
    20  	defer os.RemoveAll(tmpDir)
    21  	fname := filepath.Join(tmpDir, "test_put_get.txt.gz")
    22  	originalContents := "123,test1\n456,test2\n"
    23  
    24  	var b bytes.Buffer
    25  	gzw := gzip.NewWriter(&b)
    26  	gzw.Write([]byte(originalContents))
    27  	gzw.Close()
    28  	if err := os.WriteFile(fname, b.Bytes(), readWriteFileMode); err != nil {
    29  		t.Fatal("could not write to gzip file")
    30  	}
    31  	putDir, err := os.MkdirTemp("", "put")
    32  	if err != nil {
    33  		t.Error(err)
    34  	}
    35  
    36  	info := execResponseStageInfo{
    37  		Location:     putDir,
    38  		LocationType: "LOCAL_FS",
    39  	}
    40  	localUtil := new(localUtil)
    41  	localCli, err := localUtil.createClient(&info, false)
    42  	if err != nil {
    43  		t.Error(err)
    44  	}
    45  	uploadMeta := fileMetadata{
    46  		name:              "data1.txt.gz",
    47  		stageLocationType: "LOCAL_FS",
    48  		noSleepingTime:    true,
    49  		parallel:          4,
    50  		client:            localCli,
    51  		stageInfo:         &info,
    52  		dstFileName:       "data1.txt.gz",
    53  		srcFileName:       path.Join(tmpDir, "/test_put_get.txt.gz"),
    54  		overwrite:         true,
    55  		options: &SnowflakeFileTransferOptions{
    56  			MultiPartThreshold: dataSizeThreshold,
    57  		},
    58  	}
    59  	uploadMeta.realSrcFileName = uploadMeta.srcFileName
    60  	err = localUtil.uploadOneFileWithRetry(&uploadMeta)
    61  	if err != nil {
    62  		t.Error(err)
    63  	}
    64  	if uploadMeta.resStatus != uploaded {
    65  		t.Fatalf("failed to upload file")
    66  	}
    67  
    68  	uploadMeta.overwrite = false
    69  	err = localUtil.uploadOneFileWithRetry(&uploadMeta)
    70  	if err != nil {
    71  		t.Error(err)
    72  	}
    73  	if uploadMeta.resStatus != skipped {
    74  		t.Fatal("overwrite is false. should have skipped")
    75  	}
    76  	fileStream, _ := os.Open(fname)
    77  	ctx := WithFileStream(context.Background(), fileStream)
    78  	uploadMeta.srcStream = getFileStream(ctx)
    79  
    80  	err = localUtil.uploadOneFileWithRetry(&uploadMeta)
    81  	if err != nil {
    82  		t.Error(err)
    83  	}
    84  	if uploadMeta.resStatus != skipped {
    85  		t.Fatalf("overwrite is false. should have skipped")
    86  	}
    87  	uploadMeta.overwrite = true
    88  	err = localUtil.uploadOneFileWithRetry(&uploadMeta)
    89  	if err != nil {
    90  		t.Error(err)
    91  	}
    92  	if uploadMeta.resStatus != uploaded {
    93  		t.Fatalf("failed to upload file")
    94  	}
    95  
    96  	uploadMeta.realSrcStream = uploadMeta.srcStream
    97  	err = localUtil.uploadOneFileWithRetry(&uploadMeta)
    98  	if err != nil {
    99  		t.Error(err)
   100  	}
   101  	if uploadMeta.resStatus != uploaded {
   102  		t.Fatalf("failed to upload file")
   103  	}
   104  }
   105  
   106  func TestDownloadLocalFile(t *testing.T) {
   107  	tmpDir, err := os.MkdirTemp("", "local_put")
   108  	if err != nil {
   109  		t.Error(err)
   110  	}
   111  	defer os.RemoveAll(tmpDir)
   112  	fname := filepath.Join(tmpDir, "test_put_get.txt.gz")
   113  	originalContents := "123,test1\n456,test2\n"
   114  
   115  	var b bytes.Buffer
   116  	gzw := gzip.NewWriter(&b)
   117  	gzw.Write([]byte(originalContents))
   118  	gzw.Close()
   119  	if err := os.WriteFile(fname, b.Bytes(), readWriteFileMode); err != nil {
   120  		t.Fatal("could not write to gzip file")
   121  	}
   122  	putDir, err := os.MkdirTemp("", "put")
   123  	if err != nil {
   124  		t.Error(err)
   125  	}
   126  
   127  	info := execResponseStageInfo{
   128  		Location:     tmpDir,
   129  		LocationType: "LOCAL_FS",
   130  	}
   131  	localUtil := new(localUtil)
   132  	localCli, err := localUtil.createClient(&info, false)
   133  	if err != nil {
   134  		t.Error(err)
   135  	}
   136  	downloadMeta := fileMetadata{
   137  		name:              "test_put_get.txt.gz",
   138  		stageLocationType: "LOCAL_FS",
   139  		noSleepingTime:    true,
   140  		client:            localCli,
   141  		stageInfo:         &info,
   142  		dstFileName:       "test_put_get.txt.gz",
   143  		overwrite:         true,
   144  		srcFileName:       "test_put_get.txt.gz",
   145  		localLocation:     putDir,
   146  		options: &SnowflakeFileTransferOptions{
   147  			MultiPartThreshold: dataSizeThreshold,
   148  		},
   149  	}
   150  	err = localUtil.downloadOneFile(&downloadMeta)
   151  	if err != nil {
   152  		t.Error(err)
   153  	}
   154  	if downloadMeta.resStatus != downloaded {
   155  		t.Fatalf("failed to get file in local storage")
   156  	}
   157  
   158  	downloadMeta.srcFileName = "test_put_get.txt.gz"
   159  	err = localUtil.downloadOneFile(&downloadMeta)
   160  	if err != nil {
   161  		t.Error(err)
   162  	}
   163  	if downloadMeta.resStatus != downloaded {
   164  		t.Fatalf("failed to get file in local storage")
   165  	}
   166  
   167  	downloadMeta.srcFileName = "local://test_put_get.txt.gz"
   168  	err = localUtil.downloadOneFile(&downloadMeta)
   169  	if err == nil {
   170  		t.Error("file name is invalid. should have returned an error")
   171  	}
   172  }