diff --git a/NEXT_CHANGELOG.md b/NEXT_CHANGELOG.md index 57d93cb0fc..8fa642cc50 100644 --- a/NEXT_CHANGELOG.md +++ b/NEXT_CHANGELOG.md @@ -7,6 +7,7 @@ ### CLI * An explicitly selected profile (`--profile` or a bundle's `workspace.profile`) now takes precedence over auth environment variables (`DATABRICKS_HOST`, `DATABRICKS_TOKEN`, etc.) instead of being silently shadowed by them; env vars still fill auth fields the profile leaves empty ([#5096](https://github.com/databricks/cli/issues/5096)). +* `databricks ssh connect` now fails fast with an actionable error when `--base-environment` refers to a serverless environment on version 5 or newer, which is not supported, instead of hanging until the startup timeout ([#5825](https://github.com/databricks/cli/pull/5825)). ### Bundles diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index 9fc20aa56d..6429d371d8 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -52,6 +52,10 @@ const ( sshServerTaskKey = "start_ssh_server" serverlessEnvironmentKey = "ssh_tunnel_serverless" minEnvironmentVersion = 4 + // Serverless environment_version 5+ is not supported by the SSH tunnel: the driver + // proxy does not route to the server's port on that runtime, so the health check + // never passes and connect times out. Reject known-v5 base environments up front. + maxSupportedEnvironmentVersion = 4 ) // acceleratorProvisioningNotice maps a GPU accelerator type to the upfront notice @@ -271,6 +275,18 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt cmdio.LogString(ctx, fmt.Sprintf("Connecting to %s...", sessionID)) } + // Reject an unsupported serverless environment version before any workspace side + // effects (secret scope, key upload, notebook import) so the user gets an immediate, + // actionable error instead of a multi-minute health-check timeout. The env.yaml path + // form can't be checked here, so warn that its version is unverified. + if opts.BaseEnvironment != "" && !opts.ProxyMode { + if strings.HasPrefix(opts.BaseEnvironment, "/") { + cmdio.LogString(ctx, fmt.Sprintf("WARNING: cannot verify the serverless environment version of base environment %q (specified as an env.yaml path); ssh connect only supports environment_version <= %d, and a newer version will fail to connect", opts.BaseEnvironment, maxSupportedEnvironmentVersion)) + } else if err := validateBaseEnvironmentVersion(ctx, client, opts.BaseEnvironment); err != nil { + return err + } + } + if opts.IDE != "" && !opts.ProxyMode { if err := vscode.CheckIDECommand(opts.IDE); err != nil { return err @@ -692,6 +708,44 @@ func resolveBaseEnvironment(ctx context.Context, client *databricks.WorkspaceCli } } +// validateBaseEnvironmentVersion rejects a base environment whose serverless runtime is +// unsupported (environment_version > maxSupportedEnvironmentVersion). It only applies to the +// display-name and resource-ID forms, whose version is carried in the base-environment +// listing; the env.yaml path form (leading "/") is not checkable here and is warned about +// separately by the caller. It is fail-open: any inability to determine the version (list +// error, no matching entry, nil spec, empty/unparseable version) returns nil so a valid +// connect is never blocked on a transient error or an uncheckable input. +func validateBaseEnvironmentVersion(ctx context.Context, client *databricks.WorkspaceClient, input string) error { + if strings.HasPrefix(input, "/") { + return nil + } + + envs, err := client.Environments.ListWorkspaceBaseEnvironmentsAll(ctx, environments.ListWorkspaceBaseEnvironmentsRequest{}) + if err != nil { + return nil + } + + isResourceID := strings.HasPrefix(input, "workspace-base-environments/") + for _, e := range envs { + matched := e.Name == input + if !isResourceID { + matched = e.DisplayName == input + } + if !matched || e.Spec == nil { + continue + } + version, err := strconv.Atoi(e.Spec.EnvironmentVersion) + if err != nil { + return nil + } + if version > maxSupportedEnvironmentVersion { + return fmt.Errorf("base environment %q uses serverless environment version %d, which is not supported by ssh connect (supported: version <= %d) — use a base environment created with environment_version %d", input, version, maxSupportedEnvironmentVersion, maxSupportedEnvironmentVersion) + } + return nil + } + return nil +} + // shellSingleQuote wraps s in single quotes for safe inclusion in a shell // command, escaping any embedded single quotes. func shellSingleQuote(s string) string { diff --git a/experimental/ssh/internal/client/client_internal_test.go b/experimental/ssh/internal/client/client_internal_test.go index d2bfde5783..57f61da8c8 100644 --- a/experimental/ssh/internal/client/client_internal_test.go +++ b/experimental/ssh/internal/client/client_internal_test.go @@ -1,6 +1,7 @@ package client import ( + "errors" "strings" "testing" "time" @@ -90,6 +91,83 @@ func TestResolveBaseEnvironment(t *testing.T) { }) } +func TestValidateBaseEnvironmentVersion(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + + listReturns := func(m *mocks.MockWorkspaceClient, envs []environments.WorkspaceBaseEnvironment, err error) { + m.GetMockEnvironmentsAPI().EXPECT(). + ListWorkspaceBaseEnvironmentsAll(mock.Anything, environments.ListWorkspaceBaseEnvironmentsRequest{}). + Return(envs, err) + } + specV5 := &environments.EnvironmentSpec{EnvironmentVersion: "5"} + specV4 := &environments.EnvironmentSpec{EnvironmentVersion: "4"} + + t.Run("display name v5 rejected", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + listReturns(m, []environments.WorkspaceBaseEnvironment{ + {DisplayName: "my-env", Name: "workspace-base-environments/dbe_mine", Spec: specV5}, + }, nil) + err := validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "my-env") + require.Error(t, err) + assert.Contains(t, err.Error(), `base environment "my-env" uses serverless environment version 5`) + }) + + t.Run("resource ID v5 rejected", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + listReturns(m, []environments.WorkspaceBaseEnvironment{ + {DisplayName: "my-env", Name: "workspace-base-environments/dbe_mine", Spec: specV5}, + }, nil) + err := validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "workspace-base-environments/dbe_mine") + require.Error(t, err) + assert.Contains(t, err.Error(), "serverless environment version 5") + }) + + t.Run("display name v4 allowed", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + listReturns(m, []environments.WorkspaceBaseEnvironment{ + {DisplayName: "my-env", Name: "workspace-base-environments/dbe_mine", Spec: specV4}, + }, nil) + assert.NoError(t, validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "my-env")) + }) + + // env.yaml path form is not checkable and must not hit the API. + t.Run("path form skipped", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + assert.NoError(t, validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "/Workspace/path/to/env.yaml")) + }) + + // Fail-open cases: the version can't be determined, so connect proceeds. + t.Run("list error fails open", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + listReturns(m, nil, errors.New("boom")) + assert.NoError(t, validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "my-env")) + }) + + t.Run("no match fails open", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + listReturns(m, []environments.WorkspaceBaseEnvironment{ + {DisplayName: "other", Name: "workspace-base-environments/dbe_other", Spec: specV5}, + }, nil) + assert.NoError(t, validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "my-env")) + }) + + t.Run("nil spec fails open", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + listReturns(m, []environments.WorkspaceBaseEnvironment{ + {DisplayName: "my-env", Name: "workspace-base-environments/dbe_mine", Spec: nil}, + }, nil) + assert.NoError(t, validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "my-env")) + }) + + t.Run("unparseable version fails open", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + listReturns(m, []environments.WorkspaceBaseEnvironment{ + {DisplayName: "my-env", Name: "workspace-base-environments/dbe_mine", Spec: &environments.EnvironmentSpec{EnvironmentVersion: ""}}, + }, nil) + assert.NoError(t, validateBaseEnvironmentVersion(ctx, m.WorkspaceClient, "my-env")) + }) +} + func TestDescribeRunFailureIncludesMessageTraceAndURL(t *testing.T) { ctx := cmdio.MockDiscard(t.Context()) m := mocks.NewMockWorkspaceClient(t)