Fix Databricks Connect run/debug interpreter split-brain

packages/databricks-vscode/src/run/DatabricksDebugConfigurationProvider.ts

@@ -13,6 +13,8 @@ export interface DatabricksPythonDebugConfiguration extends DebugConfiguration {
     env?: Record<string, any>;
     isInternalDatabricksRun?: boolean;
     console?: "integratedTerminal" | "externalTerminal" | "internalConsole";
+    /** Path to the python interpreter debugpy should use */
+    python?: string;
 }
 
 function isTest(debugConfiguration: DebugConfiguration) {

packages/databricks-vscode/src/run/RunCommands.test.ts

@@ -0,0 +1,167 @@
+import * as assert from "assert";
+import {commands, ExtensionContext, Uri} from "vscode";
+import {anything, instance, mock, verify, when} from "ts-mockito";
+import {RunCommands} from "./RunCommands";
+import {ConnectionManager} from "../configuration/ConnectionManager";
+import {MsPythonExtensionWrapper} from "../language/MsPythonExtensionWrapper";
+import {FeatureManager, FeatureState} from "../feature-manager/FeatureManager";
+import {CustomWhenContext} from "../vscode-objs/CustomWhenContext";
+import {WorkspaceFolderManager} from "../vscode-objs/WorkspaceFolderManager";
+import {Telemetry} from "../telemetry";
+
+function featureState(available: boolean, executable?: string): FeatureState {
+    return {
+        available,
+        steps: new Map([
+            [
+                "checkPythonEnvironment",
+                {
+                    id: "checkPythonEnvironment",
+                    available,
+                    title: "Active Environment: .venv",
+                    message: executable,
+                },
+            ],
+        ]),
+    };
+}
+
+describe(__filename, () => {
+    let featureManagerMock: FeatureManager;
+    let pythonExtensionMock: MsPythonExtensionWrapper;
+    let connectionManagerMock: ConnectionManager;
+    let runCommands: RunCommands;
+
+    beforeEach(() => {
+        featureManagerMock = mock<FeatureManager>(FeatureManager);
+        pythonExtensionMock = mock(MsPythonExtensionWrapper);
+        connectionManagerMock = mock(ConnectionManager);
+        when(connectionManagerMock.state).thenReturn("CONNECTED");
+        runCommands = new RunCommands(
+            instance(connectionManagerMock),
+            instance(mock(WorkspaceFolderManager)),
+            instance(pythonExtensionMock),
+            instance(featureManagerMock),
+            {subscriptions: []} as unknown as ExtensionContext,
+            instance(mock(CustomWhenContext)),
+            instance(mock(Telemetry))
+        );
+    });
+
+    describe("createDbconnectDebugConfig", () => {
+        it("should pin debugpy to the verified python executable", () => {
+            const config = runCommands.createDbconnectDebugConfig(
+                Uri.file("/project/src/main.py"),
+                "/project/.venv/bin/python"
+            );
+            assert.strictEqual(config.type, "debugpy");
+            assert.strictEqual(config.databricks, true);
+            assert.strictEqual(
+                config.program,
+                Uri.file("/project/src/main.py").fsPath
+            );
+            assert.strictEqual(config.python, "/project/.venv/bin/python");
+        });
+    });
+
+    describe("checkDbconnectEnabled", () => {
+        it("should pass when the verified interpreter is still selected", async () => {
+            when(
+                featureManagerMock.isEnabled("environment.dependencies")
+            ).thenResolve(featureState(true, "/project/.venv/bin/python"));
+            when(pythonExtensionMock.getPythonExecutable()).thenResolve(
+                "/project/.venv/bin/python"
+            );
+
+            const result = await runCommands["checkDbconnectEnabled"]();
+
+            assert.strictEqual(result, true);
+            verify(
+                featureManagerMock.isEnabled("environment.dependencies", true)
+            ).never();
+        });
+
+        it("should re-verify when the selected interpreter changed", async () => {
+            when(
+                featureManagerMock.isEnabled("environment.dependencies")
+            ).thenResolve(featureState(true, "/old/python"));
+            when(
+                featureManagerMock.isEnabled("environment.dependencies", true)
+            ).thenResolve(featureState(true, "/project/.venv/bin/python"));
+            when(pythonExtensionMock.getPythonExecutable()).thenResolve(
+                "/project/.venv/bin/python"
+            );
+
+            const result = await runCommands["checkDbconnectEnabled"]();
+
+            assert.strictEqual(result, true);
+            verify(
+                featureManagerMock.isEnabled("environment.dependencies", true)
+            ).once();
+        });
+
+        it("should not re-verify while disconnected", async () => {
+            when(connectionManagerMock.state).thenReturn("DISCONNECTED");
+            when(
+                featureManagerMock.isEnabled("environment.dependencies")
+            ).thenResolve(featureState(true, "/old/python"));
+            when(pythonExtensionMock.getPythonExecutable()).thenResolve(
+                "/project/.venv/bin/python"
+            );
+
+            const result = await runCommands["checkDbconnectEnabled"]();
+
+            assert.strictEqual(result, true);
+            verify(
+                featureManagerMock.isEnabled("environment.dependencies", true)
+            ).never();
+        });
+
+        it("should not re-verify when the feature is unavailable", async () => {
+            when(featureManagerMock.isEnabled(anything())).thenResolve(
+                featureState(false, undefined)
+            );
+            when(pythonExtensionMock.getPythonExecutable()).thenResolve(
+                undefined
+            );
+            // Stub the setup command flow that runs for unavailable features.
+            const registration = commands.registerCommand(
+                "databricks.environment.setup",
+                () => false
+            );
+            try {
+                const result = await runCommands["checkDbconnectEnabled"]();
+                assert.strictEqual(result, false);
+                verify(
+                    featureManagerMock.isEnabled(
+                        "environment.dependencies",
+                        true
+                    )
+                ).never();
+            } finally {
+                registration.dispose();
+            }
+        });
+
+        it("should proceed after a successful in-flow setup", async () => {
+            // Unavailable on the first check; the setup flow fixes it, so the
+            // re-check after setup reports available and the launch proceeds.
+            when(
+                featureManagerMock.isEnabled("environment.dependencies")
+            ).thenResolve(
+                featureState(false, undefined),
+                featureState(true, "/project/.venv/bin/python")
+            );
+            const registration = commands.registerCommand(
+                "databricks.environment.setup",
+                () => undefined
+            );
+            try {
+                const result = await runCommands["checkDbconnectEnabled"]();
+                assert.strictEqual(result, true);
+            } finally {
+                registration.dispose();
+            }
+        });
+    });
+});

packages/databricks-vscode/src/run/RunCommands.ts

@@ -5,7 +5,7 @@ import {LocalUri} from "../sync/SyncDestination";
 import {DatabricksPythonDebugConfiguration} from "./DatabricksDebugConfigurationProvider";
 import {MsPythonExtensionWrapper} from "../language/MsPythonExtensionWrapper";
 import path from "path";
-import {FeatureManager} from "../feature-manager/FeatureManager";
+import {FeatureManager, FeatureState} from "../feature-manager/FeatureManager";
 import {
     escapeExecutableForTerminal,
     escapePathArgument,
@@ -113,13 +113,46 @@ export class RunCommands {
     }
 
     private async checkDbconnectEnabled() {
-        const featureState = await this.featureManager.isEnabled(
+        let featureState = await this.featureManager.isEnabled(
             "environment.dependencies"
         );
+        if (
+            featureState.available &&
+            // Re-checks block on an established connection, so don't force
+            // one while disconnected.
+            this.connection.state === "CONNECTED" &&
+            (await this.isPythonEnvironmentStale(featureState))
+        ) {
+            // The Python extension doesn't always notify us about interpreter
+            // changes, so the cached state can refer to a previously selected
+            // interpreter. Re-check before launching anything with it.
+            featureState = await this.featureManager.isEnabled(
+                "environment.dependencies",
+                true
+            );
+        }
         if (featureState.available) {
             return true;
         }
-        return await commands.executeCommand("databricks.environment.setup");
+        // Run the setup flow, then re-check: a successful setup should let the
+        // launch proceed instead of aborting and making the user re-trigger.
+        await commands.executeCommand("databricks.environment.setup");
+        return (await this.featureManager.isEnabled("environment.dependencies"))
+            .available;
+    }
+
+    private async isPythonEnvironmentStale(featureState: FeatureState) {
+        // For an accepted checkPythonEnvironment step the message holds
+        // the path of the verified python executable.
+        const verifiedExecutable = featureState.steps.get(
+            "checkPythonEnvironment"
+        )?.message;
+        const currentExecutable =
+            await this.pythonExtension.getPythonExecutable();
+        return (
+            verifiedExecutable !== undefined &&
+            verifiedExecutable !== currentExecutable
+        );
     }
 
     private getTargetResource(resource?: Uri): Uri | undefined {
@@ -129,27 +162,63 @@ export class RunCommands {
         return resource;
     }
 
-    async debugFileUsingDbconnect(resource?: Uri) {
+    /**
+     * Shared preflight for the run and debug Databricks Connect launchers:
+     * verifies the environment, resolves the target file, and resolves the
+     * python interpreter. Run and debug must use the *same* interpreter, so it
+     * is resolved here in one place. Returns undefined (after surfacing an
+     * error message) when any step fails.
+     */
+    private async resolveDbconnectLaunch(
+        resource?: Uri
+    ): Promise<{targetResource: Uri; executable: string} | undefined> {
         if (!(await this.checkDbconnectEnabled())) {
-            return;
+            return undefined;
         }
-
         const targetResource = this.getTargetResource(resource);
         if (!targetResource) {
-            window.showErrorMessage("Open a file to run");
-            return;
+            window.showErrorMessage("Open a file to run on Databricks");
+            return undefined;
         }
-        const config: DatabricksPythonDebugConfiguration = {
+        const executable = await this.pythonExtension.getPythonExecutable();
+        if (!executable) {
+            window.showErrorMessage("No python executable found");
+            return undefined;
+        }
+        return {targetResource, executable};
+    }
+
+    createDbconnectDebugConfig(
+        targetResource: Uri,
+        executable: string
+    ): DatabricksPythonDebugConfiguration {
+        return {
             program: targetResource.fsPath,
             type: "debugpy",
             name: "Databricks Connect",
             request: "launch",
             databricks: true,
             console: "integratedTerminal",
             env: {...process.env},
+            // Pin debugpy to the interpreter we have verified. Without this
+            // debugpy falls back to the Python extension's selected
+            // interpreter for the workspace folder, which can differ from the
+            // environment used by "run" and by the verification checks.
+            python: executable,
         };
+    }
 
-        debug.startDebugging(
+    async debugFileUsingDbconnect(resource?: Uri) {
+        const launch = await this.resolveDbconnectLaunch(resource);
+        if (!launch) {
+            return;
+        }
+        const config = this.createDbconnectDebugConfig(
+            launch.targetResource,
+            launch.executable
+        );
+
+        await debug.startDebugging(
             this.workspaceFolderManager.activeWorkspaceFolder,
             config
         );
@@ -161,21 +230,11 @@ export class RunCommands {
     }
 
     async runFileUsingDbconnect(resource?: Uri) {
-        if (!(await this.checkDbconnectEnabled())) {
-            return;
-        }
-
-        const targetResource = this.getTargetResource(resource);
-        if (!targetResource) {
-            window.showErrorMessage("Open a file to run");
-            return;
-        }
-
-        const executable = await this.pythonExtension.getPythonExecutable();
-        if (!executable) {
-            window.showErrorMessage("No python executable found");
+        const launch = await this.resolveDbconnectLaunch(resource);
+        if (!launch) {
             return;
         }
+        const {targetResource, executable} = launch;
 
         const terminal = window.activeTerminal ?? window.createTerminal();
         const bootstrapPath = this.context.asAbsolutePath(