diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 906d0161e..b74035f8b 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -17,6 +17,7 @@ import {ClusterModel} from "./cluster/ClusterModel"; import {ClusterCommands} from "./cluster/ClusterCommands"; import {ConfigurationDataProvider} from "./ui/configuration-view/ConfigurationDataProvider"; import {composePythonSetupEntry} from "./ui/configuration-view/pythonSetupEntry"; +import {routeEnvironmentSetup} from "./language/pythonSetupRouting"; import {COPY_COMMAND_IDS} from "./ui/configuration-view/copyActions"; import {AiToolsManager} from "./aitools/AiToolsManager"; import {AiToolsCommands} from "./aitools/AiToolsCommands"; @@ -911,7 +912,10 @@ export async function activate( pythonExtensionWrapper, environmentDependenciesInstaller, configureAutocomplete, - packageManagerTelemetry + packageManagerTelemetry, + // Constructed lazily (first isEnabled call is well after + // pythonSetupEnvironment is wired), so this reference is safe. + () => pythonSetupEnvironment.isVisible() ) ); // uv-native Python environment setup (python-setup). Constructed always, @@ -1157,8 +1161,16 @@ export async function activate( context.subscriptions.push( telemetry.registerCommand( "databricks.environment.setup", - environmentCommands.setup, - environmentCommands + // Route to the uv-native flow when it is the active surface for the + // project, else the legacy checklist. Every trigger surface (status + // bar, config-view rows, palette, the run/debug gate) funnels + // through this command, so they all dispatch here. + (stepId?: string) => + routeEnvironmentSetup( + pythonSetupEnvironment, + environmentCommands, + stepId + ) ), telemetry.registerCommand( "databricks.environment.refresh", diff --git a/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.test.ts b/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.test.ts new file mode 100644 index 000000000..23a8c930e --- /dev/null +++ b/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.test.ts @@ -0,0 +1,86 @@ +import * as assert from "assert"; +import type {Disposable} from "vscode"; +import {EnvironmentDependenciesVerifier} from "./EnvironmentDependenciesVerifier"; +import {ConnectionManager} from "../configuration/ConnectionManager"; +import {MsPythonExtensionWrapper} from "./MsPythonExtensionWrapper"; +import {EnvironmentDependenciesInstaller} from "./EnvironmentDependenciesInstaller"; +import {ConfigureAutocomplete} from "./ConfigureAutocomplete"; +import {PackageManagerTelemetry} from "./PackageManagerTelemetry"; + +// ConnectionManager/MsPythonExtensionWrapper expose their VS Code Events as +// instance properties (`emitter.event`), which ts-mockito can't stub — calling +// them throws "not a function" during construction. Hand-rolled stubs give the +// constructor real no-op event subscriptions and let us drive the one handler +// under test directly. +const noopEvent = () => ({dispose() {}}) as Disposable; + +describe(__filename, () => { + let showCalls: unknown[]; + + function makeVerifier(isUvActive: () => Promise) { + const connectionManager = { + serverless: false, + cluster: undefined, + onDidChangeCluster: noopEvent, + onDidChangeState: noopEvent, + } as unknown as ConnectionManager; + + const pythonExtension = { + pythonEnvironment: Promise.resolve({ + version: {major: 3, minor: 10, micro: 0}, + environment: {name: ".venv"}, + executable: {uri: {fsPath: "/project/.venv/bin/python"}}, + }), + getPythonExecutable: async () => "/project/.venv/bin/python", + // databricks-connect missing: the legacy path would offer to install + // it on an interpreter change. + getPackageDetailsFromEnvironment: async () => undefined, + onDidChangePythonExecutable: noopEvent, + } as unknown as MsPythonExtensionWrapper; + + const installer = { + show: (advertisement?: boolean) => { + showCalls.push(advertisement); + return Promise.resolve(); + }, + onDidTryInstallation: noopEvent, + } as unknown as EnvironmentDependenciesInstaller; + + const configureAutocomplete = { + shouldSetupBuiltins: async () => false, + onDidUpdate: noopEvent, + } as unknown as ConfigureAutocomplete; + + const packageManagerTelemetry = + {} as unknown as PackageManagerTelemetry; + + return new EnvironmentDependenciesVerifier( + connectionManager, + pythonExtension, + installer, + configureAutocomplete, + packageManagerTelemetry, + isUvActive + ); + } + + beforeEach(() => { + showCalls = []; + }); + + it("offers the legacy install prompt on interpreter change when uv is not active", async () => { + const verifier = makeVerifier(async () => false); + + await verifier["onInterpreterChanged"](); + + assert.deepStrictEqual(showCalls, [true]); + }); + + it("suppresses the legacy install prompt on interpreter change when uv is active", async () => { + const verifier = makeVerifier(async () => true); + + await verifier["onInterpreterChanged"](); + + assert.deepStrictEqual(showCalls, []); + }); +}); diff --git a/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts b/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts index 30040ceef..522a70f50 100644 --- a/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts +++ b/packages/databricks-vscode/src/language/EnvironmentDependenciesVerifier.ts @@ -20,7 +20,12 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier { private readonly pythonExtension: MsPythonExtensionWrapper, private readonly installer: EnvironmentDependenciesInstaller, private readonly configureAutocomplete: ConfigureAutocomplete, - private readonly packageManagerTelemetry: PackageManagerTelemetry + private readonly packageManagerTelemetry: PackageManagerTelemetry, + // Whether the uv-native flow is the active surface for this project. When + // it is, this legacy checklist stays evaluated (other consumers read its + // state) but must not auto-prompt, or the user gets both flows' prompts. + // Defaults to "never active" so the legacy construction is unchanged. + private readonly isUvActive: () => Promise = async () => false ) { super([ "checkCluster", @@ -43,14 +48,10 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier { await this.checkWorkspaceHasUc(); } }, this), - this.pythonExtension.onDidChangePythonExecutable(async () => { - await this.checkPythonEnvironment(); - const depsCheck = await this.checkEnvironmentDependencies(); - if (!depsCheck.available && depsCheck.action) { - await depsCheck.action(true); - } - await this.checkBuiltins(); - }, this), + this.pythonExtension.onDidChangePythonExecutable( + () => this.onInterpreterChanged(), + this + ), this.installer.onDidTryInstallation(async () => { await this.checkEnvironmentDependencies(); await this.checkBuiltins(); @@ -61,6 +62,22 @@ export class EnvironmentDependenciesVerifier extends MultiStepAccessVerifier { ); } + private async onInterpreterChanged() { + await this.checkPythonEnvironment(); + const depsCheck = await this.checkEnvironmentDependencies(); + // Suppress the legacy auto-install prompt when the uv-native flow owns + // this project — it drives its own setup, so a second prompt here is a + // duplicate. + if ( + !depsCheck.available && + depsCheck.action && + !(await this.isUvActive()) + ) { + await depsCheck.action(true); + } + await this.checkBuiltins(); + } + promptForAttachingCluster(msg: string) { return async () => { await commands.executeCommand( diff --git a/packages/databricks-vscode/src/language/pythonSetupRouting.test.ts b/packages/databricks-vscode/src/language/pythonSetupRouting.test.ts new file mode 100644 index 000000000..548064d4f --- /dev/null +++ b/packages/databricks-vscode/src/language/pythonSetupRouting.test.ts @@ -0,0 +1,47 @@ +import * as assert from "assert"; +import {routeEnvironmentSetup} from "./pythonSetupRouting"; + +describe(__filename, () => { + function makeUv(visible: boolean) { + const calls = {setup: 0}; + return { + uv: { + isVisible: () => Promise.resolve(visible), + setup: async () => { + calls.setup++; + }, + }, + calls, + }; + } + + it("routes to the uv flow when it is the active surface", async () => { + const {uv, calls} = makeUv(true); + const legacyCalls: (string | undefined)[] = []; + const legacy = { + setup: async (stepId?: string) => { + legacyCalls.push(stepId); + }, + }; + + await routeEnvironmentSetup(uv, legacy, "checkPythonEnvironment"); + + assert.strictEqual(calls.setup, 1); + assert.deepStrictEqual(legacyCalls, []); + }); + + it("routes to the legacy checklist (with the step id) when uv is not active", async () => { + const {uv, calls} = makeUv(false); + const legacyCalls: (string | undefined)[] = []; + const legacy = { + setup: async (stepId?: string) => { + legacyCalls.push(stepId); + }, + }; + + await routeEnvironmentSetup(uv, legacy, "checkPythonEnvironment"); + + assert.strictEqual(calls.setup, 0); + assert.deepStrictEqual(legacyCalls, ["checkPythonEnvironment"]); + }); +}); diff --git a/packages/databricks-vscode/src/language/pythonSetupRouting.ts b/packages/databricks-vscode/src/language/pythonSetupRouting.ts new file mode 100644 index 000000000..3c7abe4b5 --- /dev/null +++ b/packages/databricks-vscode/src/language/pythonSetupRouting.ts @@ -0,0 +1,38 @@ +/** + * The slice of the uv-native setup flow ({@link + * ../python-setup/controllers/PythonSetupEnvironmentSetup}) that the setup + * command router needs. Kept as a narrow structural interface so the router + * carries no dependency on the controller/gateway layers and stays + * unit-testable. + */ +export interface UvPythonSetup { + /** Whether the uv-native flow is the active surface for the current project. */ + isVisible(): Promise; + /** Run the uv-native setup (re-entrancy-guarded); resolves when it settles. */ + setup(): Promise; +} + +/** The legacy checklist setup, as the router invokes it. */ +export interface LegacyEnvironmentSetup { + setup(stepId?: string): Promise; +} + +/** + * Back the `databricks.environment.setup` command: run the uv-native flow when + * it is the active surface for the current project, otherwise the legacy + * checklist. Routing here means every surface that funnels through that command + * (status bar, config-view rows, palette, the run/debug gate) reaches the right + * flow without per-surface branching. `stepId` is a legacy-only affordance and + * is ignored by the uv flow, which has no per-step entry points. + */ +export async function routeEnvironmentSetup( + uv: UvPythonSetup, + legacy: LegacyEnvironmentSetup, + stepId?: string +): Promise { + if (await uv.isVisible()) { + await uv.setup(); + return; + } + await legacy.setup(stepId); +} diff --git a/packages/databricks-vscode/src/run/RunCommands.test.ts b/packages/databricks-vscode/src/run/RunCommands.test.ts index 77cc20aa8..aa79460a9 100644 --- a/packages/databricks-vscode/src/run/RunCommands.test.ts +++ b/packages/databricks-vscode/src/run/RunCommands.test.ts @@ -1,6 +1,6 @@ import * as assert from "assert"; import {commands, ExtensionContext, Uri} from "vscode"; -import {anything, instance, mock, verify, when} from "ts-mockito"; +import {instance, mock, verify, when} from "ts-mockito"; import {RunCommands} from "./RunCommands"; import {ConnectionManager} from "../configuration/ConnectionManager"; import {MsPythonExtensionWrapper} from "../language/MsPythonExtensionWrapper"; @@ -119,10 +119,16 @@ describe(__filename, () => { ).never(); }); - it("should not re-verify when the feature is unavailable", async () => { - when(featureManagerMock.isEnabled(anything())).thenResolve( - featureState(false, undefined) - ); + it("should not run the staleness re-check when the feature is unavailable", async () => { + // Unavailable on the first (cached) check, so the staleness re-check + // is skipped; setup runs and the forced post-setup check still + // reports unavailable, so the launch aborts. + when( + featureManagerMock.isEnabled("environment.dependencies") + ).thenResolve(featureState(false, undefined)); + when( + featureManagerMock.isEnabled("environment.dependencies", true) + ).thenResolve(featureState(false, undefined)); when(pythonExtensionMock.getPythonExecutable()).thenResolve( undefined ); @@ -134,12 +140,6 @@ describe(__filename, () => { try { const result = await runCommands["checkDbconnectEnabled"](); assert.strictEqual(result, false); - verify( - featureManagerMock.isEnabled( - "environment.dependencies", - true - ) - ).never(); } finally { registration.dispose(); } @@ -147,13 +147,14 @@ describe(__filename, () => { 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. + // forced 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") - ); + ).thenResolve(featureState(false, undefined)); + when( + featureManagerMock.isEnabled("environment.dependencies", true) + ).thenResolve(featureState(true, "/project/.venv/bin/python")); const registration = commands.registerCommand( "databricks.environment.setup", () => undefined diff --git a/packages/databricks-vscode/src/run/RunCommands.ts b/packages/databricks-vscode/src/run/RunCommands.ts index d23f020fb..008ee8d94 100644 --- a/packages/databricks-vscode/src/run/RunCommands.ts +++ b/packages/databricks-vscode/src/run/RunCommands.ts @@ -138,9 +138,17 @@ export class RunCommands { } // Run the setup flow, then re-check: a successful setup should let the // launch proceed instead of aborting and making the user re-trigger. + // Force a fresh check rather than reading the cache: when the setup + // command routes to the uv flow, it adopts the interpreter and the + // legacy state refreshes only on the (async) interpreter-change event, + // which may not have landed yet. await commands.executeCommand("databricks.environment.setup"); - return (await this.featureManager.isEnabled("environment.dependencies")) - .available; + return ( + await this.featureManager.isEnabled( + "environment.dependencies", + true + ) + ).available; } private async isPythonEnvironmentStale(featureState: FeatureState) {