diff --git a/packages/databricks-vscode/package.json b/packages/databricks-vscode/package.json index 95b475619..f6ac1d841 100644 --- a/packages/databricks-vscode/package.json +++ b/packages/databricks-vscode/package.json @@ -1583,7 +1583,7 @@ }, "databricks.experiments.optInto": { "type": "array", - "default": [], + "default": ["environment.pythonSetup"], "items": { "enum": [ "views.cluster", diff --git a/packages/databricks-vscode/src/extension.ts b/packages/databricks-vscode/src/extension.ts index 9ebe2c05d..6cf2046f9 100644 --- a/packages/databricks-vscode/src/extension.ts +++ b/packages/databricks-vscode/src/extension.ts @@ -16,6 +16,7 @@ import {ClusterListDataProvider} from "./cluster/ClusterListDataProvider"; 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 {AiToolsManager} from "./aitools/AiToolsManager"; import {AiToolsCommands} from "./aitools/AiToolsCommands"; import {RunCommands} from "./run/RunCommands"; @@ -51,7 +52,11 @@ import { import {PythonSetupManagerDetector} from "./python-setup/utils/PythonSetupManagerDetector"; import {PythonSetupCliClient} from "./python-setup/gateways/PythonSetupCliClient"; import {PythonSetupEnvironmentSetup} from "./python-setup/controllers/PythonSetupEnvironmentSetup"; -import {makePythonSetupDeps} from "./python-setup/controllers/pythonSetupDeps"; +import { + makePythonSetupDeps, + resolveComputeFrom, +} from "./python-setup/controllers/pythonSetupDeps"; +import {PythonSetupDriftManager} from "./python-setup/controllers/PythonSetupDriftManager"; import {resolveCliPath} from "./python-setup/utils/setupLocalArgs"; import { isPythonSetupEnabled, @@ -1014,15 +1019,112 @@ export async function activate( pythonSetupEnvironment.setup, pythonSetupEnvironment ), - // Re-run affordance on the "Python environment ready" row. Delegates to - // the same setup handler (re-entrancy-guarded); a distinct id lets the - // menu show a "Re-run Python setup" title instead of the initial one. + // Re-run affordance shared by the "Python environment ready" row and the + // drifted row. Delegates to the same setup handler (re-entrancy-guarded); + // a distinct command id lets the menu show a "Re-run Python setup" title + // and gives re-runs their own COMMAND_EXECUTION telemetry. telemetry.registerCommand( "databricks.environment.rerunPythonEnv", pythonSetupEnvironment.setup, pythonSetupEnvironment ) ); + // Drives the config-view row's "out of date" state: on compute/open/setup + // triggers it silently resolves the selected compute's env key via a CLI + // dry-run and compares it against the last successful setup's key. Every + // resolution path is fail-safe (returns undefined => "unknown", no drift) and + // never surfaces UI. + const pythonSetupDrift = new PythonSetupDriftManager({ + // Reuse the exact feature+greenfield gate the row is shown under. + isVisible: () => pythonSetupEnvironment.isVisible(), + getPersistedEnvKey: () => + stateStorage.get("databricks.pythonSetup.setupState")?.envKey, + // Cheap, synchronous compute identity (no CLI). A cluster's env key is + // derived from its Spark version, so include it: a runtime-state change + // (RUNNING -> TERMINATED) keeps the descriptor stable and is skipped, + // while a DBR edit changes it and re-checks. undefined means nothing + // comparable is attached (drift is then meaningless). + getComputeDescriptor: () => { + const cluster = connectionManager.cluster; + if (cluster) { + return `cluster:${cluster.id}:${cluster.sparkVersion}`; + } + if (connectionManager.serverless) { + const version = connectionManager.serverlessVersion; + return version === undefined + ? undefined + : `serverless:${version}`; + } + return undefined; + }, + resolveCurrentEnvKey: async (token) => { + // activeProjectUri throws when no project is active; degrade to + // "unknown" rather than letting it reject into the drift check. + let root: string | undefined; + try { + root = workspaceFolderManager.activeProjectUri.fsPath; + } catch { + return undefined; + } + const resolution = resolveComputeFrom({ + serverless: connectionManager.serverless, + cluster: connectionManager.cluster + ? {id: connectionManager.cluster.id} + : undefined, + serverlessVersion: connectionManager.serverlessVersion, + }); + if (resolution.status !== "ok") { + return undefined; + } + try { + const result = await pythonSetupClient.run( + { + mode: "default", + dryRun: true, + compute: resolution.compute, + }, + {cwd: root, token} + ); + return result.compute?.envKey; + } catch { + return undefined; + } + }, + recordDrift: (report) => telemetry.recordPythonSetupDrift(report), + }); + context.subscriptions.push( + pythonSetupDrift, + // Compute target changed (cluster attach/detach/switch). + connectionManager.onDidChangeCluster(() => + pythonSetupDrift.check("computeChange") + ), + // Serverless enable/disable and connection churn flow through state. + connectionManager.onDidChangeState(() => + pythonSetupDrift.check("computeChange") + ), + // Re-picking the serverless version while serverless is already selected + // fires neither onDidChangeCluster nor onDidChangeState -- it only writes + // the `serverlessVersion` config key -- so watch that key directly, or a + // v4 -> v2 switch would silently miss drift. + configModel.onDidChangeKey("serverlessVersion")(async () => + pythonSetupDrift.check("computeChange") + ), + // A completed setup updates the persisted state; re-evaluate so a + // successful re-run clears the badge promptly. + pythonSetupEnvironment.onDidChangeState(() => + pythonSetupDrift.check("setupCompleted") + ) + ); + // The "workspace open" trigger: evaluate once now that everything is wired. + pythonSetupDrift.check("workspaceOpen"); + + // The config-view entry combines the setup controller's readiness with the + // drift manager's `drifted` signal. + const pythonSetupEntry = composePythonSetupEntry( + pythonSetupEnvironment, + pythonSetupDrift + ); + context.subscriptions.push(pythonSetupEntry); const environmentCommands = new EnvironmentCommands( featureManager, @@ -1143,7 +1245,7 @@ export async function activate( featureManager, workspaceFolderManager, aiToolsManager, - pythonSetupEnvironment + pythonSetupEntry ); const configurationView = window.createTreeView("configurationView", { treeDataProvider: configurationDataProvider, diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupDriftManager.test.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupDriftManager.test.ts new file mode 100644 index 000000000..c38b89ffc --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupDriftManager.test.ts @@ -0,0 +1,184 @@ +import {expect} from "chai"; +import {CancellationLike} from "../gateways/PythonSetupCliClient"; +import { + PythonSetupDriftDeps, + PythonSetupDriftManager, +} from "./PythonSetupDriftManager"; + +function makeDeps(over: Partial = {}): { + deps: PythonSetupDriftDeps; + recorded: unknown[]; + calls: {resolve: number}; +} { + const recorded: unknown[] = []; + const calls = {resolve: 0}; + const deps: PythonSetupDriftDeps = { + isVisible: async () => true, + getPersistedEnvKey: () => "serverless/serverless-v4", + getComputeDescriptor: () => "cluster:c1", + // eslint-disable-next-line @typescript-eslint/no-unused-vars + resolveCurrentEnvKey: async (_token: CancellationLike) => { + calls.resolve++; + return "dbr/15.4.x-scala2.12"; + }, + recordDrift: (r) => recorded.push(r), + ...over, + }; + return {deps, recorded, calls}; +} + +describe("PythonSetupDriftManager", () => { + it("flags drift and reports telemetry when keys differ", async () => { + const {deps, recorded} = makeDeps(); + const m = new PythonSetupDriftManager(deps); + let fired = 0; + m.onDidChangeState(() => fired++); + + await m.evaluate("computeChange"); + + expect(m.drifted).to.be.true; + expect(fired).to.equal(1); + expect(recorded).to.deep.equal([ + { + trigger: "computeChange", + fromEnvKey: "serverless/serverless-v4", + toEnvKey: "dbr/15.4.x-scala2.12", + }, + ]); + m.dispose(); + }); + + it("does not flag drift when the keys match", async () => { + const {deps, recorded} = makeDeps({ + resolveCurrentEnvKey: async () => "serverless/serverless-v4", + }); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("workspaceOpen"); + expect(m.drifted).to.be.false; + expect(recorded).to.have.length(0); + m.dispose(); + }); + + it("is a no-op when not visible", async () => { + const {deps, recorded} = makeDeps({isVisible: async () => false}); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("workspaceOpen"); + expect(m.drifted).to.be.false; + expect(recorded).to.have.length(0); + m.dispose(); + }); + + it("does not flag drift when there is no persisted state", async () => { + const {deps} = makeDeps({getPersistedEnvKey: () => undefined}); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("workspaceOpen"); + expect(m.drifted).to.be.false; + m.dispose(); + }); + + it("leaves the flag unchanged when the current key is unknown", async () => { + // Start drifted, then a later check can't resolve the key: stay drifted. + const {deps} = makeDeps(); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("computeChange"); + expect(m.drifted).to.be.true; + + (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey = + async () => undefined; + await m.evaluate("workspaceOpen"); + expect(m.drifted).to.be.true; + m.dispose(); + }); + + it("clears drift once the keys match again", async () => { + const {deps} = makeDeps(); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("computeChange"); + expect(m.drifted).to.be.true; + + (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey = + async () => "serverless/serverless-v4"; + await m.evaluate("setupCompleted"); + expect(m.drifted).to.be.false; + m.dispose(); + }); + + it("reports the same mismatch only once until it clears", async () => { + const {deps, recorded} = makeDeps(); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("computeChange"); + await m.evaluate("workspaceOpen"); // same mismatch, no new telemetry + expect(recorded).to.have.length(1); + + // Clears, then a mismatch recurs on a DIFFERENT compute -> reported + // again (a different descriptor is required, since a compute-change with + // the same identity is skipped as a no-op runtime-state transition). + (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey = + async () => "serverless/serverless-v4"; + await m.evaluate("setupCompleted"); + (deps as {getComputeDescriptor: unknown}).getComputeDescriptor = () => + "cluster:c2"; + (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey = + async () => "dbr/15.4.x-scala2.12"; + await m.evaluate("computeChange"); + expect(recorded).to.have.length(2); + m.dispose(); + }); + + it("skips the dry-run when a compute-change leaves the identity unchanged", async () => { + // Same descriptor across two compute-change triggers models a cluster + // runtime-state transition (RUNNING -> TERMINATED): the env key cannot + // have changed, so the CLI dry-run must not run a second time. + const {deps, calls} = makeDeps(); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("computeChange"); + expect(calls.resolve).to.equal(1); + await m.evaluate("computeChange"); + expect(calls.resolve).to.equal(1); + m.dispose(); + }); + + it("clears drift when no comparable compute is attached", async () => { + // Start drifted, then compute is detached: drift is meaningless, so the + // stale flag must clear rather than linger. + const {deps} = makeDeps(); + const m = new PythonSetupDriftManager(deps); + await m.evaluate("computeChange"); + expect(m.drifted).to.be.true; + + (deps as {getComputeDescriptor: unknown}).getComputeDescriptor = () => + undefined; + await m.evaluate("computeChange"); + expect(m.drifted).to.be.false; + m.dispose(); + }); + + it("stays silent and leaves the flag unchanged when a dep rejects", async () => { + // A rejecting dep must resolve quietly to "unknown" -- no throw, no + // unhandled rejection, and the drift flag is left as-is. + const {deps, recorded} = makeDeps({ + isVisible: async () => { + throw new Error("network down"); + }, + }); + const m = new PythonSetupDriftManager(deps); + + // Starts not drifted: a rejection leaves it false. + await m.evaluate("workspaceOpen"); + expect(m.drifted).to.be.false; + expect(recorded).to.have.length(0); + + // Now start drifted, then a rejecting dep must not retract the flag. + (deps as {isVisible: unknown}).isVisible = async () => true; + await m.evaluate("computeChange"); + expect(m.drifted).to.be.true; + + (deps as {resolveCurrentEnvKey: unknown}).resolveCurrentEnvKey = + async () => { + throw new Error("dry-run failed"); + }; + await m.evaluate("workspaceOpen"); + expect(m.drifted).to.be.true; + m.dispose(); + }); +}); diff --git a/packages/databricks-vscode/src/python-setup/controllers/PythonSetupDriftManager.ts b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupDriftManager.ts new file mode 100644 index 000000000..a3fe7df9e --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/controllers/PythonSetupDriftManager.ts @@ -0,0 +1,187 @@ +import {CancellationTokenSource, Disposable, Event, EventEmitter} from "vscode"; +import {CancellationLike} from "../gateways/PythonSetupCliClient"; +import {PythonSetupDrift} from "../../telemetry/pythonSetupExtensions"; +import {PythonSetupDriftTrigger} from "../../telemetry/constants"; +import {isDrifted} from "../utils/driftDetection"; + +export interface PythonSetupDriftDeps { + isVisible: () => Promise; + getPersistedEnvKey: () => string | undefined; + /** + * A cheap, synchronous descriptor of the currently selected compute's + * IDENTITY (e.g. `"cluster::"`, `"serverless:v5"`), or + * `undefined` when no comparable compute is attached (nothing selected, or + * serverless with no chosen version). Unlike {@link resolveCurrentEnvKey} + * this never spawns the CLI. It lets the manager (a) skip the dry-run when a + * compute-change trigger fires but the identity is unchanged -- a cluster + * runtime-state transition rather than a switch -- and (b) clear a stale + * drift flag when nothing comparable is attached. + */ + getComputeDescriptor: () => string | undefined; + resolveCurrentEnvKey: ( + token: CancellationLike + ) => Promise; + recordDrift: (report: PythonSetupDrift) => void; +} + +/** + * Watches for compute drift: when the selected compute's environment key no + * longer matches the one recorded by the last successful setup, exposes a + * `drifted` flag (and fires `onDidChangeState`) that the config-view row renders + * as an "out of date -- re-run setup" affordance. + * + * The check is deliberately passive: it runs a silent CLI `--dry-run` (no + * progress UI, no prompt, no error surface), is gated by `isVisible` and the + * presence of a persisted state, is debounced against rapid compute switches, + * and treats any inability to resolve the current key as "unknown" -- never a + * false alarm. To avoid needless dry-runs it skips a compute-change check whose + * compute identity is unchanged (a runtime-state transition, not a switch), and + * it clears drift outright when no comparable compute is attached. + */ +export class PythonSetupDriftManager implements Disposable { + private _drifted = false; + /** `${from}->${to}` of the last reported mismatch, to dedupe telemetry. */ + private lastReported: string | undefined; + /** Compute descriptor evaluated last, to skip no-op compute-change checks. */ + private lastComputeDescriptor: string | undefined; + private generation = 0; + private debounceTimer: ReturnType | undefined; + private inFlight: CancellationTokenSource | undefined; + + private readonly stateEmitter = new EventEmitter(); + readonly onDidChangeState: Event = this.stateEmitter.event; + + constructor( + private readonly deps: PythonSetupDriftDeps, + private readonly debounceMs: number = 500 + ) {} + + get drifted(): boolean { + return this._drifted; + } + + /** Debounced entry point for triggers (compute change, open, setup done). */ + check(trigger: PythonSetupDriftTrigger): void { + if (this.debounceTimer !== undefined) { + clearTimeout(this.debounceTimer); + } + this.debounceTimer = setTimeout(() => { + this.debounceTimer = undefined; + void this.evaluate(trigger); + }, this.debounceMs); + } + + /** + * The awaitable core. Public so it is unit-testable directly; production + * code reaches it through the debounced {@link check}. + */ + async evaluate(trigger: PythonSetupDriftTrigger): Promise { + const myGeneration = ++this.generation; + + // Cancel any dry-run still running for a superseded trigger. + this.inFlight?.cancel(); + this.inFlight?.dispose(); + const source = new CancellationTokenSource(); + this.inFlight = source; + + try { + const visible = await this.deps.isVisible(); + + // A newer trigger started while we awaited: drop this stale result + // so an out-of-order early return cannot retract a fresher flag. + if (myGeneration !== this.generation) { + return; + } + if (!visible) { + this.setDrifted(false); + return; + } + const persisted = this.deps.getPersistedEnvKey(); + if (persisted === undefined) { + this.setDrifted(false); + return; + } + const descriptor = this.deps.getComputeDescriptor(); + // No comparable compute attached (detached, or serverless with no + // chosen version): drift is meaningless -- you cannot be drifted from + // nothing -- so clear any stale flag instead of leaving it set. + if (descriptor === undefined) { + this.lastComputeDescriptor = undefined; + this.setDrifted(false); + return; + } + // A compute-change trigger whose resolved identity is unchanged is a + // runtime-state transition (e.g. a cluster going RUNNING -> + // TERMINATED), not a compute switch. The environment key is derived + // from the identity, so it cannot have changed: skip the dry-run. + // workspaceOpen / setupCompleted always re-evaluate -- the first + // check must run, and a completed setup moves the persisted baseline. + if ( + trigger === "computeChange" && + descriptor === this.lastComputeDescriptor + ) { + return; + } + this.lastComputeDescriptor = descriptor; + const current = await this.deps.resolveCurrentEnvKey(source.token); + + // A newer trigger started while we awaited: drop this stale result. + if (myGeneration !== this.generation) { + return; + } + // Could not resolve the current key -> unknown. Leave the flag as-is + // rather than clearing (a transient network/auth failure must not + // silently retract a real drift warning). + if (current === undefined) { + return; + } + + const drifted = isDrifted(persisted, current); + this.setDrifted(drifted); + + if (drifted) { + const mismatch = `${persisted}->${current}`; + if (this.lastReported !== mismatch) { + this.lastReported = mismatch; + this.deps.recordDrift({ + trigger, + fromEnvKey: persisted, + toEnvKey: current, + }); + } + } + } catch { + // Any failure resolving the current state (e.g. isVisible or the + // dry-run rejecting) is treated as "unknown": stay silent and leave + // the drift flag untouched -- never surface UI, never a false alarm, + // never retract a real warning. Same fail-safe direction as the + // `current === undefined` branch above. + } finally { + if (this.inFlight === source) { + source.dispose(); + this.inFlight = undefined; + } + } + } + + private setDrifted(value: boolean): void { + if (!value) { + // Reset the telemetry dedupe latch so a recurrence is reported again. + this.lastReported = undefined; + } + if (value === this._drifted) { + return; + } + this._drifted = value; + this.stateEmitter.fire(); + } + + dispose(): void { + if (this.debounceTimer !== undefined) { + clearTimeout(this.debounceTimer); + } + this.inFlight?.cancel(); + this.inFlight?.dispose(); + this.stateEmitter.dispose(); + } +} diff --git a/packages/databricks-vscode/src/python-setup/utils/driftDetection.test.ts b/packages/databricks-vscode/src/python-setup/utils/driftDetection.test.ts new file mode 100644 index 000000000..fc2b776c2 --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/utils/driftDetection.test.ts @@ -0,0 +1,27 @@ +import {expect} from "chai"; +import {isDrifted} from "./driftDetection"; + +describe("isDrifted", () => { + it("is true when both keys are known and differ", () => { + expect(isDrifted("serverless/serverless-v4", "dbr/15.4.x-scala2.12")).to + .be.true; + }); + + it("is false when the keys are equal", () => { + expect( + isDrifted("serverless/serverless-v5", "serverless/serverless-v5") + ).to.be.false; + }); + + it("is false (fail-safe) when the current key is unknown", () => { + expect(isDrifted("serverless/serverless-v5", undefined)).to.be.false; + }); + + it("is false (fail-safe) when there is no persisted key", () => { + expect(isDrifted(undefined, "dbr/15.4.x-scala2.12")).to.be.false; + }); + + it("is false when both are unknown", () => { + expect(isDrifted(undefined, undefined)).to.be.false; + }); +}); diff --git a/packages/databricks-vscode/src/python-setup/utils/driftDetection.ts b/packages/databricks-vscode/src/python-setup/utils/driftDetection.ts new file mode 100644 index 000000000..0919573bb --- /dev/null +++ b/packages/databricks-vscode/src/python-setup/utils/driftDetection.ts @@ -0,0 +1,20 @@ +/** + * Decide whether the local environment has drifted from the selected compute. + * + * Drift means we know both the environment key we last provisioned against + * (`persistedEnvKey`, from `databricks.pythonSetup.setupState`) and the key the + * currently selected compute would resolve to (`currentEnvKey`), and they + * differ. Anything unknown — no prior setup, or a compute whose key could not be + * resolved — is deliberately NOT drift: absence of a clear signal must never + * raise a false alarm (see the design's fail-safe rule). + */ +export function isDrifted( + persistedEnvKey: string | undefined, + currentEnvKey: string | undefined +): boolean { + return ( + persistedEnvKey !== undefined && + currentEnvKey !== undefined && + persistedEnvKey !== currentEnvKey + ); +} diff --git a/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.test.ts b/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.test.ts index 8d02e5a04..c4695906d 100644 --- a/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.test.ts +++ b/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.test.ts @@ -75,6 +75,25 @@ describe("buildSetupLocalArgs", () => { }); expect(args.slice(-2)).to.deep.equal(["--output", "json"]); }); + + it("adds --dry-run when the invocation is a dry run", () => { + const args = buildSetupLocalArgs({ + mode: "default", + compute: {kind: "serverless", version: "5"}, + dryRun: true, + }); + expect(args).to.include("--dry-run"); + // Still requests machine-readable output last. + expect(args.slice(-2)).to.deep.equal(["--output", "json"]); + }); + + it("omits --dry-run by default", () => { + const args = buildSetupLocalArgs({ + mode: "default", + compute: {kind: "serverless", version: "5"}, + }); + expect(args).to.not.include("--dry-run"); + }); }); describe("resolveCliPath", () => { diff --git a/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.ts b/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.ts index ca6a5bf4e..cc3f38e63 100644 --- a/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.ts +++ b/packages/databricks-vscode/src/python-setup/utils/setupLocalArgs.ts @@ -13,6 +13,13 @@ import {PythonSetupMode} from "../models/PythonSetupResult"; */ export interface SetupLocalInvocation { mode: PythonSetupMode; + /** + * When true, pass `--dry-run`: the CLI resolves compute and reports the + * environment key without provisioning or writing to disk. Used by drift + * detection to read the authoritative `compute.envKey` for the selected + * compute. + */ + dryRun?: boolean; compute: | {kind: "cluster"; clusterId: string} | {kind: "serverless"; version: string}; @@ -41,6 +48,9 @@ export function buildSetupLocalArgs(inv: SetupLocalInvocation): string[] { if (inv.mode === "constraints-only") { args.push("--constraints-only"); } + if (inv.dryRun) { + args.push("--dry-run"); + } if (inv.constraintSourceUrl) { args.push("--constraint-source-url", inv.constraintSourceUrl); } diff --git a/packages/databricks-vscode/src/telemetry/constants.ts b/packages/databricks-vscode/src/telemetry/constants.ts index 6d41f10ca..37a335c84 100644 --- a/packages/databricks-vscode/src/telemetry/constants.ts +++ b/packages/databricks-vscode/src/telemetry/constants.ts @@ -27,6 +27,7 @@ export enum Events { PYTHON_ENV_SETUP_DETECTED = "python_env.setup.detected", PYTHON_ENV_SETUP_ATTEMPT = "python_env.setup.attempt", PYTHON_ENV_SETUP_RESULT = "python_env.setup.result", + PYTHON_ENV_DRIFT = "python_env.drift", AITOOLS_INSTALL = "aitoolsInstall", AITOOLS_UPDATE = "aitoolsUpdate", AITOOLS_UNINSTALL = "aitoolsUninstall", @@ -153,6 +154,12 @@ export type PythonSetupFailurePhase = | "adopt" | "persist"; +/** How a drift check was triggered. */ +export type PythonSetupDriftTrigger = + | "computeChange" + | "workspaceOpen" + | "setupCompleted"; + /** Documentation about all of the properties and metrics of the event. */ type EventDescription = {[K in keyof T]?: {comment?: string}}; @@ -554,6 +561,31 @@ export class EventTypes { // spawn and interpreter adoption. ...getDurationProperty(), }; + [Events.PYTHON_ENV_DRIFT]: EventType<{ + trigger: PythonSetupDriftTrigger; + fromEnvKey: string; + toEnvKey: string; + }> = { + comment: + "Emitted when the selected compute's environment key no longer matches the one the " + + "local .venv was provisioned against (from databricks.pythonSetup.setupState). Reported " + + "once per newly-detected distinct mismatch, not on every trigger. Categorical data only.", + trigger: { + comment: + "What prompted the check: computeChange | workspaceOpen | setupCompleted", + }, + fromEnvKey: { + comment: + 'The recorded environment key (e.g. "serverless/serverless-v4", ' + + '"dbr/15.4.x-scala2.12"). Constrained to those shapes before emission; anything ' + + 'else becomes "other". Never a cluster id or name', + }, + toEnvKey: { + comment: + "The environment key the currently selected compute resolves to, same closed " + + 'vocabulary as fromEnvKey (else "other")', + }, + }; } /** diff --git a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts index 065f6745a..ffdb7a51c 100644 --- a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts +++ b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.test.ts @@ -402,4 +402,35 @@ describe(__filename, () => { expect(() => reportResult({outcome: "ok"})).to.not.throw(); expect(telemetry.isTelemetryEnabled).to.equal(false); }); + + describe("recordPythonSetupDrift", () => { + it("emits python_env.drift with the trigger and sanitized keys", () => { + const {telemetry, events} = makeTelemetry(); + telemetry.recordPythonSetupDrift({ + trigger: "computeChange", + fromEnvKey: "serverless/serverless-v4", + toEnvKey: "dbr/15.4.x-scala2.12", + }); + const drift = events.find((e) => e.name === "python_env.drift"); + expect(drift, "a drift event was recorded").to.not.be.undefined; + expect(drift!.props["event.trigger"]).to.equal("computeChange"); + expect(drift!.props["event.fromEnvKey"]).to.equal( + "serverless/serverless-v4" + ); + expect(drift!.props["event.toEnvKey"]).to.equal( + "dbr/15.4.x-scala2.12" + ); + }); + + it("collapses an unrecognized env key to 'other'", () => { + const {telemetry, events} = makeTelemetry(); + telemetry.recordPythonSetupDrift({ + trigger: "workspaceOpen", + fromEnvKey: "serverless/serverless-v5", + toEnvKey: "0710-secret-cluster-id", + }); + const drift = events.find((e) => e.name === "python_env.drift")!; + expect(drift.props["event.toEnvKey"]).to.equal("other"); + }); + }); }); diff --git a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts index cf8842074..cbff7be0b 100644 --- a/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts +++ b/packages/databricks-vscode/src/telemetry/pythonSetupExtensions.ts @@ -2,6 +2,7 @@ import {Events, Telemetry} from "."; import { ComputeType, PrimaryManager, + PythonSetupDriftTrigger, PythonSetupErrorCode, PythonSetupFailurePhase, PythonSetupMode, @@ -35,6 +36,15 @@ export interface PythonSetupAttempt { trigger: PythonSetupRunTrigger; } +/** A detected drift, reduced to the categorical fields we report. */ +export interface PythonSetupDrift { + trigger: PythonSetupDriftTrigger; + /** The recorded environment key the .venv was provisioned against. */ + fromEnvKey: string; + /** The environment key the currently selected compute resolves to. */ + toEnvKey: string; +} + /** How a setup run ended, reduced to the categorical fields we report. */ export interface PythonSetupOutcomeReport { outcome: PythonSetupOutcome; @@ -181,6 +191,13 @@ declare module "." { * legacy checklist and the uv-native entry mutually exclusively. */ recordPythonSetupNoCompute(): void; + + /** + * Record a detected compute drift. Emitted once per newly-detected + * distinct mismatch by {@link PythonSetupDriftManager}; both keys are + * constrained to the categorical envKey vocabulary before emission. + */ + recordPythonSetupDrift(report: PythonSetupDrift): void; } } @@ -263,3 +280,20 @@ Telemetry.prototype.recordPythonSetupNoCompute = function () { // start(), which always stamps an elapsed time. this.recordEvent(Events.PYTHON_ENV_SETUP_RESULT, {outcome: "no_compute"}); }; + +Telemetry.prototype.recordPythonSetupDrift = function ( + report: PythonSetupDrift +): void { + this.recordEvent(Events.PYTHON_ENV_DRIFT, { + trigger: report.trigger, + // Both keys are copied from CLI/persisted JSON; constrain them to the + // closed envKey vocabulary so an unexpected runtime string (or a cluster + // id that slipped in) collapses to "other" rather than leaking + // high-cardinality / identifying content. + // categoricalEnvKey only returns undefined for an undefined input; both + // fields are required non-null strings, so the results are always + // defined (asserted here to satisfy the string-typed schema). + fromEnvKey: categoricalEnvKey(report.fromEnvKey)!, + toEnvKey: categoricalEnvKey(report.toEnvKey)!, + }); +}; diff --git a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.test.ts b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.test.ts index fd2e9b4bf..7132935b4 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.test.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.test.ts @@ -17,10 +17,12 @@ const PYTHON_SETUP_ENTRY_ID = "ENVIRONMENT_PYTHON_SETUP"; function stubPythonSetup(opts: { visible: boolean; ready: boolean; + drifted?: boolean; }): PythonSetupEntry { return { isVisible: async () => opts.visible, ready: opts.ready, + drifted: opts.drifted ?? false, // Minimal Event: registering a listener returns a no-op Disposable. onDidChangeState: () => ({dispose() {}}), }; diff --git a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts index aa35be279..deb297f36 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/EnvironmentComponent.ts @@ -8,6 +8,7 @@ import {buildPythonSetupEntry, PythonSetupEntry} from "./pythonSetupEntry"; const ENVIRONMENT_COMPONENT_ID = "ENVIRONMENT"; const PYTHON_SETUP_COMMAND = "databricks.environment.setupPythonEnv"; +const PYTHON_SETUP_RERUN_COMMAND = "databricks.environment.rerunPythonEnv"; const getItemContext = (key: string, available: boolean) => `databricks.environment.${key}.${available ? "success" : "error"}`; @@ -43,8 +44,9 @@ export class EnvironmentComponent extends BaseComponent { const pythonSetup = this.pythonSetup; if (pythonSetup && (await pythonSetup.isVisible())) { return buildPythonSetupEntry( - {ready: pythonSetup.ready}, - PYTHON_SETUP_COMMAND + {ready: pythonSetup.ready, drifted: pythonSetup.drifted}, + PYTHON_SETUP_COMMAND, + PYTHON_SETUP_RERUN_COMMAND ); } const environmentState = await this.featureManager.isEnabled( diff --git a/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.test.ts b/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.test.ts index d04d92620..830284edc 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.test.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.test.ts @@ -1,35 +1,125 @@ import {expect} from "chai"; -import {ThemeColor, ThemeIcon} from "vscode"; -import {buildPythonSetupEntry} from "./pythonSetupEntry"; +import {EventEmitter, ThemeColor, ThemeIcon} from "vscode"; +import { + buildPythonSetupEntry, + composePythonSetupEntry, +} from "./pythonSetupEntry"; describe("buildPythonSetupEntry", () => { const COMMAND = "databricks.environment.setupPythonEnv"; + const RERUN = "databricks.environment.rerunPythonEnv"; it("renders a run CTA when setup is not yet ready", () => { - const [item] = buildPythonSetupEntry({ready: false}, COMMAND); - + const [item] = buildPythonSetupEntry( + {ready: false, drifted: false}, + COMMAND, + RERUN + ); expect(item.command?.command).to.equal(COMMAND); expect((item.iconPath as ThemeIcon).id).to.equal("rocket"); - // Not-done reads as an error (red), consistent with the sibling - // checklist entries, rather than the green debug-start color. expect((item.iconPath as ThemeIcon).color).to.deep.equal( new ThemeColor("errorForeground") ); - // The label invites the user to run setup. expect(String(item.label)).to.match(/set up/i); }); it("renders a ready status line (check icon) once setup succeeded", () => { - const [item] = buildPythonSetupEntry({ready: true}, COMMAND); - + const [item] = buildPythonSetupEntry( + {ready: true, drifted: false}, + COMMAND, + RERUN + ); expect((item.iconPath as ThemeIcon).id).to.equal("check"); - // Still actionable (re-run), but presented as done. expect(item.command?.command).to.equal(COMMAND); }); - it("returns exactly one entry (mutually exclusive with the checklist)", () => { - expect(buildPythonSetupEntry({ready: false}, COMMAND)).to.have.length( - 1 + it("renders an out-of-date state that re-runs setup when drifted", () => { + const [item] = buildPythonSetupEntry( + {ready: true, drifted: true}, + COMMAND, + RERUN + ); + expect((item.iconPath as ThemeIcon).id).to.equal("warning"); + expect(item.command?.command).to.equal(RERUN); + expect(String(item.label)).to.match(/drifted/i); + }); + + it("drift takes precedence even when not ready this session", () => { + const [item] = buildPythonSetupEntry( + {ready: false, drifted: true}, + COMMAND, + RERUN + ); + expect((item.iconPath as ThemeIcon).id).to.equal("warning"); + expect(item.command?.command).to.equal(RERUN); + }); + + it("gives the drifted row a distinct id so VS Code rebinds its command", () => { + // The drifted state points at a different command than ready/set-up; if + // it reused the same tree-item id, VS Code would not reliably rebind the + // command on refresh and the re-run click would be inert. + const [drifted] = buildPythonSetupEntry( + {ready: true, drifted: true}, + COMMAND, + RERUN + ); + const [ready] = buildPythonSetupEntry( + {ready: true, drifted: false}, + COMMAND, + RERUN ); + expect(drifted.id).to.not.equal(ready.id); + }); + + it("returns exactly one entry (mutually exclusive with the checklist)", () => { + expect( + buildPythonSetupEntry( + {ready: false, drifted: false}, + COMMAND, + RERUN + ) + ).to.have.length(1); + }); +}); + +describe("composePythonSetupEntry", () => { + function fakeSetup() { + const e = new EventEmitter(); + return { + _e: e, + ready: false, + isVisible: async () => true, + onDidChangeState: e.event, + }; + } + function fakeDrift() { + const e = new EventEmitter(); + return {_e: e, drifted: false, onDidChangeState: e.event}; + } + + it("forwards ready, drifted and isVisible from the sources", async () => { + const setup = fakeSetup(); + const drift = fakeDrift(); + const entry = composePythonSetupEntry(setup, drift); + + setup.ready = true; + drift.drifted = true; + expect(entry.ready).to.be.true; + expect(entry.drifted).to.be.true; + expect(await entry.isVisible()).to.be.true; + entry.dispose(); + }); + + it("fires onDidChangeState when either source changes", () => { + const setup = fakeSetup(); + const drift = fakeDrift(); + const entry = composePythonSetupEntry(setup, drift); + let fired = 0; + entry.onDidChangeState(() => fired++); + + setup._e.fire(); + drift._e.fire(); + expect(fired).to.equal(2); + entry.dispose(); }); }); diff --git a/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.ts b/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.ts index 079182c87..020b1a550 100644 --- a/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.ts +++ b/packages/databricks-vscode/src/ui/configuration-view/pythonSetupEntry.ts @@ -1,7 +1,15 @@ -import {Event, ThemeColor, ThemeIcon} from "vscode"; +import {Disposable, Event, EventEmitter, ThemeColor, ThemeIcon} from "vscode"; import {ConfigurationTreeItem} from "./types"; const PYTHON_SETUP_ENTRY_ID = "ENVIRONMENT_PYTHON_SETUP"; +// The drifted row deliberately uses a DISTINCT tree-item id from the ready/set-up +// row. VS Code does not reliably rebind a tree node's `command` when an item +// keeps the same `id` but swaps to a different command across a refresh: the +// label/icon update but clicks still fire (or fail to fire) the old binding. The +// ready and set-up states share one command (setupPythonEnv), but the drifted +// state points at a different command (rerunPythonEnv, for its own telemetry), so +// it must be a separate node — otherwise the "re-run" click is silently inert. +const PYTHON_SETUP_DRIFTED_ENTRY_ID = "ENVIRONMENT_PYTHON_SETUP_DRIFTED"; /** * The slice of the setup orchestrator the config view needs to render its @@ -15,23 +23,50 @@ export interface PythonSetupEntry { isVisible(): Promise; /** True once a setup has completed successfully this session. */ readonly ready: boolean; - /** Fires when {@link ready} changes, so the view can refresh. */ + /** + * True when the selected compute no longer matches the recorded setup state + * (see {@link PythonSetupDriftManager}); renders the "out of date" state. + */ + readonly drifted: boolean; + /** Fires when {@link ready} or {@link drifted} changes, so the view refreshes. */ readonly onDidChangeState: Event; } /** * Build the single Python Environment child row for the uv-native setup. * - * Pure over its inputs so the label/icon/command wiring is unit-testable. Not - * ready → a run call-to-action (rocket); ready → a done status line (check). - * Either way the row runs `commandId`, so a ready environment can be re-run. - * Returns a one-element array to slot directly into `getChildren`, underscoring - * that this entry is mutually exclusive with the legacy checklist. + * Three states, in precedence order: + * - drifted -> an "out of date" warning that re-runs setup (rerunCommandId); + * - ready -> a done status line (check) that can still be re-run; + * - neither -> a run call-to-action (rocket). + * Drift wins over ready: a stale environment is the more urgent thing to show, + * and its action (re-run) is what resolves it. */ export function buildPythonSetupEntry( - state: {ready: boolean}, - commandId: string + state: {ready: boolean; drifted: boolean}, + commandId: string, + rerunCommandId: string ): ConfigurationTreeItem[] { + if (state.drifted) { + return [ + { + id: PYTHON_SETUP_DRIFTED_ENTRY_ID, + label: "Python environment is drifted", + tooltip: + "The selected compute no longer matches your Python " + + "environment. Re-run setup to align it.", + contextValue: "databricks.environment.pythonSetup.drifted", + iconPath: new ThemeIcon( + "warning", + new ThemeColor("errorForeground") + ), + command: { + title: "Re-run Python setup", + command: rerunCommandId, + }, + }, + ]; + } return [ { id: PYTHON_SETUP_ENTRY_ID, @@ -51,3 +86,41 @@ export function buildPythonSetupEntry( }, ]; } + +/** + * Combine the setup controller's `ready` state and the drift manager's + * `drifted` state into the single {@link PythonSetupEntry} the config view + * consumes, merging both change events so a change in either refreshes the row. + * Returns a Disposable that tears down the merged emitter and its subscriptions. + */ +export function composePythonSetupEntry( + setup: { + isVisible(): Promise; + readonly ready: boolean; + readonly onDidChangeState: Event; + }, + drift: { + readonly drifted: boolean; + readonly onDidChangeState: Event; + } +): PythonSetupEntry & Disposable { + const emitter = new EventEmitter(); + const subs = [ + setup.onDidChangeState(() => emitter.fire()), + drift.onDidChangeState(() => emitter.fire()), + ]; + return { + isVisible: () => setup.isVisible(), + get ready() { + return setup.ready; + }, + get drifted() { + return drift.drifted; + }, + onDidChangeState: emitter.event, + dispose() { + subs.forEach((s) => s.dispose()); + emitter.dispose(); + }, + }; +}