diff --git a/CHANGES.md b/CHANGES.md index 2940a8c6e708..b9accf9cf88f 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -83,6 +83,7 @@ * (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)). * (Python, Java, Go) Added `--dataflow_container_base_repository` (Python, Go) and `--dataflowContainerBaseRepository` (Java) to set the repository (e.g. `docker.example.com/my-prefix/cloud-dataflow/v1beta3`) from which the Dataflow runner pulls its default SDK container images, including the default images of cross-language Apache Beam environments. Defaults to `gcr.io/cloud-dataflow/v1beta3` ([#40569](https://github.com/apache/beam/issues/40569)). +* (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 diff --git a/sdks/go/pkg/beam/options/gcpopts/options.go b/sdks/go/pkg/beam/options/gcpopts/options.go index 2acb73dbd5e4..b5df80e2de06 100644 --- a/sdks/go/pkg/beam/options/gcpopts/options.go +++ b/sdks/go/pkg/beam/options/gcpopts/options.go @@ -22,8 +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" "golang.org/x/oauth2/google" ) @@ -34,8 +36,73 @@ 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") +} + +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 +} + // 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..35050d1010e6 --- /dev/null +++ b/sdks/go/pkg/beam/options/gcpopts/universe_test.go @@ -0,0 +1,120 @@ +// 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" + "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 }) + + // 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) { + 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/dataflow.go b/sdks/go/pkg/beam/runners/dataflow/dataflow.go index 87226404fcd7..20fb4fae0053 100644 --- a/sdks/go/pkg/beam/runners/dataflow/dataflow.go +++ b/sdks/go/pkg/beam/runners/dataflow/dataflow.go @@ -553,7 +553,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/runners/dataflow/dataflowlib/client_test.go b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go new file mode 100644 index 000000000000..63e3791f9bf2 --- /dev/null +++ b/sdks/go/pkg/beam/runners/dataflow/dataflowlib/client_test.go @@ -0,0 +1,91 @@ +// 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 ( + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "encoding/json" + "encoding/pem" + "os" + "path/filepath" + "testing" +) + +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") + if err := os.WriteFile(credFile, content, 0600); err != nil { + t.Fatalf("Failed writing fake credential file: %v", err) + } + t.Setenv("GOOGLE_APPLICATION_CREDENTIALS", credFile) + + tests := []struct { + name string + endpoint string + env string + want string + }{ + { + // 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/", + env: "creds.example.com", + want: "http://localhost:8080/", + }, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + t.Setenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN", test.env) + + client, err := NewClient(t.Context(), 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..841e17498769 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,22 @@ 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 != "" { + 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 new file mode 100644 index 000000000000..6b31646d6601 --- /dev/null +++ b/sdks/go/pkg/beam/util/gcsx/client_test.go @@ -0,0 +1,53 @@ +// 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 ( + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/apache/beam/sdks/v2/go/pkg/beam/options/gcpopts" +) + +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() + + orig := *gcpopts.GcsEndpoint + *gcpopts.GcsEndpoint = srv.URL + "/storage/v1/" + t.Cleanup(func() { *gcpopts.GcsEndpoint = orig }) + + ctx := t.Context() + 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 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 0986851c16d0..73300fcc6dc1 100644 --- a/sdks/go/pkg/beam/util/gcsx/gcs.go +++ b/sdks/go/pkg/beam/util/gcsx/gcs.go @@ -25,6 +25,7 @@ import ( "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" "google.golang.org/api/option" ) @@ -33,14 +34,36 @@ var userAgent = option.WithUserAgent("(GPN:Beam)") // 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()...) + // 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 != "" { + 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 } // Upload writes the given content to GCS. If the specified bucket does not