diff --git a/experimental/ssh/internal/client/client.go b/experimental/ssh/internal/client/client.go index b9fc604af9..3c4ed62155 100644 --- a/experimental/ssh/internal/client/client.go +++ b/experimental/ssh/internal/client/client.go @@ -45,11 +45,6 @@ const ( sshServerTaskKey = "start_ssh_server" serverlessEnvironmentKey = "ssh_tunnel_serverless" minEnvironmentVersion = 4 - - VSCodeOption = "vscode" - VSCodeCommand = "code" - CursorOption = "cursor" - CursorCommand = "cursor" ) type ClientOptions struct { @@ -117,8 +112,8 @@ func (o *ClientOptions) Validate() error { if o.ConnectionName != "" && !connectionNameRegex.MatchString(o.ConnectionName) { return fmt.Errorf("connection name %q must consist of letters, numbers, dashes, and underscores", o.ConnectionName) } - if o.IDE != "" && o.IDE != VSCodeOption && o.IDE != CursorOption { - return fmt.Errorf("invalid IDE value: %q, expected %q or %q", o.IDE, VSCodeOption, CursorOption) + if o.IDE != "" && o.IDE != vscode.VSCodeOption && o.IDE != vscode.CursorOption { + return fmt.Errorf("invalid IDE value: %q, expected %q or %q", o.IDE, vscode.VSCodeOption, vscode.CursorOption) } if o.EnvironmentVersion > 0 && o.EnvironmentVersion < minEnvironmentVersion { return fmt.Errorf("environment version must be >= %d, got %d", minEnvironmentVersion, o.EnvironmentVersion) @@ -212,6 +207,32 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt return errors.New("either --cluster or --name must be provided") } + if opts.IDE != "" && !opts.ProxyMode { + if err := vscode.CheckIDECommand(opts.IDE); err != nil { + return err + } + } + + // Check and update IDE settings for serverless mode, where we must set up + // desired server ports (or socket connection mode) for the connection to go through + // (as the majority of the localhost ports on the remote side are blocked by iptable rules). + // Plus the platform (always linux), and extensions (python and jupyter), to make the initial experience smoother. + if opts.IDE != "" && opts.IsServerlessMode() && !opts.ProxyMode && !opts.SkipSettingsCheck && cmdio.IsPromptSupported(ctx) { + err := vscode.CheckAndUpdateSettings(ctx, opts.IDE, opts.ConnectionName) + if err != nil { + cmdio.LogString(ctx, fmt.Sprintf("Failed to update IDE settings: %v", err)) + cmdio.LogString(ctx, vscode.GetManualInstructions(opts.IDE, opts.ConnectionName)) + cmdio.LogString(ctx, "Use --skip-settings-check to bypass IDE settings verification.") + shouldProceed, promptErr := cmdio.AskYesOrNo(ctx, "Do you want to proceed with the connection?") + if promptErr != nil { + return fmt.Errorf("failed to prompt user: %w", promptErr) + } + if !shouldProceed { + return errors.New("aborted: IDE settings need to be updated manually, user declined to proceed") + } + } + } + // Only check cluster state for dedicated clusters if !opts.IsServerlessMode() { err := checkClusterState(ctx, client, opts.ClusterID, opts.AutoStartCluster) @@ -242,26 +263,6 @@ func Run(ctx context.Context, client *databricks.WorkspaceClient, opts ClientOpt cmdio.LogString(ctx, "Using SSH key: "+keyPath) cmdio.LogString(ctx, fmt.Sprintf("Secrets scope: %s, key name: %s", secretScopeName, opts.ClientPublicKeyName)) - // Check and update IDE settings for serverless mode, where we must set up - // desired server ports (or socket connection mode) for the connection to go through - // (as the majority of the localhost ports on the remote side are blocked by iptable rules). - // Plus the platform (always linux), and extensions (python and jupyter), to make the initial experience smoother. - if opts.IDE != "" && opts.IsServerlessMode() && !opts.ProxyMode && !opts.SkipSettingsCheck && cmdio.IsPromptSupported(ctx) { - err = vscode.CheckAndUpdateSettings(ctx, opts.IDE, opts.ConnectionName) - if err != nil { - cmdio.LogString(ctx, fmt.Sprintf("Failed to update IDE settings: %v", err)) - cmdio.LogString(ctx, vscode.GetManualInstructions(opts.IDE, opts.ConnectionName)) - cmdio.LogString(ctx, "Use --skip-settings-check to bypass IDE settings verification.") - shouldProceed, promptErr := cmdio.AskYesOrNo(ctx, "Do you want to proceed with the connection?") - if promptErr != nil { - return fmt.Errorf("failed to prompt user: %w", promptErr) - } - if !shouldProceed { - return errors.New("aborted: IDE settings need to be updated manually, user declined to proceed") - } - } - } - var userName string var serverPort int var clusterID string @@ -330,7 +331,6 @@ func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, k if err != nil { return fmt.Errorf("failed to get current user: %w", err) } - databricksUserName := currentUser.UserName // Ensure SSH config entry exists configPath, err := sshconfig.GetMainConfigPath(ctx) @@ -343,23 +343,7 @@ func runIDE(ctx context.Context, client *databricks.WorkspaceClient, userName, k return fmt.Errorf("failed to ensure SSH config entry: %w", err) } - ideCommand := VSCodeCommand - if opts.IDE == CursorOption { - ideCommand = CursorCommand - } - - // Construct the remote SSH URI - // Format: ssh-remote+@ /Workspace/Users// - remoteURI := fmt.Sprintf("ssh-remote+%s@%s", userName, connectionName) - remotePath := fmt.Sprintf("/Workspace/Users/%s/", databricksUserName) - - cmdio.LogString(ctx, fmt.Sprintf("Launching %s with remote URI: %s and path: %s", opts.IDE, remoteURI, remotePath)) - - ideCmd := exec.CommandContext(ctx, ideCommand, "--remote", remoteURI, remotePath) - ideCmd.Stdout = os.Stdout - ideCmd.Stderr = os.Stderr - - return ideCmd.Run() + return vscode.LaunchIDE(ctx, opts.IDE, connectionName, userName, currentUser.UserName) } func ensureSSHConfigEntry(ctx context.Context, configPath, hostName, userName, keyPath string, serverPort int, clusterID string, opts ClientOptions) error { diff --git a/experimental/ssh/internal/vscode/run.go b/experimental/ssh/internal/vscode/run.go new file mode 100644 index 0000000000..373a54d8f4 --- /dev/null +++ b/experimental/ssh/internal/vscode/run.go @@ -0,0 +1,81 @@ +package vscode + +import ( + "context" + "fmt" + "os" + "os/exec" + + "github.com/databricks/cli/libs/cmdio" +) + +// Options as they can be set via --ide flag. +const ( + VSCodeOption = "vscode" + CursorOption = "cursor" +) + +type ideDescriptor struct { + Option string + Command string + Name string + InstallURL string + AppName string +} + +var vsCodeIDE = ideDescriptor{ + Option: VSCodeOption, + Command: "code", + Name: "VS Code", + InstallURL: "https://code.visualstudio.com/", + AppName: "Code", +} + +var cursorIDE = ideDescriptor{ + Option: CursorOption, + Command: "cursor", + Name: "Cursor", + InstallURL: "https://cursor.com/", + AppName: "Cursor", +} + +func getIDE(option string) ideDescriptor { + if option == CursorOption { + return cursorIDE + } + return vsCodeIDE +} + +// CheckIDECommand verifies the IDE CLI command is available on PATH. +func CheckIDECommand(option string) error { + ide := getIDE(option) + + if _, err := exec.LookPath(ide.Command); err != nil { + return fmt.Errorf( + "%q command not found on PATH. To fix this:\n"+ + "1. Install %s from %s\n"+ + "2. Open the Command Palette (Cmd+Shift+P / Ctrl+Shift+P) and run \"Shell Command: Install '%s' command\"\n"+ + "3. Restart your terminal", + ide.Command, ide.Name, ide.InstallURL, ide.Command, + ) + } + return nil +} + +// LaunchIDE launches the IDE with a remote SSH connection. +func LaunchIDE(ctx context.Context, ideOption, connectionName, userName, databricksUserName string) error { + ide := getIDE(ideOption) + + // Construct the remote SSH URI + // Format: ssh-remote+@ /Workspace/Users// + remoteURI := fmt.Sprintf("ssh-remote+%s@%s", userName, connectionName) + remotePath := fmt.Sprintf("/Workspace/Users/%s/", databricksUserName) + + cmdio.LogString(ctx, fmt.Sprintf("Launching %s with remote URI: %s and path: %s", ideOption, remoteURI, remotePath)) + + ideCmd := exec.CommandContext(ctx, ide.Command, "--remote", remoteURI, remotePath) + ideCmd.Stdout = os.Stdout + ideCmd.Stderr = os.Stderr + + return ideCmd.Run() +} diff --git a/experimental/ssh/internal/vscode/run_test.go b/experimental/ssh/internal/vscode/run_test.go new file mode 100644 index 0000000000..f50fa4f93d --- /dev/null +++ b/experimental/ssh/internal/vscode/run_test.go @@ -0,0 +1,90 @@ +package vscode + +import ( + "os" + "path/filepath" + "runtime" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestCheckIDECommand_Missing(t *testing.T) { + // Override PATH to ensure commands are not found + t.Setenv("PATH", t.TempDir()) + + tests := []struct { + name string + ide string + wantErrMsg string + }{ + { + name: "missing vscode command", + ide: VSCodeOption, + wantErrMsg: `"code" command not found on PATH`, + }, + { + name: "missing cursor command", + ide: CursorOption, + wantErrMsg: `"cursor" command not found on PATH`, + }, + { + name: "vscode error contains install instructions", + ide: VSCodeOption, + wantErrMsg: "https://code.visualstudio.com/", + }, + { + name: "cursor error contains install instructions", + ide: CursorOption, + wantErrMsg: "https://cursor.com/", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := CheckIDECommand(tt.ide) + require.Error(t, err) + assert.Contains(t, err.Error(), tt.wantErrMsg) + }) + } +} + +func TestCheckIDECommand_Found(t *testing.T) { + tmpDir := t.TempDir() + t.Setenv("PATH", tmpDir) + + tests := []struct { + name string + ide string + command string + }{ + { + name: "vscode command found", + ide: VSCodeOption, + command: "code", + }, + { + name: "cursor command found", + ide: CursorOption, + command: "cursor", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Create a fake executable in the temp directory. + // On Windows, exec.LookPath requires a known extension (e.g. .exe). + command := tt.command + if runtime.GOOS == "windows" { + command += ".exe" + } + fakePath := filepath.Join(tmpDir, command) + err := os.WriteFile(fakePath, []byte("#!/bin/sh\n"), 0o755) + require.NoError(t, err) + + err = CheckIDECommand(tt.ide) + assert.NoError(t, err) + }) + } +} diff --git a/experimental/ssh/internal/vscode/settings.go b/experimental/ssh/internal/vscode/settings.go index d8930f4ae4..fa71cc70a2 100644 --- a/experimental/ssh/internal/vscode/settings.go +++ b/experimental/ssh/internal/vscode/settings.go @@ -26,19 +26,8 @@ const ( remotePlatformKey = "remote.SSH.remotePlatform" defaultExtensionsKey = "remote.SSH.defaultExtensions" listenOnSocketKey = "remote.SSH.remoteServerListenOnSocket" - vscodeIDE = "vscode" - cursorIDE = "cursor" - vscodeName = "VS Code" - cursorName = "Cursor" ) -func getIDEName(ide string) string { - if ide == cursorIDE { - return cursorName - } - return vscodeName -} - type missingSettings struct { portRange bool platform bool @@ -118,7 +107,7 @@ func CheckAndUpdateSettings(ctx context.Context, ide, connectionName string) err return fmt.Errorf("failed to save settings: %w", err) } - cmdio.LogString(ctx, fmt.Sprintf("Updated %s settings for '%s'", getIDEName(ide), connectionName)) + cmdio.LogString(ctx, fmt.Sprintf("Updated %s settings for '%s'", getIDE(ide).Name, connectionName)) return nil } @@ -128,10 +117,7 @@ func getDefaultSettingsPath(ctx context.Context, ide string) (string, error) { return "", fmt.Errorf("failed to get home directory: %w", err) } - appName := "Code" - if ide == cursorIDE { - appName = "Cursor" - } + appName := getIDE(ide).AppName var settingsDir string switch runtime.GOOS { @@ -249,7 +235,7 @@ func settingsMessage(connectionName string, missing *missingSettings) string { func promptUserForUpdate(ctx context.Context, ide, connectionName string, missing *missingSettings) (bool, error) { question := fmt.Sprintf( "The following settings will be applied to %s for '%s':\n%s\nApply these settings?", - getIDEName(ide), connectionName, settingsMessage(connectionName, missing)) + getIDE(ide).Name, connectionName, settingsMessage(connectionName, missing)) return cmdio.AskYesOrNo(ctx, question) } @@ -286,7 +272,7 @@ func handleMissingFile(ctx context.Context, ide, connectionName, settingsPath st return fmt.Errorf("failed to save settings: %w", err) } - cmdio.LogString(ctx, fmt.Sprintf("Created %s settings at %s", getIDEName(ide), filepath.ToSlash(settingsPath))) + cmdio.LogString(ctx, fmt.Sprintf("Created %s settings at %s", getIDE(ide).Name, filepath.ToSlash(settingsPath))) return nil } @@ -366,5 +352,5 @@ func GetManualInstructions(ide, connectionName string) string { } return fmt.Sprintf( "To ensure the remote connection works as expected, manually add these settings to your %s settings.json:\n%s", - getIDEName(ide), settingsMessage(connectionName, missing)) + getIDE(ide).Name, settingsMessage(connectionName, missing)) } diff --git a/experimental/ssh/internal/vscode/settings_test.go b/experimental/ssh/internal/vscode/settings_test.go index 6bffdef84d..38c7688420 100644 --- a/experimental/ssh/internal/vscode/settings_test.go +++ b/experimental/ssh/internal/vscode/settings_test.go @@ -55,7 +55,7 @@ func TestGetDefaultSettingsPath_VSCode_Linux(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "HOME", "/home/testuser") - path, err := getDefaultSettingsPath(ctx, vscodeIDE) + path, err := getDefaultSettingsPath(ctx, VSCodeOption) require.NoError(t, err) assert.Equal(t, "/home/testuser/.config/Code/User/settings.json", path) } @@ -68,7 +68,7 @@ func TestGetDefaultSettingsPath_Cursor_Linux(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "HOME", "/home/testuser") - path, err := getDefaultSettingsPath(ctx, cursorIDE) + path, err := getDefaultSettingsPath(ctx, CursorOption) require.NoError(t, err) assert.Equal(t, "/home/testuser/.config/Cursor/User/settings.json", path) } @@ -81,7 +81,7 @@ func TestGetDefaultSettingsPath_VSCode_Darwin(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "HOME", "/Users/testuser") - path, err := getDefaultSettingsPath(ctx, vscodeIDE) + path, err := getDefaultSettingsPath(ctx, VSCodeOption) require.NoError(t, err) assert.Equal(t, "/Users/testuser/Library/Application Support/Code/User/settings.json", path) } @@ -94,7 +94,7 @@ func TestGetDefaultSettingsPath_Cursor_Darwin(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "HOME", "/Users/testuser") - path, err := getDefaultSettingsPath(ctx, cursorIDE) + path, err := getDefaultSettingsPath(ctx, CursorOption) require.NoError(t, err) assert.Equal(t, "/Users/testuser/Library/Application Support/Cursor/User/settings.json", path) } @@ -107,7 +107,7 @@ func TestGetDefaultSettingsPath_VSCode_Windows(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "APPDATA", `C:\Users\testuser\AppData\Roaming`) - path, err := getDefaultSettingsPath(ctx, vscodeIDE) + path, err := getDefaultSettingsPath(ctx, VSCodeOption) require.NoError(t, err) assert.Equal(t, `C:\Users\testuser\AppData\Roaming\Code\User\settings.json`, path) } @@ -120,7 +120,7 @@ func TestGetDefaultSettingsPath_Cursor_Windows(t *testing.T) { ctx := t.Context() ctx = env.Set(ctx, "APPDATA", `C:\Users\testuser\AppData\Roaming`) - path, err := getDefaultSettingsPath(ctx, cursorIDE) + path, err := getDefaultSettingsPath(ctx, CursorOption) require.NoError(t, err) assert.Equal(t, `C:\Users\testuser\AppData\Roaming\Cursor\User\settings.json`, path) } @@ -536,7 +536,7 @@ func TestMissingSettings_IsEmpty(t *testing.T) { } func TestGetManualInstructions_VSCode(t *testing.T) { - instructions := GetManualInstructions(vscodeIDE, "test-conn") + instructions := GetManualInstructions(VSCodeOption, "test-conn") assert.Contains(t, instructions, "VS Code") assert.Contains(t, instructions, "test-conn")