From 18db58203b34a8716f937f1f1e156634490b2802 Mon Sep 17 00:00:00 2001 From: Akshay Singla Date: Thu, 1 Oct 2026 03:56:46 +0000 Subject: [PATCH] Fix sandbox SSH while stopping --- cmd/sandbox/ssh.go | 18 ++++++- cmd/sandbox/ssh_test.go | 114 ++++++++++++++++++++++++++++++++++++++++ cmd/sandbox/start.go | 31 +++++++---- 3 files changed, 150 insertions(+), 13 deletions(-) diff --git a/cmd/sandbox/ssh.go b/cmd/sandbox/ssh.go index aacae607044..7fbe7df1c20 100644 --- a/cmd/sandbox/ssh.go +++ b/cmd/sandbox/ssh.go @@ -218,13 +218,27 @@ func verifyKeyRegistered(ctx context.Context, api *sandboxAPI, keyPath string) e return fmt.Errorf("your sandbox SSH key (%s) is not registered with this workspace — run `databricks sandbox register` to re-register it", want) } -// ensureRunning brings the named sandbox to Running with its own -// spinner — caller must not have one open. +// ensureRunning owns the lifecycle transition so SSH does not hang inside +// connection setup. The caller must not have a spinner open. func ensureRunning(ctx context.Context, api *sandboxAPI, id, currentStatus string) (*sandboxEntry, error) { + if strings.EqualFold(currentStatus, "running") { + return &sandboxEntry{SandboxID: id, Status: currentStatus}, nil + } + s := spin(ctx, "Starting "+cmdio.Bold(ctx, id)+"…") defer s.Close() var sb *sandboxEntry + if strings.EqualFold(currentStatus, "stopping") { + stopped, err := waitForStopped(ctx, api, s, id) + if err != nil { + s.fail("Failed to start " + id) + return nil, err + } + currentStatus = stopped.Status + s.Update("Starting " + id + "…") + } + if strings.EqualFold(currentStatus, "stopped") { updated, err := api.start(ctx, id) if err != nil { diff --git a/cmd/sandbox/ssh_test.go b/cmd/sandbox/ssh_test.go index 801a9f1b977..31f53eee6b4 100644 --- a/cmd/sandbox/ssh_test.go +++ b/cmd/sandbox/ssh_test.go @@ -1,11 +1,125 @@ package sandbox import ( + "encoding/json" + "net/http" + "net/http/httptest" + "sync/atomic" "testing" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) +func TestEnsureRunningLifecycle(t *testing.T) { + type requestStep struct { + method string + path string + status string + } + + for _, tc := range []struct { + name string + initialStatus string + steps []requestStep + wantStatus string + wantErr string + wantStarts int32 + }{ + { + name: "stopping waits for stopped before starting", + initialStatus: "Stopping", + steps: []requestStep{ + {method: http.MethodGet, path: sandboxPath("test-id"), status: "Stopped"}, + {method: http.MethodPost, path: sandboxPath("test-id") + "/start", status: "Pending"}, + {method: http.MethodGet, path: sandboxPath("test-id"), status: "Running"}, + }, + wantStatus: "Running", + wantStarts: 1, + }, + { + name: "stopped starts normally", + initialStatus: "Stopped", + steps: []requestStep{ + {method: http.MethodPost, path: sandboxPath("test-id") + "/start", status: "Pending"}, + {method: http.MethodGet, path: sandboxPath("test-id"), status: "Running"}, + }, + wantStatus: "Running", + wantStarts: 1, + }, + { + name: "running is a no-op", + initialStatus: "Running", + wantStatus: "Running", + }, + { + name: "stopped after start remains an error", + initialStatus: "Stopped", + steps: []requestStep{ + {method: http.MethodPost, path: sandboxPath("test-id") + "/start", status: "Pending"}, + {method: http.MethodGet, path: sandboxPath("test-id"), status: "Stopped"}, + }, + wantErr: `sandbox test-id reached unexpected state "Stopped" while starting`, + wantStarts: 1, + }, + } { + t.Run(tc.name, func(t *testing.T) { + var requestCount atomic.Int32 + var startCount atomic.Int32 + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/.well-known/databricks-config" { + assert.NoError(t, json.NewEncoder(w).Encode(map[string]string{ + "oidc_endpoint": "https://workspace.example.test/oidc", + })) + return + } + + index := int(requestCount.Add(1) - 1) + if index >= len(tc.steps) { + t.Errorf("unexpected request %s %s", r.Method, r.URL.Path) + http.Error(w, "unexpected request", http.StatusInternalServerError) + return + } + + step := tc.steps[index] + assert.Equal(t, step.method, r.Method) + assert.Equal(t, step.path, r.URL.Path) + if r.Method == http.MethodPost && r.URL.Path == sandboxPath("test-id")+"/start" { + startCount.Add(1) + } + assert.NoError(t, json.NewEncoder(w).Encode(sandboxEntry{ + SandboxID: "test-id", + Status: step.status, + })) + })) + t.Cleanup(server.Close) + + workspaceClient, err := databricks.NewWorkspaceClient(&databricks.Config{ + Host: server.URL, + Token: "test-token", + }) + require.NoError(t, err) + api, err := newSandboxAPI(workspaceClient) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + got, err := ensureRunning(ctx, api, "test-id", tc.initialStatus) + if tc.wantErr != "" { + require.EqualError(t, err, tc.wantErr) + assert.Nil(t, got) + } else { + require.NoError(t, err) + require.NotNil(t, got) + assert.Equal(t, tc.wantStatus, got.Status) + } + assert.Equal(t, int32(len(tc.steps)), requestCount.Load()) + assert.Equal(t, tc.wantStarts, startCount.Load()) + }) + } +} + func TestBuildSSHArgsBaseFlags(t *testing.T) { got := buildSSHArgs("happy-panda-1234", "gw.example.test", "2222", "/keys/id", nil) // Sanity: the destination is the last arg when no extras. diff --git a/cmd/sandbox/start.go b/cmd/sandbox/start.go index ec69cbd55c1..ba48a42efd9 100644 --- a/cmd/sandbox/start.go +++ b/cmd/sandbox/start.go @@ -93,12 +93,19 @@ Example: return cmd } -// waitForRunning polls the sandbox until it reaches Running or the timeout -// elapses. The spinner is updated with the elapsed seconds each poll so the -// user can tell progress is happening, and a transition to an unexpected -// terminal state (Stopped / Terminated / Failed) short-circuits with a -// useful error rather than waiting out the full timeout. func waitForRunning(ctx context.Context, api *sandboxAPI, s *spinner, id string) (*sandboxEntry, error) { + return waitForState(ctx, api, s, id, "Running", "Starting", "Stopped", "Terminated", "Failed") +} + +// The API rejects start requests until teardown finishes, so STOPPING needs +// its own polling phase. +func waitForStopped(ctx context.Context, api *sandboxAPI, s *spinner, id string) (*sandboxEntry, error) { + return waitForState(ctx, api, s, id, "Stopped", "Stopping", "Terminated", "Failed") +} + +// Centralizing lifecycle polling keeps timeout, cancellation, and terminal +// state handling consistent across start and SSH. +func waitForState(ctx context.Context, api *sandboxAPI, s *spinner, id, targetStatus, operation string, unexpectedStatuses ...string) (*sandboxEntry, error) { start := time.Now() deadline := start.Add(startWaitTimeout) for { @@ -106,16 +113,18 @@ func waitForRunning(ctx context.Context, api *sandboxAPI, s *spinner, id string) if err != nil { return nil, fmt.Errorf("polling status of %s: %w", id, err) } - switch strings.ToLower(sb.Status) { - case "running": + if strings.EqualFold(sb.Status, targetStatus) { return sb, nil - case "stopped", "terminated", "failed": - return nil, fmt.Errorf("sandbox %s reached unexpected state %q while starting", id, sb.Status) + } + for _, unexpectedStatus := range unexpectedStatuses { + if strings.EqualFold(sb.Status, unexpectedStatus) { + return nil, fmt.Errorf("sandbox %s reached unexpected state %q while %s", id, sb.Status, strings.ToLower(operation)) + } } elapsed := time.Since(start).Round(time.Second) - s.Update(fmt.Sprintf("Starting %s… (%s)", id, elapsed)) + s.Update(fmt.Sprintf("%s %s… (%s)", operation, id, elapsed)) if time.Now().After(deadline) { - return nil, fmt.Errorf("sandbox %s did not reach Running within %s (last seen %s)", id, startWaitTimeout, sb.Status) + return nil, fmt.Errorf("sandbox %s did not reach %s within %s (last seen %s)", id, targetStatus, startWaitTimeout, sb.Status) } select { case <-ctx.Done():