Skip to content

Commit d88b4e1

Browse files
caroline-dbcaroline-db
committed
air: make dry-run workspace validation strict
1 parent b5edcb6 commit d88b4e1

8 files changed

Lines changed: 185 additions & 330 deletions

File tree

‎cmd/air/run.go‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -46,8 +46,7 @@ func warnExperimentalContainers(ctx context.Context, cfg *runConfig) {
4646
}
4747
}
4848

49-
// validateDryRunInWorkspace marks transient workspace access failures incomplete;
50-
// configuration, authentication, and permission errors are returned directly.
49+
// validateDryRunInWorkspace keeps caller errors distinct from service failures.
5150
func validateDryRunInWorkspace(ctx context.Context, cmd *cobra.Command, args []string, cfg *runConfig, idempotencyToken string, timeout time.Duration) (context.Context, error) {
5251
if !cmdctx.HasWorkspaceClient(ctx) {
5352
if err := root.MustWorkspaceClient(cmd, args); err != nil {
@@ -59,10 +58,11 @@ func validateDryRunInWorkspace(ctx context.Context, cmd *cobra.Command, args []s
5958
w := cmdctx.WorkspaceClient(ctx)
6059
_, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg)
6160
if err != nil {
62-
if !validationCouldNotComplete(err) {
61+
unavailable, retryable := classifyValidationFailure(err)
62+
if !unavailable {
6363
return ctx, err
6464
}
65-
return ctx, asIncompleteConfigValidation(err)
65+
return ctx, asValidationUnavailable(err, retryable)
6666
}
6767

6868
validationCtx, cancel := context.WithTimeout(ctx, timeout)
@@ -146,20 +146,21 @@ The path must be a separate argument: cobra reserves -h as a boolean, so
146146
if dryRun {
147147
idempotencyToken := effectiveIdempotencyToken(idempotencyKey, cfg)
148148
ctx, validationErr := validateDryRunInWorkspace(ctx, cmd, args, cfg, idempotencyToken, validationTimeout)
149-
if validationErr != nil && !configValidationIncomplete(validationErr) {
150-
return validationErr
151-
}
152149
if err := validateIdempotencyToken(idempotencyToken); err != nil {
153150
return err
154151
}
155152
if validationErr != nil {
156-
cmdio.LogString(ctx, fmt.Sprintf("Warning: workspace validation could not be completed (%s); only local validation was performed.", validationErr))
153+
unavailable, ok := validationUnavailable(validationErr)
154+
if !ok {
155+
return validationErr
156+
}
157+
kind := "PERMANENT"
158+
if unavailable.retryable {
159+
kind = "TRANSIENT"
160+
}
161+
return renderError(ctx, cmd, "VALIDATION_UNAVAILABLE", kind, unavailable.retryable, validationErr)
157162
}
158163
if root.OutputType(cmd) == flags.OutputText {
159-
if validationErr != nil {
160-
cmdio.LogString(ctx, fmt.Sprintf("Dry run: local validation passed for %q; not submitting.", cfg.ExperimentName))
161-
return nil
162-
}
163164
cmdio.LogString(ctx, fmt.Sprintf("Dry run: configuration for %q is valid; not submitting.", cfg.ExperimentName))
164165
return nil
165166
}

‎cmd/air/run_test.go‎

Lines changed: 44 additions & 90 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,18 @@ package aircmd
33
import (
44
"bytes"
55
"encoding/json"
6-
"io"
76
"net/http"
87
"net/http/httptest"
98
"strings"
109
"sync/atomic"
1110
"testing"
1211
"time"
1312

13+
"github.com/databricks/cli/cmd/root"
1414
"github.com/databricks/cli/libs/cmdctx"
1515
"github.com/databricks/cli/libs/cmdio"
1616
"github.com/databricks/cli/libs/flags"
17+
"github.com/databricks/cli/libs/testserver"
1718
"github.com/stretchr/testify/assert"
1819
"github.com/stretchr/testify/require"
1920
)
@@ -137,30 +138,22 @@ type dryRunRequest struct {
137138
body map[string]any
138139
}
139140

140-
// dryRunServer records workspace reads and serves one validation response.
141-
func dryRunServer(t *testing.T, status int, response string, requests *[]dryRunRequest) *httptest.Server {
141+
func dryRunServer(t *testing.T, response any, requests *[]dryRunRequest) *testserver.Server {
142142
t.Helper()
143-
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
144-
request := dryRunRequest{method: r.Method, path: r.URL.Path}
145-
if r.URL.Path == "/.well-known/databricks-config" {
146-
*requests = append(*requests, request)
147-
w.WriteHeader(http.StatusNotFound)
148-
return
149-
}
150-
if r.URL.Path == validateConfigPath {
151-
if err := json.NewDecoder(r.Body).Decode(&request.body); err != nil {
152-
t.Errorf("decode validation request: %v", err)
143+
server := testserver.New(t)
144+
if requests != nil {
145+
server.RequestCallback = func(req *testserver.Request) {
146+
request := dryRunRequest{method: req.Method, path: req.URL.Path}
147+
if req.URL.Path == validateConfigPath {
148+
require.NoError(t, json.Unmarshal(req.Body, &request.body))
153149
}
154150
*requests = append(*requests, request)
155-
w.WriteHeader(status)
156-
_, _ = w.Write([]byte(response))
157-
return
158151
}
159-
*requests = append(*requests, request)
160-
_, _ = w.Write([]byte(`{"userName":"u@example.com","workspace_id":1}`))
161-
}))
162-
t.Cleanup(srv.Close)
163-
return srv
152+
}
153+
server.Handle("POST", validateConfigPath, func(_ testserver.Request) any { return response })
154+
testserver.AddDefaultHandlers(server)
155+
t.Cleanup(server.Close)
156+
return server
164157
}
165158

166159
type dryRunResult struct {
@@ -189,23 +182,9 @@ func runDryRunCmd(t *testing.T, timeout time.Duration, output flags.Output, srvU
189182

190183
func TestRunDryRunValidatesBackendWithoutMutation(t *testing.T) {
191184
var requests []dryRunRequest
192-
srv := dryRunServer(t, http.StatusOK, `{}`, &requests)
193-
config := minimalConfig + `
194-
environment:
195-
dependencies:
196-
- numpy
197-
code_source:
198-
type: snapshot
199-
snapshot:
200-
root_path: .
201-
max_retries: 2
202-
usage_policy_name: team-policy
203-
permissions:
204-
- group_name: users
205-
level: CAN_VIEW
206-
`
207-
208-
result := runDryRunCmd(t, dryRunValidationTimeout, flags.OutputText, srv.URL, config, "--dry-run", "--idempotency-key", "flag-token")
185+
srv := dryRunServer(t, validateConfigResponse{}, &requests)
186+
187+
result := runDryRunCmd(t, dryRunValidationTimeout, flags.OutputText, srv.URL, minimalConfig, "--dry-run", "--idempotency-key", "flag-token")
209188
require.NoError(t, result.err)
210189
assert.Contains(t, result.stderr, `Dry run: configuration for "my-run" is valid; not submitting.`)
211190

@@ -226,80 +205,55 @@ permissions:
226205
assert.Contains(t, validations[0].body, "run_options")
227206
task := validations[0].body["task"].(map[string]any)
228207
commandPath := task["deployments"].([]any)[0].(map[string]any)["command_path"].(string)
229-
assert.True(t, strings.HasPrefix(commandPath, "/Workspace/Users/u@example.com/.air/cli_launch/my-run/my-run_"))
208+
assert.True(t, strings.HasPrefix(commandPath, "/Workspace/Users/tester@databricks.com/.air/cli_launch/my-run/my-run_"))
230209
assert.True(t, strings.HasSuffix(commandPath, "/command.sh"))
231210
runOptions := validations[0].body["run_options"].(map[string]any)
232211
assert.Equal(t, "flag-token", runOptions["idempotency_token"])
233212
}
234213

235214
func TestRunDryRunConfigValidationHasDeadline(t *testing.T) {
236215
requestCanceled := make(chan struct{}, 1)
237-
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
238-
switch r.URL.Path {
239-
case "/.well-known/databricks-config":
240-
w.WriteHeader(http.StatusNotFound)
241-
case "/api/2.0/preview/scim/v2/Me":
242-
_, _ = w.Write([]byte(`{"userName":"u@example.com"}`))
243-
case validateConfigPath:
244-
_, _ = io.Copy(io.Discard, r.Body)
245-
<-r.Context().Done()
246-
requestCanceled <- struct{}{}
247-
default:
248-
http.NotFound(w, r)
249-
}
250-
}))
216+
srv := testserver.New(t)
217+
srv.Handle("POST", validateConfigPath, func(req testserver.Request) any {
218+
<-req.Context.Done()
219+
requestCanceled <- struct{}{}
220+
return nil
221+
})
222+
testserver.AddDefaultHandlers(srv)
251223
t.Cleanup(srv.Close)
252224

253225
result := runDryRunCmd(t, 100*time.Millisecond, flags.OutputText, srv.URL, minimalConfig, "--dry-run")
254-
require.NoError(t, result.err)
255-
assert.Contains(t, result.stderr, "only local validation was performed.")
256-
assert.Contains(t, result.stderr, `Dry run: local validation passed for "my-run"; not submitting.`)
226+
require.Error(t, result.err)
227+
assert.Contains(t, result.err.Error(), "config validation unavailable: failed to validate config:")
228+
assert.Contains(t, result.err.Error(), "context deadline exceeded")
229+
assert.NotContains(t, result.stderr, "Dry run: configuration")
257230
select {
258231
case <-requestCanceled:
259232
case <-time.After(time.Second):
260233
require.Fail(t, "dry-run validation request context was not canceled")
261234
}
262235
}
263236

264-
func TestRunDryRunFailsOnAuthenticationFailure(t *testing.T) {
265-
var requests []dryRunRequest
266-
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
267-
requests = append(requests, dryRunRequest{method: r.Method, path: r.URL.Path})
268-
if r.URL.Path == "/.well-known/databricks-config" {
269-
w.WriteHeader(http.StatusNotFound)
270-
return
271-
}
272-
w.WriteHeader(http.StatusUnauthorized)
273-
_, _ = w.Write([]byte(`{"error_code":"UNAUTHENTICATED","message":"log in"}`))
274-
}))
275-
t.Cleanup(srv.Close)
276-
277-
result := runDryRunCmd(t, dryRunValidationTimeout, flags.OutputText, srv.URL, minimalConfig, "--dry-run")
278-
require.Error(t, result.err)
279-
assert.Contains(t, result.err.Error(), "failed to resolve current user")
280-
assert.NotContains(t, result.stderr, "only local validation was performed")
281-
for _, request := range requests {
282-
assert.NotEqual(t, validateConfigPath, request.path)
283-
}
284-
}
285-
286-
func TestRunDryRunIncompleteKeepsJSONOutputContract(t *testing.T) {
287-
var requests []dryRunRequest
288-
srv := dryRunServer(t, http.StatusInternalServerError, `{"error_code":"INTERNAL_ERROR","message":"backend failed"}`, &requests)
237+
func TestRunDryRunUnavailableUsesJSONErrorEnvelope(t *testing.T) {
238+
srv := dryRunServer(t, testserver.Response{
239+
StatusCode: http.StatusInternalServerError,
240+
Body: map[string]string{"error_code": "INTERNAL_ERROR", "message": "backend failed"},
241+
}, nil)
289242
result := runDryRunCmd(t, dryRunValidationTimeout, flags.OutputJSON, srv.URL, minimalConfig, "--dry-run")
290-
require.NoError(t, result.err)
291-
var got struct {
292-
Data runResult `json:"data"`
293-
}
243+
require.ErrorIs(t, result.err, root.ErrAlreadyPrinted)
244+
var got errorEnvelope
294245
require.NoError(t, json.Unmarshal([]byte(result.stdout), &got))
295-
assert.Equal(t, "DRY_RUN_OK", got.Data.Status)
296-
assert.True(t, got.Data.DryRun)
297-
assert.Contains(t, result.stderr, "only local validation was performed")
246+
assert.Equal(t, "VALIDATION_UNAVAILABLE", got.Error.Code)
247+
assert.Equal(t, "TRANSIENT", got.Error.Kind)
248+
assert.True(t, got.Error.Retryable)
249+
assert.Contains(t, got.Error.Message, "config validation unavailable")
250+
assert.NotContains(t, result.stdout, "DRY_RUN_OK")
251+
assert.Empty(t, result.stderr)
298252
}
299253

300-
func TestRunDryRunValidatesEffectiveIdempotencyTokenBeforeLocalFallback(t *testing.T) {
254+
func TestRunDryRunValidatesIdempotencyTokenAfterWorkspaceValidation(t *testing.T) {
301255
var requests []dryRunRequest
302-
srv := dryRunServer(t, http.StatusOK, `{}`, &requests)
256+
srv := dryRunServer(t, validateConfigResponse{}, &requests)
303257
tooLong := strings.Repeat("x", 65)
304258

305259
result := runDryRunCmd(t, dryRunValidationTimeout, flags.OutputText, srv.URL, minimalConfig, "--dry-run", "--idempotency-key", tooLong)

‎cmd/air/runsubmit_test.go‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,13 @@ func stubValidateConfig(server *testserver.Server) {
3030
})
3131
}
3232

33+
func TestEffectiveIdempotencyToken(t *testing.T) {
34+
cfg := &runConfig{IdempotencyToken: new("from-config")}
35+
assert.Equal(t, "from-flag", effectiveIdempotencyToken("from-flag", cfg))
36+
assert.Equal(t, "from-config", effectiveIdempotencyToken("", cfg))
37+
assert.NotEmpty(t, effectiveIdempotencyToken("", &runConfig{}))
38+
}
39+
3340
func TestDlRuntimeImage(t *testing.T) {
3441
ctx := t.Context()
3542
// A config runtime version wins and is used bare.
@@ -377,6 +384,66 @@ func TestSubmitWorkload(t *testing.T) {
377384
assert.Empty(t, d.Compute.ProvisionedCapacityId)
378385
}
379386

387+
func TestSubmitWorkloadValidationOutcome(t *testing.T) {
388+
tests := []struct {
389+
name string
390+
response any
391+
wantError string
392+
wantSubmit bool
393+
}{
394+
{
395+
name: "field error",
396+
response: validateConfigResponse{Errors: []configFieldError{{
397+
Path: "task.deployments[0].compute.accelerator_count", Message: "must be positive", Code: "INVALID_COUNT",
398+
}}},
399+
wantError: "accelerator_count: must be positive",
400+
},
401+
{
402+
name: "service unavailable",
403+
response: testserver.Response{StatusCode: 503, Body: map[string]string{"error_code": "TEMPORARILY_UNAVAILABLE", "message": "try again"}},
404+
wantSubmit: true,
405+
},
406+
}
407+
408+
for _, tt := range tests {
409+
t.Run(tt.name, func(t *testing.T) {
410+
server := testserver.New(t)
411+
t.Cleanup(server.Close)
412+
var uploads, submissions atomic.Int32
413+
server.Handle("POST", validateConfigPath, func(req testserver.Request) any { return tt.response })
414+
server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any {
415+
uploads.Add(1)
416+
return req.Workspace.WorkspaceFilesImportFile(req.Vars["path"], req.Body, true)
417+
})
418+
server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
419+
submissions.Add(1)
420+
return jobs.SubmitRunResponse{RunId: 777}
421+
})
422+
testserver.AddDefaultHandlers(server)
423+
424+
w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token", RetryTimeoutSeconds: 1})
425+
require.NoError(t, err)
426+
cfgPath := writeConfigFile(t, "run.yaml", minimalConfig)
427+
cfg, err := loadRunConfig(cfgPath)
428+
require.NoError(t, err)
429+
430+
_, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false)
431+
if tt.wantError != "" {
432+
require.ErrorContains(t, err, tt.wantError)
433+
} else {
434+
require.NoError(t, err)
435+
}
436+
if tt.wantSubmit {
437+
assert.Positive(t, uploads.Load())
438+
assert.Equal(t, int32(1), submissions.Load())
439+
} else {
440+
assert.Zero(t, uploads.Load())
441+
assert.Zero(t, submissions.Load())
442+
}
443+
})
444+
}
445+
}
446+
380447
func TestSubmitWorkloadStagingErrorPreventsSubmit(t *testing.T) {
381448
server := testserver.New(t)
382449
t.Cleanup(server.Close)

0 commit comments

Comments
 (0)