Support global Python env

.devcontainer/devcontainer.json

@@ -5,22 +5,27 @@
     // Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile
     "image": "mcr.microsoft.com/devcontainers/typescript-node:1-22-bookworm",
     "containerEnv": {
-        "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0"
+        "COREPACK_ENABLE_DOWNLOAD_PROMPT": "0",
+        "CLI_ARCH": "linux_amd64" // Set the default CLI_ARCH to linux_amd64, for cli download script
     },
     // Features to add to the dev container. More info: https://containers.dev/features.
     "features": {
         "ghcr.io/devcontainers/features/docker-in-docker:2": {},
         "ghcr.io/devcontainers/features/git:1": {},
         "ghcr.io/devcontainers/features/github-cli:1": {},
+        "ghcr.io/devcontainers/features/python:1": {
+            "version": "3.12"
+        },
+        "ghcr.io/devcontainers/features/azure-cli:1": {},
         "ghcr.io/devcontainers-extra/features/act:1": {},
         "ghcr.io/devcontainers-extra/features/actionlint:1": {}
     },
     // Use 'forwardPorts' to make a list of ports inside the container available locally.
     // "forwardPorts": [],
     // Use 'postCreateCommand' to run commands after the container is created.
-    "postCreateCommand": "sudo corepack enable yarn && sudo corepack install && yarn install",
+    "postCreateCommand": "sudo corepack enable yarn && sudo corepack install",
     // Use 'postStartCommand' to run commands after the container is started.
-    "postStartCommand": "yarn -v",
+    "postStartCommand": "yarn install && yarn run build && yarn workspace databricks run package:cli:fetch",
     // Configure tool-specific properties.
     "customizations": {
         "vscode": {
@@ -41,12 +46,17 @@
                 }
             },
             "extensions": [
+                "ms-python.python",
+                "ms-toolsai.jupyter",
                 "dbaeumer.vscode-eslint",
                 "connor4312.esbuild-problem-matchers",
                 "ms-vscode.extension-test-runner",
                 "amodio.tsl-problem-matcher",
                 "esbenp.prettier-vscode",
-                "Tobermory.es6-string-html"
+                "Tobermory.es6-string-html",
+                "github.vscode-github-actions",
+                "redhat.vscode-xml",
+                "redhat.vscode-yaml"
             ]
         }
     },

.vscode/settings.json

@@ -9,10 +9,10 @@
         "dist": true // set this to false to include "dist" folder in search results
     },
     // Turn off tsc task auto detection since we have the necessary tasks as npm scripts
-    "typescript.tsc.autoDetect": "off",
+    "ts.tsc.autoDetect": "off",
     "editor.formatOnSave": true,
     "editor.defaultFormatter": "esbenp.prettier-vscode",
     "[javascript]": {
         "editor.defaultFormatter": "esbenp.prettier-vscode"
     }
-}
+}

README.md

@@ -11,32 +11,28 @@ This repository contains the source code for Databricks extensions for VSCode.
 
 Currently, we have the following packages:
 
