diff --git a/sdks/go/pkg/beam/options/gcpopts/options.go b/sdks/go/pkg/beam/options/gcpopts/options.go index b5df80e2de06..d75585047e18 100644 --- a/sdks/go/pkg/beam/options/gcpopts/options.go +++ b/sdks/go/pkg/beam/options/gcpopts/options.go @@ -22,11 +22,11 @@ import ( "os" "os/exec" "strings" - "sync" "cloud.google.com/go/storage" - "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" "github.com/apache/beam/sdks/v2/go/pkg/beam/log" + "github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcpx" + "github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcsx" "golang.org/x/oauth2/google" ) @@ -49,24 +49,7 @@ var ( // option, which is how the value reaches remote workers, where flags are not // parsed. Returns an empty string if no override is configured. func GetGcsEndpoint() string { - if *GcsEndpoint != "" { - return *GcsEndpoint - } - return runtime.GlobalOptions.Get("gcs_endpoint") -} - -const universeDomainEnvVar = "GOOGLE_CLOUD_UNIVERSE_DOMAIN" - -// findDefaultCredentials is a variable to allow overriding in tests. -var findDefaultCredentials = google.FindDefaultCredentials - -// cachedUniverseDomain caches the universe domain resolved from the default -// credentials, since resolving it may require a metadata server round trip -// and clients are created frequently. -var cachedUniverseDomain struct { - mu sync.Mutex - resolved bool - value string + return gcsx.GetGcsEndpoint() } // GetUniverseDomain returns the Google Cloud universe domain, first from the @@ -79,28 +62,7 @@ var cachedUniverseDomain struct { // option.WithUniverseDomain to support non-default universes such as Trusted // Partner Cloud (TPC). func GetUniverseDomain(ctx context.Context) string { - if ud := os.Getenv(universeDomainEnvVar); ud != "" { - return ud - } - - cachedUniverseDomain.mu.Lock() - defer cachedUniverseDomain.mu.Unlock() - if cachedUniverseDomain.resolved { - return cachedUniverseDomain.value - } - creds, err := findDefaultCredentials(ctx) - if err != nil { - // Don't cache: credentials may become available later, and client - // construction surfaces the credential error itself. - return "" - } - ud, err := creds.GetUniverseDomain() - if err != nil { - return "" - } - cachedUniverseDomain.resolved = true - cachedUniverseDomain.value = ud - return ud + return gcpx.GetUniverseDomain(ctx) } // GetProject returns the project, if non empty and exits otherwise. diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go index 841e17498769..ee1dd1b8cc03 100644 --- a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go +++ b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go @@ -27,7 +27,7 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" "github.com/apache/beam/sdks/v2/go/pkg/beam/log" pipepb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/pipeline_v1" - "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" + "github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcpx" df "google.golang.org/api/dataflow/v1b3" "google.golang.org/api/option" "google.golang.org/protobuf/proto" @@ -325,7 +325,7 @@ func NewClient(ctx context.Context, endpoint string) (*df.Service, error) { } // The client ignores the credentials' universe domain, and requests fail // if it differs from the configured one (googleapis.com by default). - if ud := gcpopts.GetUniverseDomain(ctx); ud != "" { + if ud := gcpx.GetUniverseDomain(ctx); ud != "" { opts = append(opts, option.WithUniverseDomain(ud)) } return df.NewService(ctx, opts...) diff --git a/sdks/go/pkg/beam/util/gcpx/gcp.go b/sdks/go/pkg/beam/util/gcpx/gcp.go new file mode 100644 index 000000000000..e8b5af04ae06 --- /dev/null +++ b/sdks/go/pkg/beam/util/gcpx/gcp.go @@ -0,0 +1,73 @@ +// Licensed to the Apache Software Foundation (ASF) under one or more +// contributor license agreements. See the NOTICE file distributed with +// this work for additional information regarding copyright ownership. +// The ASF licenses this file to You under the Apache License, Version 2.0 +// (the "License"); you may not use this file except in compliance with +// the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// Package gcpx contains utilities for working with Google Cloud Platform (GCP). +package gcpx + +import ( + "context" + "os" + "sync" + + "golang.org/x/oauth2/google" +) + +const universeDomainEnvVar = "GOOGLE_CLOUD_UNIVERSE_DOMAIN" + +// findDefaultCredentials is a variable to allow overriding in tests. +var findDefaultCredentials = google.FindDefaultCredentials + +// cachedUniverseDomain caches the universe domain resolved from the default +// credentials, since resolving it may require a metadata server round trip +// and clients are created frequently. +var cachedUniverseDomain struct { + mu sync.Mutex + resolved bool + value string +} + +// GetUniverseDomain returns the Google Cloud universe domain, first from the +// GOOGLE_CLOUD_UNIVERSE_DOMAIN environment variable, then falling back to the +// universe domain of the default credentials. Returns an empty string if it +// cannot be determined. +// +// Google API clients don't infer the universe domain from the credentials, +// and requests fail if the two differ, so clients should pass this value via +// option.WithUniverseDomain to support non-default universes such as Trusted +// Partner Cloud (TPC). +func GetUniverseDomain(ctx context.Context) string { + if ud := os.Getenv(universeDomainEnvVar); ud != "" { + return ud + } + + cachedUniverseDomain.mu.Lock() + defer cachedUniverseDomain.mu.Unlock() + if cachedUniverseDomain.resolved { + return cachedUniverseDomain.value + } + creds, err := findDefaultCredentials(ctx) + if err != nil { + // Don't cache: credentials may become available later, and client + // construction surfaces the credential error itself. + return "" + } + ud, err := creds.GetUniverseDomain() + if err != nil { + return "" + } + cachedUniverseDomain.resolved = true + cachedUniverseDomain.value = ud + return ud +} diff --git a/sdks/go/pkg/beam/options/gcpopts/universe_test.go b/sdks/go/pkg/beam/util/gcpx/gcp_test.go similarity index 77% rename from sdks/go/pkg/beam/options/gcpopts/universe_test.go rename to sdks/go/pkg/beam/util/gcpx/gcp_test.go index 35050d1010e6..e8ba73085cdb 100644 --- a/sdks/go/pkg/beam/options/gcpopts/universe_test.go +++ b/sdks/go/pkg/beam/util/gcpx/gcp_test.go @@ -13,14 +13,13 @@ // See the License for the specific language governing permissions and // limitations under the License. -package gcpopts +package gcpx import ( "context" "errors" "testing" - "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" "golang.org/x/oauth2/google" ) @@ -30,6 +29,7 @@ func TestGetUniverseDomain(t *testing.T) { env string credsUD string credsErr error + udErr error want string wantCalls int // credential lookups across two calls }{ @@ -51,6 +51,12 @@ func TestGetUniverseDomain(t *testing.T) { want: "", wantCalls: 2, }, + { + name: "universe domain provider errors are not cached", + udErr: errors.New("metadata server error"), + want: "", + wantCalls: 2, + }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { @@ -64,7 +70,7 @@ func TestGetUniverseDomain(t *testing.T) { return nil, test.credsErr } return &google.Credentials{ - UniverseDomainProvider: func() (string, error) { return test.credsUD, nil }, + UniverseDomainProvider: func() (string, error) { return test.credsUD, test.udErr }, }, nil } t.Cleanup(func() { findDefaultCredentials = orig }) @@ -96,25 +102,3 @@ func resetUniverseDomainCache(t *testing.T) { reset() t.Cleanup(reset) } - -func TestGetGcsEndpoint(t *testing.T) { - const flagValue = "https://flag.example.com/storage/v1/" - const optionValue = "https://option.example.com/storage/v1/" - - origFlag := *GcsEndpoint - t.Cleanup(func() { - *GcsEndpoint = origFlag - runtime.GlobalOptions.Set("gcs_endpoint", "") - }) - // On workers, the value arrives as a pipeline option rather than a flag. - runtime.GlobalOptions.Set("gcs_endpoint", optionValue) - - *GcsEndpoint = "" - if got := GetGcsEndpoint(); got != optionValue { - t.Errorf("GetGcsEndpoint() with flag unset = %q, want pipeline option %q", got, optionValue) - } - *GcsEndpoint = flagValue - if got := GetGcsEndpoint(); got != flagValue { - t.Errorf("GetGcsEndpoint() with flag set = %q, want flag %q", got, flagValue) - } -} diff --git a/sdks/go/pkg/beam/util/gcsx/client_test.go b/sdks/go/pkg/beam/util/gcsx/client_test.go index 6b31646d6601..3ec62a3a734e 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -16,14 +16,42 @@ package gcsx import ( + "flag" "net/http" "net/http/httptest" "strings" "testing" - "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" ) +func TestGetGcsEndpoint(t *testing.T) { + const flagValue = "https://flag.example.com/storage/v1/" + const optionValue = "https://option.example.com/storage/v1/" + + origCommandLine := flag.CommandLine + flag.CommandLine = flag.NewFlagSet(t.Name(), flag.ContinueOnError) + t.Cleanup(func() { + flag.CommandLine = origCommandLine + runtime.GlobalOptions.Set("gcs_endpoint", "") + }) + // On workers, the value arrives as a pipeline option rather than a flag. + runtime.GlobalOptions.Set("gcs_endpoint", optionValue) + + if got := GetGcsEndpoint(); got != optionValue { + t.Errorf("GetGcsEndpoint() with flag unregistered = %q, want pipeline option %q", got, optionValue) + } + + gcsEndpoint := flag.String("gcs_endpoint", "", "") + if got := GetGcsEndpoint(); got != optionValue { + t.Errorf("GetGcsEndpoint() with flag unset = %q, want pipeline option %q", got, optionValue) + } + *gcsEndpoint = flagValue + if got := GetGcsEndpoint(); got != flagValue { + t.Errorf("GetGcsEndpoint() with flag set = %q, want flag %q", got, flagValue) + } +} + func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { var gotPath string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -33,9 +61,8 @@ func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { })) defer srv.Close() - orig := *gcpopts.GcsEndpoint - *gcpopts.GcsEndpoint = srv.URL + "/storage/v1/" - t.Cleanup(func() { *gcpopts.GcsEndpoint = orig }) + runtime.GlobalOptions.Set("gcs_endpoint", srv.URL+"/storage/v1/") + t.Cleanup(func() { runtime.GlobalOptions.Set("gcs_endpoint", "") }) ctx := t.Context() client, err := NewUnauthenticatedClient(ctx) diff --git a/sdks/go/pkg/beam/util/gcsx/gcs.go b/sdks/go/pkg/beam/util/gcsx/gcs.go index 73300fcc6dc1..9bbc14567e67 100644 --- a/sdks/go/pkg/beam/util/gcsx/gcs.go +++ b/sdks/go/pkg/beam/util/gcsx/gcs.go @@ -18,14 +18,16 @@ package gcsx import ( "context" + "flag" "fmt" "io" "net/url" "path" "cloud.google.com/go/storage" + "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" "github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors" - "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" + "github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcpx" "google.golang.org/api/option" ) @@ -43,7 +45,7 @@ func NewClient(ctx context.Context, scope string) (*storage.Client, error) { opts = append(opts, endpointOptions()...) // storage.NewClient ignores the credentials' universe domain, and requests // fail if it differs from the configured one (googleapis.com by default). - if ud := gcpopts.GetUniverseDomain(ctx); ud != "" { + if ud := gcpx.GetUniverseDomain(ctx); ud != "" { opts = append(opts, option.WithUniverseDomain(ud)) } return storage.NewClient(ctx, opts...) @@ -60,12 +62,27 @@ func NewUnauthenticatedClient(ctx context.Context) (*storage.Client, error) { // endpointOptions returns the client options for the --gcs_endpoint override, // if any. func endpointOptions() []option.ClientOption { - if ep := gcpopts.GetGcsEndpoint(); ep != "" { + if ep := GetGcsEndpoint(); ep != "" { return []option.ClientOption{option.WithEndpoint(ep)} } return nil } +// GetGcsEndpoint returns the GCS endpoint override from the --gcs_endpoint +// flag. If the flag is unset or not registered, it falls back to the +// "gcs_endpoint" pipeline option, which is how the value reaches remote +// workers, where flags are not parsed. Returns an empty string if no override +// is configured. +// +// Looking up the flag by name avoids importing gcpopts, which registers global +// flags such as --region and --project that can conflict with user binaries. +func GetGcsEndpoint() string { + if f := flag.Lookup("gcs_endpoint"); f != nil && f.Value.String() != "" { + return f.Value.String() + } + return runtime.GlobalOptions.Get("gcs_endpoint") +} + // Upload writes the given content to GCS. If the specified bucket does not // exist, it is created first. Returns the full path of the object. func Upload(ctx context.Context, client *storage.Client, project, bucket, object string, r io.Reader) (string, error) {