From 5a7acefe7ffb607fb25621430e1b6822518518d8 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 20 Oct 2026 10:42:40 -0400 Subject: [PATCH 1/6] [Go] Support GCS endpoint override and non-default universes Add a --gcs_endpoint flag. gcsx clients now pass the endpoint and the universe domain, taken from GOOGLE_CLOUD_UNIVERSE_DOMAIN or the default credentials, to storage.NewClient. This allows GCS access in Trusted Partner Cloud (TPC). --- sdks/go/pkg/beam/options/gcpopts/options.go | 18 +++ sdks/go/pkg/beam/runners/dataflow/dataflow.go | 2 +- sdks/go/pkg/beam/util/gcsx/client_test.go | 148 ++++++++++++++++++ sdks/go/pkg/beam/util/gcsx/gcs.go | 80 +++++++++- 4 files changed, 245 insertions(+), 3 deletions(-) create mode 100644 sdks/go/pkg/beam/util/gcsx/client_test.go diff --git a/sdks/go/pkg/beam/options/gcpopts/options.go b/sdks/go/pkg/beam/options/gcpopts/options.go index 2acb73dbd5e4..df66b20ef0cc 100644 --- a/sdks/go/pkg/beam/options/gcpopts/options.go +++ b/sdks/go/pkg/beam/options/gcpopts/options.go @@ -24,6 +24,7 @@ import ( "strings" "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" "golang.org/x/oauth2/google" ) @@ -34,8 +35,25 @@ var ( // Region is the GCP region where the job should be run. Region = flag.String("region", "", "GCP Region (required)") + + // GcsEndpoint overrides the Google Cloud Storage API endpoint. + GcsEndpoint = flag.String("gcs_endpoint", "", "The URL for the GCS API, "+ + "e.g. https://storage.googleapis.com/storage/v1/ (optional). If not set, the "+ + "endpoint is derived from the universe domain, which is read from "+ + "GOOGLE_CLOUD_UNIVERSE_DOMAIN or the default credentials.") ) +// GetGcsEndpoint returns the GCS endpoint override from the --gcs_endpoint +// flag. If the flag is unset, 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. +func GetGcsEndpoint() string { + if *GcsEndpoint != "" { + return *GcsEndpoint + } + return runtime.GlobalOptions.Get("gcs_endpoint") +} + // GetProject returns the project, if non empty and exits otherwise. // Convenience function. func GetProject(ctx context.Context) string { diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflow.go b/sdks/go/pkg/beam/runners/dataflow/dataflow.go index ce5e99b0a8de..a22293601b52 100644 --- a/sdks/go/pkg/beam/runners/dataflow/dataflow.go +++ b/sdks/go/pkg/beam/runners/dataflow/dataflow.go @@ -481,7 +481,7 @@ func checkSoftDeletePolicyEnabled(ctx context.Context, bucketName string, locati log.Warnf(ctx, "Error parsing bucket name: %v", err) return } - client, err := storage.NewClient(ctx) + client, err := gcsx.NewClient(ctx, storage.ScopeReadOnly) if err != nil { log.Warnf(ctx, "Error creating GCS client: %v", err) return diff --git a/sdks/go/pkg/beam/util/gcsx/client_test.go b/sdks/go/pkg/beam/util/gcsx/client_test.go new file mode 100644 index 000000000000..1997966c5067 --- /dev/null +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -0,0 +1,148 @@ +// 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 gcsx + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "cloud.google.com/go/storage" + "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" + "golang.org/x/oauth2/google" +) + +// 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) +} + +// fakeCredentials overrides findDefaultCredentials for the duration of the +// test, and returns a pointer to the number of times it was called. +func fakeCredentials(t *testing.T, universe string, err error) *int { + t.Helper() + calls := 0 + orig := findDefaultCredentials + findDefaultCredentials = func(ctx context.Context, scopes ...string) (*google.Credentials, error) { + calls++ + if err != nil { + return nil, err + } + return &google.Credentials{ + UniverseDomainProvider: func() (string, error) { return universe, nil }, + }, nil + } + t.Cleanup(func() { findDefaultCredentials = orig }) + return &calls +} + +func setGcsEndpointFlag(t *testing.T, endpoint string) { + t.Helper() + orig := *gcpopts.GcsEndpoint + *gcpopts.GcsEndpoint = endpoint + t.Cleanup(func() { *gcpopts.GcsEndpoint = orig }) +} + +func TestUniverseDomain_EnvVarTakesPrecedence(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, "env.example.com") + calls := fakeCredentials(t, "creds.example.com", nil) + + if got, want := universeDomain(context.Background(), storage.ScopeReadWrite), "env.example.com"; got != want { + t.Errorf("universeDomain() = %q, want %q", got, want) + } + if *calls != 0 { + t.Errorf("credentials looked up %d times, want 0", *calls) + } +} + +func TestUniverseDomain_FallsBackToCredentials(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, "") + calls := fakeCredentials(t, "creds.example.com", nil) + + for i := 0; i < 2; i++ { + if got, want := universeDomain(context.Background(), storage.ScopeReadWrite), "creds.example.com"; got != want { + t.Errorf("universeDomain() = %q, want %q", got, want) + } + } + if *calls != 1 { + t.Errorf("credentials looked up %d times, want 1 (cached)", *calls) + } +} + +func TestUniverseDomain_CredentialsError(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, "") + calls := fakeCredentials(t, "", errors.New("no credentials")) + + for i := 0; i < 2; i++ { + if got := universeDomain(context.Background(), storage.ScopeReadWrite); got != "" { + t.Errorf("universeDomain() = %q, want empty", got) + } + } + if *calls != 2 { + t.Errorf("credentials looked up %d times, want 2 (errors not cached)", *calls) + } +} + +func TestEndpointOptions(t *testing.T) { + setGcsEndpointFlag(t, "") + if got := endpointOptions(); len(got) != 0 { + t.Errorf("endpointOptions() with no endpoint = %v, want none", got) + } + setGcsEndpointFlag(t, "https://storage.example.com/storage/v1/") + if got := endpointOptions(); len(got) != 1 { + t.Errorf("endpointOptions() with endpoint returned %d options, want 1", len(got)) + } +} + +func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { + var gotPath string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + gotPath = r.URL.Path + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"name": "some-bucket"}`)) + })) + defer srv.Close() + setGcsEndpointFlag(t, srv.URL+"/storage/v1/") + + ctx := context.Background() + client, err := NewUnauthenticatedClient(ctx) + if err != nil { + t.Fatalf("NewUnauthenticatedClient() failed: %v", err) + } + defer client.Close() + + if _, err := client.Bucket("some-bucket").Attrs(ctx); err != nil { + t.Fatalf("Attrs() failed: %v", err) + } + if !strings.HasPrefix(gotPath, "/storage/v1/b/some-bucket") { + t.Errorf("request path = %q, want prefix %q", gotPath, "/storage/v1/b/some-bucket") + } +} diff --git a/sdks/go/pkg/beam/util/gcsx/gcs.go b/sdks/go/pkg/beam/util/gcsx/gcs.go index 0986851c16d0..34269a8b237d 100644 --- a/sdks/go/pkg/beam/util/gcsx/gcs.go +++ b/sdks/go/pkg/beam/util/gcsx/gcs.go @@ -21,26 +21,102 @@ import ( "fmt" "io" "net/url" + "os" "path" + "sync" "cloud.google.com/go/storage" "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" ) var userAgent = option.WithUserAgent("(GPN:Beam)") +const ( + defaultUniverseDomain = "googleapis.com" + 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 GCS clients are created frequently. +var cachedUniverseDomain struct { + mu sync.Mutex + resolved bool + value string +} + // NewClient creates a new GCS client with default application credentials, and supplied // OAuth scope. The OAuth scopes are defined in https://pkg.go.dev/cloud.google.com/go/storage#pkg-constants. // Sets the user agent to Beam. +// +// The endpoint can be overridden with the --gcs_endpoint flag. In non-default +// universes (e.g. Trusted Partner Cloud), the universe domain is taken from +// GOOGLE_CLOUD_UNIVERSE_DOMAIN or, if unset, from the default credentials. func NewClient(ctx context.Context, scope string) (*storage.Client, error) { - return storage.NewClient(ctx, option.WithScopes(scope), userAgent) + opts := []option.ClientOption{option.WithScopes(scope), userAgent} + opts = append(opts, endpointOptions()...) + if ud := universeDomain(ctx, scope); ud != "" && ud != defaultUniverseDomain { + opts = append(opts, option.WithUniverseDomain(ud)) + } + return storage.NewClient(ctx, opts...) } // NewUnauthenticatedClient creates a new GCS client without authentication. // Sets the user agent to Beam. func NewUnauthenticatedClient(ctx context.Context) (*storage.Client, error) { - return storage.NewClient(ctx, option.WithoutAuthentication(), userAgent) + opts := []option.ClientOption{option.WithoutAuthentication(), userAgent} + opts = append(opts, endpointOptions()...) + return storage.NewClient(ctx, opts...) +} + +// endpointOptions returns the client options for the --gcs_endpoint override, +// if any. +func endpointOptions() []option.ClientOption { + if ep := gcpopts.GetGcsEndpoint(); ep != "" { + return []option.ClientOption{option.WithEndpoint(ep)} + } + return nil +} + +// universeDomain returns the universe domain to use for GCS clients. +// +// storage.NewClient only reads the universe domain from +// option.WithUniverseDomain or GOOGLE_CLOUD_UNIVERSE_DOMAIN; it ignores the +// credentials' universe domain, and requests fail if the latter differs from +// the configured one (googleapis.com by default). Resolve it explicitly to +// support non-default universes such as Trusted Partner Cloud (TPC). +// +// Returns an empty string if it cannot be determined, leaving the decision to +// the client library. +func universeDomain(ctx context.Context, scope string) 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, scope) + if err != nil { + // Don't cache: credentials may become available later, and + // storage.NewClient 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 From f89e7d9a2189ba8b7114a762c29b5ba3b793d3ba Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 20 Oct 2026 10:49:45 -0400 Subject: [PATCH 2/6] [Go] Support non-default universes in the Dataflow client Replace the deprecated df.New with df.NewService, passing --dataflow_endpoint via option.WithEndpoint and the universe domain from GOOGLE_CLOUD_UNIVERSE_DOMAIN or the default credentials. Share the universe domain resolution with the GCS client via gcpopts. --- sdks/go/pkg/beam/options/gcpopts/options.go | 53 +++++++ .../pkg/beam/options/gcpopts/universe_test.go | 132 ++++++++++++++++++ .../dataflow/dataflowlib/client_test.go | 85 +++++++++++ .../beam/runners/dataflow/dataflowlib/job.go | 26 ++-- sdks/go/pkg/beam/util/gcsx/client_test.go | 79 ----------- sdks/go/pkg/beam/util/gcsx/gcs.go | 59 +------- 6 files changed, 288 insertions(+), 146 deletions(-) create mode 100644 sdks/go/pkg/beam/options/gcpopts/universe_test.go create mode 100644 sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go diff --git a/sdks/go/pkg/beam/options/gcpopts/options.go b/sdks/go/pkg/beam/options/gcpopts/options.go index df66b20ef0cc..d48426632413 100644 --- a/sdks/go/pkg/beam/options/gcpopts/options.go +++ b/sdks/go/pkg/beam/options/gcpopts/options.go @@ -22,6 +22,7 @@ import ( "os" "os/exec" "strings" + "sync" "cloud.google.com/go/storage" "github.com/apache/beam/sdks/v2/go/pkg/beam/core/runtime" @@ -54,6 +55,58 @@ func GetGcsEndpoint() string { return runtime.GlobalOptions.Get("gcs_endpoint") } +const ( + // DefaultUniverseDomain is the universe domain of the public Google Cloud. + DefaultUniverseDomain = "googleapis.com" + 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, "https://www.googleapis.com/auth/cloud-platform") + 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 +} + // GetProject returns the project, if non empty and exits otherwise. // Convenience function. func GetProject(ctx context.Context) string { diff --git a/sdks/go/pkg/beam/options/gcpopts/universe_test.go b/sdks/go/pkg/beam/options/gcpopts/universe_test.go new file mode 100644 index 000000000000..e9f27081701a --- /dev/null +++ b/sdks/go/pkg/beam/options/gcpopts/universe_test.go @@ -0,0 +1,132 @@ +// 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 ( + "context" + "errors" + "os" + "path/filepath" + "testing" + + "golang.org/x/oauth2/google" +) + +// 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) +} + +// fakeCredentials overrides findDefaultCredentials for the duration of the +// test, and returns a pointer to the number of times it was called. +func fakeCredentials(t *testing.T, universe string, err error) *int { + t.Helper() + calls := 0 + orig := findDefaultCredentials + findDefaultCredentials = func(ctx context.Context, scopes ...string) (*google.Credentials, error) { + calls++ + if err != nil { + return nil, err + } + return &google.Credentials{ + UniverseDomainProvider: func() (string, error) { return universe, nil }, + }, nil + } + t.Cleanup(func() { findDefaultCredentials = orig }) + return &calls +} + +func TestGetUniverseDomain_EnvVarTakesPrecedence(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, "env.example.com") + calls := fakeCredentials(t, "creds.example.com", nil) + + if got, want := GetUniverseDomain(context.Background()), "env.example.com"; got != want { + t.Errorf("GetUniverseDomain() = %q, want %q", got, want) + } + if *calls != 0 { + t.Errorf("credentials looked up %d times, want 0", *calls) + } +} + +func TestGetUniverseDomain_FallsBackToCredentials(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, "") + calls := fakeCredentials(t, "creds.example.com", nil) + + for i := 0; i < 2; i++ { + if got, want := GetUniverseDomain(context.Background()), "creds.example.com"; got != want { + t.Errorf("GetUniverseDomain() = %q, want %q", got, want) + } + } + if *calls != 1 { + t.Errorf("credentials looked up %d times, want 1 (cached)", *calls) + } +} + +func TestGetUniverseDomain_CredentialsError(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, "") + calls := fakeCredentials(t, "", errors.New("no credentials")) + + for i := 0; i < 2; i++ { + if got := GetUniverseDomain(context.Background()); got != "" { + t.Errorf("GetUniverseDomain() = %q, want empty", got) + } + } + if *calls != 2 { + t.Errorf("credentials looked up %d times, want 2 (errors not cached)", *calls) + } +} + +func TestGetUniverseDomain_FromCredentialFile(t *testing.T) { + resetUniverseDomainCache(t) + t.Setenv(universeDomainEnvVar, "") + + credFile := filepath.Join(t.TempDir(), "creds.json") + content := []byte(`{ + "type": "service_account", + "project_id": "test-project", + "universe_domain": "file.example.com" + }`) + if err := os.WriteFile(credFile, content, 0600); err != nil { + t.Fatalf("Failed writing fake credential file: %v", err) + } + t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", credFile) + + if got, want := GetUniverseDomain(context.Background()), "file.example.com"; got != want { + t.Errorf("GetUniverseDomain() = %q, want %q", got, want) + } +} + +func TestGetGcsEndpoint(t *testing.T) { + orig := *GcsEndpoint + t.Cleanup(func() { *GcsEndpoint = orig }) + + *GcsEndpoint = "https://storage.example.com/storage/v1/" + if got, want := GetGcsEndpoint(), "https://storage.example.com/storage/v1/"; got != want { + t.Errorf("GetGcsEndpoint() = %q, want %q", got, want) + } +} diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go new file mode 100644 index 000000000000..1ed0c02619e8 --- /dev/null +++ b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go @@ -0,0 +1,85 @@ +// 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 dataflowlib + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +// setupFakeCredentials points application default credentials at a fake +// authorized user credential file, so clients can be constructed without +// real credentials. No token is fetched until a request is made. +func setupFakeCredentials(t *testing.T) { + t.Helper() + credFile := filepath.Join(t.TempDir(), "creds.json") + content := []byte(`{ + "type": "authorized_user", + "client_id": "fake-client-id", + "client_secret": "fake-client-secret", + "refresh_token": "fake-refresh-token" + }`) + if err := os.WriteFile(credFile, content, 0600); err != nil { + t.Fatalf("Failed writing fake credential file: %v", err) + } + t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", credFile) +} + +func TestNewClient_BasePath(t *testing.T) { + tests := []struct { + name string + endpoint string + universe string + want string + }{ + { + name: "default", + want: "https://dataflow.googleapis.com/", + }, + { + name: "endpoint override", + endpoint: "https://dataflow.example.com/", + want: "https://dataflow.example.com/", + }, + { + name: "universe domain", + universe: "example.com", + want: "https://dataflow.example.com/", + }, + { + name: "endpoint takes precedence over universe domain", + endpoint: "http://localhost:8080/", + universe: "example.com", + want: "http://localhost:8080/", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + setupFakeCredentials(t) + t.Setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN", test.universe) + + client, err := NewClient(context.Background(), test.endpoint) + if err != nil { + t.Fatalf("NewClient(%q) failed: %v", test.endpoint, err) + } + if got := client.BasePath; got != test.want { + t.Errorf("NewClient(%q).BasePath = %q, want %q", test.endpoint, got, test.want) + } + }) + } +} diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go index 27a0a0d0f224..fd6d14c3465e 100644 --- a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go +++ b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go @@ -27,8 +27,9 @@ 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" - "golang.org/x/oauth2/google" + "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" df "google.golang.org/api/dataflow/v1b3" + "google.golang.org/api/option" "google.golang.org/protobuf/proto" // Importing to get the side effect of the remote execution hook. See init(). @@ -312,20 +313,23 @@ func currentStateMessage(currentState, jobID string) (bool, string, error) { // NewClient creates a new dataflow client with default application credentials // and CloudPlatformScope. The Dataflow endpoint is optionally overridden. +// +// In non-default universes (e.g. Trusted Partner Cloud), the universe domain +// is taken from GOOGLE_CLOUD_UNIVERSE_DOMAIN or, if unset, from the default +// credentials. func NewClient(ctx context.Context, endpoint string) (*df.Service, error) { - cl, err := google.DefaultClient(ctx, df.CloudPlatformScope) - if err != nil { - return nil, err - } - client, err := df.New(cl) - if err != nil { - return nil, err - } + opts := []option.ClientOption{option.WithScopes(df.CloudPlatformScope)} if endpoint != "" { log.Infof(ctx, "Dataflow endpoint override: %s", endpoint) - client.BasePath = endpoint + opts = append(opts, option.WithEndpoint(endpoint)) + } + // 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 != "" && ud != gcpopts.DefaultUniverseDomain { + log.Infof(ctx, "Using Google Cloud universe domain: %s", ud) + opts = append(opts, option.WithUniverseDomain(ud)) } - return client, nil + return df.NewService(ctx, opts...) } // GetRunningJobByName gets a Dataflow job running by its name and returns an diff --git a/sdks/go/pkg/beam/util/gcsx/client_test.go b/sdks/go/pkg/beam/util/gcsx/client_test.go index 1997966c5067..b9a0d1771b14 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -17,50 +17,14 @@ package gcsx import ( "context" - "errors" "net/http" "net/http/httptest" "strings" "testing" - "cloud.google.com/go/storage" "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" - "golang.org/x/oauth2/google" ) -// 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) -} - -// fakeCredentials overrides findDefaultCredentials for the duration of the -// test, and returns a pointer to the number of times it was called. -func fakeCredentials(t *testing.T, universe string, err error) *int { - t.Helper() - calls := 0 - orig := findDefaultCredentials - findDefaultCredentials = func(ctx context.Context, scopes ...string) (*google.Credentials, error) { - calls++ - if err != nil { - return nil, err - } - return &google.Credentials{ - UniverseDomainProvider: func() (string, error) { return universe, nil }, - }, nil - } - t.Cleanup(func() { findDefaultCredentials = orig }) - return &calls -} - func setGcsEndpointFlag(t *testing.T, endpoint string) { t.Helper() orig := *gcpopts.GcsEndpoint @@ -68,49 +32,6 @@ func setGcsEndpointFlag(t *testing.T, endpoint string) { t.Cleanup(func() { *gcpopts.GcsEndpoint = orig }) } -func TestUniverseDomain_EnvVarTakesPrecedence(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, "env.example.com") - calls := fakeCredentials(t, "creds.example.com", nil) - - if got, want := universeDomain(context.Background(), storage.ScopeReadWrite), "env.example.com"; got != want { - t.Errorf("universeDomain() = %q, want %q", got, want) - } - if *calls != 0 { - t.Errorf("credentials looked up %d times, want 0", *calls) - } -} - -func TestUniverseDomain_FallsBackToCredentials(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, "") - calls := fakeCredentials(t, "creds.example.com", nil) - - for i := 0; i < 2; i++ { - if got, want := universeDomain(context.Background(), storage.ScopeReadWrite), "creds.example.com"; got != want { - t.Errorf("universeDomain() = %q, want %q", got, want) - } - } - if *calls != 1 { - t.Errorf("credentials looked up %d times, want 1 (cached)", *calls) - } -} - -func TestUniverseDomain_CredentialsError(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, "") - calls := fakeCredentials(t, "", errors.New("no credentials")) - - for i := 0; i < 2; i++ { - if got := universeDomain(context.Background(), storage.ScopeReadWrite); got != "" { - t.Errorf("universeDomain() = %q, want empty", got) - } - } - if *calls != 2 { - t.Errorf("credentials looked up %d times, want 2 (errors not cached)", *calls) - } -} - func TestEndpointOptions(t *testing.T) { setGcsEndpointFlag(t, "") if got := endpointOptions(); len(got) != 0 { diff --git a/sdks/go/pkg/beam/util/gcsx/gcs.go b/sdks/go/pkg/beam/util/gcsx/gcs.go index 34269a8b237d..4521712bcad5 100644 --- a/sdks/go/pkg/beam/util/gcsx/gcs.go +++ b/sdks/go/pkg/beam/util/gcsx/gcs.go @@ -21,36 +21,16 @@ import ( "fmt" "io" "net/url" - "os" "path" - "sync" "cloud.google.com/go/storage" "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" ) var userAgent = option.WithUserAgent("(GPN:Beam)") -const ( - defaultUniverseDomain = "googleapis.com" - 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 GCS clients are created frequently. -var cachedUniverseDomain struct { - mu sync.Mutex - resolved bool - value string -} - // NewClient creates a new GCS client with default application credentials, and supplied // OAuth scope. The OAuth scopes are defined in https://pkg.go.dev/cloud.google.com/go/storage#pkg-constants. // Sets the user agent to Beam. @@ -61,7 +41,9 @@ var cachedUniverseDomain struct { func NewClient(ctx context.Context, scope string) (*storage.Client, error) { opts := []option.ClientOption{option.WithScopes(scope), userAgent} opts = append(opts, endpointOptions()...) - if ud := universeDomain(ctx, scope); ud != "" && ud != defaultUniverseDomain { + // 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 != "" && ud != gcpopts.DefaultUniverseDomain { opts = append(opts, option.WithUniverseDomain(ud)) } return storage.NewClient(ctx, opts...) @@ -84,41 +66,6 @@ func endpointOptions() []option.ClientOption { return nil } -// universeDomain returns the universe domain to use for GCS clients. -// -// storage.NewClient only reads the universe domain from -// option.WithUniverseDomain or GOOGLE_CLOUD_UNIVERSE_DOMAIN; it ignores the -// credentials' universe domain, and requests fail if the latter differs from -// the configured one (googleapis.com by default). Resolve it explicitly to -// support non-default universes such as Trusted Partner Cloud (TPC). -// -// Returns an empty string if it cannot be determined, leaving the decision to -// the client library. -func universeDomain(ctx context.Context, scope string) 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, scope) - if err != nil { - // Don't cache: credentials may become available later, and - // storage.NewClient 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 9072f14edcb18b7dfa09b71e14d6b5b2b1b8bf2e Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 20 Oct 2026 10:58:18 -0400 Subject: [PATCH 3/6] Add CHANGES.md --- CHANGES.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGES.md b/CHANGES.md index d06940d12cf4..b40afb3ec58f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -82,6 +82,7 @@ * (Python) Added `--aws_credentials_secret_manager` so AWS credentials (`--aws_access_key_id`, `--aws_secret_access_key`, `--aws_session_token` and their `--s3_*` overrides) can be secret specs resolved through a secret manager such as `GoogleCloudSecretManager`, instead of plain-text values ([#40528](https://github.com/apache/beam/issues/40528)). * (Python) Added `AwsSecret` (`AwsSecretsManager` / `type:awssecret`) to retrieve secrets from AWS Secrets Manager, configured from `AwsOptions` (`--aws_*`) with optional per-secret connection overrides ([#40528](https://github.com/apache/beam/issues/40528)). * (Go) Added `schema.RegisterLogicalTypeConversion` and `schema.RegisterLogicalTypeCoder` to register portable logical types in one call, and Go types for the standard `beam:logical_type:date:v1` and `beam:logical_type:micros_instant:v1` logical types, so schemas from other SDKs with these types can be used in cross-language pipelines ([#39684](https://github.com/apache/beam/issues/39684)). +* (Go) GCS and Dataflow clients now support non-default universes such as Trusted Partner Cloud (TPC), using the universe domain from `GOOGLE_CLOUD_UNIVERSE_DOMAIN` or the default credentials. Added a `--gcs_endpoint` pipeline option to override the GCS endpoint ([#40569](https://github.com/apache/beam/issues/40569)). ## Breaking Changes From acbfd090bbaac67b6ee8ac18b31bfb120f36d318 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 20 Oct 2026 11:08:47 -0400 Subject: [PATCH 4/6] [Go] Simplify universe domain handling and dedup tests --- sdks/go/pkg/beam/options/gcpopts/options.go | 8 +- .../pkg/beam/options/gcpopts/universe_test.go | 166 ++++++++---------- .../dataflow/dataflowlib/client_test.go | 65 ++++--- .../beam/runners/dataflow/dataflowlib/job.go | 3 +- sdks/go/pkg/beam/util/gcsx/client_test.go | 27 +-- sdks/go/pkg/beam/util/gcsx/gcs.go | 2 +- 6 files changed, 122 insertions(+), 149 deletions(-) diff --git a/sdks/go/pkg/beam/options/gcpopts/options.go b/sdks/go/pkg/beam/options/gcpopts/options.go index d48426632413..b5df80e2de06 100644 --- a/sdks/go/pkg/beam/options/gcpopts/options.go +++ b/sdks/go/pkg/beam/options/gcpopts/options.go @@ -55,11 +55,7 @@ func GetGcsEndpoint() string { return runtime.GlobalOptions.Get("gcs_endpoint") } -const ( - // DefaultUniverseDomain is the universe domain of the public Google Cloud. - DefaultUniverseDomain = "googleapis.com" - universeDomainEnvVar = "GOOGLE_CLOUD_UNIVERSE_DOMAIN" -) +const universeDomainEnvVar = "GOOGLE_CLOUD_UNIVERSE_DOMAIN" // findDefaultCredentials is a variable to allow overriding in tests. var findDefaultCredentials = google.FindDefaultCredentials @@ -92,7 +88,7 @@ func GetUniverseDomain(ctx context.Context) string { if cachedUniverseDomain.resolved { return cachedUniverseDomain.value } - creds, err := findDefaultCredentials(ctx, "https://www.googleapis.com/auth/cloud-platform") + creds, err := findDefaultCredentials(ctx) if err != nil { // Don't cache: credentials may become available later, and client // construction surfaces the credential error itself. diff --git a/sdks/go/pkg/beam/options/gcpopts/universe_test.go b/sdks/go/pkg/beam/options/gcpopts/universe_test.go index e9f27081701a..686d1646c9a0 100644 --- a/sdks/go/pkg/beam/options/gcpopts/universe_test.go +++ b/sdks/go/pkg/beam/options/gcpopts/universe_test.go @@ -18,13 +18,69 @@ package gcpopts import ( "context" "errors" - "os" - "path/filepath" "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, + }, + } + 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 }) + + for i := 0; i < 2; i++ { + if got := GetUniverseDomain(context.Background()); 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) { @@ -39,94 +95,24 @@ func resetUniverseDomainCache(t *testing.T) { t.Cleanup(reset) } -// fakeCredentials overrides findDefaultCredentials for the duration of the -// test, and returns a pointer to the number of times it was called. -func fakeCredentials(t *testing.T, universe string, err error) *int { - t.Helper() - calls := 0 - orig := findDefaultCredentials - findDefaultCredentials = func(ctx context.Context, scopes ...string) (*google.Credentials, error) { - calls++ - if err != nil { - return nil, err - } - return &google.Credentials{ - UniverseDomainProvider: func() (string, error) { return universe, nil }, - }, nil - } - t.Cleanup(func() { findDefaultCredentials = orig }) - return &calls -} - -func TestGetUniverseDomain_EnvVarTakesPrecedence(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, "env.example.com") - calls := fakeCredentials(t, "creds.example.com", nil) - - if got, want := GetUniverseDomain(context.Background()), "env.example.com"; got != want { - t.Errorf("GetUniverseDomain() = %q, want %q", got, want) - } - if *calls != 0 { - t.Errorf("credentials looked up %d times, want 0", *calls) - } -} - -func TestGetUniverseDomain_FallsBackToCredentials(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, "") - calls := fakeCredentials(t, "creds.example.com", nil) - - for i := 0; i < 2; i++ { - if got, want := GetUniverseDomain(context.Background()), "creds.example.com"; got != want { - t.Errorf("GetUniverseDomain() = %q, want %q", got, want) - } - } - if *calls != 1 { - t.Errorf("credentials looked up %d times, want 1 (cached)", *calls) - } -} - -func TestGetUniverseDomain_CredentialsError(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, "") - calls := fakeCredentials(t, "", errors.New("no credentials")) - - for i := 0; i < 2; i++ { - if got := GetUniverseDomain(context.Background()); got != "" { - t.Errorf("GetUniverseDomain() = %q, want empty", got) - } - } - if *calls != 2 { - t.Errorf("credentials looked up %d times, want 2 (errors not cached)", *calls) - } -} - -func TestGetUniverseDomain_FromCredentialFile(t *testing.T) { - resetUniverseDomainCache(t) - t.Setenv(universeDomainEnvVar, "") - - credFile := filepath.Join(t.TempDir(), "creds.json") - content := []byte(`{ - "type": "service_account", - "project_id": "test-project", - "universe_domain": "file.example.com" - }`) - if err := os.WriteFile(credFile, content, 0600); err != nil { - t.Fatalf("Failed writing fake credential file: %v", err) - } - t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", credFile) - - if got, want := GetUniverseDomain(context.Background()), "file.example.com"; got != want { - t.Errorf("GetUniverseDomain() = %q, want %q", got, want) - } -} - func TestGetGcsEndpoint(t *testing.T) { - orig := *GcsEndpoint - t.Cleanup(func() { *GcsEndpoint = orig }) - - *GcsEndpoint = "https://storage.example.com/storage/v1/" - if got, want := GetGcsEndpoint(), "https://storage.example.com/storage/v1/"; got != want { - t.Errorf("GetGcsEndpoint() = %q, want %q", got, want) + 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/runners/dataflow/dataflowlib/client_test.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go index 1ed0c02619e8..117fb12fe643 100644 --- a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go +++ b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go @@ -17,61 +17,68 @@ package dataflowlib import ( "context" + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "encoding/json" + "encoding/pem" "os" "path/filepath" "testing" ) -// setupFakeCredentials points application default credentials at a fake -// authorized user credential file, so clients can be constructed without -// real credentials. No token is fetched until a request is made. -func setupFakeCredentials(t *testing.T) { - t.Helper() +func TestNewClient_BasePath(t *testing.T) { + // Fake service account credentials in a non-default universe. The auth + // library parses the key when the client is created (non-default universes + // use self-signed JWTs), so it must be a valid key. No token is fetched + // until a request is made. + key, err := rsa.GenerateKey(rand.Reader, 2048) + if err != nil { + t.Fatalf("Failed generating key: %v", err) + } + der, err := x509.MarshalPKCS8PrivateKey(key) + if err != nil { + t.Fatalf("Failed marshaling key: %v", err) + } + content, err := json.Marshal(map[string]string{ + "type": "service_account", + "project_id": "test-project", + "client_email": "test@test-project.iam.gserviceaccount.com", + "private_key_id": "test-key-id", + "private_key": string(pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: der})), + "universe_domain": "creds.example.com", + }) + if err != nil { + t.Fatalf("Failed marshaling credentials: %v", err) + } credFile := filepath.Join(t.TempDir(), "creds.json") - content := []byte(`{ - "type": "authorized_user", - "client_id": "fake-client-id", - "client_secret": "fake-client-secret", - "refresh_token": "fake-refresh-token" - }`) if err := os.WriteFile(credFile, content, 0600); err != nil { t.Fatalf("Failed writing fake credential file: %v", err) } t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", credFile) -} -func TestNewClient_BasePath(t *testing.T) { tests := []struct { name string endpoint string - universe string + env string want string }{ { - name: "default", - want: "https://dataflow.googleapis.com/", - }, - { - name: "endpoint override", - endpoint: "https://dataflow.example.com/", - want: "https://dataflow.example.com/", - }, - { - name: "universe domain", - universe: "example.com", - want: "https://dataflow.example.com/", + // The client library doesn't read the credentials' universe domain, + // so this verifies that NewClient passes it through. + name: "universe domain from credentials", + want: "https://dataflow.creds.example.com/", }, { name: "endpoint takes precedence over universe domain", endpoint: "http://localhost:8080/", - universe: "example.com", + env: "creds.example.com", want: "http://localhost:8080/", }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) { - setupFakeCredentials(t) - t.Setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN", test.universe) + t.Setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN", test.env) client, err := NewClient(context.Background(), test.endpoint) if err != nil { diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go index fd6d14c3465e..841e17498769 100644 --- a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go +++ b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/job.go @@ -325,8 +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 != "" && ud != gcpopts.DefaultUniverseDomain { - log.Infof(ctx, "Using Google Cloud universe domain: %s", ud) + if ud := gcpopts.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 b9a0d1771b14..f0bd1ef44d97 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -25,24 +25,6 @@ import ( "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" ) -func setGcsEndpointFlag(t *testing.T, endpoint string) { - t.Helper() - orig := *gcpopts.GcsEndpoint - *gcpopts.GcsEndpoint = endpoint - t.Cleanup(func() { *gcpopts.GcsEndpoint = orig }) -} - -func TestEndpointOptions(t *testing.T) { - setGcsEndpointFlag(t, "") - if got := endpointOptions(); len(got) != 0 { - t.Errorf("endpointOptions() with no endpoint = %v, want none", got) - } - setGcsEndpointFlag(t, "https://storage.example.com/storage/v1/") - if got := endpointOptions(); len(got) != 1 { - t.Errorf("endpointOptions() with endpoint returned %d options, want 1", len(got)) - } -} - func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { var gotPath string srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -51,7 +33,10 @@ func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { w.Write([]byte(`{"name": "some-bucket"}`)) })) defer srv.Close() - setGcsEndpointFlag(t, srv.URL+"/storage/v1/") + + orig := *gcpopts.GcsEndpoint + *gcpopts.GcsEndpoint = srv.URL + "/storage/v1/" + t.Cleanup(func() { *gcpopts.GcsEndpoint = orig }) ctx := context.Background() client, err := NewUnauthenticatedClient(ctx) @@ -63,7 +48,7 @@ func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { if _, err := client.Bucket("some-bucket").Attrs(ctx); err != nil { t.Fatalf("Attrs() failed: %v", err) } - if !strings.HasPrefix(gotPath, "/storage/v1/b/some-bucket") { - t.Errorf("request path = %q, want prefix %q", gotPath, "/storage/v1/b/some-bucket") + if want := "/storage/v1/b/some-bucket"; !strings.HasPrefix(gotPath, want) { + t.Errorf("request path = %q, want prefix %q", gotPath, want) } } diff --git a/sdks/go/pkg/beam/util/gcsx/gcs.go b/sdks/go/pkg/beam/util/gcsx/gcs.go index 4521712bcad5..73300fcc6dc1 100644 --- a/sdks/go/pkg/beam/util/gcsx/gcs.go +++ b/sdks/go/pkg/beam/util/gcsx/gcs.go @@ -43,7 +43,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 != "" && ud != gcpopts.DefaultUniverseDomain { + if ud := gcpopts.GetUniverseDomain(ctx); ud != "" { opts = append(opts, option.WithUniverseDomain(ud)) } return storage.NewClient(ctx, opts...) From f5afce3ebc00c3d87bbe9992ac89f07e125fb8cf Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 20 Oct 2026 21:36:04 -0400 Subject: [PATCH 5/6] Use t.Context() in new tests --- sdks/go/pkg/beam/options/gcpopts/universe_test.go | 2 +- sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go | 3 +-- sdks/go/pkg/beam/util/gcsx/client_test.go | 3 +-- 3 files changed, 3 insertions(+), 5 deletions(-) diff --git a/sdks/go/pkg/beam/options/gcpopts/universe_test.go b/sdks/go/pkg/beam/options/gcpopts/universe_test.go index 686d1646c9a0..b8069ed6e566 100644 --- a/sdks/go/pkg/beam/options/gcpopts/universe_test.go +++ b/sdks/go/pkg/beam/options/gcpopts/universe_test.go @@ -70,7 +70,7 @@ func TestGetUniverseDomain(t *testing.T) { t.Cleanup(func() { findDefaultCredentials = orig }) for i := 0; i < 2; i++ { - if got := GetUniverseDomain(context.Background()); got != test.want { + if got := GetUniverseDomain(t.Context()); got != test.want { t.Errorf("GetUniverseDomain() = %q, want %q", got, test.want) } } diff --git a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go index 117fb12fe643..63e3791f9bf2 100644 --- a/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go +++ b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go @@ -16,7 +16,6 @@ package dataflowlib import ( - "context" "crypto/rand" "crypto/rsa" "crypto/x509" @@ -80,7 +79,7 @@ func TestNewClient_BasePath(t *testing.T) { t.Run(test.name, func(t *testing.T) { t.Setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN", test.env) - client, err := NewClient(context.Background(), test.endpoint) + client, err := NewClient(t.Context(), test.endpoint) if err != nil { t.Fatalf("NewClient(%q) failed: %v", test.endpoint, err) } diff --git a/sdks/go/pkg/beam/util/gcsx/client_test.go b/sdks/go/pkg/beam/util/gcsx/client_test.go index f0bd1ef44d97..6b31646d6601 100644 --- a/sdks/go/pkg/beam/util/gcsx/client_test.go +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -16,7 +16,6 @@ package gcsx import ( - "context" "net/http" "net/http/httptest" "strings" @@ -38,7 +37,7 @@ func TestNewUnauthenticatedClient_UsesGcsEndpoint(t *testing.T) { *gcpopts.GcsEndpoint = srv.URL + "/storage/v1/" t.Cleanup(func() { *gcpopts.GcsEndpoint = orig }) - ctx := context.Background() + ctx := t.Context() client, err := NewUnauthenticatedClient(ctx) if err != nil { t.Fatalf("NewUnauthenticatedClient() failed: %v", err) From a72498943359e9359bd9e4cd958edd698f04d85f Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Tue, 20 Oct 2026 21:38:38 -0400 Subject: [PATCH 6/6] Add a comment to explain the loop in test. --- sdks/go/pkg/beam/options/gcpopts/universe_test.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdks/go/pkg/beam/options/gcpopts/universe_test.go b/sdks/go/pkg/beam/options/gcpopts/universe_test.go index b8069ed6e566..35050d1010e6 100644 --- a/sdks/go/pkg/beam/options/gcpopts/universe_test.go +++ b/sdks/go/pkg/beam/options/gcpopts/universe_test.go @@ -69,6 +69,8 @@ func TestGetUniverseDomain(t *testing.T) { } 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)