--   [databricks-vscode](https://github.com/databricks/databricks-vscode/tree/main/packages/databricks-vscode)
-    The VSCode extension for Databricks published to the VSCode marketplace.
--   [databricks-vscode-types](https://github.com/databricks/databricks-vscode/tree/main/packages/databricks-vscode-types)
-    Type definition of the public API of the VSCode extension.
+- [databricks-vscode](https://github.com/databricks/databricks-vscode/tree/main/packages/databricks-vscode)
+  The VSCode extension for Databricks published to the VSCode marketplace.
+- [databricks-vscode-types](https://github.com/databricks/databricks-vscode/tree/main/packages/databricks-vscode-types)
+  Type definition of the public API of the VSCode extension.
 
 ### Getting Started
 
-Prepare yarn:
+Start the [devcontainer](.devcontainer/devcontainer.json) or use the system of your choice with Node.js 22+ installed and prepare yarn:
 
-```
+```sh
 npm install -g yarn@2
 yarn install
-```
-
-Prepare Databricks JavaScript SDK:
-
-```
-yarn run install:sdk
+yarn run build
 ```
 
 Prepare Databricks CLI:
 
-```
+```sh
 yarn workspace databricks run package:cli:fetch
 ```
 
+Then open the [code workspace](https://code.visualstudio.com/docs/editing/workspaces/workspaces).
 After that you are ready to build and test the `databricks-vscode` extension.
 
 ### Found an issue?

packages/databricks-vscode/.vscode/extensions.json

@@ -2,9 +2,13 @@
     // See http://go.microsoft.com/fwlink/?LinkId=827846
     // for the documentation about the extensions.json format
     "recommendations": [
+        "ms-python.python",
+        "ms-toolsai.jupyter",
         "dbaeumer.vscode-eslint",
         "amodio.tsl-problem-matcher",
         "esbenp.prettier-vscode",
-        "Tobermory.es6-string-html"
+        "Tobermory.es6-string-html",
+        "redhat.vscode-xml",
+        "redhat.vscode-yaml"
     ]
-}
+}

packages/databricks-vscode/src/language/EnvironmentCommands.test.ts

@@ -0,0 +1,110 @@
+import assert from "assert";
+import {Uri} from "vscode";
+import {Environment} from "./MsPythonExtensionApi";
+import {environmentName} from "../utils/environmentUtils";
+
+function makeEnvironment(overrides: Partial<Environment> = {}): Environment {
+    return {
+        id: "env-1",
+        path: "/usr/bin/python3",
+        executable: {
+            uri: Uri.file("/usr/bin/python3"),
+            bitness: "64-bit",
+            sysPrefix: "/usr",
+        },
+        environment: {
+            type: "VirtualEnvironment",
+            name: "test-venv",
+            folderUri: Uri.file("/home/user/test-venv"),
+            workspaceFolder: undefined,
+        },
+        version: {
+            major: 3,
+            minor: 10,
+            micro: 0,
+            release: {level: "final", serial: 0},
+            sysVersion: "3.10.0",
+        },
+        tools: [],
+        ...overrides,
+    };
+}
+
+function makeGlobalEnvironment(): Environment {
+    return makeEnvironment({
+        id: "global-python",
+        path: "/usr/bin/python3",
+        environment: undefined,
+    });
+}
+
+describe(__filename, () => {
+    it("should use path as name for global environments", () => {
+        const env = makeGlobalEnvironment();
+        const name = environmentName(env);
+        assert.strictEqual(name, "3.10.0 /usr/bin/python3");
+    });
+
+    it("should use environment name for virtual environments", () => {
+        const env = makeEnvironment();
+        const name = environmentName(env);
+        assert.strictEqual(name, "3.10.0 test-venv");
+    });
+
+    it("should handle environments without version", () => {
+        const env = makeGlobalEnvironment();
+        (env as any).version = undefined;
+        const name = environmentName(env);
+        assert.strictEqual(name, "/usr/bin/python3");
+    });
+
+    it("should identify global environment as having no environment property", () => {
+        const globalEnv = makeGlobalEnvironment();
+        assert.strictEqual(globalEnv.environment, undefined);
+        assert.ok(globalEnv.version);
+        assert.ok(globalEnv.executable.uri);
+    });
+
+    it("should identify virtual environment as having an environment property", () => {
+        const venvEnv = makeEnvironment();
+        assert.ok(venvEnv.environment);
+        assert.strictEqual(venvEnv.environment!.type, "VirtualEnvironment");
+        assert.strictEqual(venvEnv.environment!.name, "test-venv");
+    });
+
+    it("should distinguish global from non-global by environment field", () => {
+        const globalEnv = makeGlobalEnvironment();
+        const venvEnv = makeEnvironment();
+        assert.strictEqual(!globalEnv.environment, true);
+        assert.strictEqual(!venvEnv.environment, false);
+    });
+
+    it("should create pick item with 'Global' description for global environments", () => {
+        const env = makeGlobalEnvironment();
+        const isGlobal = !env.environment;
+        const item = {
+            label: environmentName(env),
+            description: isGlobal ? "Global" : env.environment?.type,
+            detail: env.path,
+            path: env.path,
+            isGlobal,
+        };
+        assert.strictEqual(item.description, "Global");
+        assert.strictEqual(item.isGlobal, true);
+        assert.strictEqual(item.path, "/usr/bin/python3");
+    });
+
+    it("should create pick item with env type for virtual environments", () => {
+        const env = makeEnvironment();
+        const isGlobal = !env.environment;
+        const item = {
+            label: environmentName(env),
+            description: isGlobal ? "Global" : env.environment?.type,
+            detail: env.path,
+            path: env.path,
+            isGlobal,
+        };
+        assert.strictEqual(item.description, "VirtualEnvironment");
+        assert.strictEqual(item.isGlobal, false);
+    });
+});

packages/databricks-vscode/src/language/EnvironmentCommands.ts

@@ -80,7 +80,22 @@ export class EnvironmentCommands {
         if (environments.length > 0) {
             await this.showEnvironmentsQuickPick(environments);
         } else {
-            await this.pythonExtension.createPythonEnvironment();
+            const createNewLabel = "$(add) Create new environment";
+            const useExistingLabel = "$(globe) Use existing Python";
+            const selectedPick = await window.showQuickPick(
+                [
+                    {label: createNewLabel, alwaysShow: true},
+                    {label: useExistingLabel, alwaysShow: true},
+                ],
+                {title: "Select Python Environment"}
+            );
+            if (selectedPick) {
+                if (selectedPick.label === createNewLabel) {
+                    await this.pythonExtension.createPythonEnvironment();
+                } else if (selectedPick.label === useExistingLabel) {
+                    await this.showPythonInterpreterQuickPick();
+                }
+            }
         }
     }
 
@@ -93,10 +108,12 @@ export class EnvironmentCommands {
             })
         );
         const createNewLabel = "$(add) Create new environment";
+        const useExistingLabel = "$(globe) Use existing Python";
         const usePythonExtensionLabel =
             "$(gear) Use Python Extension to setup environments";
         const staticPicks: QuickPickItem[] = [
             {label: createNewLabel, alwaysShow: true},
+            {label: useExistingLabel, alwaysShow: true},
             {label: usePythonExtensionLabel, alwaysShow: true},
         ];
         const selectedPick = await window.showQuickPick(
@@ -106,6 +123,8 @@ export class EnvironmentCommands {
         if (selectedPick) {
             if (selectedPick.label === createNewLabel) {
                 await this.pythonExtension.createPythonEnvironment();
+            } else if (selectedPick.label === useExistingLabel) {
+                await this.showPythonInterpreterQuickPick();
             } else if (selectedPick.label === usePythonExtensionLabel) {
                 await this.pythonExtension.selectPythonInterpreter();
             } else if (selectedPick.path) {
@@ -116,6 +135,79 @@ export class EnvironmentCommands {
         }
     }
 
+    private async showPythonInterpreterQuickPick() {
+        const environments =
+            await this.pythonExtension.getAllKnownEnvironments();
+        if (environments.length === 0) {
+            window.showInformationMessage(
+                "No Python interpreters found. Install Python or create a virtual environment."
+            );
+            return;
+        }
+
+        type EnvPickItem = QuickPickItem & {
+            path: string;
+            isGlobal: boolean;
+        };
+        const items: EnvPickItem[] = environments.map((env) => {
+            const isGlobal = !env.environment;
+            return {
+                label: environmentName(env),
+                description: isGlobal ? "Global" : env.environment?.type,
+                detail: env.path,
+                path: env.path,
+                isGlobal,
+            };
+        });
+
+        const activeEnvPath =
+            this.pythonExtension.api.environments.getActiveEnvironmentPath();
+
+        const quickPick = window.createQuickPick<EnvPickItem>();
+        quickPick.title = "Select Python Interpreter";
+        quickPick.items = items;
+        quickPick.canSelectMany = false;
+
+        if (activeEnvPath) {
+            const currentItem = items.find(
+                (item) => item.path === activeEnvPath.path
+            );
+            if (currentItem) {
+                quickPick.activeItems = [currentItem];
+            }
+        }
+
+        quickPick.show();
+
+        return new Promise<void>((resolve) => {
+            quickPick.onDidAccept(async () => {
+                const selected = quickPick.selectedItems[0];
+                quickPick.dispose();
+                if (selected) {
+                    if (selected.isGlobal) {
+                        const confirm = await window.showWarningMessage(
+                            "You selected a global Python interpreter. Installing packages like databricks-connect into a global Python may affect other applications. Consider using a virtual environment instead.",
+                            {modal: true},
+                            "Use Global Python"
+                        );
+                        if (confirm !== "Use Global Python") {
+                            resolve();
+                            return;
+                        }
+                    }
+                    await this.pythonExtension.api.environments.updateActiveEnvironmentPath(
+                        selected.path
+                    );
+                }
+                resolve();
+            });
+            quickPick.onDidHide(() => {
+                quickPick.dispose();
+                resolve();
+            });
+        });
+    }
+
     async reinstallDBConnect(cluster?: Cluster) {
         const state = await this.featureManager.isEnabled(
             "environment.dependencies"