SSH: enable databricks extension on the remote side (#4764)

## Changes
The extension changes are in a separate PR -
databricks/databricks-vscode#1861
Here we set our extension id to the IDE settings
And also set special env var that will tell the extension to initialize
in a special "remote" mode, wher it only sets the active venv and does
nothing else (for now)

## Why
venv activation is one of the main friction points

## 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/internal/client/ssh-server-bootstrap.py

@@ -92,6 +92,7 @@ def run_ssh_server():
     os.environ["DATABRICKS_TOKEN"] = ctx.apiToken
     os.environ["DATABRICKS_CLUSTER_ID"] = ctx.clusterId
     os.environ["DATABRICKS_VIRTUAL_ENV"] = sys.executable
+    os.environ["DATABRICKS_REMOTE_ENV"] = "1"
     python_path = os.path.dirname(sys.executable)
     os.environ["PATH"] = f"{python_path}:{os.environ['PATH']}"
     if os.environ.get("VIRTUAL_ENV") is None:

experimental/ssh/internal/vscode/settings.go

@@ -22,6 +22,7 @@ const (
 	remotePlatform       = "linux"
 	pythonExtension      = "ms-python.python"
 	jupyterExtension     = "ms-toolsai.jupyter"
+	databricksExtension  = "databricks.databricks"
 	serverPickPortsKey   = "remote.SSH.serverPickPortsFromRange"
 	remotePlatformKey    = "remote.SSH.remotePlatform"
 	defaultExtensionsKey = "remote.SSH.defaultExtensions"
@@ -178,7 +179,7 @@ func hasCorrectListenOnSocket(v hujson.Value) bool {
 }
 
 func getMissingExtensions(v hujson.Value) []string {
-	required := []string{pythonExtension, jupyterExtension}
+	required := []string{pythonExtension, jupyterExtension, databricksExtension}
 	found := v.Find(jsonPtr(defaultExtensionsKey))
 	if found == nil {
 		return required
@@ -244,7 +245,7 @@ func handleMissingFile(ctx context.Context, ide, connectionName, settingsPath st
 		portRange:      true,
 		platform:       true,
 		listenOnSocket: true,
-		extensions:     []string{pythonExtension, jupyterExtension},
+		extensions:     []string{pythonExtension, jupyterExtension, databricksExtension},
 	}
 	shouldCreate, err := promptUserForUpdate(ctx, ide, connectionName, missing)
 	if err != nil {
@@ -348,7 +349,7 @@ func GetManualInstructions(ide, connectionName string) string {
 		portRange:      true,
 		platform:       true,
 		listenOnSocket: true,
-		extensions:     []string{pythonExtension, jupyterExtension},
+		extensions:     []string{pythonExtension, jupyterExtension, databricksExtension},
 	}
 	return fmt.Sprintf(
 		"To ensure the remote connection works as expected, manually add these settings to your %s settings.json:\n%s",

experimental/ssh/internal/vscode/settings_test.go

@@ -199,7 +199,7 @@ func TestValidateSettings_Complete(t *testing.T) {
 		"remote.SSH.serverPickPortsFromRange": {"test-conn": "29500-29505"},
 		"remote.SSH.remotePlatform": {"test-conn": "linux"},
 		"remote.SSH.remoteServerListenOnSocket": true,
-		"remote.SSH.defaultExtensions": ["ms-python.python", "ms-toolsai.jupyter"]
+		"remote.SSH.defaultExtensions": ["ms-python.python", "ms-toolsai.jupyter", "databricks.databricks"]
 	}`)
 
 	missing := validateSettings(v, "test-conn")
@@ -213,7 +213,7 @@ func TestValidateSettings_Missing(t *testing.T) {
 	assert.False(t, missing.isEmpty())
 	assert.True(t, missing.portRange)
 	assert.True(t, missing.platform)
-	assert.Equal(t, []string{"ms-python.python", "ms-toolsai.jupyter"}, missing.extensions)
+	assert.Equal(t, []string{"ms-python.python", "ms-toolsai.jupyter", "databricks.databricks"}, missing.extensions)
 }
 
 func TestValidateSettings_IncorrectValues(t *testing.T) {
@@ -227,15 +227,15 @@ func TestValidateSettings_IncorrectValues(t *testing.T) {
 	assert.False(t, missing.isEmpty())
 	assert.True(t, missing.portRange)
 	assert.True(t, missing.platform)
-	assert.Equal(t, []string{"ms-toolsai.jupyter"}, missing.extensions)
+	assert.Equal(t, []string{"ms-toolsai.jupyter", "databricks.databricks"}, missing.extensions)
 }
 
 func TestValidateSettings_DuplicateExtensionsNotReported(t *testing.T) {
 	v := parseTestValue(t, `{
 		"remote.SSH.serverPickPortsFromRange": {"test-conn": "29500-29505"},
 		"remote.SSH.remotePlatform": {"test-conn": "linux"},
 		"remote.SSH.remoteServerListenOnSocket": true,
-		"remote.SSH.defaultExtensions": ["ms-python.python", "ms-python.python", "ms-toolsai.jupyter"]
+		"remote.SSH.defaultExtensions": ["ms-python.python", "ms-python.python", "ms-toolsai.jupyter", "databricks.databricks"]
 	}`)
 
 	missing := validateSettings(v, "test-conn")
@@ -246,7 +246,7 @@ func TestValidateSettings_MissingConnection(t *testing.T) {
 	v := parseTestValue(t, `{
 		"remote.SSH.serverPickPortsFromRange": {"other-conn": "29500-29505"},
 		"remote.SSH.remotePlatform": {"other-conn": "linux"},
-		"remote.SSH.defaultExtensions": ["ms-python.python", "ms-toolsai.jupyter"]
+		"remote.SSH.defaultExtensions": ["ms-python.python", "ms-toolsai.jupyter", "databricks.databricks"]
 	}`)
 
 	// Validating for a different connection should show port and platform as missing
@@ -273,7 +273,7 @@ func TestUpdateSettings_PreserveExistingConnections(t *testing.T) {
 	missing := &missingSettings{
 		portRange:  true,
 		platform:   true,
-		extensions: []string{"ms-python.python", "ms-toolsai.jupyter"},
+		extensions: []string{"ms-python.python", "ms-toolsai.jupyter", "databricks.databricks"},
 	}
 
 	err := updateSettings(&v, "conn-c", missing)
@@ -307,10 +307,11 @@ func TestUpdateSettings_PreserveExistingConnections(t *testing.T) {
 
 	// Check that extensions were merged
 	exts := findStringSlice(t, v, jsonPtr(defaultExtensionsKey))
-	assert.Len(t, exts, 3)
+	assert.Len(t, exts, 4)
 	assert.Contains(t, exts, "other.extension")
 	assert.Contains(t, exts, "ms-python.python")
 	assert.Contains(t, exts, "ms-toolsai.jupyter")
+	assert.Contains(t, exts, "databricks.databricks")
 }
 
 func TestUpdateSettings_NewConnection(t *testing.T) {
@@ -319,7 +320,7 @@ func TestUpdateSettings_NewConnection(t *testing.T) {
 	missing := &missingSettings{
 		portRange:  true,
 		platform:   true,
-		extensions: []string{"ms-python.python", "ms-toolsai.jupyter"},
+		extensions: []string{"ms-python.python", "ms-toolsai.jupyter", "databricks.databricks"},
 	}
 
 	err := updateSettings(&v, "new-conn", missing)
@@ -334,9 +335,10 @@ func TestUpdateSettings_NewConnection(t *testing.T) {
 	assert.Equal(t, "linux", val)
 
 	exts := findStringSlice(t, v, jsonPtr(defaultExtensionsKey))
-	assert.Len(t, exts, 2)
+	assert.Len(t, exts, 3)
 	assert.Contains(t, exts, "ms-python.python")
 	assert.Contains(t, exts, "ms-toolsai.jupyter")
+	assert.Contains(t, exts, "databricks.databricks")
 }
 
 func TestUpdateSettings_GlobalExtensions(t *testing.T) {
@@ -544,6 +546,7 @@ func TestGetManualInstructions_VSCode(t *testing.T) {
 	assert.Contains(t, instructions, "linux")
 	assert.Contains(t, instructions, "ms-python.python")
 	assert.Contains(t, instructions, "ms-toolsai.jupyter")
+	assert.Contains(t, instructions, "databricks.databricks")
 	assert.Contains(t, instructions, "remote.SSH.serverPickPortsFromRange")
 	assert.Contains(t, instructions, "remote.SSH.remotePlatform")
 	assert.Contains(t, instructions, "remote.SSH.defaultExtensions")