Use host metadata to determine GCP SA token requirement

Port of Python SDK PR #1322. Add requiresGcpSaAccessToken() which
checks host metadata to determine if a GCP SA access token is needed.
For workspace hosts (metadata has workspace_id), the SA token is
skipped. For account hosts or when metadata is unavailable, falls
back to checking AccountID. GoogleDefaultCredentials and
GoogleCredentials now use this metadata-based decision instead of
always attempting SA token creation.

Co-authored-by: Isaac

Author: hectorcast-db

config/auth_gcp_google_credentials.go

@@ -34,15 +34,21 @@ func (c GoogleCredentials) Configure(ctx context.Context, cfg *Config) (credenti
 	if err != nil {
 		return nil, fmt.Errorf("could not obtain OIDC token from JSON: %w", err)
 	}
-	// Obtain token source for creating Google Cloud Platform token.
+	opts := cacheOptions(cfg)
+	if !requiresGcpSaAccessToken(ctx, cfg) {
+		logger.Infof(ctx, "Using Google Credentials for Workspace")
+		visitor := refreshableVisitor(inner, opts...)
+		return credentials.CredentialsProviderFn(visitor), nil
+	}
+	// Account-level host: obtain token source for creating Google Cloud Platform token.
 	creds, err := google.CredentialsFromJSON(ctx, json,
 		"https://www.googleapis.com/auth/cloud-platform",
 		"https://www.googleapis.com/auth/compute")
 	if err != nil {
 		return nil, fmt.Errorf("could not obtain OAuth2 token from JSON: %w", err)
 	}
-	logger.Infof(ctx, "Using Google Credentials")
-	visitor := serviceToServiceVisitorWithFallback(inner, creds.TokenSource, "X-Databricks-GCP-SA-Access-Token", cacheOptions(cfg)...)
+	logger.Infof(ctx, "Using Google Credentials for Account")
+	visitor := serviceToServiceVisitorWithFallback(inner, creds.TokenSource, "X-Databricks-GCP-SA-Access-Token", opts...)
 	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }
 

config/auth_gcp_google_id.go

@@ -29,7 +29,12 @@ func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (c
 		return nil, err
 	}
 	opts := cacheOptions(cfg)
-	// Always attempt to create SA token source for the secondary header.
+	if !requiresGcpSaAccessToken(ctx, cfg) {
+		logger.Infof(ctx, "Using Google Default Application Credentials for Workspace")
+		visitor := refreshableVisitor(inner, opts...)
+		return credentials.CredentialsProviderFn(visitor), nil
+	}
+	// Account-level host: attempt to create SA token source for the secondary header.
 	// If it fails, fall back to refreshableVisitor with a warning.
 	platform, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{
 		TargetPrincipal: cfg.GoogleServiceAccount,
@@ -43,7 +48,7 @@ func (c GoogleDefaultCredentials) Configure(ctx context.Context, cfg *Config) (c
 		visitor := refreshableVisitor(inner, opts...)
 		return credentials.CredentialsProviderFn(visitor), nil
 	}
-	logger.Infof(ctx, "Using Google Default Application Credentials")
+	logger.Infof(ctx, "Using Google Default Application Credentials for Account")
 	visitor := serviceToServiceVisitorWithFallback(inner, platform, "X-Databricks-GCP-SA-Access-Token", opts...)
 	return credentials.NewOAuthCredentialsProvider(visitor, inner.Token), nil
 }

config/auth_gcp_helpers.go

@@ -0,0 +1,26 @@
+package config
+
+import (
+	"context"
+
+	"github.com/databricks/databricks-sdk-go/logger"
+)
+
+// requiresGcpSaAccessToken determines whether a GCP SA access token is needed
+// for the X-Databricks-GCP-SA-Access-Token header. It uses host metadata to
+// determine if the host is an account-level host (no workspace_id), and falls
+// back to checking AccountID when metadata is unavailable.
+func requiresGcpSaAccessToken(ctx context.Context, cfg *Config) bool {
+	if cfg.Host == "" {
+		return cfg.AccountID != ""
+	}
+	if err := cfg.EnsureResolved(); err != nil {
+		return cfg.AccountID != ""
+	}
+	meta, err := getHostMetadata(ctx, cfg.CanonicalHostName(), cfg.refreshClient)
+	if err != nil {
+		logger.Debugf(ctx, "Failed to fetch host metadata for GCP SA check: %v", err)
+		return cfg.AccountID != ""
+	}
+	return meta.WorkspaceID == ""
+}

