Repository: snowflakedb/gosnowflake
Pull request #519 (2022-01-04 21:49:16)

SNOW-522730 Fix S3 GET in Nested Stage
### Description
Fix error where nested files in s3 stages were not returning the proper file.

### 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/bind_uploader.go b/bind_uploader.go
index 1d9a242..775cf5f 100644
--- a/bind_uploader.go
+++ b/bind_uploader.go
@@ -13,8 +13,8 @@ import (
 )
 
 const (
-	bindStageName   = "SYSTEM$BIND"
-	createStageStmt = "CREATE OR REPLACE TEMPORARY STAGE " + bindStageName +
+	bindStageName            = "SYSTEM$BIND"
+	createTemporaryStageStmt = "CREATE OR REPLACE TEMPORARY STAGE " + bindStageName +
 		" file_format=" + "(type=csv field_optionally_enclosed_by='\"')"
 
 	// size (in bytes) of max input stream (10MB default) as per JDBC specs
@@ -91,7 +91,7 @@ func (bu *bindUploader) createStageIfNeeded() error {
 	if bu.arrayBindStage != "" {
 		return nil
 	}
-	data, err := bu.sc.exec(bu.ctx, createStageStmt, false, false, false, []driver.NamedValue{})
+	data, err := bu.sc.exec(bu.ctx, createTemporaryStageStmt, false, false, false, []driver.NamedValue{})
 	if err != nil {
 		newThreshold := "0"
 		bu.sc.cfg.Params[sessionArrayBindStageThreshold] = &newThreshold
diff --git a/s3_util.go b/s3_util.go
index 5b44cce..bbd1f47 100644
--- a/s3_util.go
+++ b/s3_util.go
@@ -202,11 +202,7 @@ func (util *snowflakeS3Util) nativeDownloadFile(
 	meta *fileMetadata,
 	fullDstFileName string,
 	maxConcurrency int64) error {
-	s3loc, err := util.extractBucketNameAndPath(meta.stageInfo.Location)
-	if err != nil {
-		return err
-	}
-	s3path := s3loc.s3Path + strings.TrimLeft(meta.dstFileName, "/")
+	s3Obj, _ := util.getS3Object(meta, meta.srcFileName)
 	client, ok := meta.client.(*s3.Client)
 	if !ok {
 		return &SnowflakeError{
@@ -223,8 +219,8 @@ func (util *snowflakeS3Util) nativeDownloadFile(
 		u.Concurrency = int(maxConcurrency)
 	})
 	if _, err = downloader.Download(context.Background(), f, &s3.GetObjectInput{
-		Bucket: &s3loc.bucketName,
-		Key:    &s3path,
+		Bucket: s3Obj.Bucket,
+		Key:    s3Obj.Key,
 	}); err != nil {
 		var ae smithy.APIError
 		if errors.As(err, &ae) {


Test patch:
diff --git a/put_get_test.go b/put_get_test.go
index 8c259e2..63881b7 100644
--- a/put_get_test.go
+++ b/put_get_test.go
@@ -17,6 +17,8 @@ import (
 	"testing"
 )
 
+const createStageStmt = "CREATE OR REPLACE STAGE %v URL = '%v' CREDENTIALS = (%v)"
+
 func TestPutError(t *testing.T) {
 	if isWindows {
 		t.Skip("permission model is different")
@@ -53,10 +55,10 @@ func TestPutError(t *testing.T) {
 			RaisePutGetError: false,
 		},
 	}
-	if err := fta.execute(); err != nil {
+	if err = fta.execute(); err != nil {
 		t.Fatal(err)
 	}
-	if _, err := fta.result(); err != nil {
+	if _, err = fta.result(); err != nil {
 		t.Fatal(err)
 	}
 
@@ -66,10 +68,10 @@ func TestPutError(t *testing.T) {
 			RaisePutGetError: true,
 		},
 	}
-	if err := fta.execute(); err != nil {
+	if err = fta.execute(); err != nil {
 		t.Fatal(err)
 	}
-	if _, err := fta.result(); err == nil {
+	if _, err = fta.result(); err == nil {
 		t.Fatalf("should raise permission error")
 	}
 }
@@ -108,23 +110,31 @@ func cleanupPut(dbt *DBTest, td *tcPutGetData) {
 	dbt.mustExec("drop warehouse " + td.warehouse)
 }
 
-func createTestData(dbt *DBTest) (*tcPutGetData, error) {
+func getAWSCredentials() (string, string, string, error) {
 	keyID, ok := os.LookupEnv("AWS_ACCESS_KEY_ID")
 	if !ok {
-		return nil, fmt.Errorf("key id invalid")
+		return "", "", "", fmt.Errorf("key id invalid")
 	}
 	secretKey, ok := os.LookupEnv("AWS_SECRET_ACCESS_KEY")
 	if !ok {
-		return nil, fmt.Errorf("secret key invalid")
+		return keyID, "", "", fmt.Errorf("secret key invalid")
 	}
 	bucket, present := os.LookupEnv("SF_AWS_USER_BUCKET")
 	if !present {
 		usr, err := usr.Current()
 		if err != nil {
-			return nil, err
+			return keyID, secretKey, "", err
 		}
 		bucket = fmt.Sprintf("sfc-dev1-regression/%v/reg", usr.Username)
 	}
+	return keyID, secretKey, bucket, nil
+}
+
+func createTestData(dbt *DBTest) (*tcPutGetData, error) {
+	keyID, secretKey, bucket, err := getAWSCredentials()
+	if err != nil {
+		return nil, err
+	}
 	uniqueName := randomString(10)
 	database := fmt.Sprintf("%v_db", uniqueName)
 	wh := fmt.Sprintf("%v_wh", uniqueName)
@@ -143,7 +153,7 @@ func createTestData(dbt *DBTest) (*tcPutGetData, error) {
 		bucket,
 	}
 
-	if _, err := dbt.db.Exec("use role sysadmin"); err != nil {
+	if _, err = dbt.db.Exec("use role sysadmin"); err != nil {
 		return nil, err
 	}
 	dbt.mustExec(fmt.Sprintf(
@@ -275,7 +285,6 @@ func TestPutOverwrite(t *testing.T) {
 	f.Close()
 
 	runTests(t, dsn, func(dbt *DBTest) {
-		var err error
 		if _, err = dbt.db.Exec("use role sysadmin"); err != nil {
 			t.Skip("snowflake admin account not accessible")
 		}
@@ -366,7 +375,7 @@ func testPutGet(t *testing.T, isStream bool) {
 	gzw := gzip.NewWriter(&b)
 	gzw.Write([]byte(originalContents))
 	gzw.Close()
-	if err := ioutil.WriteFile(fname, b.Bytes(), os.ModePerm); err != nil {
+	if err = ioutil.WriteFile(fname, b.Bytes(), os.ModePerm); err != nil {
 		t.Fatal("could not write to gzip file")
 	}
 
@@ -400,7 +409,7 @@ func testPutGet(t *testing.T, isStream bool) {
 
 		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 {
+			if err = rows.Scan(&s0, &s1, &s2, &s3, &s4, &s5, &s6, &s7); err != nil {
 				t.Fatal(err)
 			}
 		}
@@ -422,7 +431,7 @@ func testPutGet(t *testing.T, isStream bool) {
 		rows = dbt.mustQuery(sqlText)
 		defer rows.Close()
 		for rows.Next() {
-			if err := rows.Scan(&s0, &s1, &s2, &s3); err != nil {
+			if err = rows.Scan(&s0, &s1, &s2, &s3); err != nil {
 				t.Error(err)
 			}
 			if !strings.HasPrefix(s0, "data_") {
diff --git a/put_get_with_aws_test.go b/put_get_with_aws_test.go
index 89c30d2..2ada3d3 100644
--- a/put_get_with_aws_test.go
+++ b/put_get_with_aws_test.go
@@ -8,10 +8,13 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/url"
 	"os"
 	"path/filepath"
+	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
@@ -260,3 +263,111 @@ func TestPretendToPutButList(t *testing.T) {
 		t.Fatal("list buckets should fail")
 	}
 }
+
+func TestPutGetAWSStage(t *testing.T) {
+	if runningOnGithubAction() && !runningOnAWS() {
+		t.Skip("skipping non aws environment")
+	}
+
+	tmpDir, err := ioutil.TempDir("", "put_get")
+	if err != nil {
+		t.Error(err)
+	}
+	defer os.RemoveAll(tmpDir)
+	name := "test_put_get.txt.gz"
+	fname := filepath.Join(tmpDir, name)
+	originalContents := "123,test1\n456,test2\n"
+	stageName := "test_put_get_stage_" + 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) {
+		var createStageQuery string
+		keyID, secretKey, _, err := getAWSCredentials()
+		if err != nil {
+			t.Skip("snowflake admin account not accessible")
+		}
+		createStageQuery = fmt.Sprintf(createStageStmt,
+			stageName,
+			"s3://"+stageName,
+			fmt.Sprintf("AWS_KEY_ID='%v' AWS_SECRET_KEY='%v'", keyID, secretKey))
+		dbt.mustExec(createStageQuery)
+
+		defer dbt.mustExec("DROP STAGE IF EXISTS " + stageName)
+
+		sql := "put 'file://%v' @~/%v auto_compress=false"
+		sqlText := fmt.Sprintf(sql, strings.ReplaceAll(fname, "\\", "\\\\"), stageName)
+		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)
+		}
+
+		sql = fmt.Sprintf("get @~/%v 'file://%v'", stageName, tmpDir)
+		sqlText = strings.ReplaceAll(sql, "\\", "\\\\")
+		rows = dbt.mustQuery(sqlText)
+		defer rows.Close()
+		for rows.Next() {
+			if err = rows.Scan(&s0, &s1, &s2, &s3); err != nil {
+				t.Error(err)
+			}
+
+			if strings.Compare(s0, name) != 0 {
+				t.Error("a file was not downloaded by GET")
+			}
+			if v, err := strconv.Atoi(s1); err != nil || v != 41 {
+				t.Error("did not return the right file size")
+			}
+			if s2 != "DOWNLOADED" {
+				t.Error("did not return DOWNLOADED status")
+			}
+			if s3 != "" {
+				t.Errorf("returned %v", s3)
+			}
+		}
+
+		files, err := filepath.Glob(filepath.Join(tmpDir, "*"))
+		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)
+		}
+		var contents string
+		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])
+			}
+		}
+
+		if contents != originalContents {
+			t.Error("output is different from the original file")
+		}
+	})
+}