From 67ad1ad3412931327bc26b3d5778170a04036b66 Mon Sep 17 00:00:00 2001 From: Xiaoyi Zhang Date: Wed, 19 Aug 2026 02:00:17 +0000 Subject: [PATCH] Add oauth-m2m-gcp auth type for GCP access token passthrough Authenticates the request identity with a Databricks OAuth service principal (Authorization header) and passes a Google Cloud access token through the X-Databricks-GCP-SA-Access-Token header, so GCP account-level provisioning APIs (e.g. workspace / VPC-endpoint / customer-managed-key creation) can be called with a Databricks-governed identity when SSO is enabled and Google ID token auth is disabled. Reuses the existing serviceToServiceVisitor: the Databricks OAuth token (extracted into databricksOAuthTokenSource) is the primary, and a Google cloud-platform access token (from google_credentials or google_service_account) is the secondary. Selected explicitly via AuthType "oauth-m2m-gcp", since it combines the oauth and google credential groups that the single-auth-method check otherwise rejects. Co-authored-by: Isaac --- NEXT_CHANGELOG.md | 4 + README.md | 20 +++++ config/auth_default.go | 1 + config/auth_gcp_oauth_m2m.go | 111 ++++++++++++++++++++++++++ config/auth_gcp_oauth_m2m_test.go | 125 ++++++++++++++++++++++++++++++ config/auth_m2m.go | 22 +++++- config/config.go | 8 +- internal/auth_gcp_m2m_test.go | 70 +++++++++++++++++ 8 files changed, 354 insertions(+), 7 deletions(-) create mode 100644 config/auth_gcp_oauth_m2m.go create mode 100644 config/auth_gcp_oauth_m2m_test.go create mode 100644 internal/auth_gcp_m2m_test.go diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index d3af4fbb2..033fb0a51 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -6,10 +6,14 @@ ### New Features and Improvements +* Added the `oauth-m2m-gcp` authentication type, which authenticates with a Databricks OAuth service-principal token and passes a Google Cloud access token through the `X-Databricks-GCP-SA-Access-Token` header. This lets GCP account-level provisioning APIs be called with a Databricks-governed identity when SSO is enabled and Google ID token auth is disabled. It must be selected explicitly via `AuthType: "oauth-m2m-gcp"`. + ### Bug Fixes ### Documentation +* Documented the `oauth-m2m-gcp` authentication type in the GCP section of the README. + ### Internal Changes ### API Changes diff --git a/README.md b/README.md index f26ee44bf..6ab56480b 100644 --- a/README.md +++ b/README.md @@ -348,6 +348,26 @@ w, err := databricks.NewWorkspaceClient(&databricks.Config{ }) ``` +### GCP service account access token passthrough with Databricks OAuth + +When single sign-on (SSO) is enabled on a GCP account, Google ID token authentication (`google-id`) is disabled. To call GCP account-level APIs that provision Google Cloud resources (for example workspace or VPC-endpoint creation) in that case, you can authenticate the request identity with a Databricks OAuth service principal while still passing a Google Cloud access token through for provisioning. + +Set `AuthType: "oauth-m2m-gcp"`. This must be set explicitly, because the mode combines the OAuth (`ClientID`/`ClientSecret`) and Google (`GoogleCredentials`/`GoogleServiceAccount`) credential groups, which are otherwise rejected together. Provide: + +- `Host`, `ClientID` and `ClientSecret` for the Databricks OAuth service principal — used for the `Authorization` header; and +- `GoogleCredentials` or `GoogleServiceAccount` for the Google Cloud access token — sent in the `X-Databricks-GCP-SA-Access-Token` header. + +```go +w, err := databricks.NewAccountClient(&databricks.Config{ + Host: askFor("Host:"), + AccountID: askFor("Account ID:"), + ClientID: askFor("Databricks OAuth Client ID:"), + ClientSecret: askFor("Databricks OAuth Client Secret:"), + GoogleServiceAccount: askFor("Google Service Account:"), + AuthType: "oauth-m2m-gcp", +}) +``` + ### Overriding `.databrickscfg` For [Databricks native authentication](#databricks-native-authentication), you can override the default behavior in `*databricks.Config` for using `.databrickscfg` as follows: diff --git a/config/auth_default.go b/config/auth_default.go index fc99d3c41..63d778344 100644 --- a/config/auth_default.go +++ b/config/auth_default.go @@ -117,6 +117,7 @@ func (c *DefaultCredentials) Configure(ctx context.Context, cfg *Config) (creden AzureClientSecretCredentials{}, AzureCliCredentials{}, // Google strategies. + GcpM2mCredentials{}, GoogleCredentials{}, GoogleDefaultCredentials{}, ) diff --git a/config/auth_gcp_oauth_m2m.go b/config/auth_gcp_oauth_m2m.go new file mode 100644 index 000000000..f09edb61c --- /dev/null +++ b/config/auth_gcp_oauth_m2m.go @@ -0,0 +1,111 @@ +package config + +import ( + "context" + "fmt" + + "github.com/databricks/databricks-sdk-go/config/credentials" + "github.com/databricks/databricks-sdk-go/config/experimental/auth" + "github.com/databricks/databricks-sdk-go/config/experimental/auth/authconv" + "github.com/databricks/databricks-sdk-go/logger" + "golang.org/x/oauth2/google" + "google.golang.org/api/impersonate" +) + +// gcpServiceAccountAccessTokenHeader carries a Google Cloud OAuth access token +// (cloud-platform scope) that Databricks uses to provision GCP resources on the +// caller's behalf. It is a passthrough credential, independent of the identity +// in the Authorization header. +const gcpServiceAccountAccessTokenHeader = "X-Databricks-GCP-SA-Access-Token" + +// GcpM2mCredentials authenticates the request identity with a Databricks OAuth +// service-principal (M2M) token, and additionally attaches a Google Cloud +// access token in the X-Databricks-GCP-SA-Access-Token header so Databricks can +// provision GCP resources on the caller's behalf. +// +// This is the supported way to call GCP account-level provisioning APIs when +// SSO is enabled: identity comes from a Databricks-governed service principal +// (Authorization header), while the Google access token supplies only the GCP +// resource credential. +// +// It must be selected explicitly via auth_type = "oauth-m2m-gcp". The mode +// combines the "oauth" (client_id/client_secret) and "google" +// (google_credentials/google_service_account) config groups, which the default +// single-auth-method conflict check rejects unless an auth type is set. +type GcpM2mCredentials struct { + // googleTokenSource, when non-nil, is used as the Google Cloud access token + // source instead of deriving one from the config. Test-only seam. + googleTokenSource auth.TokenSource +} + +func (c GcpM2mCredentials) Name() string { + return "oauth-m2m-gcp" +} + +func (c GcpM2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) { + if !cfg.IsGcp() || cfg.ClientID == "" || cfg.ClientSecret == "" { + return nil, nil + } + if c.googleTokenSource == nil && cfg.GoogleCredentials == "" && cfg.GoogleServiceAccount == "" { + return nil, nil + } + + primary, err := databricksOAuthTokenSource(ctx, cfg) + if err != nil { + return nil, err + } + + secondary := c.googleTokenSource + if secondary == nil { + secondary, err = googleAccessTokenSource(ctx, cfg) + if err != nil { + return nil, err + } + } + + logger.Infof(ctx, "Using Databricks OAuth (M2M) with GCP service account access token passthrough") + + // Google token sources cache internally; disable async refresh to avoid + // redundant work. The Databricks primary refreshes synchronously on expiry. + opts := append(cacheOptions(cfg), auth.WithAsyncRefresh(false)) + // secondaryOptional is false: this mode exists to attach the GCP access + // token, so a failure to obtain it must fail the request rather than + // silently drop the header. + visitor := serviceToServiceVisitor(primary, secondary, gcpServiceAccountAccessTokenHeader, false, opts...) + return newVisitorOAuthCredentials(visitor, auth.NewCachedTokenSource(primary, opts...)), nil +} + +// googleAccessTokenSource returns a token source for a Google Cloud OAuth +// access token (cloud-platform scope). It prefers a service-account JSON key +// (google_credentials) and falls back to service-account impersonation +// (google_service_account). +func googleAccessTokenSource(ctx context.Context, cfg *Config) (auth.TokenSource, error) { + scopes := []string{ + "https://www.googleapis.com/auth/cloud-platform", + "https://www.googleapis.com/auth/compute", + } + switch { + case cfg.GoogleCredentials != "": + jsonBytes, err := readCredentials(cfg.GoogleCredentials) + if err != nil { + return nil, fmt.Errorf("could not read GoogleCredentials. "+ + "Make sure the file exists, or the JSON content is valid: %w", err) + } + creds, err := google.CredentialsFromJSON(ctx, jsonBytes, scopes...) + if err != nil { + return nil, fmt.Errorf("could not obtain GCP access token from JSON: %w", err) + } + return authconv.AuthTokenSource(creds.TokenSource), nil + case cfg.GoogleServiceAccount != "": + platform, err := impersonate.CredentialsTokenSource(ctx, impersonate.CredentialsConfig{ + TargetPrincipal: cfg.GoogleServiceAccount, + Scopes: scopes, + }) + if err != nil { + return nil, fmt.Errorf("could not create GCP SA access token source: %w", err) + } + return authconv.AuthTokenSource(platform), nil + default: + return nil, fmt.Errorf("oauth-m2m-gcp requires google_credentials or google_service_account to be set") + } +} diff --git a/config/auth_gcp_oauth_m2m_test.go b/config/auth_gcp_oauth_m2m_test.go new file mode 100644 index 000000000..4f18349a0 --- /dev/null +++ b/config/auth_gcp_oauth_m2m_test.go @@ -0,0 +1,125 @@ +package config + +import ( + "context" + "net/http" + "testing" + + "github.com/databricks/databricks-sdk-go/common/environment" + "github.com/databricks/databricks-sdk-go/config/experimental/auth" + "github.com/databricks/databricks-sdk-go/httpclient/fixtures" + "github.com/stretchr/testify/require" + "golang.org/x/oauth2" +) + +func staticAuthTokenSource(token string) auth.TokenSource { + return auth.TokenSourceFn(func(ctx context.Context) (*oauth2.Token, error) { + return &oauth2.Token{TokenType: "Bearer", AccessToken: token}, nil + }) +} + +// The core behavior: identity token in Authorization, Google access token in +// the X-Databricks-GCP-SA-Access-Token passthrough header. The Google source is +// injected so the test is hermetic; the Databricks token is minted against a +// mocked account token endpoint. +func TestGcpM2m_SetsBothHeaders(t *testing.T) { + cfg := &Config{ + Host: "https://accounts.gcp.databricks.com", + AccountID: "abc", + Cloud: environment.CloudGCP, + ClientID: "b", + ClientSecret: "c", + GoogleServiceAccount: "sa@proj.iam.gserviceaccount.com", + AuthType: "oauth-m2m-gcp", + ConfigFile: "/dev/null", + HTTPTransport: fixtures.MappingTransport{ + "POST /oidc/accounts/abc/v1/token": { + Response: oauth2.Token{ + TokenType: "Bearer", + AccessToken: "db-oauth", + }, + }, + }, + } + require.NoError(t, cfg.EnsureResolved()) + + strat := GcpM2mCredentials{googleTokenSource: staticAuthTokenSource("gcp-access")} + provider, err := strat.Configure(context.Background(), cfg) + require.NoError(t, err) + require.NotNil(t, provider) + + req, err := http.NewRequest("GET", "http://localhost", nil) + require.NoError(t, err) + require.NoError(t, provider.SetHeaders(req)) + + require.Equal(t, "Bearer db-oauth", req.Header.Get("Authorization")) + require.Equal(t, "gcp-access", req.Header.Get("X-Databricks-GCP-SA-Access-Token")) +} + +// Combining oauth (client_id/client_secret) and google +// (google_service_account) config groups is rejected by the single-auth-method +// conflict check unless an explicit auth_type is set. +func TestGcpM2m_RequiresExplicitAuthType(t *testing.T) { + newCfg := func(authType string) *Config { + return &Config{ + Host: "https://accounts.gcp.databricks.com", + AccountID: "abc", + Cloud: environment.CloudGCP, + ClientID: "b", + ClientSecret: "c", + GoogleServiceAccount: "sa@proj.iam.gserviceaccount.com", + AuthType: authType, + ConfigFile: "/dev/null", + } + } + + err := newCfg("").EnsureResolved() + require.Error(t, err) + require.Contains(t, err.Error(), "more than one authorization method") + + require.NoError(t, newCfg("oauth-m2m-gcp").EnsureResolved()) +} + +// The strategy is inert (returns a nil provider so the chain moves on) unless it +// is on GCP with Databricks OAuth SP credentials and a Google source. +func TestGcpM2m_ConfigureSkips(t *testing.T) { + cases := []struct { + name string + cfg *Config + }{ + { + name: "not gcp", + cfg: &Config{ + Host: "https://foo.cloud.databricks.com", ClientID: "b", ClientSecret: "c", + GoogleServiceAccount: "sa@proj.iam.gserviceaccount.com", ConfigFile: "/dev/null", + }, + }, + { + name: "missing client secret", + cfg: &Config{ + Cloud: environment.CloudGCP, Host: "https://accounts.gcp.databricks.com", + ClientID: "b", GoogleServiceAccount: "sa@proj.iam.gserviceaccount.com", ConfigFile: "/dev/null", + }, + }, + { + name: "no google source", + cfg: &Config{ + Cloud: environment.CloudGCP, Host: "https://accounts.gcp.databricks.com", + ClientID: "b", ClientSecret: "c", ConfigFile: "/dev/null", + }, + }, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + provider, err := GcpM2mCredentials{}.Configure(context.Background(), tc.cfg) + require.NoError(t, err) + require.Nil(t, provider) + }) + } +} + +func TestGoogleAccessTokenSource_RequiresASource(t *testing.T) { + _, err := googleAccessTokenSource(context.Background(), &Config{}) + require.Error(t, err) + require.Contains(t, err.Error(), "requires google_credentials or google_service_account") +} diff --git a/config/auth_m2m.go b/config/auth_m2m.go index e62a52f14..e8a70ca18 100644 --- a/config/auth_m2m.go +++ b/config/auth_m2m.go @@ -22,6 +22,24 @@ func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials if cfg.ClientID == "" || cfg.ClientSecret == "" { return nil, nil } + ts, err := databricksOAuthTokenSource(ctx, cfg) + if err != nil { + return nil, err + } + return credentials.NewOAuthCredentialsProviderFromTokenSource( + auth.NewCachedTokenSource(ts, cacheOptions(cfg)...), + ), nil +} + +// databricksOAuthTokenSource returns a token source that mints Databricks OAuth +// access tokens for the configured service principal (client_id/client_secret) +// via the client-credentials grant. +// +// The returned source is uncached: it performs a network call on each Token() +// and does not itself dedupe. Callers must wrap it in auth.NewCachedTokenSource +// (as M2mCredentials does), or pass it to serviceToServiceVisitor, which caches +// internally. +func databricksOAuthTokenSource(ctx context.Context, cfg *Config) (auth.TokenSource, error) { endpoints, err := cfg.getOidcEndpoints(ctx) if err != nil { return nil, fmt.Errorf("oidc: %w", err) @@ -52,7 +70,5 @@ func (c M2mCredentials) Configure(ctx context.Context, cfg *Config) (credentials return ccfg.Token(ctx) }) - return credentials.NewOAuthCredentialsProviderFromTokenSource( - auth.NewCachedTokenSource(auth.NewRetryingTokenSource(ts), cacheOptions(cfg)...), - ), nil + return auth.NewRetryingTokenSource(ts), nil } diff --git a/config/config.go b/config/config.go index 8ba92ab87..e5396ecb6 100644 --- a/config/config.go +++ b/config/config.go @@ -146,8 +146,8 @@ type Config struct { // will disable the fallback mechanism. OAuthCallbackPort int `name:"oauth_callback_port" env:"DATABRICKS_OAUTH_CALLBACK_PORT" auth:"-"` - GoogleServiceAccount string `name:"google_service_account" env:"DATABRICKS_GOOGLE_SERVICE_ACCOUNT" auth:"google" auth_types:"google-id"` - GoogleCredentials string `name:"google_credentials" env:"GOOGLE_CREDENTIALS" auth:"google,sensitive" auth_types:"google-credentials"` + GoogleServiceAccount string `name:"google_service_account" env:"DATABRICKS_GOOGLE_SERVICE_ACCOUNT" auth:"google" auth_types:"google-id,oauth-m2m-gcp"` + GoogleCredentials string `name:"google_credentials" env:"GOOGLE_CREDENTIALS" auth:"google,sensitive" auth_types:"google-credentials,oauth-m2m-gcp"` // Azure Resource Manager ID for Azure Databricks workspace, which is exhanged for a Host AzureResourceID string `name:"azure_workspace_resource_id" env:"DATABRICKS_AZURE_RESOURCE_ID" auth:"azure" auth_types:"azure-cli,azure-msi"` @@ -176,8 +176,8 @@ type Config struct { // versions of Go SDK. AzureLoginAppID string `name:"azure_login_app_id" env:"DATABRICKS_AZURE_LOGIN_APP_ID" auth:"azure"` - ClientID string `name:"client_id" env:"DATABRICKS_CLIENT_ID" auth:"oauth" auth_types:"oauth-m2m"` - ClientSecret string `name:"client_secret" env:"DATABRICKS_CLIENT_SECRET" auth:"oauth,sensitive" auth_types:"oauth-m2m"` + ClientID string `name:"client_id" env:"DATABRICKS_CLIENT_ID" auth:"oauth" auth_types:"oauth-m2m,oauth-m2m-gcp"` + ClientSecret string `name:"client_secret" env:"DATABRICKS_CLIENT_SECRET" auth:"oauth,sensitive" auth_types:"oauth-m2m,oauth-m2m-gcp"` // Scopes is a list of OAuth scopes to request when authenticating. // diff --git a/internal/auth_gcp_m2m_test.go b/internal/auth_gcp_m2m_test.go new file mode 100644 index 000000000..2ce7c7de8 --- /dev/null +++ b/internal/auth_gcp_m2m_test.go @@ -0,0 +1,70 @@ +package internal + +import ( + "net/http" + "testing" + + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/stretchr/testify/require" +) + +// headerCapturingTransport records the headers of the most recent outgoing +// request and delegates to an inner RoundTripper. +type headerCapturingTransport struct { + inner http.RoundTripper + last http.Header +} + +func (t *headerCapturingTransport) RoundTrip(req *http.Request) (*http.Response, error) { + t.last = req.Header.Clone() + return t.inner.RoundTrip(req) +} + +// TestMwsAccAccountOAuthM2MGcpAuth verifies the oauth-m2m-gcp strategy end to +// end on a GCP account: a Databricks OAuth service-principal token authenticates +// the account API call (Authorization header), and a Google Cloud access token +// is attached in the X-Databricks-GCP-SA-Access-Token passthrough header on the +// real outgoing request. +// +// Gated on DATABRICKS_GOOGLE_SERVICE_ACCOUNT, which is only set in GCP CI, so +// the test is skipped on other clouds. +func TestMwsAccAccountOAuthM2MGcpAuth(t *testing.T) { + ctx, _ := accountTest(t) + t.Log(GetEnvOrSkipTest(t, "CLOUD_ENV")) + + host := GetEnvOrSkipTest(t, "DATABRICKS_HOST") + accountID := GetEnvOrSkipTest(t, "DATABRICKS_ACCOUNT_ID") + clientID := GetEnvOrSkipTest(t, "TEST_DATABRICKS_CLIENT_ID") + clientSecret := GetEnvOrSkipTest(t, "TEST_DATABRICKS_CLIENT_SECRET") + // The Google service account whose cloud-platform access token is passed + // through. Only present in GCP CI, so this skips elsewhere. + googleServiceAccount := GetEnvOrSkipTest(t, "DATABRICKS_GOOGLE_SERVICE_ACCOUNT") + + recorder := &headerCapturingTransport{inner: http.DefaultTransport} + + accCfg := &databricks.Config{ + Host: host, + AccountID: accountID, + ClientID: clientID, + ClientSecret: clientSecret, + GoogleServiceAccount: googleServiceAccount, + AuthType: "oauth-m2m-gcp", + HTTPTransport: recorder, + } + + accClient, err := databricks.NewAccountClient(accCfg) + require.NoError(t, err) + + // A real account API call proves the Databricks OAuth identity authenticates + // and that minting the Google access token does not break the request. + it := accClient.ServicePrincipals.List(ctx, iam.ListAccountServicePrincipalsRequest{}) + _, err = it.Next(ctx) + require.NoError(t, err) + + // Both credentials must have been attached to the actual outgoing request. + require.NotNil(t, recorder.last, "no request was captured") + require.Contains(t, recorder.last.Get("Authorization"), "Bearer ") + require.NotEmpty(t, recorder.last.Get("X-Databricks-GCP-SA-Access-Token"), + "expected the GCP SA access token passthrough header to be set") +}