config/auth_gcp_helpers_test.go

@@ -0,0 +1,100 @@
+package config
+
+import (
+	"context"
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/httpclient/fixtures"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestRequiresGcpSaAccessToken_WorkspaceFromMetadata(t *testing.T) {
+	noopLoader := mockLoader(func(cfg *Config) error { return nil })
+	cfg := &Config{
+		Host:    testHMHost,
+		Loaders: []Loader{noopLoader},
+		HTTPTransport: fixtures.SliceTransport{
+			{
+				Method:       "GET",
+				Resource:     "/.well-known/databricks-config",
+				ReuseRequest: true,
+				Status:       200,
+				Response:     `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `", "workspace_id": "` + testHMWorkspaceID + `"}`,
+			},
+		},
+	}
+	result := requiresGcpSaAccessToken(context.Background(), cfg)
+	assert.False(t, result, "workspace host should not require GCP SA token")
+}
+
+func TestRequiresGcpSaAccessToken_AccountFromMetadata(t *testing.T) {
+	noopLoader := mockLoader(func(cfg *Config) error { return nil })
+	cfg := &Config{
+		Host:    testHMHost,
+		Loaders: []Loader{noopLoader},
+		HTTPTransport: fixtures.SliceTransport{
+			{
+				Method:       "GET",
+				Resource:     "/.well-known/databricks-config",
+				ReuseRequest: true,
+				Status:       200,
+				Response:     `{"oidc_endpoint": "` + testHMHost + `/oidc", "account_id": "` + testHMAccountID + `"}`,
+			},
+		},
+	}
+	result := requiresGcpSaAccessToken(context.Background(), cfg)
+	assert.True(t, result, "account host (no workspace_id) should require GCP SA token")
+}
+
+func TestRequiresGcpSaAccessToken_MetadataError_FallsBackToAccountID(t *testing.T) {
+	noopLoader := mockLoader(func(cfg *Config) error { return nil })
+	cfg := &Config{
+		Host:      testHMHost,
+		AccountID: testHMAccountID,
+		Loaders:   []Loader{noopLoader},
+		HTTPTransport: fixtures.SliceTransport{
+			{
+				Method:       "GET",
+				Resource:     "/.well-known/databricks-config",
+				ReuseRequest: true,
+				Status:       500,
+				Response:     `{"error": "internal error"}`,
+			},
+		},
+	}
+	result := requiresGcpSaAccessToken(context.Background(), cfg)
+	assert.True(t, result, "with account_id set and metadata error, should require SA token")
+}
+
+func TestRequiresGcpSaAccessToken_MetadataError_NoAccountID(t *testing.T) {
+	noopLoader := mockLoader(func(cfg *Config) error { return nil })
+	cfg := &Config{
+		Host:    testHMHost,
+		Loaders: []Loader{noopLoader},
+		HTTPTransport: fixtures.SliceTransport{
+			{
+				Method:       "GET",
+				Resource:     "/.well-known/databricks-config",
+				ReuseRequest: true,
+				Status:       500,
+				Response:     `{"error": "internal error"}`,
+			},
+		},
+	}
+	result := requiresGcpSaAccessToken(context.Background(), cfg)
+	assert.False(t, result, "without account_id and metadata error, should not require SA token")
+}
+
+func TestRequiresGcpSaAccessToken_NoHost_WithAccountID(t *testing.T) {
+	cfg := &Config{
+		AccountID: testHMAccountID,
+	}
+	result := requiresGcpSaAccessToken(context.Background(), cfg)
+	assert.True(t, result, "no host with account_id should require SA token")
+}
+
+func TestRequiresGcpSaAccessToken_NoHost_NoAccountID(t *testing.T) {
+	cfg := &Config{}
+	result := requiresGcpSaAccessToken(context.Background(), cfg)
+	assert.False(t, result, "no host and no account_id should not require SA token")
+}