From 2a6e07d768566537c06b118f0604798d09fa9239 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Thu, 22 Oct 2026 14:50:21 +0000 Subject: [PATCH 1/4] remove functions from gcpopts and move to gcsx --- sdks/go/pkg/beam/options/gcpopts/options.go | 45 +--------- .../pkg/beam/options/gcpopts/universe_test.go | 75 +---------------- .../beam/runners/dataflow/dataflowlib/job.go | 4 +- sdks/go/pkg/beam/util/gcsx/client_test.go | 83 ++++++++++++++++++- sdks/go/pkg/beam/util/gcsx/gcs.go | 73 +++++++++++++++- 5 files changed, 157 insertions(+), 123 deletions(-) diff --git a/sdks/go/pkg/beam/options/gcpopts/options.go b/sdks/go/pkg/beam/options/gcpopts/options.go index b5df80e2de06..cd2c6efefdfe 100644 --- a/sdks/go/pkg/beam/options/gcpopts/options.go +++ b/sdks/go/pkg/beam/options/gcpopts/options.go @@ -22,11 +22,10 @@ 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/gcsx" "golang.org/x/oauth2/google" ) @@ -49,24 +48,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 +61,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 gcsx.GetUniverseDomain(ctx) } // GetProject returns the project, if non empty and exits otherwise. diff --git a/sdks/go/pkg/beam/options/gcpopts/universe_test.go b/sdks/go/pkg/beam/options/gcpopts/universe_test.go index 35050d1010e6..8954453abc48 100644 --- a/sdks/go/pkg/beam/options/gcpopts/universe_test.go +++ b/sdks/go/pkg/beam/options/gcpopts/universe_test.go @@ -16,85 +16,16 @@ package gcpopts import ( - "context" - "errors" "testing" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" - "golang.org/x/oauth2/google" ) func TestGetUniverseDomain(t *testing.T) { - tests := []struct { - name string - env string - credsUD string - credsErr error - want string - wantCalls int // credential lookups across two calls - }{ - { - name: "env var takes precedence", - env: "env.example.com", - credsUD: "creds.example.com", - want: "env.example.com", - }, - { - name: "falls back to credentials and caches", - credsUD: "creds.example.com", - want: "creds.example.com", - wantCalls: 1, - }, - { - name: "credential errors are not cached", - credsErr: errors.New("no credentials"), - want: "", - wantCalls: 2, - }, + t.Setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN", "env.example.com") + if got, want := GetUniverseDomain(t.Context()), "env.example.com"; got != want { + t.Errorf("GetUniverseDomain() = %q, want %q", got, want) } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, test.env) - calls := 0 - orig := findDefaultCredentials - findDefaultCredentials = func(context.Context, ...string) (*google.Credentials, error) { - calls++ - if test.credsErr != nil { - return nil, test.credsErr - } - return &google.Credentials{ - UniverseDomainProvider: func() (string, error) { return test.credsUD, nil }, - }, nil - } - t.Cleanup(func() { findDefaultCredentials = orig }) - - // Call twice so wantCalls can verify caching: a successful lookup - // is reused by the second call, while a failed one is retried. - for i := 0; i < 2; i++ { - if got := GetUniverseDomain(t.Context()); got != test.want { - t.Errorf("GetUniverseDomain() = %q, want %q", got, test.want) - } - } - if calls != test.wantCalls { - t.Errorf("credentials looked up %d times, want %d", calls, test.wantCalls) - } - }) - } -} - -// resetUniverseDomainCache clears the cached universe domain and restores it -// when the test finishes. -func resetUniverseDomainCache(t *testing.T) { - t.Helper() - reset := func() { - cachedUniverseDomain.mu.Lock() - cachedUniverseDomain.resolved = false - cachedUniverseDomain.value = "" - cachedUniverseDomain.mu.Unlock() - } - reset() - t.Cleanup(reset) } func TestGetGcsEndpoint(t *testing.T) { diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go index 841e17498769..4c8c9e7237cd 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/gcsx" 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 := gcsx.GetUniverseDomain(ctx); ud != "" { opts = append(opts, option.WithUniverseDomain(ud)) } return df.NewService(ctx, opts...) diff --git a/sdks/go/pkg/beam/util/gcsx/client_test.go b/sdks/go/pkg/beam/util/gcsx/client_test.go index 6b31646d6601..08be9cc0e18f 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -16,12 +16,15 @@ package gcsx import ( + "context" + "errors" "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" + "golang.org/x/oauth2/google" ) func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { @@ -33,9 +36,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) @@ -51,3 +53,76 @@ func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { t.Errorf("request path = %q, want prefix %q", gotPath, want) } } + +func TestGetUniverseDomain(t *testing.T) { + tests := []struct { + name string + env string + credsUD string + credsErr error + want string + wantCalls int // credential lookups across two calls + }{ + { + name: "env var takes precedence", + env: "env.example.com", + credsUD: "creds.example.com", + want: "env.example.com", + }, + { + name: "falls back to credentials and caches", + credsUD: "creds.example.com", + want: "creds.example.com", + wantCalls: 1, + }, + { + name: "credential errors are not cached", + credsErr: errors.New("no credentials"), + want: "", + wantCalls: 2, + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, test.env) + calls := 0 + orig := findDefaultCredentials + findDefaultCredentials = func(context.Context, ...string) (*google.Credentials, error) { + calls++ + if test.credsErr != nil { + return nil, test.credsErr + } + return &google.Credentials{ + UniverseDomainProvider: func() (string, error) { return test.credsUD, nil }, + }, nil + } + t.Cleanup(func() { findDefaultCredentials = orig }) + + // Call twice so wantCalls can verify caching: a successful lookup + // is reused by the second call, while a failed one is retried. + for i := 0; i < 2; i++ { + if got := GetUniverseDomain(t.Context()); got != test.want { + t.Errorf("GetUniverseDomain() = %q, want %q", got, test.want) + } + } + if calls != test.wantCalls { + t.Errorf("credentials looked up %d times, want %d", calls, test.wantCalls) + } + }) + } +} + +// resetUniverseDomainCache clears the cached universe domain and restores it +// when the test finishes. +func resetUniverseDomainCache(t *testing.T) { + t.Helper() + reset := func() { + cachedUniverseDomain.mu.Lock() + cachedUniverseDomain.resolved = false + cachedUniverseDomain.value = "" + cachedUniverseDomain.mu.Unlock() + } + reset() + t.Cleanup(reset) +} diff --git a/sdks/go/pkg/beam/util/gcsx/gcs.go b/sdks/go/pkg/beam/util/gcsx/gcs.go index 73300fcc6dc1..172776b2eb1f 100644 --- a/sdks/go/pkg/beam/util/gcsx/gcs.go +++ b/sdks/go/pkg/beam/util/gcsx/gcs.go @@ -18,14 +18,18 @@ package gcsx import ( "context" + "flag" "fmt" "io" "net/url" + "os" "path" + "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/internal/errors" - "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" + "golang.org/x/oauth2/google" "google.golang.org/api/option" ) @@ -43,7 +47,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 := GetUniverseDomain(ctx); ud != "" { opts = append(opts, option.WithUniverseDomain(ud)) } return storage.NewClient(ctx, opts...) @@ -60,12 +64,75 @@ 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") +} + +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 +} + // 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) { From 71ac52f6a06758e1b430dff17447ca8b9bbf7e08 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Thu, 22 Oct 2026 19:42:47 +0000 Subject: [PATCH 2/4] switch to a separate package design --- sdks/go/pkg/beam/options/gcpopts/options.go | 3 +- .../beam/runners/dataflow/dataflowlib/job.go | 4 +- sdks/go/pkg/beam/util/gcpx/gcp.go | 73 ++++++++++++ sdks/go/pkg/beam/util/gcpx/gcp_test.go | 104 ++++++++++++++++++ sdks/go/pkg/beam/util/gcsx/client_test.go | 76 ------------- sdks/go/pkg/beam/util/gcsx/gcs.go | 54 +-------- 6 files changed, 183 insertions(+), 131 deletions(-) create mode 100644 sdks/go/pkg/beam/util/gcpx/gcp.go create mode 100644 sdks/go/pkg/beam/util/gcpx/gcp_test.go diff --git a/sdks/go/pkg/beam/options/gcpopts/options.go b/sdks/go/pkg/beam/options/gcpopts/options.go index cd2c6efefdfe..d75585047e18 100644 --- a/sdks/go/pkg/beam/options/gcpopts/options.go +++ b/sdks/go/pkg/beam/options/gcpopts/options.go @@ -25,6 +25,7 @@ import ( "cloud.google.com/go/storage" "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" ) @@ -61,7 +62,7 @@ func GetGcsEndpoint() string { // option.WithUniverseDomain to support non-default universes such as Trusted // Partner Cloud (TPC). func GetUniverseDomain(ctx context.Context) string { - return gcsx.GetUniverseDomain(ctx) + 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 4c8c9e7237cd..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/util/gcsx" + "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 := gcsx.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/util/gcpx/gcp_test.go b/sdks/go/pkg/beam/util/gcpx/gcp_test.go new file mode 100644 index 000000000000..e8ba73085cdb --- /dev/null +++ b/sdks/go/pkg/beam/util/gcpx/gcp_test.go @@ -0,0 +1,104 @@ +// 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 + +import ( + "context" + "errors" + "testing" + + "golang.org/x/oauth2/google" +) + +func TestGetUniverseDomain(t *testing.T) { + tests := []struct { + name string + env string + credsUD string + credsErr error + udErr error + want string + wantCalls int // credential lookups across two calls + }{ + { + name: "env var takes precedence", + env: "env.example.com", + credsUD: "creds.example.com", + want: "env.example.com", + }, + { + name: "falls back to credentials and caches", + credsUD: "creds.example.com", + want: "creds.example.com", + wantCalls: 1, + }, + { + name: "credential errors are not cached", + credsErr: errors.New("no credentials"), + 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) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, test.env) + calls := 0 + orig := findDefaultCredentials + findDefaultCredentials = func(context.Context, ...string) (*google.Credentials, error) { + calls++ + if test.credsErr != nil { + return nil, test.credsErr + } + return &google.Credentials{ + UniverseDomainProvider: func() (string, error) { return test.credsUD, test.udErr }, + }, nil + } + t.Cleanup(func() { findDefaultCredentials = orig }) + + // Call twice so wantCalls can verify caching: a successful lookup + // is reused by the second call, while a failed one is retried. + for i := 0; i < 2; i++ { + if got := GetUniverseDomain(t.Context()); got != test.want { + t.Errorf("GetUniverseDomain() = %q, want %q", got, test.want) + } + } + if calls != test.wantCalls { + t.Errorf("credentials looked up %d times, want %d", calls, test.wantCalls) + } + }) + } +} + +// resetUniverseDomainCache clears the cached universe domain and restores it +// when the test finishes. +func resetUniverseDomainCache(t *testing.T) { + t.Helper() + reset := func() { + cachedUniverseDomain.mu.Lock() + cachedUniverseDomain.resolved = false + cachedUniverseDomain.value = "" + cachedUniverseDomain.mu.Unlock() + } + reset() + t.Cleanup(reset) +} diff --git a/sdks/go/pkg/beam/util/gcsx/client_test.go b/sdks/go/pkg/beam/util/gcsx/client_test.go index 08be9cc0e18f..af15c20da740 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -16,15 +16,12 @@ package gcsx import ( - "context" - "errors" "net/http" "net/http/httptest" "strings" "testing" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" - "golang.org/x/oauth2/google" ) func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { @@ -53,76 +50,3 @@ func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { t.Errorf("request path = %q, want prefix %q", gotPath, want) } } - -func TestGetUniverseDomain(t *testing.T) { - tests := []struct { - name string - env string - credsUD string - credsErr error - want string - wantCalls int // credential lookups across two calls - }{ - { - name: "env var takes precedence", - env: "env.example.com", - credsUD: "creds.example.com", - want: "env.example.com", - }, - { - name: "falls back to credentials and caches", - credsUD: "creds.example.com", - want: "creds.example.com", - wantCalls: 1, - }, - { - name: "credential errors are not cached", - credsErr: errors.New("no credentials"), - want: "", - wantCalls: 2, - }, - } - for _, test := range tests { - t.Run(test.name, func(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, test.env) - calls := 0 - orig := findDefaultCredentials - findDefaultCredentials = func(context.Context, ...string) (*google.Credentials, error) { - calls++ - if test.credsErr != nil { - return nil, test.credsErr - } - return &google.Credentials{ - UniverseDomainProvider: func() (string, error) { return test.credsUD, nil }, - }, nil - } - t.Cleanup(func() { findDefaultCredentials = orig }) - - // Call twice so wantCalls can verify caching: a successful lookup - // is reused by the second call, while a failed one is retried. - for i := 0; i < 2; i++ { - if got := GetUniverseDomain(t.Context()); got != test.want { - t.Errorf("GetUniverseDomain() = %q, want %q", got, test.want) - } - } - if calls != test.wantCalls { - t.Errorf("credentials looked up %d times, want %d", calls, test.wantCalls) - } - }) - } -} - -// resetUniverseDomainCache clears the cached universe domain and restores it -// when the test finishes. -func resetUniverseDomainCache(t *testing.T) { - t.Helper() - reset := func() { - cachedUniverseDomain.mu.Lock() - cachedUniverseDomain.resolved = false - cachedUniverseDomain.value = "" - cachedUniverseDomain.mu.Unlock() - } - reset() - t.Cleanup(reset) -} diff --git a/sdks/go/pkg/beam/util/gcsx/gcs.go b/sdks/go/pkg/beam/util/gcsx/gcs.go index 172776b2eb1f..9bbc14567e67 100644 --- a/sdks/go/pkg/beam/util/gcsx/gcs.go +++ b/sdks/go/pkg/beam/util/gcsx/gcs.go @@ -22,14 +22,12 @@ import ( "fmt" "io" "net/url" - "os" "path" - "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/internal/errors" - "golang.org/x/oauth2/google" + "github.com/apache/beam/sdks/v2/go/pkg/beam/util/gcpx" "google.golang.org/api/option" ) @@ -47,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 := GetUniverseDomain(ctx); ud != "" { + if ud := gcpx.GetUniverseDomain(ctx); ud != "" { opts = append(opts, option.WithUniverseDomain(ud)) } return storage.NewClient(ctx, opts...) @@ -85,54 +83,6 @@ func GetGcsEndpoint() string { 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 -} - -// 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 -} - // 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) { From 79d6200a37d1373f50c00a1bd5b16600a646968b Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Fri, 23 Oct 2026 16:02:32 +0000 Subject: [PATCH 3/4] address reviewers' comments --- .../pkg/beam/options/gcpopts/universe_test.go | 51 ------------------- sdks/go/pkg/beam/util/gcsx/client_test.go | 32 ++++++++++++ 2 files changed, 32 insertions(+), 51 deletions(-) delete mode 100644 sdks/go/pkg/beam/options/gcpopts/universe_test.go diff --git a/sdks/go/pkg/beam/options/gcpopts/universe_test.go b/sdks/go/pkg/beam/options/gcpopts/universe_test.go deleted file mode 100644 index 8954453abc48..000000000000 --- a/sdks/go/pkg/beam/options/gcpopts/universe_test.go +++ /dev/null @@ -1,51 +0,0 @@ -// 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 gcpopts - -import ( - "testing" - - "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" -) - -func TestGetUniverseDomain(t *testing.T) { - t.Setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN", "env.example.com") - if got, want := GetUniverseDomain(t.Context()), "env.example.com"; got != want { - t.Errorf("GetUniverseDomain() = %q, want %q", got, want) - } -} - -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 af15c20da740..0dd7f3d09e57 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -16,6 +16,7 @@ package gcsx import ( + "flag" "net/http" "net/http/httptest" "strings" @@ -24,6 +25,37 @@ import ( "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/" + + // On workers, the value arrives as a pipeline option rather than a flag. + runtime.GlobalOptions.Set("gcs_endpoint", optionValue) + t.Cleanup(func() { + runtime.GlobalOptions.Set("gcs_endpoint", "") + }) + + if flag.Lookup("gcs_endpoint") == nil { + if got := GetGcsEndpoint(); got != optionValue { + t.Errorf("GetGcsEndpoint() with flag unregistered = %q, want pipeline option %q", got, optionValue) + } + flag.String("gcs_endpoint", "", "") + } + origFlag := flag.Lookup("gcs_endpoint").Value.String() + t.Cleanup(func() { + _ = flag.Set("gcs_endpoint", origFlag) + }) + + _ = flag.Set("gcs_endpoint", "") + if got := GetGcsEndpoint(); got != optionValue { + t.Errorf("GetGcsEndpoint() with flag unset = %q, want pipeline option %q", got, optionValue) + } + _ = flag.Set("gcs_endpoint", 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) { From 76e78461f25f224072fd5cf35b636e416de246f7 Mon Sep 17 00:00:00 2001 From: Derrick Williams Date: Fri, 23 Oct 2026 16:43:19 +0000 Subject: [PATCH 4/4] forgot one change --- sdks/go/pkg/beam/util/gcsx/client_test.go | 22 +++++++++------------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/sdks/go/pkg/beam/util/gcsx/client_test.go b/sdks/go/pkg/beam/util/gcsx/client_test.go index 0dd7f3d09e57..3ec62a3a734e 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -29,28 +29,24 @@ func TestGetGcsEndpoint(t *testing.T) { const flagValue = "https://flag.example.com/storage/v1/" const optionValue = "https://option.example.com/storage/v1/" - // On workers, the value arrives as a pipeline option rather than a flag. - runtime.GlobalOptions.Set("gcs_endpoint", optionValue) + 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 flag.Lookup("gcs_endpoint") == nil { - if got := GetGcsEndpoint(); got != optionValue { - t.Errorf("GetGcsEndpoint() with flag unregistered = %q, want pipeline option %q", got, optionValue) - } - flag.String("gcs_endpoint", "", "") + if got := GetGcsEndpoint(); got != optionValue { + t.Errorf("GetGcsEndpoint() with flag unregistered = %q, want pipeline option %q", got, optionValue) } - origFlag := flag.Lookup("gcs_endpoint").Value.String() - t.Cleanup(func() { - _ = flag.Set("gcs_endpoint", origFlag) - }) - _ = flag.Set("gcs_endpoint", "") + gcsEndpoint := flag.String("gcs_endpoint", "", "") if got := GetGcsEndpoint(); got != optionValue { t.Errorf("GetGcsEndpoint() with flag unset = %q, want pipeline option %q", got, optionValue) } - _ = flag.Set("gcs_endpoint", flagValue) + *gcsEndpoint = flagValue if got := GetGcsEndpoint(); got != flagValue { t.Errorf("GetGcsEndpoint() with flag set = %q, want flag %q", got, flagValue) }