Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1789753: Support GCS region specific endpoint #1271

Draft
wants to merge 3 commits into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 30 additions & 2 deletions gcs_storage_client.go
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ import (
"os"
"strconv"
"strings"

"cloud.google.com/go/storage"
"google.golang.org/api/option"
)

const (
Expand All @@ -20,6 +23,7 @@ const (
gcsMetadataMatdescKey = gcsMetadataPrefix + "matdesc"
gcsMetadataEncryptionDataProp = gcsMetadataPrefix + "encryptiondata"
gcsFileHeaderDigest = "gcs-file-header-digest"
gcsRegionMeCentral2 = "me-central2"
)

type snowflakeGcsClient struct {
Expand Down Expand Up @@ -79,6 +83,7 @@ func (util *snowflakeGcsClient) getFileHeader(meta *fileMetadata, filename strin
}
return client.Do(req)
})

if err != nil {
return nil, err
}
Expand Down Expand Up @@ -402,8 +407,31 @@ func (util *snowflakeGcsClient) isTokenExpired(resp *http.Response) bool {
return resp.StatusCode == 401
}

func newGcsClient() gcsAPI {
return &http.Client{
func newGcsClient(info *execResponseStageInfo) (gcsAPI, error) {
httpClient := &http.Client{
Transport: SnowflakeTransport,
}

// TODO: SNOW-1789759 hardcoded region will be replaced in the future
endpoint := getGcsCustomEndpoint(info)
_, err := storage.NewClient(context.Background(), option.WithHTTPClient(httpClient))
if endpoint != "" {
_, err = storage.NewClient(context.Background(), option.WithHTTPClient(httpClient), option.WithEndpoint(endpoint))
}
if err != nil {
return nil, err
}

return httpClient, nil
}

func getGcsCustomEndpoint(info *execResponseStageInfo) string {
// TODO: SNOW-1789759 hardcoded region will be replaced in the future
isRegionalURLEnabled := (strings.ToLower(info.Region) == gcsRegionMeCentral2) || info.UseRegionalURL
if info.EndPoint != "" {
return fmt.Sprintf("https://%s", info.EndPoint)
} else if info.Region != "" && isRegionalURLEnabled {
return fmt.Sprintf("https://storage.%s.rep.googleapis.com", strings.ToLower(info.Region))
}
return ""
}
81 changes: 81 additions & 0 deletions gcs_storage_client_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1126,3 +1126,84 @@ func Test_snowflakeGcsClient_nativeDownloadFile(t *testing.T) {
t.Error("should have raised an error")
}
}

func TestGetGcsCustomEndpoint(t *testing.T) {
testcases := []struct {
desc string
in execResponseStageInfo
out string
}{
{
desc: "when the useRegionalURL is only enabled",
in: execResponseStageInfo{
UseRegionalURL: true,
EndPoint: "",
Region: "mockLocation",
},
out: "https://storage.mocklocation.rep.googleapis.com",
},
{
desc: "when the region is me-central2",
in: execResponseStageInfo{
UseRegionalURL: false,
EndPoint: "",
Region: "me-central2",
},
out: "https://storage.me-central2.rep.googleapis.com",
},
{
desc: "when the region is me-central2 (mixed case)",
in: execResponseStageInfo{
UseRegionalURL: false,
EndPoint: "",
Region: "ME-cEntRal2",
},
out: "https://storage.me-central2.rep.googleapis.com",
},
{
desc: "when the region is me-central2 (uppercase)",
in: execResponseStageInfo{
UseRegionalURL: false,
EndPoint: "",
Region: "ME-CENTRAL2",
},
out: "https://storage.me-central2.rep.googleapis.com",
},
{
desc: "when the endPoint is specified",
in: execResponseStageInfo{
UseRegionalURL: false,
EndPoint: "storage.specialEndPoint.rep.googleapis.com",
Region: "ME-cEntRal1",
},
out: "https://storage.specialEndPoint.rep.googleapis.com",
},
{
desc: "when both the endPoint and the useRegionalUrl are specified",
in: execResponseStageInfo{
UseRegionalURL: true,
EndPoint: "storage.specialEndPoint.rep.googleapis.com",
Region: "ME-cEntRal1",
},
out: "https://storage.specialEndPoint.rep.googleapis.com",
},
{
desc: "when both the endPoint is specified and the region is me-central2",
in: execResponseStageInfo{
UseRegionalURL: true,
EndPoint: "storage.specialEndPoint.rep.googleapis.com",
Region: "ME-CENTRAL2",
},
out: "https://storage.specialEndPoint.rep.googleapis.com",
},
}

for _, test := range testcases {
t.Run(test.desc, func(t *testing.T) {
endpoint := getGcsCustomEndpoint(&test.in)
if endpoint != test.out {
t.Errorf("failed. in: %v, expected: %v, got: %v", test.in, test.out, endpoint)
}
})
}
}
59 changes: 50 additions & 9 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ module github.com/snowflakedb/gosnowflake
go 1.21

require (
cloud.google.com/go/storage v1.48.0
github.com/99designs/keyring v1.2.2
github.com/Azure/azure-sdk-for-go/sdk/azcore v1.4.0
github.com/Azure/azure-sdk-for-go/sdk/storage/azblob v1.0.0
Expand All @@ -17,13 +18,24 @@ require (
github.com/golang-jwt/jwt/v5 v5.2.1
github.com/pkg/browser v0.0.0-20210911075715-681adbf594b8
github.com/sirupsen/logrus v1.9.3
golang.org/x/crypto v0.22.0
golang.org/x/sys v0.19.0
golang.org/x/crypto v0.29.0
golang.org/x/sys v0.27.0
google.golang.org/api v0.210.0
)

require (
cel.dev/expr v0.16.1 // indirect
cloud.google.com/go v0.116.0 // indirect
cloud.google.com/go/auth v0.11.0 // indirect
cloud.google.com/go/auth/oauth2adapt v0.2.6 // indirect
cloud.google.com/go/compute/metadata v0.5.2 // indirect
cloud.google.com/go/iam v1.2.2 // indirect
cloud.google.com/go/monitoring v1.21.2 // indirect
github.com/99designs/go-keychain v0.0.0-20191008050251-8e49817e8af4 // indirect
github.com/Azure/azure-sdk-for-go/sdk/internal v1.1.2 // indirect
github.com/GoogleCloudPlatform/opentelemetry-operations-go/detectors/gcp v1.24.1 // indirect
github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.48.1 // indirect
github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.48.1 // indirect
github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c // indirect
github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.2 // indirect
github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.5 // indirect
Expand All @@ -33,26 +45,55 @@ require (
github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.7 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.7 // indirect
github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.5 // indirect
github.com/census-instrumentation/opencensus-proto v0.4.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cncf/xds/go v0.0.0-20240905190251-b4127c9b8d78 // indirect
github.com/danieljoos/wincred v1.1.2 // indirect
github.com/dvsekhvalnov/jose2go v1.6.0 // indirect
github.com/envoyproxy/go-control-plane v0.13.0 // indirect
github.com/envoyproxy/protoc-gen-validate v1.1.0 // indirect
github.com/felixge/httpsnoop v1.0.4 // indirect
github.com/go-logr/logr v1.4.2 // indirect
github.com/go-logr/stdr v1.2.2 // indirect
github.com/goccy/go-json v0.10.2 // indirect
github.com/godbus/dbus v0.0.0-20190726142602-4481cbc300e2 // indirect
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
github.com/google/flatbuffers v24.3.25+incompatible // indirect
github.com/google/s2a-go v0.1.8 // indirect
github.com/google/uuid v1.6.0 // indirect
github.com/googleapis/enterprise-certificate-proxy v0.3.4 // indirect
github.com/googleapis/gax-go/v2 v2.14.0 // indirect
github.com/gsterjov/go-libsecret v0.0.0-20161001094733-a6f4afe4910c // indirect
github.com/jmespath/go-jmespath v0.4.0 // indirect
github.com/klauspost/compress v1.17.8 // indirect
github.com/klauspost/cpuid/v2 v2.2.7 // indirect
github.com/kr/pretty v0.3.1 // indirect
github.com/mtibben/percent v0.2.1 // indirect
github.com/pierrec/lz4/v4 v4.1.21 // indirect
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
github.com/zeebo/xxh3 v1.0.2 // indirect
go.opencensus.io v0.24.0 // indirect
go.opentelemetry.io/contrib/detectors/gcp v1.29.0 // indirect
go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc v0.54.0 // indirect
go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp v0.54.0 // indirect
go.opentelemetry.io/otel v1.29.0 // indirect
go.opentelemetry.io/otel/metric v1.29.0 // indirect
go.opentelemetry.io/otel/sdk v1.29.0 // indirect
go.opentelemetry.io/otel/sdk/metric v1.29.0 // indirect
go.opentelemetry.io/otel/trace v1.29.0 // indirect
golang.org/x/exp v0.0.0-20240222234643-814bf88cf225 // indirect
golang.org/x/mod v0.17.0 // indirect
golang.org/x/net v0.24.0 // indirect
golang.org/x/sync v0.7.0 // indirect
golang.org/x/term v0.19.0 // indirect
golang.org/x/text v0.14.0 // indirect
golang.org/x/tools v0.20.0 // indirect
golang.org/x/net v0.31.0 // indirect
golang.org/x/oauth2 v0.24.0 // indirect
golang.org/x/sync v0.9.0 // indirect
golang.org/x/term v0.26.0 // indirect
golang.org/x/text v0.20.0 // indirect
golang.org/x/time v0.8.0 // indirect
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d // indirect
golang.org/x/xerrors v0.0.0-20231012003039-104605ab7028 // indirect
gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect
google.golang.org/genproto v0.0.0-20241118233622-e639e219e697 // indirect
google.golang.org/genproto/googleapis/api v0.0.0-20241113202542-65e8d215514f // indirect
google.golang.org/genproto/googleapis/rpc v0.0.0-20241118233622-e639e219e697 // indirect
google.golang.org/grpc v1.67.2 // indirect
google.golang.org/grpc/stats/opentelemetry v0.0.0-20240907200651-3ffb98b2c93a // indirect
google.golang.org/protobuf v1.35.2 // indirect
)
Loading
Loading