add cov to build files

add some more test
snowflakedb · May 11, 2023 · b7a6e68 · b7a6e68
1 parent 6859e04
commit b7a6e68
Show file tree

Hide file tree

Showing 5 changed files with 207 additions and 23 deletions.
diff --git a/ci/scripts/test_component.sh b/ci/scripts/test_component.sh
@@ -12,4 +12,4 @@ if [[ -n "$GITHUB_WORKFLOW" ]]; then
 fi
 env | grep SNOWFLAKE | grep -v PASS | sort
 cd $TOPDIR
-go test -timeout 30m -race $COVFLAGS -v . 
+go test -timeout 30m -race -coverprofile=coverage.out -covermode=atomic $COVFLAGS -v . 
diff --git a/ci/test.bat b/ci/test.bat
@@ -42,4 +42,4 @@ echo [INFO] Database:  %SNOWFLAKE_TEST_DATABASE%
 echo [INFO] Warehouse: %SNOWFLAKE_TEST_WAREHOUSE%
 echo [INFO] Role:      %SNOWFLAKE_TEST_ROLE%
 
-go test --timeout 30m --tags=sfdebug -race -v .
+go test --timeout 30m --tags=sfdebug -race -coverprofile=coverage.out -covermode=atomic -v .
diff --git a/gcs_storage_client.go b/gcs_storage_client.go
@@ -100,7 +100,10 @@ func (util *snowflakeGcsClient) getFileHeader(meta *fileMetadata, filename strin
 		var encryptionMeta *encryptMetadata
 		if resp.Header.Get(gcsMetadataEncryptionDataProp) != "" {
 			var encryptData *encryptionData
-			json.Unmarshal([]byte(resp.Header.Get(gcsMetadataEncryptionDataProp)), encryptData)
+			err := json.Unmarshal([]byte(resp.Header.Get(gcsMetadataEncryptionDataProp)), &encryptData)
+			if err != nil {
+				logger.Error(err)
+			}
 			if encryptData != nil {
 				encryptionMeta = &encryptMetadata{
 					key: encryptData.WrappedContentKey.EncryptionKey,

diff --git a/gcs_storage_client_test.go b/gcs_storage_client_test.go
@@ -273,7 +273,7 @@ func TestUploadFileWithGcsUploadFailedWithTokenExpired(t *testing.T) {
 		mockGcsClient: &clientMock{
 			DoFunc: func(req *http.Request) (*http.Response, error) {
 				return &http.Response{
-					Status:     "400 Unauthorized",
+					Status:     "401 Unauthorized",
 					StatusCode: 401,
 				}, nil
 			},
@@ -602,7 +602,7 @@ func TestGetHeaderReturnsError(t *testing.T) {
 	}
 }
 
-func TestGetHeaderNotAllowed(t *testing.T) {
+func TestGetHeaderBadRequest(t *testing.T) {
 	info := execResponseStageInfo{
 		Location:     "gcs/teststage/users/34/",
 		LocationType: "GCS",
@@ -616,7 +616,7 @@ func TestGetHeaderNotAllowed(t *testing.T) {
 		mockGcsClient: &clientMock{
 			DoFunc: func(req *http.Request) (*http.Response, error) {
 				return &http.Response{
-					Status:     "400 Not Allowed",
+					Status:     "400 Bad Request",
 					StatusCode: 400,
 				}, nil
 			},
@@ -759,3 +759,201 @@ func TestUploadFileWithBadRequest(t *testing.T) {
 			renewPresignedURL, uploadMeta.resStatus)
 	}
 }
+
+func TestGetFileHeaderEncryptionData(t *testing.T) {
+	mockEncDataResp := "{\"EncryptionMode\":\"FullBlob\",\"WrappedContentKey\": {\"KeyId\":\"symmKey1\",\"EncryptedKey\":\"testencryptedkey12345678910==\",\"Algorithm\":\"AES_CBC_256\"},\"EncryptionAgent\": {\"Protocol\":\"1.0\",\"EncryptionAlgorithm\":\"AES_CBC_256\"},\"ContentEncryptionIV\":\"testIVkey12345678910==\",\"KeyWrappingMetadata\":{\"EncryptionLibrary\":\"Java 5.3.0\"}}"
+	mockMatDesc := "{\"queryid\":\"01abc874-0406-1bf0-0000-53b10668e056\",\"smkid\":\"92019681909886\",\"key\":\"128\"}"
+	info := execResponseStageInfo{
+		Location:     "gcs/teststage/users/34/",
+		LocationType: "GCS",
+		Creds: execResponseCredentials{
+			GcsAccessToken: "test-token-124456577",
+		},
+	}
+	meta := fileMetadata{
+		client:    info.Creds.GcsAccessToken,
+		stageInfo: &info,
+		mockGcsClient: &clientMock{
+			DoFunc: func(req *http.Request) (*http.Response, error) {
+				return &http.Response{
+					Status:     "200 OK",
+					StatusCode: 200,
+					Header: http.Header{
+						"X-Goog-Meta-Encryptiondata": []string{mockEncDataResp},
+						"Content-Length":             []string{"4256"},
+						"X-Goog-Meta-Sfc-Digest":     []string{"123456789abcdef"},
+						"X-Goog-Meta-Matdesc":        []string{mockMatDesc},
+					},
+				}, nil
+			},
+		},
+	}
+	header, err := new(snowflakeGcsClient).getFileHeader(&meta, "file.txt")
+	if err != nil {
+		t.Fatal(err)
+	}
+	expectedFileHeader := &fileHeader{
+		digest:        "123456789abcdef",
+		contentLength: 4256,
+		encryptionMetadata: &encryptMetadata{
+			key:     "testencryptedkey12345678910==",
+			iv:      "testIVkey12345678910==",
+			matdesc: mockMatDesc,
+		},
+	}
+	if header.contentLength != expectedFileHeader.contentLength || header.digest != expectedFileHeader.digest || header.encryptionMetadata.iv != expectedFileHeader.encryptionMetadata.iv || header.encryptionMetadata.key != expectedFileHeader.encryptionMetadata.key || header.encryptionMetadata.matdesc != expectedFileHeader.encryptionMetadata.matdesc {
+		t.Fatalf("unexpected file header. expected: %v, got: %v", expectedFileHeader, header)
+	}
+}
+
+func TestUploadFileToGcsNoStatus(t *testing.T) {
+	info := execResponseStageInfo{
+		Location:     "gcs-blob/storage/users/456/",
+		LocationType: "GCS",
+	}
+	encMat := snowflakeFileEncryption{
+		QueryStageMasterKey: "abCdEFO0upIT36dAxGsa0w==",
+		QueryID:             "01abc874-0406-1bf0-0000-53b10668e056",
+		SMKID:               92019681909886,
+	}
+	initialParallel := int64(100)
+	dir, err := os.Getwd()
+	if err != nil {
+		t.Error(err)
+	}
+
+	gcsCli, err := new(snowflakeGcsClient).createClient(&info, false)
+	if err != nil {
+		t.Error(err)
+	}
+	uploadMeta := fileMetadata{
+		name:               "data1.txt.gz",
+		stageLocationType:  "GCS",
+		noSleepingTime:     true,
+		parallel:           initialParallel,
+		client:             gcsCli,
+		sha256Digest:       "123456789abcdef",
+		stageInfo:          &info,
+		dstFileName:        "data1.txt.gz",
+		srcFileName:        path.Join(dir, "/test_data/put_get_1.txt"),
+		overwrite:          true,
+		dstCompressionType: compressionTypes["GZIP"],
+		encryptionMaterial: &encMat,
+		options: &SnowflakeFileTransferOptions{
+			MultiPartThreshold: dataSizeThreshold,
+		},
+		mockGcsClient: &clientMock{
+			DoFunc: func(req *http.Request) (*http.Response, error) {
+				return &http.Response{
+					Status:     "401 Unauthorized",
+					StatusCode: 401,
+				}, nil
+			},
+		},
+	}
+
+	uploadMeta.realSrcFileName = uploadMeta.srcFileName
+	fi, err := os.Stat(uploadMeta.srcFileName)
+	if err != nil {
+		t.Error(err)
+	}
+	uploadMeta.uploadSize = fi.Size()
+
+	err = new(remoteStorageUtil).uploadOneFile(&uploadMeta)
+	if err == nil {
+		t.Error("should have raised an error")
+	}
+}
+
+func TestDownloadFileFromGcsError(t *testing.T) {
+	info := execResponseStageInfo{
+		Location:     "gcs/teststage/users/34/",
+		LocationType: "GCS",
+	}
+	dir, err := os.Getwd()
+	if err != nil {
+		t.Error(err)
+	}
+
+	gcsCli, err := new(snowflakeGcsClient).createClient(&info, false)
+	if err != nil {
+		t.Error(err)
+	}
+
+	downloadMeta := fileMetadata{
+		name:              "data1.txt.gz",
+		stageLocationType: "GCS",
+		noSleepingTime:    true,
+		client:            gcsCli,
+		stageInfo:         &info,
+		dstFileName:       "data1.txt.gz",
+		overwrite:         true,
+		srcFileName:       "data1.txt.gz",
+		localLocation:     dir,
+		options: &SnowflakeFileTransferOptions{
+			MultiPartThreshold: dataSizeThreshold,
+		},
+		mockGcsClient: &clientMock{
+			DoFunc: func(req *http.Request) (*http.Response, error) {
+				return &http.Response{
+					Status:     "403 Unauthorized",
+					StatusCode: 401,
+				}, nil
+			},
+		},
+		resStatus: downloaded, // bypass file header request
+	}
+	err = new(remoteStorageUtil).downloadOneFile(&downloadMeta)
+	if err == nil {
+		t.Error("should have raised an error")
+	}
+}
+
+func TestDownloadFileWithBadRequest(t *testing.T) {
+	info := execResponseStageInfo{
+		Location:     "gcs/teststage/users/34/",
+		LocationType: "GCS",
+	}
+	dir, err := os.Getwd()
+	if err != nil {
+		t.Error(err)
+	}
+
+	gcsCli, err := new(snowflakeGcsClient).createClient(&info, false)
+	if err != nil {
+		t.Error(err)
+	}
+
+	downloadMeta := fileMetadata{
+		name:              "data1.txt.gz",
+		stageLocationType: "GCS",
+		noSleepingTime:    true,
+		client:            gcsCli,
+		stageInfo:         &info,
+		dstFileName:       "data1.txt.gz",
+		overwrite:         true,
+		srcFileName:       "data1.txt.gz",
+		localLocation:     dir,
+		options: &SnowflakeFileTransferOptions{
+			MultiPartThreshold: dataSizeThreshold,
+		},
+		mockGcsClient: &clientMock{
+			DoFunc: func(req *http.Request) (*http.Response, error) {
+				return &http.Response{
+					Status:     "400 Bad Request",
+					StatusCode: 400,
+				}, nil
+			},
+		},
+		resStatus: downloaded, // bypass file header request
+	}
+	err = new(remoteStorageUtil).downloadOneFile(&downloadMeta)
+	if err == nil {
+		t.Error("should have raised an error")
+	}
+
+	if downloadMeta.resStatus != renewPresignedURL {
+		t.Fatalf("expected %v result status, got: %v",
+			renewPresignedURL, downloadMeta.resStatus)
+	}
+}
diff --git a/storage_client.go b/storage_client.go
@@ -226,23 +226,6 @@ func (rsu *remoteStorageUtil) downloadOneFile(meta *fileMetadata) error {
 				meta.dstFileSize = fi.Size()
 			}
 			return nil
-		} else if meta.resStatus == renewPresignedURL {
-			return nil
-		} else if meta.resStatus == renewToken {
-			return nil
-		} else if meta.resStatus == needRetryWithLowerConcurrency {
-			maxConcurrency = meta.parallel - int64(retry)*meta.parallel/int64(maxRetry)
-			maxConcurrency = int64Max(defaultConcurrency, maxConcurrency)
-			meta.lastMaxConcurrency = int(maxConcurrency)
-			if !meta.noSleepingTime {
-				sleepingTime := intMin(int(math.Exp2(float64(retry))), 16)
-				time.Sleep(time.Duration(sleepingTime) * time.Second)
-			}
-		} else if meta.resStatus == needRetry {
-			if !meta.noSleepingTime {
-				sleepingTime := intMin(int(math.Exp2(float64(retry))), 16)
-				time.Sleep(time.Duration(sleepingTime) * time.Second)
-			}
 		}
 		lastErr = meta.lastError
 	}