From 67145b48ef9bfa2efaa1c99d99c962bd5ae54fc4 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 21 Oct 2026 18:47:11 +0000 Subject: [PATCH 1/5] air: trim the SCIM Me call and overlap it with config validation `air run` resolves the user with SCIM Me before it can build the command path that config validation needs, so the two calls ran back to back. - Ask SCIM Me for only `userName`; the full response also resolves groups and entitlements. - When the workspace's own OAuth token names a user, start validation on that user's path while SCIM Me is in flight. SCIM's answer still decides the launch directory: if it differs from the token, validation reruns on SCIM's path before anything is created, uploaded, or submitted. Co-authored-by: Isaac --- cmd/air/run.go | 9 +++- cmd/air/runlaunch.go | 94 +++++++++++++++++++++++++++-------- cmd/air/runlaunch_test.go | 100 ++++++++++++++++++++++++++++++++++---- cmd/air/runsubmit.go | 12 ++--- 4 files changed, 176 insertions(+), 39 deletions(-) diff --git a/cmd/air/run.go b/cmd/air/run.go index c8027b09e9a..07d3e2c44a2 100644 --- a/cmd/air/run.go +++ b/cmd/air/run.go @@ -56,7 +56,12 @@ func validateDryRunInWorkspace(ctx context.Context, cmd *cobra.Command, args []s } w := cmdctx.WorkspaceClient(ctx) - _, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) + // Path and validation failures are classified differently, so keep validation's result apart. + var validateErr error + _, _, _, err := prospectiveLaunchPaths(ctx, w, cfg, func(funcDir, commandPath string) error { + validateErr = validateDryRunConfig(ctx, w, cfg, commandPath, submittedContainers(cfg, funcDir), idempotencyToken, timeout) + return nil + }) if err != nil { unavailable, retryable := classifyValidationFailure(err) if !unavailable { @@ -65,7 +70,7 @@ func validateDryRunInWorkspace(ctx context.Context, cmd *cobra.Command, args []s return ctx, asValidationUnavailable(err, retryable) } - return ctx, validateDryRunConfig(ctx, w, cfg, commandPath, submittedContainers(cfg, funcDir), idempotencyToken, timeout) + return ctx, validateErr } func newRunCommand() *cobra.Command { diff --git a/cmd/air/runlaunch.go b/cmd/air/runlaunch.go index 5566ec98ca5..7c9f7853a78 100644 --- a/cmd/air/runlaunch.go +++ b/cmd/air/runlaunch.go @@ -2,6 +2,8 @@ package aircmd import ( "context" + "encoding/base64" + "encoding/json" "errors" "fmt" "path" @@ -11,8 +13,10 @@ import ( "github.com/databricks/cli/libs/env" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" + "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/workspace" + "golang.org/x/sync/errgroup" ) // userWorkspaceDirEnv overrides the per-user workspace directory; mirrors the @@ -21,30 +25,49 @@ const userWorkspaceDirEnv = "DATABRICKS_INTERNAL_USER_WORKSPACE_DIR" // currentUserEmail returns the authenticated user's email (works for any domain). func currentUserEmail(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { - me, err := w.CurrentUser.Me(ctx, iam.MeRequest{}) + // Only userName is needed; the full response also resolves groups and entitlements. + me, err := w.CurrentUser.Me(ctx, iam.MeRequest{Attributes: "userName"}) if err != nil { return "", fmt.Errorf("failed to resolve current user: %w", err) } return me.UserName, nil } -// userWorkspaceDir returns the user's workspace home, honoring the env override. -func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { - if override := env.Get(ctx, userWorkspaceDirEnv); override != "" { - return override, nil +// userNameFromWorkspaceToken returns the sub claim of cfg's OAuth access token when this +// workspace issued it, or "" otherwise. It is only a guess: sub need not equal the SCIM userName. +func userNameFromWorkspaceToken(ctx context.Context, cfg *config.Config) string { + token, err := cfg.GetTokenSource().Token(ctx) + if err != nil { + return "" } - override, err := profileValue(w.Config, "databricks_internal_user_workspace_dir") + parts := strings.Split(token.AccessToken, ".") + if len(parts) != 3 { + return "" + } + payload, err := base64.RawURLEncoding.DecodeString(parts[1]) if err != nil { - return "", fmt.Errorf("failed to read selected profile: %w", err) + return "" } - if override != "" { + var claims struct { + Issuer string `json:"iss"` + Subject string `json:"sub"` + } + if json.Unmarshal(payload, &claims) != nil || claims.Issuer != cfg.Host+"/oidc" { + return "" + } + return claims.Subject +} + +// userWorkspaceDirOverride returns the per-user workspace directory set by env or profile, if any. +func userWorkspaceDirOverride(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + if override := env.Get(ctx, userWorkspaceDirEnv); override != "" { return override, nil } - email, err := currentUserEmail(ctx, w) + override, err := profileValue(w.Config, "databricks_internal_user_workspace_dir") if err != nil { - return "", err + return "", fmt.Errorf("failed to read selected profile: %w", err) } - return "/Workspace/Users/" + email, nil + return override, nil } // cliLaunchDir returns a unique workspace directory for a run's launch artifacts: @@ -57,19 +80,50 @@ func cliLaunchDir(base, experiment, run string) string { return path.Join(base, ".air", "cli_launch", experiment, run+"_"+unique) } -// prospectiveLaunchPaths returns the workspace paths a submission would use -// without creating directories or uploading artifacts. -func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig) (string, string, string, error) { - base, err := userWorkspaceDir(ctx, w) - if err != nil { - return "", "", "", err - } +// prospectiveLaunchPaths returns the workspace paths a submission would use without creating +// directories or uploading artifacts, after running validate on them. +// +// The user's directory comes from SCIM Me unless overridden. The access token usually names the +// user already, so validate first runs on that guess while SCIM Me is in flight, and again on +// SCIM's answer if it differs. +func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, validate func(launchDir, commandPath string) error) (string, string, string, error) { runName := "" if cfg.MLflowRunName != nil { runName = *cfg.MLflowRunName } - launchDir := cliLaunchDir(base, cfg.ExperimentName, runName) - return base, launchDir, path.Join(launchDir, commandScriptName), nil + validated := func(base string) (string, string, string, error) { + launchDir := cliLaunchDir(base, cfg.ExperimentName, runName) + commandPath := path.Join(launchDir, commandScriptName) + return base, launchDir, commandPath, validate(launchDir, commandPath) + } + + override, err := userWorkspaceDirOverride(ctx, w) + if err != nil { + return "", "", "", err + } + if override != "" { + return validated(override) + } + + guess := userNameFromWorkspaceToken(ctx, w.Config) + var email string + var g errgroup.Group + g.Go(func() (err error) { + email, err = currentUserEmail(ctx, w) + return err + }) + var base, launchDir, commandPath string + var guessErr error + if guess != "" { + base, launchDir, commandPath, guessErr = validated("/Workspace/Users/" + guess) + } + if err := g.Wait(); err != nil { + return "", "", "", err + } + if guess == "" || email != guess { + return validated("/Workspace/Users/" + email) + } + return base, launchDir, commandPath, guessErr } // ensureExperimentDirectory creates experimentDir if it is missing, matching the diff --git a/cmd/air/runlaunch_test.go b/cmd/air/runlaunch_test.go index 7654e544964..eb12ade37f1 100644 --- a/cmd/air/runlaunch_test.go +++ b/cmd/air/runlaunch_test.go @@ -1,6 +1,10 @@ package aircmd import ( + "context" + "encoding/base64" + "encoding/json" + "errors" "os" "path/filepath" "strings" @@ -9,8 +13,11 @@ import ( "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/config" + sdkauth "github.com/databricks/databricks-sdk-go/config/experimental/auth" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/oauth2" ) func TestCliLaunchDir(t *testing.T) { @@ -34,15 +41,91 @@ func newFakeWorkspaceClient(t *testing.T) *databricks.WorkspaceClient { func TestUserWorkspaceDir(t *testing.T) { w := newFakeWorkspaceClient(t) t.Setenv(userWorkspaceDirEnv, "") - dir, err := userWorkspaceDir(t.Context(), w) - require.NoError(t, err) + dir := launchBase(t, w) assert.True(t, strings.HasPrefix(dir, "/Workspace/Users/"), dir) // The env override wins without an API call. t.Setenv(userWorkspaceDirEnv, "/Workspace/custom") - dir, err = userWorkspaceDir(t.Context(), w) + assert.Equal(t, "/Workspace/custom", launchBase(t, w)) +} + +// launchBase returns the user directory prospectiveLaunchPaths resolves, with validation stubbed out. +func launchBase(t *testing.T, w *databricks.WorkspaceClient) string { + t.Helper() + base, _, _, err := prospectiveLaunchPaths(t.Context(), w, baseRunConfig(), func(string, string) error { return nil }) + require.NoError(t, err) + return base +} + +func fakeJWT(t *testing.T, issuer, subject string) string { + payload, err := json.Marshal(map[string]string{"iss": issuer, "sub": subject}) require.NoError(t, err) - assert.Equal(t, "/Workspace/custom", dir) + enc := base64.RawURLEncoding + return enc.EncodeToString([]byte(`{"alg":"none"}`)) + "." + enc.EncodeToString(payload) + ".sig" +} + +func oauthConfig(host, accessToken string) *databricks.Config { + ts := sdkauth.TokenSourceFn(func(context.Context) (*oauth2.Token, error) { + return &oauth2.Token{AccessToken: accessToken}, nil + }) + return &databricks.Config{Host: host, Credentials: config.NewTokenSourceStrategy("test-oauth", ts)} +} + +// The token's user is only a guess that lets validation start early: SCIM Me always decides. +func TestProspectiveLaunchPathsTokenGuess(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + var attributes []string + server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any { + attributes = append(attributes, req.URL.Query().Get("attributes")) + return testserver.TestUser + }) + testserver.AddDefaultHandlers(server) + issuer := server.URL + "/oidc" + scimDir := "/Workspace/Users/" + testserver.TestUser.UserName + otherDir := "/Workspace/Users/other@example.com" + otherToken := fakeJWT(t, issuer, "other@example.com") + + for _, tc := range []struct { + name string + cfg *databricks.Config + override string + validated []string + }{ + {"token names the SCIM user", oauthConfig(server.URL, fakeJWT(t, issuer, testserver.TestUser.UserName)), "", []string{scimDir}}, + {"token names another user", oauthConfig(server.URL, otherToken), "", []string{otherDir, scimDir}}, + {"token from another issuer", oauthConfig(server.URL, fakeJWT(t, "https://accounts.test/oidc", "other@example.com")), "", []string{scimDir}}, + {"opaque token", oauthConfig(server.URL, "not-a-jwt"), "", []string{scimDir}}, + {"pat", &databricks.Config{Host: server.URL, Token: otherToken}, "", []string{scimDir}}, + {"override", oauthConfig(server.URL, otherToken), "/Workspace/custom", []string{"/Workspace/custom"}}, + } { + t.Run(tc.name, func(t *testing.T) { + t.Setenv(userWorkspaceDirEnv, tc.override) + w, err := databricks.NewWorkspaceClient(tc.cfg) + require.NoError(t, err) + + var launchDirs []string + base, launchDir, commandPath, err := prospectiveLaunchPaths(t.Context(), w, baseRunConfig(), func(launchDir, _ string) error { + launchDirs = append(launchDirs, launchDir) + if strings.HasPrefix(launchDir, otherDir+"/") { + return errors.New("validated an unconfirmed user") + } + return nil + }) + require.NoError(t, err) + require.Len(t, launchDirs, len(tc.validated)) + for i, dir := range tc.validated { + assert.True(t, strings.HasPrefix(launchDirs[i], dir+"/.air/cli_launch/"), launchDirs[i]) + } + assert.Equal(t, tc.validated[len(tc.validated)-1], base) + assert.Equal(t, launchDirs[len(launchDirs)-1], launchDir) + assert.Equal(t, launchDir+"/command.sh", commandPath) + }) + } + require.NotEmpty(t, attributes) + for _, a := range attributes { + assert.Equal(t, "userName", a) + } } func TestUserWorkspaceDirProfilePrecedence(t *testing.T) { @@ -65,8 +148,7 @@ databricks_internal_user_workspace_dir = /Workspace/second w.Config.ConfigFile = configPath w.Config.Profile = "first" t.Setenv(userWorkspaceDirEnv, "") - dir, err := userWorkspaceDir(t.Context(), w) - require.NoError(t, err) + dir := launchBase(t, w) assert.Equal(t, "/Workspace/first", dir) }) @@ -75,8 +157,7 @@ databricks_internal_user_workspace_dir = /Workspace/second w.Config.ConfigFile = configPath w.Config.Profile = "" t.Setenv(userWorkspaceDirEnv, "") - dir, err := userWorkspaceDir(t.Context(), w) - require.NoError(t, err) + dir := launchBase(t, w) assert.Equal(t, "/Workspace/second", dir) }) @@ -85,8 +166,7 @@ databricks_internal_user_workspace_dir = /Workspace/second w.Config.ConfigFile = configPath w.Config.Profile = "first" t.Setenv(userWorkspaceDirEnv, "/Workspace/env") - dir, err := userWorkspaceDir(t.Context(), w) - require.NoError(t, err) + dir := launchBase(t, w) assert.Equal(t, "/Workspace/env", dir) }) } diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index d1e94e068dc..3ae29f1038a 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -337,17 +337,15 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run defer func() { logRunEvent(ctx, cfg, snap, runID, time.Since(start), err) }() // Compute and validate the actual submission path before creating artifacts. - base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) + // Validate server-side before uploading, so invalid configs leave no artifacts. + token := effectiveIdempotencyToken(idempotencyKey, cfg) + base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg, func(funcDir, commandPath string) error { + return preflightValidate(ctx, w, cfg, commandPath, submittedContainers(cfg, funcDir), token) + }) if err != nil { return 0, "", err } containers := submittedContainers(cfg, funcDir) - token := effectiveIdempotencyToken(idempotencyKey, cfg) - - // Validate server-side before uploading, so invalid configs leave no artifacts. - if err := preflightValidate(ctx, w, cfg, commandPath, containers, token); err != nil { - return 0, "", err - } // Retain the local guard after backend validation so an unavailable endpoint // cannot let a bad key reach upload or submission. From 5f2bd15e4db22e69c2fb9fd35cd76a0b4c04519d Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 21 Oct 2026 18:56:49 +0000 Subject: [PATCH 2/5] Add changelog entry Co-authored-by: Isaac --- .nextchanges/air/7153.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 .nextchanges/air/7153.md diff --git a/.nextchanges/air/7153.md b/.nextchanges/air/7153.md new file mode 100644 index 00000000000..6aabf43bec5 --- /dev/null +++ b/.nextchanges/air/7153.md @@ -0,0 +1 @@ +* `air run` now spends less time before submitting: it fetches only the user name from SCIM and validates the config while that lookup is in flight. ([#7153](https://github.com/databricks/cli/pull/7153)) From 8163540d8d6e4a2f583ca9c202b4e02179b4a73c Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 21 Oct 2026 19:07:59 +0000 Subject: [PATCH 3/5] air: keep userWorkspaceDir and compare it with the token's guess Leaves userWorkspaceDir and its tests unchanged: the token's guess is checked against its answer, so an override is handled like any other mismatch. Co-authored-by: Isaac --- cmd/air/runlaunch.go | 47 ++++++++++++++++++++------------------- cmd/air/runlaunch_test.go | 27 ++++++++++------------ 2 files changed, 36 insertions(+), 38 deletions(-) diff --git a/cmd/air/runlaunch.go b/cmd/air/runlaunch.go index 7c9f7853a78..b48ae89b516 100644 --- a/cmd/air/runlaunch.go +++ b/cmd/air/runlaunch.go @@ -58,8 +58,8 @@ func userNameFromWorkspaceToken(ctx context.Context, cfg *config.Config) string return claims.Subject } -// userWorkspaceDirOverride returns the per-user workspace directory set by env or profile, if any. -func userWorkspaceDirOverride(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { +// userWorkspaceDir returns the user's workspace home, honoring the env override. +func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { if override := env.Get(ctx, userWorkspaceDirEnv); override != "" { return override, nil } @@ -67,7 +67,14 @@ func userWorkspaceDirOverride(ctx context.Context, w *databricks.WorkspaceClient if err != nil { return "", fmt.Errorf("failed to read selected profile: %w", err) } - return override, nil + if override != "" { + return override, nil + } + email, err := currentUserEmail(ctx, w) + if err != nil { + return "", err + } + return "/Workspace/Users/" + email, nil } // cliLaunchDir returns a unique workspace directory for a run's launch artifacts: @@ -83,9 +90,8 @@ func cliLaunchDir(base, experiment, run string) string { // prospectiveLaunchPaths returns the workspace paths a submission would use without creating // directories or uploading artifacts, after running validate on them. // -// The user's directory comes from SCIM Me unless overridden. The access token usually names the -// user already, so validate first runs on that guess while SCIM Me is in flight, and again on -// SCIM's answer if it differs. +// The access token usually names the user, so validate first runs on that guess while +// userWorkspaceDir looks up the real directory, and again if the two differ. func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, validate func(launchDir, commandPath string) error) (string, string, string, error) { runName := "" if cfg.MLflowRunName != nil { @@ -97,33 +103,28 @@ func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, return base, launchDir, commandPath, validate(launchDir, commandPath) } - override, err := userWorkspaceDirOverride(ctx, w) - if err != nil { - return "", "", "", err + guessBase := "" + if user := userNameFromWorkspaceToken(ctx, w.Config); user != "" { + guessBase = "/Workspace/Users/" + user } - if override != "" { - return validated(override) - } - - guess := userNameFromWorkspaceToken(ctx, w.Config) - var email string + var base string var g errgroup.Group g.Go(func() (err error) { - email, err = currentUserEmail(ctx, w) + base, err = userWorkspaceDir(ctx, w) return err }) - var base, launchDir, commandPath string - var guessErr error - if guess != "" { - base, launchDir, commandPath, guessErr = validated("/Workspace/Users/" + guess) + var launchDir, commandPath string + var validateErr error + if guessBase != "" { + _, launchDir, commandPath, validateErr = validated(guessBase) } if err := g.Wait(); err != nil { return "", "", "", err } - if guess == "" || email != guess { - return validated("/Workspace/Users/" + email) + if base != guessBase { + return validated(base) } - return base, launchDir, commandPath, guessErr + return base, launchDir, commandPath, validateErr } // ensureExperimentDirectory creates experimentDir if it is missing, matching the diff --git a/cmd/air/runlaunch_test.go b/cmd/air/runlaunch_test.go index eb12ade37f1..b63115c6553 100644 --- a/cmd/air/runlaunch_test.go +++ b/cmd/air/runlaunch_test.go @@ -41,20 +41,15 @@ func newFakeWorkspaceClient(t *testing.T) *databricks.WorkspaceClient { func TestUserWorkspaceDir(t *testing.T) { w := newFakeWorkspaceClient(t) t.Setenv(userWorkspaceDirEnv, "") - dir := launchBase(t, w) + dir, err := userWorkspaceDir(t.Context(), w) + require.NoError(t, err) assert.True(t, strings.HasPrefix(dir, "/Workspace/Users/"), dir) // The env override wins without an API call. t.Setenv(userWorkspaceDirEnv, "/Workspace/custom") - assert.Equal(t, "/Workspace/custom", launchBase(t, w)) -} - -// launchBase returns the user directory prospectiveLaunchPaths resolves, with validation stubbed out. -func launchBase(t *testing.T, w *databricks.WorkspaceClient) string { - t.Helper() - base, _, _, err := prospectiveLaunchPaths(t.Context(), w, baseRunConfig(), func(string, string) error { return nil }) + dir, err = userWorkspaceDir(t.Context(), w) require.NoError(t, err) - return base + assert.Equal(t, "/Workspace/custom", dir) } func fakeJWT(t *testing.T, issuer, subject string) string { @@ -71,7 +66,7 @@ func oauthConfig(host, accessToken string) *databricks.Config { return &databricks.Config{Host: host, Credentials: config.NewTokenSourceStrategy("test-oauth", ts)} } -// The token's user is only a guess that lets validation start early: SCIM Me always decides. +// The token's user is only a guess that lets validation start early: userWorkspaceDir decides. func TestProspectiveLaunchPathsTokenGuess(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) @@ -95,9 +90,8 @@ func TestProspectiveLaunchPathsTokenGuess(t *testing.T) { {"token names the SCIM user", oauthConfig(server.URL, fakeJWT(t, issuer, testserver.TestUser.UserName)), "", []string{scimDir}}, {"token names another user", oauthConfig(server.URL, otherToken), "", []string{otherDir, scimDir}}, {"token from another issuer", oauthConfig(server.URL, fakeJWT(t, "https://accounts.test/oidc", "other@example.com")), "", []string{scimDir}}, - {"opaque token", oauthConfig(server.URL, "not-a-jwt"), "", []string{scimDir}}, {"pat", &databricks.Config{Host: server.URL, Token: otherToken}, "", []string{scimDir}}, - {"override", oauthConfig(server.URL, otherToken), "/Workspace/custom", []string{"/Workspace/custom"}}, + {"override", oauthConfig(server.URL, otherToken), "/Workspace/custom", []string{otherDir, "/Workspace/custom"}}, } { t.Run(tc.name, func(t *testing.T) { t.Setenv(userWorkspaceDirEnv, tc.override) @@ -148,7 +142,8 @@ databricks_internal_user_workspace_dir = /Workspace/second w.Config.ConfigFile = configPath w.Config.Profile = "first" t.Setenv(userWorkspaceDirEnv, "") - dir := launchBase(t, w) + dir, err := userWorkspaceDir(t.Context(), w) + require.NoError(t, err) assert.Equal(t, "/Workspace/first", dir) }) @@ -157,7 +152,8 @@ databricks_internal_user_workspace_dir = /Workspace/second w.Config.ConfigFile = configPath w.Config.Profile = "" t.Setenv(userWorkspaceDirEnv, "") - dir := launchBase(t, w) + dir, err := userWorkspaceDir(t.Context(), w) + require.NoError(t, err) assert.Equal(t, "/Workspace/second", dir) }) @@ -166,7 +162,8 @@ databricks_internal_user_workspace_dir = /Workspace/second w.Config.ConfigFile = configPath w.Config.Profile = "first" t.Setenv(userWorkspaceDirEnv, "/Workspace/env") - dir := launchBase(t, w) + dir, err := userWorkspaceDir(t.Context(), w) + require.NoError(t, err) assert.Equal(t, "/Workspace/env", dir) }) } From 6db6cf9582367345ad23a4718933fde9c84098ea Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 21 Oct 2026 19:22:47 +0000 Subject: [PATCH 4/5] air: guess the user only from cached workspace tokens, and from overrides Only databricks-cli and oauth-m2m serve Token from the cache request headers use; azure-cli and similar providers would fetch a fresh token just to read it. A configured directory override is used as the guess, so override users still validate once. Co-authored-by: Isaac --- cmd/air/runlaunch.go | 35 ++++++++++++++++++++++++++--------- cmd/air/runlaunch_test.go | 15 +++++++++------ 2 files changed, 35 insertions(+), 15 deletions(-) diff --git a/cmd/air/runlaunch.go b/cmd/air/runlaunch.go index b48ae89b516..4aac773b215 100644 --- a/cmd/air/runlaunch.go +++ b/cmd/air/runlaunch.go @@ -10,6 +10,7 @@ import ( "strings" "uuid" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/env" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" @@ -36,6 +37,11 @@ func currentUserEmail(ctx context.Context, w *databricks.WorkspaceClient) (strin // userNameFromWorkspaceToken returns the sub claim of cfg's OAuth access token when this // workspace issued it, or "" otherwise. It is only a guess: sub need not equal the SCIM userName. func userNameFromWorkspaceToken(ctx context.Context, cfg *config.Config) string { + // These mint workspace tokens and serve Token from the cache request headers use; others, + // such as azure-cli, would fetch a fresh token here. + if cfg.AuthType != auth.AuthTypeDatabricksCli && cfg.AuthType != auth.AuthTypeOAuthM2M { + return "" + } token, err := cfg.GetTokenSource().Token(ctx) if err != nil { return "" @@ -58,8 +64,8 @@ func userNameFromWorkspaceToken(ctx context.Context, cfg *config.Config) string return claims.Subject } -// userWorkspaceDir returns the user's workspace home, honoring the env override. -func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { +// userWorkspaceDirOverride returns the per-user workspace directory set by env or profile, if any. +func userWorkspaceDirOverride(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { if override := env.Get(ctx, userWorkspaceDirEnv); override != "" { return override, nil } @@ -67,8 +73,13 @@ func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (strin if err != nil { return "", fmt.Errorf("failed to read selected profile: %w", err) } - if override != "" { - return override, nil + return override, nil +} + +// userWorkspaceDir returns the user's workspace home, honoring the env override. +func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { + if override, err := userWorkspaceDirOverride(ctx, w); override != "" || err != nil { + return override, err } email, err := currentUserEmail(ctx, w) if err != nil { @@ -90,8 +101,9 @@ func cliLaunchDir(base, experiment, run string) string { // prospectiveLaunchPaths returns the workspace paths a submission would use without creating // directories or uploading artifacts, after running validate on them. // -// The access token usually names the user, so validate first runs on that guess while -// userWorkspaceDir looks up the real directory, and again if the two differ. +// The directory is usually known before SCIM Me answers (an override, or the user a +// workspace-issued token names), so validate first runs on that guess while userWorkspaceDir +// looks up the real directory, and again if the two differ. func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, validate func(launchDir, commandPath string) error) (string, string, string, error) { runName := "" if cfg.MLflowRunName != nil { @@ -103,9 +115,14 @@ func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, return base, launchDir, commandPath, validate(launchDir, commandPath) } - guessBase := "" - if user := userNameFromWorkspaceToken(ctx, w.Config); user != "" { - guessBase = "/Workspace/Users/" + user + guessBase, err := userWorkspaceDirOverride(ctx, w) + if err != nil { + return "", "", "", err + } + if guessBase == "" { + if user := userNameFromWorkspaceToken(ctx, w.Config); user != "" { + guessBase = "/Workspace/Users/" + user + } } var base string var g errgroup.Group diff --git a/cmd/air/runlaunch_test.go b/cmd/air/runlaunch_test.go index b63115c6553..3b342024664 100644 --- a/cmd/air/runlaunch_test.go +++ b/cmd/air/runlaunch_test.go @@ -10,6 +10,7 @@ import ( "strings" "testing" + "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" @@ -59,11 +60,11 @@ func fakeJWT(t *testing.T, issuer, subject string) string { return enc.EncodeToString([]byte(`{"alg":"none"}`)) + "." + enc.EncodeToString(payload) + ".sig" } -func oauthConfig(host, accessToken string) *databricks.Config { +func oauthConfig(host, authType, accessToken string) *databricks.Config { ts := sdkauth.TokenSourceFn(func(context.Context) (*oauth2.Token, error) { return &oauth2.Token{AccessToken: accessToken}, nil }) - return &databricks.Config{Host: host, Credentials: config.NewTokenSourceStrategy("test-oauth", ts)} + return &databricks.Config{Host: host, Credentials: config.NewTokenSourceStrategy(authType, ts)} } // The token's user is only a guess that lets validation start early: userWorkspaceDir decides. @@ -87,11 +88,13 @@ func TestProspectiveLaunchPathsTokenGuess(t *testing.T) { override string validated []string }{ - {"token names the SCIM user", oauthConfig(server.URL, fakeJWT(t, issuer, testserver.TestUser.UserName)), "", []string{scimDir}}, - {"token names another user", oauthConfig(server.URL, otherToken), "", []string{otherDir, scimDir}}, - {"token from another issuer", oauthConfig(server.URL, fakeJWT(t, "https://accounts.test/oidc", "other@example.com")), "", []string{scimDir}}, + {"token names the SCIM user", oauthConfig(server.URL, auth.AuthTypeDatabricksCli, fakeJWT(t, issuer, testserver.TestUser.UserName)), "", []string{scimDir}}, + {"token names another user", oauthConfig(server.URL, auth.AuthTypeOAuthM2M, otherToken), "", []string{otherDir, scimDir}}, + {"token from another issuer", oauthConfig(server.URL, auth.AuthTypeDatabricksCli, fakeJWT(t, "https://accounts.test/oidc", "other@example.com")), "", []string{scimDir}}, + // Other auth types may not cache Token, so they skip the guess. + {"other auth type", oauthConfig(server.URL, auth.AuthTypeAzureCli, otherToken), "", []string{scimDir}}, {"pat", &databricks.Config{Host: server.URL, Token: otherToken}, "", []string{scimDir}}, - {"override", oauthConfig(server.URL, otherToken), "/Workspace/custom", []string{otherDir, "/Workspace/custom"}}, + {"override", oauthConfig(server.URL, auth.AuthTypeDatabricksCli, otherToken), "/Workspace/custom", []string{"/Workspace/custom"}}, } { t.Run(tc.name, func(t *testing.T) { t.Setenv(userWorkspaceDirEnv, tc.override) From ecdd14ce56ad27254f6ed781dbe09d7f8ddbf9ef Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 21 Oct 2026 20:12:58 +0000 Subject: [PATCH 5/5] air: drop the token guess; keep only the trimmed SCIM Me call The overlap saved ~30ms on top of the trimmed call but carried most of the change's complexity, so keep just `attributes=userName`. Co-authored-by: Isaac --- .nextchanges/air/7153.md | 2 +- cmd/air/run.go | 9 +--- cmd/air/runlaunch.go | 97 ++++++--------------------------------- cmd/air/runlaunch_test.go | 82 +++++---------------------------- cmd/air/runsubmit.go | 12 +++-- 5 files changed, 34 insertions(+), 168 deletions(-) diff --git a/.nextchanges/air/7153.md b/.nextchanges/air/7153.md index 6aabf43bec5..f163925bc66 100644 --- a/.nextchanges/air/7153.md +++ b/.nextchanges/air/7153.md @@ -1 +1 @@ -* `air run` now spends less time before submitting: it fetches only the user name from SCIM and validates the config while that lookup is in flight. ([#7153](https://github.com/databricks/cli/pull/7153)) +* `air run` now spends less time before submitting: it fetches only the user name from SCIM. ([#7153](https://github.com/databricks/cli/pull/7153)) diff --git a/cmd/air/run.go b/cmd/air/run.go index 07d3e2c44a2..c8027b09e9a 100644 --- a/cmd/air/run.go +++ b/cmd/air/run.go @@ -56,12 +56,7 @@ func validateDryRunInWorkspace(ctx context.Context, cmd *cobra.Command, args []s } w := cmdctx.WorkspaceClient(ctx) - // Path and validation failures are classified differently, so keep validation's result apart. - var validateErr error - _, _, _, err := prospectiveLaunchPaths(ctx, w, cfg, func(funcDir, commandPath string) error { - validateErr = validateDryRunConfig(ctx, w, cfg, commandPath, submittedContainers(cfg, funcDir), idempotencyToken, timeout) - return nil - }) + _, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) if err != nil { unavailable, retryable := classifyValidationFailure(err) if !unavailable { @@ -70,7 +65,7 @@ func validateDryRunInWorkspace(ctx context.Context, cmd *cobra.Command, args []s return ctx, asValidationUnavailable(err, retryable) } - return ctx, validateErr + return ctx, validateDryRunConfig(ctx, w, cfg, commandPath, submittedContainers(cfg, funcDir), idempotencyToken, timeout) } func newRunCommand() *cobra.Command { diff --git a/cmd/air/runlaunch.go b/cmd/air/runlaunch.go index 4aac773b215..a888d80d85d 100644 --- a/cmd/air/runlaunch.go +++ b/cmd/air/runlaunch.go @@ -2,22 +2,17 @@ package aircmd import ( "context" - "encoding/base64" - "encoding/json" "errors" "fmt" "path" "strings" "uuid" - "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/env" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" - "github.com/databricks/databricks-sdk-go/config" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/workspace" - "golang.org/x/sync/errgroup" ) // userWorkspaceDirEnv overrides the per-user workspace directory; mirrors the @@ -34,38 +29,8 @@ func currentUserEmail(ctx context.Context, w *databricks.WorkspaceClient) (strin return me.UserName, nil } -// userNameFromWorkspaceToken returns the sub claim of cfg's OAuth access token when this -// workspace issued it, or "" otherwise. It is only a guess: sub need not equal the SCIM userName. -func userNameFromWorkspaceToken(ctx context.Context, cfg *config.Config) string { - // These mint workspace tokens and serve Token from the cache request headers use; others, - // such as azure-cli, would fetch a fresh token here. - if cfg.AuthType != auth.AuthTypeDatabricksCli && cfg.AuthType != auth.AuthTypeOAuthM2M { - return "" - } - token, err := cfg.GetTokenSource().Token(ctx) - if err != nil { - return "" - } - parts := strings.Split(token.AccessToken, ".") - if len(parts) != 3 { - return "" - } - payload, err := base64.RawURLEncoding.DecodeString(parts[1]) - if err != nil { - return "" - } - var claims struct { - Issuer string `json:"iss"` - Subject string `json:"sub"` - } - if json.Unmarshal(payload, &claims) != nil || claims.Issuer != cfg.Host+"/oidc" { - return "" - } - return claims.Subject -} - -// userWorkspaceDirOverride returns the per-user workspace directory set by env or profile, if any. -func userWorkspaceDirOverride(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { +// userWorkspaceDir returns the user's workspace home, honoring the env override. +func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { if override := env.Get(ctx, userWorkspaceDirEnv); override != "" { return override, nil } @@ -73,13 +38,8 @@ func userWorkspaceDirOverride(ctx context.Context, w *databricks.WorkspaceClient if err != nil { return "", fmt.Errorf("failed to read selected profile: %w", err) } - return override, nil -} - -// userWorkspaceDir returns the user's workspace home, honoring the env override. -func userWorkspaceDir(ctx context.Context, w *databricks.WorkspaceClient) (string, error) { - if override, err := userWorkspaceDirOverride(ctx, w); override != "" || err != nil { - return override, err + if override != "" { + return override, nil } email, err := currentUserEmail(ctx, w) if err != nil { @@ -98,50 +58,19 @@ func cliLaunchDir(base, experiment, run string) string { return path.Join(base, ".air", "cli_launch", experiment, run+"_"+unique) } -// prospectiveLaunchPaths returns the workspace paths a submission would use without creating -// directories or uploading artifacts, after running validate on them. -// -// The directory is usually known before SCIM Me answers (an override, or the user a -// workspace-issued token names), so validate first runs on that guess while userWorkspaceDir -// looks up the real directory, and again if the two differ. -func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, validate func(launchDir, commandPath string) error) (string, string, string, error) { - runName := "" - if cfg.MLflowRunName != nil { - runName = *cfg.MLflowRunName - } - validated := func(base string) (string, string, string, error) { - launchDir := cliLaunchDir(base, cfg.ExperimentName, runName) - commandPath := path.Join(launchDir, commandScriptName) - return base, launchDir, commandPath, validate(launchDir, commandPath) - } - - guessBase, err := userWorkspaceDirOverride(ctx, w) +// prospectiveLaunchPaths returns the workspace paths a submission would use +// without creating directories or uploading artifacts. +func prospectiveLaunchPaths(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig) (string, string, string, error) { + base, err := userWorkspaceDir(ctx, w) if err != nil { return "", "", "", err } - if guessBase == "" { - if user := userNameFromWorkspaceToken(ctx, w.Config); user != "" { - guessBase = "/Workspace/Users/" + user - } - } - var base string - var g errgroup.Group - g.Go(func() (err error) { - base, err = userWorkspaceDir(ctx, w) - return err - }) - var launchDir, commandPath string - var validateErr error - if guessBase != "" { - _, launchDir, commandPath, validateErr = validated(guessBase) - } - if err := g.Wait(); err != nil { - return "", "", "", err - } - if base != guessBase { - return validated(base) + runName := "" + if cfg.MLflowRunName != nil { + runName = *cfg.MLflowRunName } - return base, launchDir, commandPath, validateErr + launchDir := cliLaunchDir(base, cfg.ExperimentName, runName) + return base, launchDir, path.Join(launchDir, commandScriptName), nil } // ensureExperimentDirectory creates experimentDir if it is missing, matching the diff --git a/cmd/air/runlaunch_test.go b/cmd/air/runlaunch_test.go index 3b342024664..edb4803ba65 100644 --- a/cmd/air/runlaunch_test.go +++ b/cmd/air/runlaunch_test.go @@ -1,24 +1,17 @@ package aircmd import ( - "context" - "encoding/base64" - "encoding/json" - "errors" "os" "path/filepath" "strings" "testing" - "github.com/databricks/cli/libs/auth" "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" - "github.com/databricks/databricks-sdk-go/config" - sdkauth "github.com/databricks/databricks-sdk-go/config/experimental/auth" + "github.com/databricks/databricks-sdk-go/service/iam" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" - "golang.org/x/oauth2" ) func TestCliLaunchDir(t *testing.T) { @@ -53,76 +46,23 @@ func TestUserWorkspaceDir(t *testing.T) { assert.Equal(t, "/Workspace/custom", dir) } -func fakeJWT(t *testing.T, issuer, subject string) string { - payload, err := json.Marshal(map[string]string{"iss": issuer, "sub": subject}) - require.NoError(t, err) - enc := base64.RawURLEncoding - return enc.EncodeToString([]byte(`{"alg":"none"}`)) + "." + enc.EncodeToString(payload) + ".sig" -} - -func oauthConfig(host, authType, accessToken string) *databricks.Config { - ts := sdkauth.TokenSourceFn(func(context.Context) (*oauth2.Token, error) { - return &oauth2.Token{AccessToken: accessToken}, nil - }) - return &databricks.Config{Host: host, Credentials: config.NewTokenSourceStrategy(authType, ts)} -} - -// The token's user is only a guess that lets validation start early: userWorkspaceDir decides. -func TestProspectiveLaunchPathsTokenGuess(t *testing.T) { +func TestUserWorkspaceDirRequestsOnlyUserName(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) var attributes []string server.Handle("GET", "/api/2.0/preview/scim/v2/Me", func(req testserver.Request) any { attributes = append(attributes, req.URL.Query().Get("attributes")) - return testserver.TestUser + return iam.User{UserName: "me@example.com"} }) testserver.AddDefaultHandlers(server) - issuer := server.URL + "/oidc" - scimDir := "/Workspace/Users/" + testserver.TestUser.UserName - otherDir := "/Workspace/Users/other@example.com" - otherToken := fakeJWT(t, issuer, "other@example.com") - - for _, tc := range []struct { - name string - cfg *databricks.Config - override string - validated []string - }{ - {"token names the SCIM user", oauthConfig(server.URL, auth.AuthTypeDatabricksCli, fakeJWT(t, issuer, testserver.TestUser.UserName)), "", []string{scimDir}}, - {"token names another user", oauthConfig(server.URL, auth.AuthTypeOAuthM2M, otherToken), "", []string{otherDir, scimDir}}, - {"token from another issuer", oauthConfig(server.URL, auth.AuthTypeDatabricksCli, fakeJWT(t, "https://accounts.test/oidc", "other@example.com")), "", []string{scimDir}}, - // Other auth types may not cache Token, so they skip the guess. - {"other auth type", oauthConfig(server.URL, auth.AuthTypeAzureCli, otherToken), "", []string{scimDir}}, - {"pat", &databricks.Config{Host: server.URL, Token: otherToken}, "", []string{scimDir}}, - {"override", oauthConfig(server.URL, auth.AuthTypeDatabricksCli, otherToken), "/Workspace/custom", []string{"/Workspace/custom"}}, - } { - t.Run(tc.name, func(t *testing.T) { - t.Setenv(userWorkspaceDirEnv, tc.override) - w, err := databricks.NewWorkspaceClient(tc.cfg) - require.NoError(t, err) - - var launchDirs []string - base, launchDir, commandPath, err := prospectiveLaunchPaths(t.Context(), w, baseRunConfig(), func(launchDir, _ string) error { - launchDirs = append(launchDirs, launchDir) - if strings.HasPrefix(launchDir, otherDir+"/") { - return errors.New("validated an unconfirmed user") - } - return nil - }) - require.NoError(t, err) - require.Len(t, launchDirs, len(tc.validated)) - for i, dir := range tc.validated { - assert.True(t, strings.HasPrefix(launchDirs[i], dir+"/.air/cli_launch/"), launchDirs[i]) - } - assert.Equal(t, tc.validated[len(tc.validated)-1], base) - assert.Equal(t, launchDirs[len(launchDirs)-1], launchDir) - assert.Equal(t, launchDir+"/command.sh", commandPath) - }) - } - require.NotEmpty(t, attributes) - for _, a := range attributes { - assert.Equal(t, "userName", a) - } + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + t.Setenv(userWorkspaceDirEnv, "") + + dir, err := userWorkspaceDir(t.Context(), w) + require.NoError(t, err) + assert.Equal(t, "/Workspace/Users/me@example.com", dir) + assert.Equal(t, []string{"userName"}, attributes) } func TestUserWorkspaceDirProfilePrecedence(t *testing.T) { diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index 3ae29f1038a..d1e94e068dc 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -337,15 +337,17 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run defer func() { logRunEvent(ctx, cfg, snap, runID, time.Since(start), err) }() // Compute and validate the actual submission path before creating artifacts. - // Validate server-side before uploading, so invalid configs leave no artifacts. - token := effectiveIdempotencyToken(idempotencyKey, cfg) - base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg, func(funcDir, commandPath string) error { - return preflightValidate(ctx, w, cfg, commandPath, submittedContainers(cfg, funcDir), token) - }) + base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) if err != nil { return 0, "", err } containers := submittedContainers(cfg, funcDir) + token := effectiveIdempotencyToken(idempotencyKey, cfg) + + // Validate server-side before uploading, so invalid configs leave no artifacts. + if err := preflightValidate(ctx, w, cfg, commandPath, containers, token); err != nil { + return 0, "", err + } // Retain the local guard after backend validation so an unavailable endpoint // cannot let a bad key reach upload or submission.