Skip to content

Commit bac4c39

Browse files
hectorcast-dbhectorcast-db
committed
PR Comments
1 parent e94811a commit bac4c39

File tree

7 files changed

+73
-39
lines changed

7 files changed

+73
-39
lines changed

‎NEXT_CHANGELOG.md

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,12 @@
33
## Release v0.61.0
44

55
### New Features and Improvements
6-
* Introduce support for Databricks Workload Identity Federation in GitHub workflows ([1177](https://github.com/databricks/databricks-sdk-go/pull/1177))
6+
* Introduce support for Databricks Workload Identity Federation in GitHub workflows ([1177](https://github.com/databricks/databricks-sdk-go/pull/1177)).
77
See README.md for instructions.
8+
[Breaking] Users running their worklows in GitHub Actions, which use Cloud native authentication and also have a `DATABRICKS_CLIENT_ID` and `DATABRICKS_HOST`
9+
environment variables set may see their authentication start failing due to the order in which the SDK tries different authentication methods.
10+
In such case, the `DATABRICKS_AUTH_TYPE` environment variable must be set to match the previously used authentication method.
11+
812

913
### Bug Fixes
1014

‎README.md

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -174,20 +174,17 @@ Depending on the Databricks authentication method, the SDK uses the following in
174174

175175
### Databricks native authentication
176176

177-
By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Databricks basic (username/password) authentication (`AuthType: "basic"` in `*databricks.Config`). If unsucesful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "databricks-wif"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported.
177+
By default, the Databricks SDK for Go initially tries Databricks token authentication (`AuthType: "pat"` in `*databricks.Config`). If the SDK is unsuccessful, it then tries Workload Identity Federation (WIF) based authentication(`AuthType: "databricks-wif"` in `*databricks.Config`). Currently, only GitHub provided JWT Tokens is supported.
178178

179179
- For Databricks token authentication, you must provide `Host` and `Token`; or their environment variable or `.databrickscfg` file field equivalents.
180-
- For Databricks basic authentication, you must provide `Host`, `Username`, and `Password` _(for AWS workspace-level operations)_; or `Host`, `AccountID`, `Username`, and `Password` _(for AWS, Azure, or GCP account-level operations)_; or their environment variable or `.databrickscfg` file field equivalents.
181-
- For Databricks wif authentication, you must provide `Host`, `ClientID` and `TokenAudience`; or their environment variable or `.databrickscfg` file field equivalents.
180+
- For Databricks wif authentication, you must provide `Host`, `ClientID` and `TokenAudience` _(optional)_; or their environment variable or `.databrickscfg` file field equivalents.
182181

183182
| `*databricks.Config` argument | Description | Environment variable / `.databrickscfg` file field |
184183
| ----------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------------------------------------------------- |
185184
| `Host` | _(String)_ The Databricks host URL for either the Databricks workspace endpoint or the Databricks accounts endpoint. | `DATABRICKS_HOST` / `host` |
186185
| `AccountID` | _(String)_ The Databricks account ID for the Databricks accounts endpoint. Only has effect when `Host` is either `https://accounts.cloud.databricks.com/` _(AWS)_, `https://accounts.azuredatabricks.net/` _(Azure)_, or `https://accounts.gcp.databricks.com/` _(GCP)_. | `DATABRICKS_ACCOUNT_ID` / `account_id` |
187186
| `Token` | _(String)_ The Databricks personal access token (PAT) _(AWS, Azure, and GCP)_ or Azure Active Directory (Azure AD) token _(Azure)_. | `DATABRICKS_TOKEN` / `token` |
188-
| `Username` | _(String)_ The Databricks username part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_USERNAME` / `username` |
189-
| `Password` | _(String)_ The Databricks password part of basic authentication. Only possible when `Host` is `*.cloud.databricks.com` _(AWS)_. | `DATABRICKS_PASSWORD` / `password` |
190-
| `TokenAudience` | _(String)_ The Audience for the JWT Token. | `TOKEN_AUDIENCE` / `token_audience` |
187+
| `TokenAudience` | _(String)_ When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier. | `TOKEN_AUDIENCE` / `token_audience` |
191188

192189
For example, to use Databricks token authentication:
193190

‎config/auth_databricks_wif.go

Lines changed: 20 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -11,11 +11,16 @@ import (
1111
"golang.org/x/oauth2/clientcredentials"
1212
)
1313

14+
// DatabricksWIFCredentials uses a Token Supplier to get a JWT Token and exchanges
15+
// it for a Databricks Token.
16+
//
17+
// Supported suppliers:
18+
// - GitHub OIDC
1419
type DatabricksWIFCredentials struct{}
1520

1621
// Configure implements CredentialsStrategy.
1722
func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (credentials.CredentialsProvider, error) {
18-
if cfg.Host == "" || cfg.ClientID == "" || cfg.TokenAudience == "" {
23+
if cfg.Host == "" || cfg.ClientID == "" {
1924
return nil, nil
2025
}
2126

@@ -38,12 +43,12 @@ func (d DatabricksWIFCredentials) Configure(ctx context.Context, cfg *Config) (c
3843
}
3944

4045
ts := &databricksWIFTokenSource{
41-
ctx: ctx,
42-
jwtTokenSupplicer: &supplier,
43-
audience: cfg.TokenAudience,
44-
clientId: cfg.ClientID,
45-
cfg: cfg,
46-
tokenEndpoint: endpoints.TokenEndpoint,
46+
ctx: ctx,
47+
idTokenSupplier: &supplier,
48+
audience: cfg.TokenAudience,
49+
clientId: cfg.ClientID,
50+
cfg: cfg,
51+
tokenEndpoint: endpoints.TokenEndpoint,
4752
}
4853

4954
visitor := refreshableVisitor(ts)
@@ -56,16 +61,16 @@ func (d DatabricksWIFCredentials) Name() string {
5661
}
5762

5863
type databricksWIFTokenSource struct {
59-
ctx context.Context
60-
jwtTokenSupplicer *GithubOIDCTokenSupplier
61-
tokenEndpoint string
62-
audience string
63-
clientId string
64-
cfg *Config
64+
ctx context.Context
65+
idTokenSupplier *GithubOIDCTokenSupplier
66+
tokenEndpoint string
67+
audience string
68+
clientId string
69+
cfg *Config
6570
}
6671

6772
func (d *databricksWIFTokenSource) Token() (*oauth2.Token, error) {
68-
token, err := d.jwtTokenSupplicer.GetOIDCToken(d.ctx, d.audience)
73+
token, err := d.idTokenSupplier.GetOIDCToken(d.ctx, d.audience)
6974
if err != nil {
7075
return nil, err
7176
}
@@ -87,7 +92,7 @@ func (d *databricksWIFTokenSource) Token() (*oauth2.Token, error) {
8792
},
8893
}
8994

90-
logger.Debugf(d.ctx, "Getting tokken for client %s", d.clientId)
95+
logger.Debugf(d.ctx, "Getting token for client %s", d.clientId)
9196

9297
return tsConfig.Token(d.ctx)
9398
}

‎config/auth_databricks_wif_test.go

Lines changed: 35 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -25,16 +25,6 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) {
2525
},
2626
wantErrPrefix: errPrefix("databricks-wif auth: not configured"),
2727
},
28-
{
29-
desc: "missing token audience",
30-
cfg: &Config{
31-
Host: "http://host.com/test",
32-
ClientID: "client-id",
33-
ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1",
34-
ActionsIDTokenRequestToken: "token-1337",
35-
},
36-
wantErrPrefix: errPrefix("databricks-wif auth: not configured"),
37-
},
3828
{
3929
desc: "missing host",
4030
cfg: &Config{
@@ -257,6 +247,41 @@ func TestDatabricksGithubWIFCredentials(t *testing.T) {
257247
"Authorization": "access-token test-auth-token",
258248
},
259249
},
250+
{
251+
desc: "default token audience",
252+
cfg: &Config{
253+
ClientID: "client-id",
254+
AccountID: "ac123",
255+
Host: "https://accounts.databricks.com",
256+
ActionsIDTokenRequestURL: "http://endpoint.com/test?version=1",
257+
ActionsIDTokenRequestToken: "token-1337",
258+
HTTPTransport: fixtures.MappingTransport{
259+
"GET /test?version=1": {
260+
Status: http.StatusOK,
261+
ExpectedHeaders: map[string]string{
262+
"Authorization": "Bearer token-1337",
263+
"Accept": "application/json",
264+
},
265+
Response: `{"value": "id-token-42"}`,
266+
},
267+
"POST /oidc/accounts/ac123/v1/token": {
268+
Status: http.StatusOK,
269+
ExpectedHeaders: map[string]string{
270+
"Content-Type": "application/x-www-form-urlencoded",
271+
},
272+
Response: map[string]string{
273+
"token_type": "access-token",
274+
"access_token": "test-auth-token",
275+
"refresh_token": "refresh",
276+
"expires_on": "0",
277+
},
278+
},
279+
},
280+
},
281+
wantHeaders: map[string]string{
282+
"Authorization": "access-token test-auth-token",
283+
},
284+
},
260285
}
261286

262287
for _, tc := range testCases {

‎config/config.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ type Config struct {
133133
// Environment override to return when resolving the current environment.
134134
DatabricksEnvironment *environment.DatabricksEnvironment
135135

136-
// OIDC and WIF audience for the token
136+
// When using Workload Identity Federation, the audience to specify when fetching an ID token from the ID token supplier.
137137
TokenAudience string `name:"audience" auth:"-"`
138138

139139
Loaders []Loader

‎config/oidc_github.go

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package config
33
import (
44
"context"
55
"fmt"
6-
"time"
76

87
"github.com/databricks/databricks-sdk-go/httpclient"
98
"github.com/databricks/databricks-sdk-go/logger"
@@ -27,17 +26,17 @@ func (g *GithubOIDCTokenSupplier) GetOIDCToken(ctx context.Context, audience str
2726
resp := struct { // anonymous struct to parse the response
2827
Value string `json:"value"`
2928
}{}
30-
err := g.cfg.refreshClient.Do(ctx, "GET", fmt.Sprintf("%s&audience=%s", g.cfg.ActionsIDTokenRequestURL, audience),
29+
requestUrl := g.cfg.ActionsIDTokenRequestURL
30+
if audience != "" {
31+
requestUrl = fmt.Sprintf("%s&audience=%s", requestUrl, audience)
32+
}
33+
err := g.cfg.refreshClient.Do(ctx, "GET", requestUrl,
3134
httpclient.WithRequestHeader("Authorization", fmt.Sprintf("Bearer %s", g.cfg.ActionsIDTokenRequestToken)),
3235
httpclient.WithResponseUnmarshal(&resp),
3336
)
3437
if err != nil {
3538
return "", fmt.Errorf("failed to request ID token from %s: %w", g.cfg.ActionsIDTokenRequestURL, err)
3639
}
3740

38-
// GitHub issued time is not allways in sync, and can give tokens which are not yet valid.
39-
// TODO: Remove this after Databricks API is updated to handle such cases.
40-
time.Sleep(2 * time.Second)
41-
4241
return resp.Value, nil
4342
}

‎internal/auth_test.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import (
1212
)
1313

1414
func TestUcAccWifAuth(t *testing.T) {
15+
// This test cannot be run locally. It can only be run from GitHub Workflows.
16+
_ = GetEnvOrSkipTest(t, "ACTIONS_ID_TOKEN_REQUEST_URL")
1517
ctx, a := ucacctTest(t)
1618

1719
// Create SP with access to the workspace
@@ -74,6 +76,8 @@ func TestUcAccWifAuth(t *testing.T) {
7476
}
7577

7678
func TestUcAccWifAuthWorkspace(t *testing.T) {
79+
// This test cannot be run locally. It can only be run from GitHub Workflows.
80+
_ = GetEnvOrSkipTest(t, "ACTIONS_ID_TOKEN_REQUEST_URL")
7781
ctx, a := ucacctTest(t)
7882

7983
workspaceIdEnvVar := GetEnvOrSkipTest(t, "TEST_WORKSPACE_ID")

0 commit comments

Comments
 (0)