Repository: snowflakedb/gosnowflake
Pull request #1402 (2025-05-07 05:15:57)

SNOW-2032700 Go driver support virtual-style domains(2)
### Description

SNOW-XXX Please explain the changes you made here.
- [x] Added the condition when the endpoint and the useVirtualUrl are enabled.
- [x] added testing cases.

### Checklist
- [ ] Created tests which fail without the change (if possible)
- [ ] Extended the README / documentation, if necessary


sfc-gh-snowflakedb-snyk-sa: ![](https://res.cloudinary.com/snyk/image/upload/r-d/scm-platform/snyk-pull-requests/pr-banner-default.svg)

### :tada: **Snyk checks have passed. No issues have been found so far.**

:white_check_mark: **security/snyk** check is complete. No issues have been found. [(View Details)](https://app.snyk.io/org/snowflakedb-sca-scanning-public-repo/pr-checks/877f49b8-27cd-4bdf-a06d-036a16c103ca?refs=4b8b60c31d12d31ab14788032c11fc1b7f289867&source=prChecksComment)

:white_check_mark: **license/snyk** check is complete. No issues have been found. [(View Details)](https://app.snyk.io/org/snowflakedb-sca-scanning-public-repo/pr-checks/877f49b8-27cd-4bdf-a06d-036a16c103ca/license?refs=4b8b60c31d12d31ab14788032c11fc1b7f289867&source=prChecksComment)




Patch:
diff --git a/gcs_storage_client.go b/gcs_storage_client.go
index 64b5079..596b665 100644
--- a/gcs_storage_client.go
+++ b/gcs_storage_client.go
@@ -394,10 +394,20 @@ func (util *snowflakeGcsClient) extractBucketNameAndPath(location string) *gcsLo
 func (util *snowflakeGcsClient) generateFileURL(stageInfo *execResponseStageInfo, filename string) (*url.URL, error) {
 	gcsLoc := util.extractBucketNameAndPath(stageInfo.Location)
 	fullFilePath := gcsLoc.path + filename
-	endPoint := util.getGcsCustomEndpoint(stageInfo)
-	if stageInfo.UseVirtualURL {
+	endPoint := "https://storage.googleapis.com"
+
+	// TODO: SNOW-1789759 hardcoded region will be replaced in the future
+	isRegionalURLEnabled := (strings.ToLower(stageInfo.Region) == gcsRegionMeCentral2) || stageInfo.UseRegionalURL
+	if stageInfo.EndPoint != "" {
+		endPoint = fmt.Sprintf("https://%s", stageInfo.EndPoint)
+	} else if stageInfo.UseVirtualURL {
+		bucketName := util.extractBucketNameAndPath(stageInfo.Location).bucketName
+		endPoint = fmt.Sprintf("https://%s.storage.googleapis.com", bucketName)
 		return url.Parse(endPoint + "/" + url.QueryEscape(fullFilePath))
+	} else if stageInfo.Region != "" && isRegionalURLEnabled {
+		endPoint = fmt.Sprintf("https://storage.%s.rep.googleapis.com", strings.ToLower(stageInfo.Region))
 	}
+
 	return url.Parse(endPoint + "/" + gcsLoc.bucketName + "/" + url.QueryEscape(fullFilePath))
 }
 
@@ -410,19 +420,3 @@ func newGcsClient(cfg *Config) gcsAPI {
 		Transport: getTransport(cfg),
 	}
 }
-
-func (util *snowflakeGcsClient) getGcsCustomEndpoint(info *execResponseStageInfo) string {
-	endpoint := "https://storage.googleapis.com"
-
-	// TODO: SNOW-1789759 hardcoded region will be replaced in the future
-	isRegionalURLEnabled := (strings.ToLower(info.Region) == gcsRegionMeCentral2) || info.UseRegionalURL
-	if info.EndPoint != "" {
-		endpoint = fmt.Sprintf("https://%s", info.EndPoint)
-	} else if info.UseVirtualURL {
-		bucketName := util.extractBucketNameAndPath(info.Location).bucketName
-		endpoint = fmt.Sprintf("https://%s.storage.googleapis.com", bucketName)
-	} else if info.Region != "" && isRegionalURLEnabled {
-		endpoint = fmt.Sprintf("https://storage.%s.rep.googleapis.com", strings.ToLower(info.Region))
-	}
-	return endpoint
-}


Test patch:
diff --git a/gcs_storage_client_test.go b/gcs_storage_client_test.go
index 18f4f6f..bcd56c5 100644
--- a/gcs_storage_client_test.go
+++ b/gcs_storage_client_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"net/http"
 	"net/url"
 	"os"
@@ -113,6 +114,18 @@ func TestGenerateFileURL(t *testing.T) {
 		})
 	}
 
+	for _, test := range testcases {
+		t.Run(test.location, func(t *testing.T) {
+			stageInfo := &execResponseStageInfo{}
+			stageInfo.Location = test.location
+			gcsURL, err := gcsUtil.generateFileURL(stageInfo, test.fname)
+			assertNilF(t, err, "error should be nil")
+			expectedURL, err := url.Parse("https://storage.googleapis.com/" + test.bucket + "/" + url.QueryEscape(test.filepath))
+			assertNilF(t, err, "error should be nil")
+			assertEqualE(t, gcsURL.String(), expectedURL.String(), "failed. expected: %v but got: %v", expectedURL.String(), gcsURL.String())
+		})
+	}
+
 	for _, test := range testcases {
 		t.Run(test.location, func(t *testing.T) {
 			stageInfo := &execResponseStageInfo{}
@@ -1136,131 +1149,143 @@ func Test_snowflakeGcsClient_nativeDownloadFile(t *testing.T) {
 
 func TestGetGcsCustomEndpoint(t *testing.T) {
 	testcases := []struct {
-		desc string
-		in   execResponseStageInfo
-		out  string
+		desc            string
+		in              execResponseStageInfo
+		expectedFileURL string
 	}{
 		{
 			desc: "when the endPoint is not specified and UseRegionalURL is false",
 			in: execResponseStageInfo{
 				UseRegionalURL: false,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "",
 				Region:         "WEST-1",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.googleapis.com",
+			expectedFileURL: "https://storage.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when the useRegionalURL is only enabled",
 			in: execResponseStageInfo{
 				UseRegionalURL: true,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "",
 				Region:         "mockLocation",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.mocklocation.rep.googleapis.com",
+			expectedFileURL: "https://storage.mocklocation.rep.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when the region is me-central2",
 			in: execResponseStageInfo{
 				UseRegionalURL: false,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "",
 				Region:         "me-central2",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.me-central2.rep.googleapis.com",
+			expectedFileURL: "https://storage.me-central2.rep.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when the region is me-central2 (mixed case)",
 			in: execResponseStageInfo{
 				UseRegionalURL: false,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "",
 				Region:         "ME-cEntRal2",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.me-central2.rep.googleapis.com",
+			expectedFileURL: "https://storage.me-central2.rep.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when the region is me-central2 (uppercase)",
 			in: execResponseStageInfo{
 				UseRegionalURL: false,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "",
 				Region:         "ME-CENTRAL2",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.me-central2.rep.googleapis.com",
+			expectedFileURL: "https://storage.me-central2.rep.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when the endPoint is specified",
 			in: execResponseStageInfo{
 				UseRegionalURL: false,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "storage.specialEndPoint.rep.googleapis.com",
 				Region:         "ME-cEntRal1",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.specialEndPoint.rep.googleapis.com",
+			expectedFileURL: "https://storage.specialEndPoint.rep.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when both the endPoint and the useRegionalUrl are specified",
 			in: execResponseStageInfo{
 				UseRegionalURL: true,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "storage.specialEndPoint.rep.googleapis.com",
 				Region:         "ME-cEntRal1",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.specialEndPoint.rep.googleapis.com",
+			expectedFileURL: "https://storage.specialEndPoint.rep.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when both the endPoint is specified and the region is me-central2",
 			in: execResponseStageInfo{
 				UseRegionalURL: true,
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				EndPoint:       "storage.specialEndPoint.rep.googleapis.com",
 				Region:         "ME-CENTRAL2",
 				UseVirtualURL:  false,
 			},
-			out: "https://storage.specialEndPoint.rep.googleapis.com",
+			expectedFileURL: "https://storage.specialEndPoint.rep.googleapis.com/my-travel-maps",
 		},
 		{
 			desc: "when only the useVirtualUrl is enabled",
 			in: execResponseStageInfo{
-				Location:       "my-travel-maps",
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				UseRegionalURL: false,
 				EndPoint:       "",
 				Region:         "WEST-1",
 				UseVirtualURL:  true,
 			},
-			out: "https://my-travel-maps.storage.googleapis.com",
+			expectedFileURL: "https://my-travel-maps.storage.googleapis.com",
 		},
 		{
 			desc: "when both the useRegionalURL and useVirtualUrl are enabled",
 			in: execResponseStageInfo{
-				Location:       "my-travel-maps",
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				UseRegionalURL: true,
 				EndPoint:       "",
 				Region:         "ME-CENTRAL2",
 				UseVirtualURL:  true,
 			},
-			out: "https://my-travel-maps.storage.googleapis.com",
+			expectedFileURL: "https://my-travel-maps.storage.googleapis.com",
 		},
 		{
 			desc: "when all the options are enabled",
 			in: execResponseStageInfo{
-				Location:       "my-travel-maps",
+				Location:       "my-travel-maps/mock_directory/mock_path/",
 				UseRegionalURL: true,
 				EndPoint:       "storage.specialEndPoint.rep.googleapis.com",
 				Region:         "ME-CENTRAL2",
 				UseVirtualURL:  true,
 			},
-			out: "https://storage.specialEndPoint.rep.googleapis.com",
+			expectedFileURL: "https://storage.specialEndPoint.rep.googleapis.com/my-travel-maps",
 		},
 	}
 
 	for _, test := range testcases {
 		t.Run(test.desc, func(t *testing.T) {
-			endpoint := new(snowflakeGcsClient).getGcsCustomEndpoint(&test.in)
-			if endpoint != test.out {
-				t.Errorf("failed. in: %v, expected: %v, got: %v", test.in, test.out, endpoint)
-			}
+			gcs := new(snowflakeGcsClient)
+			fileURL, err := gcs.generateFileURL(&test.in, "mock_file")
+			assertNilF(t, err, "Should not fail")
+
+			expectedURL, err := url.Parse(test.expectedFileURL + "/" + url.QueryEscape("mock_directory/mock_path/mock_file"))
+			assertNilF(t, err, "Should not fail")
+
+			assertEqualF(t, fileURL.String(), expectedURL.String(), "failed. in: %v, expected: %v, got: %v", fmt.Sprintf("%v", test.in), expectedURL.String(), fileURL.String())
 		})
 	}
 }