diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 762a61349..91392b8df 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -1559,7 +1559,7 @@ export async function activate( context.subscriptions.push(pythonSetupEntry); // Once-per-session adoption gauge: for a project with a uv-native setup on - // record, whether its managed .venv still exists. Distinct from drift above + // record, whether its managed venv still exists. Distinct from drift above // (which compares compute env keys) — drift never checks that the venv is // actually present. Measurement only, best-effort, and it derives no env key // of its own: drift already reports env-key mismatch via python_env.drift. @@ -1585,8 +1585,8 @@ export async function activate( : connectionManager.cluster ? "cluster" : "none", - venvExists: (root) => - existsSync(venvInterpreterPath(path.join(root, ".venv"))), + savedSetup: () => stateStorage.get("databricks.pythonSetup.setupState"), + venvExists: (venvDir) => existsSync(venvInterpreterPath(venvDir)), record: (report) => telemetry.recordPythonSetupAdoption(report), }); // A connect-time reading: report once the connection is CONNECTED (so the diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.test.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.test.ts index a86de60c1..de6f2512d 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.test.ts @@ -1,4 +1,5 @@ import {expect} from "chai"; +import path from "node:path"; import {PythonSetupAdoption} from "../../telemetry/pythonSetupExtensions"; import { PythonSetupAdoptionDeps, @@ -15,6 +16,7 @@ function makeDeps(over: Partial = {}): { isAttributable: () => true, isVpexActive: () => true, getTargetType: () => "serverless", + savedSetup: () => undefined, venvExists: () => true, record: (r) => recorded.push(r), ...over, @@ -39,6 +41,52 @@ describe("PythonSetupAdoptionManager", () => { ]); }); + it("checks the venv folder that the last setup saved", () => { + // A uv workspace member shares the venv at the workspace root. + const checked: string[] = []; + const {deps} = makeDeps({ + projectRoot: () => "/mono/bundles/sales", + savedSetup: () => ({ + projectRoot: "/mono/bundles/sales", + venvPath: "/mono/.venv", + }), + venvExists: (venvDir) => { + checked.push(venvDir); + return true; + }, + }); + new PythonSetupAdoptionManager(deps).report(); + expect(checked).to.deep.equal(["/mono/.venv"]); + }); + + it("does not emit while the saved setup belongs to another project", () => { + // Two bundles in one workspace folder: the state is from bundle A. + let root = "/ws/b"; + const {deps, recorded} = makeDeps({ + projectRoot: () => root, + savedSetup: () => ({projectRoot: "/ws/a", venvPath: "/ws/a/.venv"}), + }); + const manager = new PythonSetupAdoptionManager(deps); + manager.report(); + expect(recorded).to.deep.equal([]); + // Not latched: switching to A still emits. + root = "/ws/a"; + manager.report(); + expect(recorded).to.have.length(1); + }); + + it("checks /.venv when the saved state has no venv folder", () => { + const checked: string[] = []; + const {deps} = makeDeps({ + venvExists: (venvDir) => { + checked.push(venvDir); + return true; + }, + }); + new PythonSetupAdoptionManager(deps).report(); + expect(checked).to.deep.equal([path.join("/ws/project", ".venv")]); + }); + it("passes through the attached compute kind, including none", () => { const {deps, recorded} = makeDeps({getTargetType: () => "none"}); new PythonSetupAdoptionManager(deps).report(); diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.ts index 6021374be..e21624c1b 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupAdoptionManager.ts @@ -1,3 +1,5 @@ +import path from "node:path"; +import {PythonSetupState} from "../../vscode-objs/StateStorage"; import {TargetCompute} from "../../telemetry/constants"; import {PythonSetupAdoption} from "../../telemetry/pythonSetupExtensions"; @@ -8,7 +10,7 @@ export interface PythonSetupAdoptionDeps { * Whether the persisted setup state can be attributed to the active project. * `setupState` is a single workspace-scoped key with no per-project * namespacing, so in a multi-root workspace it may belong to a sibling root: - * checking a different root's `.venv` against it would emit a spurious + * checking a different root's venv against it would emit a spurious * `venvPresent: false` and inflate the denominator. False there suppresses the * reading rather than emit an untrustworthy one. */ @@ -21,15 +23,22 @@ export interface PythonSetupAdoptionDeps { isVpexActive: () => boolean; /** The compute kind attached right now (cluster / serverless / none). */ getTargetType: () => TargetCompute; - /** Whether the project's managed `.venv` interpreter exists on disk. */ - venvExists: (root: string) => boolean; + /** + * The project and venv folders in the persisted setup state; each is + * undefined in state saved by older versions. + */ + savedSetup: () => + | Pick + | undefined; + /** Whether the managed venv interpreter in `venvDir` exists on disk. */ + venvExists: (venvDir: string) => boolean; /** Emit the gauge. Wraps telemetry; must be best-effort (never throw). */ record: (report: PythonSetupAdoption) => void; } /** * Emits the once-per-session {@link PythonSetupAdoption} gauge: for a project - * that has a uv-native Python setup on record, whether its managed `.venv` is + * that has a uv-native Python setup on record, whether its managed venv is * still present and what compute is attached. Purely a measurement — it reads * state and records, never changing the flow it observes. * @@ -69,8 +78,19 @@ export class PythonSetupAdoptionManager { if (!this.deps.isVpexActive()) { return; } + // The single setupState key can hold the setup of another bundle in + // this workspace folder. Do NOT latch: switching to it still reports. + const saved = this.deps.savedSetup(); + if ( + saved?.projectRoot !== undefined && + saved.projectRoot !== root + ) { + return; + } this.deps.record({ - venvPresent: this.deps.venvExists(root), + venvPresent: this.deps.venvExists( + saved?.venvPath ?? path.join(root, ".venv") + ), currentTargetType: this.deps.getTargetType(), }); this.reported.add(root); diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts index 3f083ffd4..bce0116d8 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.test.ts @@ -2,6 +2,7 @@ import {expect} from "chai"; import path from "node:path"; import { PythonSetupEnvironmentSetup, + PythonSetupPersistedState, PythonSetupSetupDeps, } from "./PythonSetupEnvironmentSetup"; import { @@ -128,6 +129,7 @@ function makeDeps( realPath: async (p) => p, adoptInterpreter: async () => {}, openProjectFile: async () => {}, + checkProjectFileBackup: async () => {}, restoreProjectFile: async () => {}, saveState: () => {}, notify: async () => {}, @@ -160,7 +162,7 @@ describe("PythonSetupEnvironmentSetup.setup", () => { it("runs the CLI and marks ready on a successful real run", async () => { const cli = makeCli({resolve: SUCCESS_REAL_RUN}); const adopted: Array<{venvPath: string; root: string}> = []; - const saved: Array<{envKey: string; pythonVersion: string}> = []; + const saved: PythonSetupPersistedState[] = []; const succeeded: PythonSetupResult[] = []; const setup = new PythonSetupEnvironmentSetup( makeDeps({ @@ -196,6 +198,8 @@ describe("PythonSetupEnvironmentSetup.setup", () => { { envKey: SUCCESS_REAL_RUN.compute!.envKey, pythonVersion: SUCCESS_REAL_RUN.resolved!.pythonVersion, + venvPath: path.resolve("/proj", ".venv"), + projectRoot: "/proj", }, ]); // The success is announced, with the parsed result. @@ -230,34 +234,9 @@ describe("PythonSetupEnvironmentSetup.setup", () => { ]); }); - it("adopts the shared venv at the root for a uv workspace member", async () => { - const member = "/mono/bundles/sales"; - const shared = path.resolve(member, "../../.venv"); - const adopted: Array<{venvPath: string; root: string}> = []; - const venvDirs: string[] = []; - const setup = new PythonSetupEnvironmentSetup( - makeDeps({ - cli: makeCli({resolve: SUCCESS_REAL_RUN_UV_WORKSPACE_MEMBER}), - projectRoot: () => member, - adoptInterpreter: async (venvPath, root) => { - adopted.push({venvPath, root}); - }, - showSuccess: async (_r, _flags, venvDir) => { - venvDirs.push(venvDir); - }, - }) - ); - - await setup.setup(); - - expect(setup.ready).to.equal(true); - expect(adopted).to.deep.equal([{venvPath: shared, root: member}]); - expect(venvDirs).to.deep.equal([shared]); - }); - - it("counts a climbing venvPath from the real project folder", async () => { - // The CLI counts "../../.venv" from the symlink-resolved project folder, - // so the controller must not join it to the folder by text. + it("resolves a relative venvPath from the real project folder", async () => { + // The CLI counts venvPath from the symlink-resolved project folder, so + // the controller must not join it to the project folder by text. const adopted: string[] = []; const setup = new PythonSetupEnvironmentSetup( makeDeps({ @@ -314,6 +293,36 @@ describe("PythonSetupEnvironmentSetup.setup", () => { expect(adopted).to.deep.equal([abs]); }); + it("adopts and saves the shared venv at the root for a uv workspace member", async () => { + const member = "/mono/bundles/sales"; + const shared = path.resolve(member, "../../.venv"); + const adopted: Array<{venvPath: string; root: string}> = []; + const saved: PythonSetupPersistedState[] = []; + const venvDirs: string[] = []; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli: makeCli({resolve: SUCCESS_REAL_RUN_UV_WORKSPACE_MEMBER}), + projectRoot: () => member, + adoptInterpreter: async (venvPath, root) => { + adopted.push({venvPath, root}); + }, + saveState: (s) => { + saved.push(s); + }, + showSuccess: async (_r, _flags, venvDir) => { + venvDirs.push(venvDir); + }, + }) + ); + + await setup.setup(); + + expect(setup.ready).to.equal(true); + expect(adopted).to.deep.equal([{venvPath: shared, root: member}]); + expect(saved.map((s) => s.venvPath)).to.deep.equal([shared]); + expect(venvDirs).to.deep.equal([shared]); + }); + it("reports ready only for the project it was set up for, not another active project", async () => { // Set up /projA, then model the active project switching to /projB // (never provisioned). readiness must track the run's project, not leak @@ -2098,6 +2107,434 @@ describe("PythonSetupEnvironmentSetup constraint-conflict recovery", () => { expect(openedRoots).to.deep.equal(["/proj"]); }); + /** {@link conflictResult} for `/mono/bundles/sales`, a uv workspace member. */ + function memberConflictResult(): PythonSetupResult { + return { + ...conflictResult(), + backupPath: "/mono/bundles/sales/pyproject.toml.bak", + uvWorkspaceRoot: "/mono", + uvWorkspaceBackupPath: "/mono/pyproject.toml.bak", + }; + } + + async function retryMemberConflict(result: PythonSetupResult) { + const restored: Array<{root: string; backupPath: string}> = []; + const cli = makeScriptedCli([result, SUCCESS_REAL_RUN]); + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + projectRoot: () => "/mono/bundles/sales", + pickSetupPreset: async () => "full", + restoreProjectFile: async (root, backupPath) => { + restored.push({root, backupPath}); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + return {restored, calls: cli.calls}; + } + + it("restores the uv workspace root pyproject.toml too before a member retry", async () => { + // The pins live in the root file, and --no-constraints does not touch + // it, so a retry without this restore conflicts again. + const {restored, calls} = await retryMemberConflict( + memberConflictResult() + ); + + expect(restored).to.deep.equal([ + { + root: "/mono/bundles/sales", + backupPath: "/mono/bundles/sales/pyproject.toml.bak", + }, + {root: "/mono", backupPath: "/mono/pyproject.toml.bak"}, + ]); + expect(calls).to.have.length(2); + expect(calls[1].skipConstraints).to.equal(true); + }); + + it("checks every backup before it restores the first one", async () => { + // A failed root check must not leave the member file restored alone. + const restored: string[] = []; + const cli = makeScriptedCli([memberConflictResult(), SUCCESS_REAL_RUN]); + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + projectRoot: () => "/mono/bundles/sales", + pickSetupPreset: async () => "full", + checkProjectFileBackup: async (folder) => { + if (folder === "/mono") { + throw new Error("backup is gone"); + } + }, + restoreProjectFile: async (_root, backupPath) => { + restored.push(backupPath); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + let threw = false; + try { + await (retryAction as {run: () => Promise}).run(); + } catch { + threw = true; + } + + expect(threw).to.equal(true); + expect(restored).to.deep.equal([]); + expect(cli.calls).to.have.length(1); + }); + + it("ignores a stale retry of the uv workspace root project after a member run", async () => { + // The root project's result has no uvWorkspaceRoot, but its backup is the + // shared root file, which a later member run rewrote. + let activeRoot = "/mono"; + const restored: string[] = []; + const cli = makeScriptedCli([ + conflictResult(), + {...SUCCESS_REAL_RUN_UV_WORKSPACE_MEMBER, uvWorkspaceRoot: "/mono"}, + ]); + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + projectRoot: () => activeRoot, + pickSetupPreset: async () => "full", + restoreProjectFile: async (_root, backupPath) => { + restored.push(backupPath); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + activeRoot = "/mono/bundles/a"; + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + + expect(restored).to.deep.equal([]); + expect(cli.calls).to.have.length(2); + }); + + it("keeps a retry live after a setup run of an unrelated project", async () => { + let activeRoot = "/projA"; + const restored: string[] = []; + const cli = makeScriptedCli([ + conflictResult(), + SUCCESS_REAL_RUN, + SUCCESS_REAL_RUN, + ]); + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + projectRoot: () => activeRoot, + pickSetupPreset: async () => "full", + restoreProjectFile: async (_root, backupPath) => { + restored.push(backupPath); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + activeRoot = "/projB"; + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + + expect(restored).to.deep.equal(["/proj/pyproject.toml.bak"]); + expect(cli.calls).to.have.length(3); + }); + + it("restores a backup that the CLI named from the real project folder", async () => { + // A standalone project open through a symlink: the CLI reports real paths. + const restored: Array<{root: string; backupPath: string}> = []; + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli: makeScriptedCli([ + { + ...conflictResult(), + backupPath: "/real/proj/pyproject.toml.bak", + }, + SUCCESS_REAL_RUN, + ]), + projectRoot: () => "/link/proj", + realPath: async () => "/real/proj", + pickSetupPreset: async () => "full", + restoreProjectFile: async (root, backupPath) => { + restored.push({root, backupPath}); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + + expect(restored).to.deep.equal([ + {root: "/real/proj", backupPath: "/real/proj/pyproject.toml.bak"}, + ]); + }); + + it("offers the retry when only the uv workspace root was backed up", async () => { + const {restored} = await retryMemberConflict({ + ...memberConflictResult(), + backupPath: undefined, + }); + + expect(restored).to.deep.equal([ + {root: "/mono", backupPath: "/mono/pyproject.toml.bak"}, + ]); + }); + + it("offers no retry for a member when the uv workspace root was not backed up", async () => { + // The root kept its pins from an earlier run, and --no-constraints does + // not remove them, so a retry would conflict again. + let actions: PythonSetupErrorAction[] | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli: makeCli({ + resolve: { + ...memberConflictResult(), + uvWorkspaceBackupPath: undefined, + }, + }), + projectRoot: () => "/mono/bundles/sales", + pickSetupPreset: async () => "full", + showError: async (_message, _detail, a) => { + actions = a; + }, + }) + ); + + await setup.setup(); + + expect(actions?.map((a) => a.label)).to.not.include( + "Retry DB Connect setup" + ); + }); + + it("offers no retry when the reported uv workspace root does not contain the project", async () => { + let actions: PythonSetupErrorAction[] | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli: makeCli({ + resolve: { + ...memberConflictResult(), + uvWorkspaceRoot: "/elsewhere", + uvWorkspaceBackupPath: "/elsewhere/pyproject.toml.bak", + }, + }), + projectRoot: () => "/mono/bundles/sales", + pickSetupPreset: async () => "full", + showError: async (_message, _detail, a) => { + actions = a; + }, + }) + ); + + await setup.setup(); + + expect(actions?.map((a) => a.label)).to.not.include( + "Retry DB Connect setup" + ); + }); + + it("ignores a stale member retry after another member synced the same uv workspace", async () => { + // Member B's later run wrote its pins into the shared root; A's old root + // backup must not replace them. + let activeRoot = "/mono/bundles/a"; + const restored: string[] = []; + const notified: string[] = []; + const cli = makeScriptedCli([ + memberConflictResult(), + {...SUCCESS_REAL_RUN_UV_WORKSPACE_MEMBER, uvWorkspaceRoot: "/mono"}, + ]); + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + projectRoot: () => activeRoot, + pickSetupPreset: async () => "full", + restoreProjectFile: async (_root, backupPath) => { + restored.push(backupPath); + }, + notify: async (m) => { + notified.push(m); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + activeRoot = "/mono/bundles/b"; + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + + expect(restored).to.deep.equal([]); + expect(cli.calls).to.have.length(2); + // The click is not silent: the user learns to run the setup again. + expect(notified).to.have.length(1); + }); + + it("ignores a stale member retry when the workspace is open through a symlink", async () => { + // The CLI reports the real root; later runs keep the opened folder form. + let activeRoot = "/link/bundles/a"; + const restored: string[] = []; + const cli = makeScriptedCli([ + memberConflictResult(), + {...SUCCESS_REAL_RUN_UV_WORKSPACE_MEMBER, uvWorkspaceRoot: "/mono"}, + ]); + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + projectRoot: () => activeRoot, + realPath: async (p) => p.replace(/^\/link/, "/mono"), + pickSetupPreset: async () => "full", + restoreProjectFile: async (_root, backupPath) => { + restored.push(backupPath); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + activeRoot = "/link/bundles/b"; + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + + expect(restored).to.deep.equal([]); + expect(cli.calls).to.have.length(2); + }); + + it("ignores a stale member retry after a setup of the uv workspace root project", async () => { + // The root project is not a member, so its result has no + // uvWorkspaceRoot, but its run still rewrites the shared root file. + let activeRoot = "/mono/bundles/a"; + const restored: string[] = []; + const cli = makeScriptedCli([memberConflictResult(), SUCCESS_REAL_RUN]); + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli, + projectRoot: () => activeRoot, + pickSetupPreset: async () => "full", + restoreProjectFile: async (_root, backupPath) => { + restored.push(backupPath); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + activeRoot = "/mono"; + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + + expect(restored).to.deep.equal([]); + expect(cli.calls).to.have.length(2); + }); + + it("trusts a uv workspace root that contains the real project folder", async () => { + // The project is open through a symlink; the CLI reports real paths. + const restored: Array<{root: string; backupPath: string}> = []; + let retryAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli: makeScriptedCli([ + memberConflictResult(), + SUCCESS_REAL_RUN, + ]), + projectRoot: () => "/link/sales", + realPath: async () => "/mono/bundles/sales", + pickSetupPreset: async () => "full", + restoreProjectFile: async (root, backupPath) => { + restored.push({root, backupPath}); + }, + showError: async (_message, _detail, actions) => { + retryAction = actions?.find( + (a) => a.label === "Retry DB Connect setup" + ); + }, + }) + ); + + await setup.setup(); + await (retryAction as {run: () => Promise}).run(); + + expect(restored).to.deep.equal([ + { + root: "/mono/bundles/sales", + backupPath: "/mono/bundles/sales/pyproject.toml.bak", + }, + {root: "/mono", backupPath: "/mono/pyproject.toml.bak"}, + ]); + }); + + it("opens the uv workspace root pyproject.toml, which holds the pins, for a member", async () => { + const openedRoots: string[] = []; + let openAction: PythonSetupErrorAction | undefined; + const setup = new PythonSetupEnvironmentSetup( + makeDeps({ + cli: makeCli({resolve: memberConflictResult()}), + projectRoot: () => "/mono/bundles/sales", + pickSetupPreset: async () => "full", + openProjectFile: async (root) => { + openedRoots.push(root); + }, + showError: async (_message, _detail, actions) => { + openAction = actions?.find( + (a) => a.label === "Open pyproject.toml" + ); + }, + }) + ); + + await setup.setup(); + await (openAction as {run: () => Promise}).run(); + + expect(openedRoots).to.deep.equal(["/mono"]); + }); + it("records the conflict failure without offering a report", async () => { const telemetry = makeTelemetryRecorder(); const setup = new PythonSetupEnvironmentSetup( diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts index bb7d73412..149e036f6 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupEnvironmentSetup.ts @@ -29,6 +29,7 @@ import { ReportEnvironment, } from "../utils/reportSetupIssue"; import {isReauthRequiredError} from "../utils/authErrors"; +import {containsPath} from "../utils/containsPath"; import {SetupLocalInvocation} from "../utils/setupLocalArgs"; import { presetToFlags, @@ -74,6 +75,16 @@ export type ResolvedCompute = export interface PythonSetupPersistedState { envKey: string; pythonVersion: string; + /** The absolute venv folder the run provisioned. */ + venvPath: string; + /** The project folder of the run. */ + projectRoot: string; +} + +/** A pyproject.toml to restore: `/pyproject.toml` from `backupPath`. */ +interface ProjectFileBackup { + folder: string; + backupPath: string; } type ReadyLocalEnvironmentResult = PythonSetupResult & @@ -160,19 +171,27 @@ export interface PythonSetupSetupDeps { * Open the project's pyproject.toml in an editor — the "Open pyproject.toml" * button on a constraint-conflict failure, so the user can inspect and * adjust the dependencies that clashed. Takes the run's captured root so it - * targets the project the failing run mutated, even after a mid-run switch. + * targets the project the failing run mutated, even after a mid-run switch; + * for a uv workspace member it is the workspace root, which holds the pins. */ - openProjectFile: (projectRoot: string) => Promise; + openProjectFile: (folder: string) => Promise; /** - * Restore `/pyproject.toml` from the CLI's pre-merge - * `backupPath` (see {@link PythonSetupResult.backupPath}) by copying it over - * the file. First step of the "Retry DB Connect setup" recovery, whose + * Restore `/pyproject.toml` from a CLI pre-merge backup inside + * `folder` (the project, or the uv workspace root for a member) by copying it + * over the file. First step of the "Retry DB Connect setup" recovery, whose * rationale lives on {@link buildConflictRecoveryActions}; a rejection there * aborts the retry before any CLI run. */ - restoreProjectFile: ( - projectRoot: string, + restoreProjectFile: (folder: string, backupPath: string) => Promise; + + /** + * Reject a backup that {@link restoreProjectFile} would refuse or cannot + * read. The retry checks every backup first, so a bad second backup does not + * leave the first file restored alone. + */ + checkProjectFileBackup: ( + folder: string, backupPath: string ) => Promise; @@ -321,6 +340,11 @@ export class PythonSetupEnvironmentSetup implements Disposable { * "ready" line for it. */ private readonly readyRoots = new Set(); + /** + * The project folder of each setup run this session, in start order, as + * opened and symlink-resolved (the CLI reports paths in the second form). + */ + private readonly runFolders: string[][] = []; /** * True when the currently active project has been set up successfully this * session. Reads the live `projectRoot` so it tracks the active project the @@ -468,6 +492,7 @@ export class PythonSetupEnvironmentSetup implements Disposable { trigger?: PythonSetupRunTrigger ): Promise { const {cli, withProgress} = this.deps; + this.runFolders.push([cwd, await this.realPathOrSelf(cwd)]); const flags = presetToFlags(preset); const invocation: SetupLocalInvocation = { @@ -556,21 +581,22 @@ export class PythonSetupEnvironmentSetup implements Disposable { // back to. Without either, fall through to the ordinary doc-link // handling rather than offer a "retry as DB Connect" that can't work // (nothing to restore) or would loop (a run that already skipped pins). - const conflictBackupPath = + const isPinConflict = result.error?.code === "E_PROVISION_CONFLICT" && - !invocation.skipConstraints && - // Truthiness, not just `!== undefined`: a (contract-forbidden) - // empty backupPath has nothing to restore, so it must fall - // through rather than offer a Retry that could only throw. - result.backupPath - ? result.backupPath - : undefined; - const recoverableConflict = conflictBackupPath !== undefined; + !invocation.skipConstraints; + const cwdForms = isPinConflict + ? [cwd, await this.realPathOrSelf(cwd)] + : [cwd]; + const conflictBackups = isPinConflict + ? conflictBackupsOf(result, cwdForms) + : []; + const recoverableConflict = conflictBackups.length > 0; const actions = recoverableConflict ? this.buildConflictRecoveryActions( compute, cwd, - conflictBackupPath + trustedUvWorkspaceRoot(result, cwdForms), + conflictBackups ) : remediationActions.some( (candidate) => @@ -683,6 +709,8 @@ export class PythonSetupEnvironmentSetup implements Disposable { this.deps.saveState({ envKey: result.compute.envKey, pythonVersion: result.resolved.pythonVersion, + venvPath: venvDir, + projectRoot: cwd, }); // Record readiness for the project this run provisioned (the // captured cwd), not the live active project — a mid-run switch must @@ -756,27 +784,47 @@ export class PythonSetupEnvironmentSetup implements Disposable { ); } + /** + * {@link PythonSetupSetupDeps.realPath}, or the folder as given if that + * fails: these callers only compare paths, and must not break the run or + * hide its CLI error. + */ + private async realPathOrSelf(cwd: string): Promise { + try { + return await this.deps.realPath(cwd); + } catch { + return cwd; + } + } + /** * The two recovery buttons for a Full-preset constraint conflict, both * run-actions (their behavior needs the run's live compute/cwd): * * - "Retry DB Connect setup" restores pyproject.toml from the CLI's pre-merge - * `backupPath`, then re-runs as DB Connect (`--no-constraints`). The restore - * is load-bearing: the conflict fails *after* the pins were merged to disk, - * so `--no-constraints` alone would leave the conflicting pins in place and - * only skip re-adding them. It runs through {@link runGuarded} — so a click - * can't race an in-flight run, and the restore only fires when the retry - * actually runs — and no-ops once the project is ready, so a stale toast's - * Retry can't downgrade an environment a later run provisioned. A failed - * restore throws before the CLI spawns, leaving showError to log it. + * `backupPath` (for a uv workspace member, also the root pyproject.toml + * from `uvWorkspaceBackupPath`), then re-runs as DB Connect + * (`--no-constraints`). The restore is load-bearing: the conflict fails + * *after* the pins were merged to disk, so `--no-constraints` alone would + * leave the conflicting pins in place and only skip re-adding them. It runs + * through {@link runGuarded}, so a click can't race an in-flight run. It + * does nothing once the project is ready, and it asks for a new setup once + * a later run touched the same files, so a stale toast can't undo that run. + * A failed restore throws before the CLI spawns, leaving showError to log it. * - "Open pyproject.toml" opens the merged file so the user can adjust the * dependencies that clashed. */ private buildConflictRecoveryActions( compute: SetupCompute, cwd: string, - backupPath: string + uvWorkspaceRoot: string | undefined, + backups: ProjectFileBackup[] ): PythonSetupErrorAction[] { + const runsAtConflict = this.runFolders.length; + // The files the retry restores: the uv workspace root for a member, else + // the project. A later run inside this folder can have rewritten them; + // a run of the root project is inside it too. + const scope = uvWorkspaceRoot ?? cwd; return [ { label: "Retry DB Connect setup", @@ -786,10 +834,35 @@ export class PythonSetupEnvironmentSetup implements Disposable { if (this.readyRoots.has(cwd)) { return; } + if ( + this.runFolders + .slice(runsAtConflict) + .some((forms) => + forms.some((folder) => + containsPath(scope, folder) + ) + ) + ) { + return this.deps.notify( + "A later Python environment setup changed these " + + "files after this conflict. Run the setup again." + ); + } return this.runGuarded(async () => { // Restore before re-running; a failure throws here, before // any attempt is recorded or the CLI spawns. - await this.deps.restoreProjectFile(cwd, backupPath); + for (const {folder, backupPath} of backups) { + await this.deps.checkProjectFileBackup( + folder, + backupPath + ); + } + for (const {folder, backupPath} of backups) { + await this.deps.restoreProjectFile( + folder, + backupPath + ); + } await this.runResolved( compute, cwd, @@ -801,7 +874,8 @@ export class PythonSetupEnvironmentSetup implements Disposable { }, { label: "Open pyproject.toml", - run: () => this.deps.openProjectFile(cwd), + // A uv workspace member keeps its pins in the root file. + run: () => this.deps.openProjectFile(uvWorkspaceRoot ?? cwd), }, ]; } @@ -902,3 +976,62 @@ export class PythonSetupEnvironmentSetup implements Disposable { this.stateEmitter.dispose(); } } + +/** + * The result's uv workspace root, if it contains the project folder in one of + * `cwdForms` (as given, symlink-resolved). The restore writes there, so the + * root must at least be a parent of the project. + */ +function trustedUvWorkspaceRoot( + result: PythonSetupResult, + cwdForms: string[] +): string | undefined { + const root = result.uvWorkspaceRoot; + return root && cwdForms.some((cwd) => containsPath(root, cwd)) + ? root + : undefined; +} + +/** + * The pre-merge backups a conflict retry restores. Truthiness, not just + * `!== undefined`: a (contract-forbidden) empty path has nothing to restore. + */ +function conflictBackupsOf( + result: PythonSetupResult, + cwdForms: string[] +): ProjectFileBackup[] { + if (result.uvWorkspaceRoot) { + // A member's pins live in the root file, and --no-constraints leaves + // them. Without a root backup, a retry would conflict again. + const root = trustedUvWorkspaceRoot(result, cwdForms); + if (root === undefined || !result.uvWorkspaceBackupPath) { + return []; + } + const backups: ProjectFileBackup[] = [ + {folder: root, backupPath: result.uvWorkspaceBackupPath}, + ]; + const {backupPath} = result; + if (backupPath) { + backups.unshift({ + folder: backupFolder(backupPath, cwdForms), + backupPath, + }); + } + return backups; + } + return result.backupPath + ? [ + { + folder: backupFolder(result.backupPath, cwdForms), + backupPath: result.backupPath, + }, + ] + : []; +} + +/** The form of the project folder that the CLI named `backupPath` from. */ +function backupFolder(backupPath: string, cwdForms: string[]): string { + return ( + cwdForms.find((form) => containsPath(form, backupPath)) ?? cwdForms[0] + ); +} diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts index a53f4469f..1f481039b 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.test.ts @@ -1,5 +1,6 @@ import {expect} from "chai"; import { + chmod, mkdir, mkdtemp, readdir, @@ -492,11 +493,14 @@ describe("makePythonSetupDeps saveState", () => { deps.saveState({ envKey: "serverless/serverless-v5", pythonVersion: "3.12", + venvPath: "/proj/.venv", + projectRoot: "/proj", }); expect(persisted).to.have.length(1); expect(persisted[0].envKey).to.equal("serverless/serverless-v5"); expect(persisted[0].pythonVersion).to.equal("3.12"); + expect(persisted[0].venvPath).to.equal("/proj/.venv"); // A wiring-supplied ISO-8601 timestamp is added. expect(persisted[0].timestamp).to.be.a("string"); expect(Number.isNaN(Date.parse(persisted[0].timestamp))).to.equal( @@ -1169,6 +1173,77 @@ describe("makePythonSetupDeps restoreProjectFile", () => { }); }); +describe("makePythonSetupDeps checkProjectFileBackup", () => { + it("accepts a backup that exists inside the folder", async () => { + const dir = await mkdtemp(path.join(tmpdir(), "vpex-check-")); + try { + const backup = path.join(dir, "pyproject.toml.bak"); + await writeFile(backup, "original\n"); + const deps = makePythonSetupDeps(makeWiring()); + await deps.checkProjectFileBackup(dir, backup); + } finally { + await rm(dir, {recursive: true, force: true}); + } + }); + + it("accepts a backup whose name starts with two dots", async () => { + const dir = await mkdtemp(path.join(tmpdir(), "vpex-check-")); + try { + const backup = path.join(dir, "..cache.bak"); + await writeFile(backup, "original\n"); + const deps = makePythonSetupDeps(makeWiring()); + await deps.checkProjectFileBackup(dir, backup); + } finally { + await rm(dir, {recursive: true, force: true}); + } + }); + + it("rejects a backup that cannot be read", async function () { + if (process.platform === "win32" || process.getuid?.() === 0) { + // chmod does not remove read access there. + this.skip(); + } + const dir = await mkdtemp(path.join(tmpdir(), "vpex-check-")); + try { + const backup = path.join(dir, "pyproject.toml.bak"); + await writeFile(backup, "original\n"); + await chmod(backup, 0o200); + const deps = makePythonSetupDeps(makeWiring()); + let threw = false; + try { + await deps.checkProjectFileBackup(dir, backup); + } catch { + threw = true; + } + expect(threw).to.equal(true); + } finally { + await rm(dir, {recursive: true, force: true}); + } + }); + + it("rejects a missing backup or one outside the folder", async () => { + const dir = await mkdtemp(path.join(tmpdir(), "vpex-check-")); + const outside = await mkdtemp(path.join(tmpdir(), "vpex-out-")); + try { + const foreign = path.join(outside, "x.bak"); + await writeFile(foreign, "x\n"); + const deps = makePythonSetupDeps(makeWiring()); + for (const backup of [path.join(dir, "missing.bak"), foreign]) { + let threw = false; + try { + await deps.checkProjectFileBackup(dir, backup); + } catch { + threw = true; + } + expect(threw, backup).to.equal(true); + } + } finally { + await rm(dir, {recursive: true, force: true}); + await rm(outside, {recursive: true, force: true}); + } + }); +}); + describe("makePythonSetupDeps showSuccess", () => { let originalInfo: typeof window.showInformationMessage; let originalWarn: typeof window.showWarningMessage; diff --git a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts index 4fac7e515..691c39cfb 100644 --- a/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts +++ b/packages/databricks-vscode/src/python-setup/controllers/pythonSetupDeps.ts @@ -1,7 +1,8 @@ import {randomBytes} from "crypto"; import {existsSync} from "fs"; -import {copyFile, realpath, rename, rm} from "fs/promises"; +import {access, constants, copyFile, realpath, rename, rm} from "fs/promises"; import path from "path"; +import {containsPath} from "../utils/containsPath"; import {commands, ProgressLocation, Uri, window} from "vscode"; import {PackageManagerDetection} from "../../language/packageManagerDetection"; import {Telemetry} from "../../telemetry"; @@ -281,30 +282,23 @@ export function makePythonSetupDeps( Uri.file(projectRoot) ); }, - openProjectFile: async (projectRoot: string) => { + openProjectFile: async (folder: string) => { // Open the pyproject.toml the conflicting run merged into, so the // user can inspect/adjust the dependencies that clashed. A // constraint conflict always mutates disk (constraints are merged // before provisioning fails), so the file exists. await window.showTextDocument( - Uri.file(path.join(projectRoot, "pyproject.toml")) + Uri.file(path.join(folder, "pyproject.toml")) ); }, - restoreProjectFile: async (projectRoot: string, backupPath: string) => { + checkProjectFileBackup: async (folder: string, backupPath: string) => { + assertBackupInside(folder, backupPath); + await access(backupPath, constants.R_OK); + }, + restoreProjectFile: async (folder: string, backupPath: string) => { // Restore the CLI's pre-merge backup over pyproject.toml (the seam's // doc covers why the DB Connect retry needs this). - const root = path.resolve(projectRoot); - // backupPath comes from the CLI result: refuse anything resolving - // outside the project so a malformed/unexpected path can't copy an - // arbitrary file over pyproject.toml. Lexical containment (path is - // from the trusted local CLI, so symlink canonicalization is not - // warranted). - const rel = path.relative(root, path.resolve(backupPath)); - if (rel === "" || rel.startsWith("..") || path.isAbsolute(rel)) { - throw new Error( - `Refusing to restore pyproject.toml from a backup outside the project: ${backupPath}` - ); - } + const root = assertBackupInside(folder, backupPath); // Write atomically: copy to a temp sibling on the same filesystem, // then rename over pyproject.toml, so an interrupted or failed copy // can never leave the project file truncated. Clean the temp up if @@ -507,3 +501,21 @@ export function makePythonSetupDeps( ), }; } + +/** + * Throw unless `backupPath` is inside `folder`; return the resolved folder. + * backupPath comes from the CLI result, so a malformed/unexpected path must + * not copy an arbitrary file over pyproject.toml. Lexical containment (the + * path is from the trusted local CLI, so symlink canonicalization is not + * warranted). + */ +function assertBackupInside(folder: string, backupPath: string): string { + const root = path.resolve(folder); + const backup = path.resolve(backupPath); + if (backup === root || !containsPath(root, backup)) { + throw new Error( + `Refusing to restore pyproject.toml from a backup outside ${root}: ${backupPath}` + ); + } + return root; +} diff --git a/packages/databricks-vscode/src/python-setup/utils/containsPath.test.ts b/packages/databricks-vscode/src/python-setup/utils/containsPath.test.ts new file mode 100644 index 000000000..00e36bb2c --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/utils/containsPath.test.ts @@ -0,0 +1,24 @@ +import {expect} from "chai"; +import path from "node:path"; +import {containsPath} from "./containsPath"; + +describe("containsPath", () => { + const root = path.resolve("/mono"); + + it("contains the folder itself and paths inside it", () => { + expect(containsPath(root, root)).to.equal(true); + expect(containsPath(root, path.join(root, "a", "b"))).to.equal(true); + }); + + it("contains a name that starts with two dots", () => { + expect(containsPath(root, path.join(root, "..cache.bak"))).to.equal( + true + ); + }); + + it("does not contain the parent, a sibling, or another root", () => { + expect(containsPath(root, path.dirname(root))).to.equal(false); + expect(containsPath(root, path.resolve("/mono-2/x"))).to.equal(false); + expect(containsPath(root, path.join(root, "..", "x"))).to.equal(false); + }); +}); diff --git a/packages/databricks-vscode/src/python-setup/utils/containsPath.ts b/packages/databricks-vscode/src/python-setup/utils/containsPath.ts new file mode 100644 index 000000000..4e4255ec4 --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/utils/containsPath.ts @@ -0,0 +1,11 @@ +import path from "node:path"; + +/** Whether `child` is `parent` or a path inside it, compared by text. */ +export function containsPath(parent: string, child: string): boolean { + const rel = path.relative(parent, child); + return ( + rel !== ".." && + !rel.startsWith(".." + path.sep) && + !path.isAbsolute(rel) + ); +} diff --git a/packages/databricks-vscode/src/telemetry/README.md b/packages/databricks-vscode/src/telemetry/README.md index c8bec1d5a..a8a4e06e4 100644 --- a/packages/databricks-vscode/src/telemetry/README.md +++ b/packages/databricks-vscode/src/telemetry/README.md @@ -131,8 +131,13 @@ root would emit a spurious `venvPresent: false` and inflate the denominator. So gauge is skipped when the workspace has more than one root, rather than record an untrustworthy reading. -The one-root guard is a heuristic, not proof of provenance — the key records no root, -so a rare edge (a multi-root workspace reduced to one root mid-session, leaving the +Within one root, the state records the project folder of its setup (in state saved +by newer versions), and the gauge skips a reading for another project. It checks the +venv folder that the state records: for a uv workspace member, the shared venv at the +workspace root. + +The one-root guard is a heuristic, not proof of provenance — older state records no +root, so a rare edge (a multi-root workspace reduced to one root mid-session, leaving the prior root's key) can still mis-attribute. Eliminating that needs the per-project storage schema, deferred (the drift detector shares the single-key limitation and does not even guard multi-root). Multi-root Databricks workspaces are uncommon, so diff --git a/packages/databricks-vscode/src/telemetry/constants.ts b/packages/databricks-vscode/src/telemetry/constants.ts index e46d94e77..073f7b0de 100644 --- a/packages/databricks-vscode/src/telemetry/constants.ts +++ b/packages/databricks-vscode/src/telemetry/constants.ts @@ -731,7 +731,8 @@ export class EventTypes { "still physically exists, which drift never checks. Categorical/boolean data only.", venvPresent: { comment: - "Whether the project's managed .venv interpreter still exists on disk. False means " + + "Whether the project's managed venv interpreter (for a uv workspace member, the " + + "shared venv at the root) still exists on disk. False means " + "the environment was provisioned once but is now gone (deleted or never restored)", }, currentTargetType: { diff --git a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts index f62995ced..3599e1b15 100644 --- a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts +++ b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts @@ -65,7 +65,10 @@ export interface PythonSetupDrift { * kind attached when the reading was taken. Both categorical/boolean. */ export interface PythonSetupAdoption { - /** Whether the project's managed `.venv` interpreter still exists on disk. */ + /** + * Whether the project's managed venv interpreter still exists on disk (for a + * uv workspace member, the shared venv at the root). + */ venvPresent: boolean; /** The compute kind attached at the time of the reading. */ currentTargetType: TargetCompute; diff --git a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts index bb0f8809b..50f6c1dad 100644 --- a/packages/databricks-vscode/src/vscode-objs/StateStorage.ts +++ b/packages/databricks-vscode/src/vscode-objs/StateStorage.ts @@ -15,11 +15,18 @@ import {StoredFavoriteRef} from "../ui/unity-catalog/types"; * - `pythonVersion`: the interpreter minor version that was provisioned (e.g. * `3.12`). * - `timestamp`: ISO-8601 time of the successful setup. + * - `venvPath`: the absolute venv folder. A uv workspace member shares the + * venv at the workspace root. + * - `projectRoot`: the project folder of the setup. + * + * State saved by older versions has no `venvPath` or `projectRoot`. */ export interface PythonSetupState { envKey: string; pythonVersion: string; timestamp: string; + venvPath?: string; + projectRoot?: string; } /* eslint-disable @typescript-eslint/naming-convention */