This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f1ec3ac0fac [Go] Support non-default universes for GCS and Dataflow
clients (#40686)
f1ec3ac0fac is described below
commit f1ec3ac0fac702d0f96e730193b756ccdbe6f8ee
Author: Shunping Huang <[email protected]>
AuthorDate: Tue Oct 20 22:46:48 2026 -0400
[Go] Support non-default universes for GCS and Dataflow clients (#40686)
---
CHANGES.md | 1 +
sdks/go/pkg/beam/options/gcpopts/options.go | 67 ++++++++++++
sdks/go/pkg/beam/options/gcpopts/universe_test.go | 120 +++++++++++++++++++++
sdks/go/pkg/beam/runners/dataflow/dataflow.go | 2 +-
.../runners/dataflow/dataflowlib/client_test.go | 91 ++++++++++++++++
.../pkg/beam/runners/dataflow/dataflowlib/job.go | 25 +++--
sdks/go/pkg/beam/util/gcsx/client_test.go | 53 +++++++++
sdks/go/pkg/beam/util/gcsx/gcs.go | 27 ++++-
8 files changed, 372 insertions(+), 14 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 2940a8c6e70..b9accf9cf88 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 2acb73dbd5e..b5df80e2de0 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 00000000000..35050d1010e
--- /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 87226404fcd..20fb4fae005 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 00000000000..63e3791f9bf
--- /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": "[email protected]",
+ "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 27a0a0d0f22..841e1749876 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 00000000000..6b31646d660
--- /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 0986851c16d..73300fcc6dc 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
