SSH: bump serverless environment version and expose a flag for it (#4663)

## Changes
Version 5 has been deployed recently, and SGC bumped their minimum to 4,
breaking us.

## Why
<!-- Why are these changes needed? Provide the context that the reviewer
might be missing.
For example, were there any decisions behind the change that are not
reflected in the code itself? -->

## Tests
<!-- How have you tested the changes? -->

<!-- If your PR needs to be included in the release notes for next
release,
add a separate entry in NEXT_CHANGELOG.md as part of your PR. -->

Author: ilia-db

experimental/ssh/cmd/connect.go

@@ -35,6 +35,7 @@ the SSH server and handling the connection proxy.
 	var userKnownHostsFile string
 	var liteswap string
 	var skipSettingsCheck bool
+	var environmentVersion int
 
 	cmd.Flags().StringVar(&clusterID, "cluster", "", "Databricks cluster ID (for dedicated clusters)")
 	cmd.Flags().DurationVar(&shutdownDelay, "shutdown-delay", defaultShutdownDelay, "Delay before shutting down the server after the last client disconnects")
@@ -67,6 +68,9 @@ the SSH server and handling the connection proxy.
 	cmd.Flags().BoolVar(&skipSettingsCheck, "skip-settings-check", false, "Skip checking and updating IDE settings")
 	cmd.Flags().MarkHidden("skip-settings-check")
 
+	cmd.Flags().IntVar(&environmentVersion, "environment-version", defaultEnvironmentVersion, "Environment version for serverless compute")
+	cmd.Flags().MarkHidden("environment-version")
+
 	cmd.PreRunE = func(cmd *cobra.Command, args []string) error {
 		// CLI in the proxy mode is executed by the ssh client and can't prompt for input
 		if proxyMode {
@@ -101,6 +105,7 @@ the SSH server and handling the connection proxy.
 			UserKnownHostsFile:   userKnownHostsFile,
 			Liteswap:             liteswap,
 			SkipSettingsCheck:    skipSettingsCheck,
+			EnvironmentVersion:   environmentVersion,
 			AdditionalArgs:       args,
 		}
 		if err := opts.Validate(); err != nil {

experimental/ssh/cmd/constants.go

@@ -3,10 +3,11 @@ package ssh
 import "time"
 
 const (
-	defaultServerPort      = 7772
-	defaultMaxClients      = 10
-	defaultShutdownDelay   = 10 * time.Minute
-	defaultHandoverTimeout = 30 * time.Minute
+	defaultServerPort         = 7772
+	defaultMaxClients         = 10
+	defaultShutdownDelay      = 10 * time.Minute
+	defaultHandoverTimeout    = 30 * time.Minute
+	defaultEnvironmentVersion = 4
 
 	serverTimeout        = 24 * time.Hour
 	taskStartupTimeout   = 10 * time.Minute

experimental/ssh/internal/client/client.go

@@ -44,6 +44,7 @@ var connectionNameRegex = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9_-]*$`)
 const (
 	sshServerTaskKey         = "start_ssh_server"
 	serverlessEnvironmentKey = "ssh_tunnel_serverless"
+	minEnvironmentVersion    = 4
 
 	VSCodeOption  = "vscode"
 	VSCodeCommand = "code"
@@ -98,6 +99,8 @@ type ClientOptions struct {
 	Liteswap string
 	// If true, skip checking and updating IDE settings.
 	SkipSettingsCheck bool
+	// Environment version for serverless compute.
+	EnvironmentVersion int
 }
 
 func (o *ClientOptions) Validate() error {
@@ -117,6 +120,9 @@ func (o *ClientOptions) Validate() error {
 	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.EnvironmentVersion > 0 && o.EnvironmentVersion < minEnvironmentVersion {
+		return fmt.Errorf("environment version must be >= %d, got %d", minEnvironmentVersion, o.EnvironmentVersion)
+	}
 	return nil
 }
 
@@ -182,6 +188,10 @@ func (o *ClientOptions) ToProxyCommand() (string, error) {
 		proxyCommand += " --liteswap=" + o.Liteswap
 	}
 
+	if o.EnvironmentVersion > 0 {
+		proxyCommand += " --environment-version=" + strconv.Itoa(o.EnvironmentVersion)
+	}
+
 	return proxyCommand, nil
 }
 
@@ -508,7 +518,7 @@ func submitSSHTunnelJob(ctx context.Context, client *databricks.WorkspaceClient,
 			{
 				EnvironmentKey: serverlessEnvironmentKey,
 				Spec: &compute.Environment{
-					EnvironmentVersion: "3",
+					EnvironmentVersion: strconv.Itoa(max(opts.EnvironmentVersion, minEnvironmentVersion)),
 				},
 			},
 		}

experimental/ssh/internal/client/client_test.go

@@ -76,6 +76,15 @@ func TestValidate(t *testing.T) {
 			name: "valid IDE cursor",
 			opts: client.ClientOptions{ClusterID: "abc-123", IDE: "cursor"},
 		},
+		{
+			name:    "environment version too low",
+			opts:    client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 3},
+			wantErr: "environment version must be >= 4, got 3",
+		},
+		{
+			name: "valid environment version",
+			opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 4},
+		},
 	}
 
 	for _, tt := range tests {
@@ -140,6 +149,11 @@ func TestToProxyCommand(t *testing.T) {
 			opts: client.ClientOptions{ClusterID: "abc-123", Liteswap: "test-env"},
 			want: quoted + " ssh connect --proxy --cluster=abc-123 --auto-start-cluster=false --shutdown-delay=0s --liteswap=test-env",
 		},
+		{
+			name: "with environment version",
+			opts: client.ClientOptions{ClusterID: "abc-123", EnvironmentVersion: 4},
+			want: quoted + " ssh connect --proxy --cluster=abc-123 --auto-start-cluster=false --shutdown-delay=0s --environment-version=4",
+		},
 	}
 
 	for _, tt := range tests {