diff --git a/packages/databricks-vscode/src/run/DatabricksDebugConfigurationProvider.ts b/packages/databricks-vscode/src/run/DatabricksDebugConfigurationProvider.ts index 76b2adc5f..ff08684fb 100644 --- a/packages/databricks-vscode/src/run/DatabricksDebugConfigurationProvider.ts +++ b/packages/databricks-vscode/src/run/DatabricksDebugConfigurationProvider.ts @@ -13,6 +13,8 @@ export interface DatabricksPythonDebugConfiguration extends DebugConfiguration { env?: Record; isInternalDatabricksRun?: boolean; console?: "integratedTerminal" | "externalTerminal" | "internalConsole"; + /** Path to the python interpreter debugpy should use */ + python?: string; } function isTest(debugConfiguration: DebugConfiguration) { diff --git a/packages/databricks-vscode/src/run/RunCommands.test.ts b/packages/databricks-vscode/src/run/RunCommands.test.ts new file mode 100644 index 000000000..db3d2a3dd --- /dev/null +++ b/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); + 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(); + } + }); + }); +}); diff --git a/packages/databricks-vscode/src/run/RunCommands.ts b/packages/databricks-vscode/src/run/RunCommands.ts index a4f2e869e..194e58132 100644 --- a/packages/databricks-vscode/src/run/RunCommands.ts +++ b/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,17 +162,37 @@ 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", @@ -147,9 +200,25 @@ export class RunCommands { 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(