Skip to content

Commit 315f77d

Browse files
rugpanovrugpanov
authored
feat(python-setup): GA the uv-native environment setup (#2124)
## Why The uv-native "Set up Python environment" flow shipped behind an opt-in feature flag (`environment.pythonSetup` in `databricks.experiments.optInto`) because its CLI command, `environments setup-local`, was only available in custom CLI builds. That command now ships in the bundled CLI (v1.12.1), so the flag — and the temporary `cliPathOverride` escape hatch — have served their purpose. This makes the flow generally available. ## What - **Remove the opt-in flag.** Drop `environment.pythonSetup` from the `databricks.experiments.optInto` enum, delete the `PYTHON_SETUP_FEATURE_ID` constant, and construct the `FeatureManager` with an empty disabled list. - **Remove the temporary override.** Delete the `databricks.experiments.cliPathOverride` setting, its `WorkspaceConfigs` getter, and `resolveCliPath`; the CLI client always uses the bundled CLI. - **Visibility is now purely project-based.** The uv flow shows only for uv-suitable projects (`isUvSetupSuitable`) — clean uv/greenfield projects with no competing manager. Projects driven by pip/poetry/conda keep the legacy checklist via `routeEnvironmentSetup`, exactly as before. - **Serverless-version prompt scoped to uv projects.** `ConnectionCommands.selectServerless` now gates the version prompt on uv-suitability (injected `isUvSetupVisible`) rather than the flag, so enabling serverless on a non-uv project keeps the plain, version-less enable — the old flow for pip users is unchanged. - **Tests.** Delete the flag-coupling drift test; update the gate/deps/args/connection tests to the always-on behavior and add a regression test pinning the pip-path (plain serverless enable, no version prompt). ## Backward compatibility - Users who already added `environment.pythonSetup` to `experiments.optInto` — harmless leftover; the value is simply no longer read. - Users who set `experiments.cliPathOverride` — the setting is removed (VS Code flags it as unknown, benign) and the bundled CLI is used, which now ships the command. - No persisted-state or telemetry-schema changes. ## Verification - `yarn --cwd packages/databricks-vscode build` — clean - `yarn --cwd packages/databricks-vscode test:lint` — clean - `yarn --cwd packages/databricks-vscode test:unit` — 829 passing, 0 failing This pull request and its description were written by Isaac.
1 parent 253b366 commit 315f77d

17 files changed

Lines changed: 205 additions & 477 deletions

‎packages/databricks-vscode/package.json‎

Lines changed: 2 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1797,24 +1797,17 @@
17971797
"items": {
17981798
"enum": [
17991799
"views.cluster",
1800-
"views.workspace",
1801-
"environment.pythonSetup"
1800+
"views.workspace"
18021801
],
18031802
"enumDescriptions": [
18041803
"Show cluster view in the explorer.",
1805-
"Show workspace browser in the explorer.",
1806-
"Enable the uv-native one-click Python environment setup for Databricks Connect."
1804+
"Show workspace browser in the explorer."
18071805
],
18081806
"type": "string"
18091807
},
18101808
"uniqueItems": true,
18111809
"description": "Opt into experimental features."
18121810
},
1813-
"databricks.experiments.cliPathOverride": {
1814-
"type": "string",
1815-
"default": "",
1816-
"description": "Absolute path to a custom Databricks CLI that provides `environments setup-local`. Temporary experimental setting used while the command ships only in custom CLI builds; leave empty to use the bundled CLI. Will be removed once the command is generally available."
1817-
},
18181811
"databricks.wsfs.rearrangeCells": {
18191812
"type": "boolean",
18201813
"default": true,

‎packages/databricks-vscode/src/configuration/ConnectionCommands.test.ts‎

Lines changed: 112 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -83,16 +83,26 @@ describe(__filename, () => {
8383

8484
describe("attachClusterQuickPick return contract", () => {
8585
let originalCreateQuickPick: typeof window.createQuickPick;
86+
let originalShowQuickPick: typeof window.showQuickPick;
8687
let fakePick: FakeComputeQuickPick;
8788
let attachCalls: string[];
8889
let commands: ConnectionCommands;
90+
// The serverless-version sub-picker uses window.showQuickPick (distinct
91+
// from the compute picker's createQuickPick), so it is stubbed
92+
// separately; each test sets what the user "picks" here.
93+
let versionPick: {version: string} | undefined;
8994

9095
beforeEach(() => {
9196
originalCreateQuickPick = window.createQuickPick;
9297
fakePick = new FakeComputeQuickPick();
9398
(window as unknown as {createQuickPick: unknown}).createQuickPick =
9499
() => fakePick;
95100

101+
originalShowQuickPick = window.showQuickPick;
102+
versionPick = undefined;
103+
(window as unknown as {showQuickPick: unknown}).showQuickPick =
104+
async () => versionPick;
105+
96106
attachCalls = [];
97107
const connectionManager = {
98108
workspaceClient: {},
@@ -115,11 +125,14 @@ describe(__filename, () => {
115125
clusterModel as never,
116126
{} as never,
117127
{} as never,
118-
{} as never
128+
{} as never,
129+
() => Promise.resolve(false)
119130
);
120131
});
121132

122133
afterEach(() => {
134+
(window as unknown as {showQuickPick: unknown}).showQuickPick =
135+
originalShowQuickPick;
123136
(window as unknown as {createQuickPick: unknown}).createQuickPick =
124137
originalCreateQuickPick;
125138
});
@@ -130,6 +143,9 @@ describe(__filename, () => {
130143
cluster: {id},
131144
}) as unknown as QuickPickItem;
132145

146+
const serverlessItem = () =>
147+
({label: "$(cloud) Serverless"}) as unknown as QuickPickItem;
148+
133149
it("resolves to the attached cluster and attaches it exactly once", async () => {
134150
const resultP = commands.attachClusterQuickPickCommand()();
135151
await fakePick.accept([clusterItem("c1")]);
@@ -191,7 +207,8 @@ describe(__filename, () => {
191207
} as never,
192208
{} as never,
193209
{} as never,
194-
{} as never
210+
{} as never,
211+
() => Promise.resolve(false)
195212
);
196213

197214
const resultP = cmds.attachClusterQuickPickCommand()();
@@ -224,13 +241,105 @@ describe(__filename, () => {
224241
} as never,
225242
{} as never,
226243
{} as never,
227-
{} as never
244+
{} as never,
245+
() => Promise.resolve(false)
228246
);
229247

230248
const resultP = cmds.attachClusterQuickPickCommand()();
231249
await fakePick.accept([clusterItem("c1")]);
232250

233251
assert.equal(await resultP, undefined);
234252
});
253+
254+
it("keeps the plain, version-less serverless enable for a non-uv-suitable project", async () => {
255+
// The pip/poetry/conda flow is unchanged at GA: selecting serverless
256+
// enables it directly, with no version sub-prompt and no serverless
257+
// target returned for the (inapplicable) uv setup.
258+
const enableCalls: Array<string | undefined> = [];
259+
const cmds = new ConnectionCommands(
260+
{} as never,
261+
{
262+
workspaceClient: {},
263+
databricksWorkspace: {userName: "me"},
264+
attachCluster: async () => {},
265+
enableServerless: async (version?: string) => {
266+
enableCalls.push(version);
267+
},
268+
} as never,
269+
{
270+
refresh() {},
271+
onDidChange() {
272+
return {dispose() {}};
273+
},
274+
roots: [],
275+
} as never,
276+
{} as never,
277+
{} as never,
278+
{} as never,
279+
// Project is not uv-suitable, so the uv setup is not the active
280+
// surface and no version should be requested.
281+
() => Promise.resolve(false)
282+
);
283+
284+
const resultP = cmds.attachClusterQuickPickCommand()();
285+
await fakePick.accept([serverlessItem()]);
286+
287+
assert.deepEqual(enableCalls, [undefined]);
288+
assert.equal(await resultP, undefined);
289+
});
290+
291+
const makeRecordingServerlessCommands = (
292+
enableCalls: Array<string | undefined>
293+
) =>
294+
new ConnectionCommands(
295+
{} as never,
296+
{
297+
workspaceClient: {},
298+
databricksWorkspace: {userName: "me"},
299+
attachCluster: async () => {},
300+
enableServerless: async (version?: string) => {
301+
enableCalls.push(version);
302+
},
303+
} as never,
304+
{
305+
refresh() {},
306+
onDidChange() {
307+
return {dispose() {}};
308+
},
309+
roots: [],
310+
} as never,
311+
{} as never,
312+
{} as never,
313+
{} as never,
314+
// uv-suitable project: selecting serverless resolves a version
315+
// first, so the version sub-picker runs.
316+
() => Promise.resolve(true)
317+
);
318+
319+
it("prompts for a version and enables it for a uv-suitable project", async () => {
320+
versionPick = {version: "5"};
321+
const enableCalls: Array<string | undefined> = [];
322+
const cmds = makeRecordingServerlessCommands(enableCalls);
323+
324+
const resultP = cmds.attachClusterQuickPickCommand()();
325+
await fakePick.accept([serverlessItem()]);
326+
327+
// The confirmed version is enabled and returned as the serverless
328+
// target for the uv setup to provision against.
329+
assert.deepEqual(enableCalls, ["5"]);
330+
assert.deepEqual(await resultP, {kind: "serverless", version: "5"});
331+
});
332+
333+
it("makes no compute change when the version picker is dismissed on a uv-suitable project", async () => {
334+
versionPick = undefined; // user escaped the version picker
335+
const enableCalls: Array<string | undefined> = [];
336+
const cmds = makeRecordingServerlessCommands(enableCalls);
337+
338+
const resultP = cmds.attachClusterQuickPickCommand()();
339+
await fakePick.accept([serverlessItem()]);
340+
341+
assert.deepEqual(enableCalls, []);
342+
assert.equal(await resultP, undefined);
343+
});
235344
});
236345
});

‎packages/databricks-vscode/src/configuration/ConnectionCommands.ts‎

Lines changed: 19 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -24,10 +24,7 @@ import {
2424
} from "../ui/configuration-view/AuthTypeComponent";
2525
import {ManualLoginSource} from "../telemetry/constants";
2626
import {onError} from "../utils/onErrorDecorator";
27-
import {
28-
isPythonSetupEnabled,
29-
resolveServerlessVersion,
30-
} from "../python-setup/utils/serverlessVersionResolver";
27+
import {resolveServerlessVersion} from "../python-setup/utils/serverlessVersionResolver";
3128
import {pickServerlessVersion} from "../python-setup/utils/serverlessVersionPicker";
3229
import {collectServerlessVersionObservations} from "../python-setup/utils/serverlessVersionObservations";
3330
import type {SetupCompute} from "../python-setup/controllers/PythonSetupEnvironmentSetup";
@@ -98,7 +95,15 @@ export class ConnectionCommands implements Disposable {
9895
private readonly clusterModel: ClusterModel,
9996
private readonly configModel: ConfigModel,
10097
private readonly cli: CliWrapper,
101-
private readonly workspaceFolderManager: WorkspaceFolderManager
98+
private readonly workspaceFolderManager: WorkspaceFolderManager,
99+
/**
100+
* Whether the uv-native Python setup is the active surface for the
101+
* current project (i.e. the project is uv-suitable). Only then does
102+
* enabling serverless prompt for an environment version to record for
103+
* that setup; a project driven by a competing manager keeps the plain,
104+
* version-less serverless enable.
105+
*/
106+
private readonly isUvSetupVisible: () => Promise<boolean>
102107
) {}
103108

104109
/**
@@ -294,18 +299,19 @@ export class ConnectionCommands implements Disposable {
294299
}
295300

296301
/**
297-
* Enable serverless compute. When the uv-native python-setup feature is
298-
* opted into, first ask the user to confirm the serverless environment
299-
* version (ranked from the project's bundle) and persist it with the
300-
* selection, so setup need not re-prompt. If they dismiss the version
301-
* picker, no compute change is made. With the feature off this is the
302-
* plain, unchanged serverless enable.
302+
* Enable serverless compute. When the uv-native python-setup is the active
303+
* surface for this project, first ask the user to confirm the serverless
304+
* environment version (ranked from the project's bundle) and persist it with
305+
* the selection, so setup need not re-prompt. If they dismiss the version
306+
* picker, no compute change is made. For a project the uv setup does not fit
307+
* (a competing manager is driving it) this is the plain, unchanged serverless
308+
* enable.
303309
*
304310
* Returns the confirmed version, or `undefined` if the picker was dismissed
305-
* or the feature is off (serverless enabled but version-less).
311+
* or the project is not uv-suitable (serverless enabled but version-less).
306312
*/
307313
private async selectServerless(): Promise<string | undefined> {
308-
if (!isPythonSetupEnabled()) {
314+
if (!(await this.isUvSetupVisible())) {
309315
await this.connectionManager.enableServerless();
310316
return undefined;
311317
}

‎packages/databricks-vscode/src/extension.ts‎

Lines changed: 15 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -47,11 +47,7 @@ import {CustomWhenContext} from "./vscode-objs/CustomWhenContext";
4747
import {StateStorage} from "./vscode-objs/StateStorage";
4848
import path from "node:path";
4949
import {existsSync} from "node:fs";
50-
import {
51-
FeatureId,
52-
FeatureManager,
53-
PYTHON_SETUP_FEATURE_ID,
54-
} from "./feature-manager/FeatureManager";
50+
import {FeatureId, FeatureManager} from "./feature-manager/FeatureManager";
5551
import {PythonSetupManagerDetector} from "./python-setup/utils/PythonSetupManagerDetector";
5652
import {PythonSetupCliClient} from "./python-setup/gateways/PythonSetupCliClient";
5753
import {PythonSetupEnvironmentSetup} from "./python-setup/controllers/PythonSetupEnvironmentSetup";
@@ -63,11 +59,7 @@ import {PythonSetupDriftManager} from "./python-setup/controllers/PythonSetupDri
6359
import {PythonSetupAdoptionManager} from "./python-setup/controllers/PythonSetupAdoptionManager";
6460
import {SetupCompute} from "./python-setup/controllers/PythonSetupEnvironmentSetup";
6561
import {venvInterpreterPath} from "./python-setup/utils/venvInterpreterPath";
66-
import {resolveCliPath} from "./python-setup/utils/setupLocalArgs";
67-
import {
68-
isPythonSetupEnabled,
69-
makeServerlessVersionPrompt,
70-
} from "./python-setup/utils/serverlessVersionResolver";
62+
import {makeServerlessVersionPrompt} from "./python-setup/utils/serverlessVersionResolver";
7163
import {collectPackageManagerSignals} from "./language/packageManagerSignals";
7264
import {EnvironmentDependenciesVerifier} from "./language/EnvironmentDependenciesVerifier";
7365
import {MsPythonExtensionWrapper} from "./language/MsPythonExtensionWrapper";
@@ -900,13 +892,7 @@ export async function activate(
900892
connectionManager,
901893
pythonExtensionWrapper
902894
);
903-
// python-setup ships disabled by default: the CLI's `environments
904-
// setup-local` command is available only in custom CLI builds for now, so
905-
// the whole flow stays hidden until a user opts in via
906-
// `databricks.experiments.optInto` (see PYTHON_SETUP_FEATURE_ID).
907-
const featureManager = new FeatureManager<FeatureId>([
908-
PYTHON_SETUP_FEATURE_ID,
909-
]);
895+
const featureManager = new FeatureManager<FeatureId>([]);
910896
featureManager.registerFeature(
911897
"environment.dependencies",
912898
() =>
@@ -921,9 +907,9 @@ export async function activate(
921907
() => pythonSetupEnvironment.isVisible()
922908
)
923909
);
924-
// uv-native Python environment setup (python-setup). Constructed always,
925-
// but inert unless the user opts in: the detector/gate keep the entry hidden
926-
// otherwise, so this changes nothing for existing users.
910+
// uv-native Python environment setup (python-setup). The detector/gate keep
911+
// the entry visible only for uv-suitable projects; projects driven by a
912+
// competing manager (pip/poetry/conda) fall back to the legacy checklist.
927913
const pythonSetupDetector = new PythonSetupManagerDetector(
928914
async (projectRoot) =>
929915
collectPackageManagerSignals(
@@ -932,11 +918,7 @@ export async function activate(
932918
)
933919
);
934920
const pythonSetupClient = new PythonSetupCliClient(
935-
() =>
936-
resolveCliPath({
937-
override: workspaceConfigs.pythonSetupCliPathOverride,
938-
bundled: cli.cliPath,
939-
}),
921+
() => cli.cliPath,
940922
() => {
941923
// Overlay the extension's workspace auth onto the ambient
942924
// environment, so the CLI provisions against the workspace we are
@@ -958,9 +940,9 @@ export async function activate(
958940
};
959941
}
960942
);
961-
// Created lazily on first setup output so a non-opted-in user never gets an
962-
// empty "Databricks Python Environment Setup" entry in the Output dropdown
963-
// (the feature is otherwise fully inert for them).
943+
// Created lazily on first setup output so a user who never runs setup does
944+
// not get an empty "Databricks Python Environment Setup" entry in the Output
945+
// dropdown.
964946
let pythonSetupLogChannel: OutputChannel | undefined;
965947
const getPythonSetupLogChannel = () => {
966948
if (pythonSetupLogChannel === undefined) {
@@ -984,7 +966,6 @@ export async function activate(
984966
return undefined;
985967
}
986968
},
987-
isEnabled: isPythonSetupEnabled,
988969
detect: (projectRoot) => pythonSetupDetector.detect(projectRoot),
989970
attachedCompute: () => ({
990971
serverless: connectionManager.serverless,
@@ -1355,7 +1336,11 @@ export async function activate(
13551336
clusterModel,
13561337
configModel,
13571338
cli,
1358-
workspaceFolderManager
1339+
workspaceFolderManager,
1340+
// Only prompt for a serverless environment version when the uv-native
1341+
// setup is the active surface for this project; a project driven by a
1342+
// competing manager keeps the plain, version-less serverless enable.
1343+
() => pythonSetupEnvironment.isVisible()
13591344
);
13601345

13611346
context.subscriptions.push(

‎packages/databricks-vscode/src/feature-manager/FeatureManager.ts‎

Lines changed: 1 addition & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -8,23 +8,7 @@ import {Loggers} from "../logger";
88

99
export type FeatureEnableAction = (...args: any[]) => Promise<void>;
1010

11-
/**
12-
* Feature id for the uv-native Python environment setup (python-setup).
13-
*
14-
* This exact string is the single source of truth used in two coupled places,
15-
* and they must stay identical or the feature can never unlock:
16-
* - the {@link FeatureManager} `disabledFeatures` entry that hides it by
17-
* default (see extension.ts), and
18-
* - the value a user adds to `databricks.experiments.optInto` to opt in
19-
* (see the enum in package.json), which {@link FeatureManager} matches
20-
* verbatim against the disabled id to decide whether to unlock.
21-
* Exporting it as a constant keeps those two from drifting apart.
22-
*/
23-
export const PYTHON_SETUP_FEATURE_ID = "environment.pythonSetup";
24-
25-
export type FeatureId =
26-
| "environment.dependencies"
27-
| typeof PYTHON_SETUP_FEATURE_ID;
11+
export type FeatureId = "environment.dependencies";
2812

2913
export interface FeatureState {
3014
available: boolean;

0 commit comments

Comments
 (0)