Skip to content

Commit 5a85750

Browse files
ilia-dbilia-db
authored
Prompt users to setup virtual environment (#1169)
## Changes This PR adds "Local Python Environment" UI element to our configuration panel. It's only visible when we detect any unresolved problems: - No cluster is connected (needed for the dbconnect setup) - Workspace doesn't have Unity Catalog enabled (needed for the dbconnect setup) - Virtual environment is not enabled - Databricks Connect is not installed The problems are presented as an item in the "Python Environment" tree node (last entry in the configuration view). Each item with "run" icon is clickable and leads to relevant action to resolve the problem. Majority or the actions we've already had for the "Databricks Connect" button (bottom left corner, this PR doesn't change that). One change here is that we force virtual environment activation before installing dbconnect package. There are two items that are visible in the "success" state: active environment and dbconnect version. Both of them have gear icons that lead to actions to change the environment or reinstall dbconnect (see screenshots below). For now we don't have any additional logic for the newly created project (through bundle init), when you create them and open in the VSCode, we will also show the "Local Python Environment" UI that will prompt users to enable the environment and install necessary dependencies into it. For managing virtual environments themselves we use python extension commands. <img width="460" alt="Screenshot 2024-04-22 at 16 53 11" src="https://github.com/databricks/databricks-vscode/assets/148094031/2be9026f-374b-40e0-83b0-14bde240fbcb"> <img width="460" alt="Screenshot 2024-04-22 at 16 52 22" src="https://github.com/databricks/databricks-vscode/assets/148094031/3f0528b4-b0a7-4b81-86c4-f4e4e4a33a97"> <img width="460" alt="Screenshot 2024-04-22 at 16 52 04" src="https://github.com/databricks/databricks-vscode/assets/148094031/4dcc1024-f451-415f-a1d0-a13e9073f8dc">
1 parent b279467 commit 5a85750

23 files changed

+796
-733
lines changed

‎packages/databricks-vscode/package.json‎

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -290,6 +290,27 @@
290290
"enablement": "databricks.context.activated",
291291
"category": "Databricks"
292292
},
293+
{
294+
"command": "databricks.environment.refresh",
295+
"title": "Refresh python environment status",
296+
"icon": "$(refresh)",
297+
"enablement": "databricks.context.activated && databricks.context.loggedIn",
298+
"category": "Databricks"
299+
},
300+
{
301+
"command": "databricks.environment.selectPythonInterpreter",
302+
"title": "Change Python environment",
303+
"icon": "$(gear)",
304+
"enablement": "databricks.context.activated && databricks.context.loggedIn",
305+
"category": "Databricks"
306+
},
307+
{
308+
"command": "databricks.environment.reinstallDBConnect",
309+
"title": "Reinstall Databricks Connect",
310+
"icon": "$(gear)",
311+
"enablement": "databricks.context.activated && databricks.context.loggedIn",
312+
"category": "Databricks"
313+
},
293314
{
294315
"command": "databricks.bundle.variable.openFile",
295316
"title": "Override bundle variables",
@@ -488,6 +509,24 @@
488509
{
489510
"command": "databricks.utils.copy",
490511
"when": "view == dabsResourceExplorerView || view == configurationView"
512+
},
513+
{
514+
"command": "databricks.environment.refresh",
515+
"when": "view == configurationView && viewItem =~ /^databricks.environment.root.(success|error)$/",
516+
"group": "inline@0",
517+
"icon": "$(refresh)"
518+
},
519+
{
520+
"command": "databricks.environment.selectPythonInterpreter",
521+
"when": "view == configurationView && viewItem =~ /^databricks.environment.checkPythonEnvironment.success$/",
522+
"group": "inline@0",
523+
"icon": "$(gear)"
524+
},
525+
{
526+
"command": "databricks.environment.reinstallDBConnect",
527+
"when": "view == configurationView && viewItem =~ /^databricks.environment.checkEnvironmentDependencies.success$/",
528+
"group": "inline@0",
529+
"icon": "$(gear)"
491530
}
492531
],
493532
"editor/title": [
@@ -773,7 +812,6 @@
773812
"default": [],
774813
"items": {
775814
"enum": [
776-
"notebooks.dbconnect",
777815
"views.cluster",
778816
"views.workspace"
779817
],

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ export class ConnectionCommands implements Disposable {
119119
}
120120

121121
attachClusterQuickPickCommand() {
122-
return async () => {
122+
return async (title?: string) => {
123123
const workspaceClient = this.connectionManager.workspaceClient;
124124
const me = this.connectionManager.databricksWorkspace?.userName;
125125
if (!workspaceClient || !me) {
@@ -130,6 +130,7 @@ export class ConnectionCommands implements Disposable {
130130
const quickPick = window.createQuickPick<
131131
ClusterItem | QuickPickItem
132132
>();
133+
quickPick.title = title;
133134
quickPick.keepScrollPosition = true;
134135
quickPick.busy = true;
135136

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

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -177,10 +177,6 @@ export class ConnectionManager implements Disposable {
177177
}
178178

179179
get cluster(): Cluster | undefined {
180-
if (this.state !== "CONNECTED") {
181-
return;
182-
}
183-
184180
return this._clusterManager?.cluster;
185181
}
186182

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

Lines changed: 41 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -36,16 +36,15 @@ import {CustomWhenContext} from "./vscode-objs/CustomWhenContext";
3636
import {StateStorage} from "./vscode-objs/StateStorage";
3737
import path from "node:path";
3838
import {FeatureId, FeatureManager} from "./feature-manager/FeatureManager";
39-
import {DbConnectAccessVerifier} from "./language/DbConnectAccessVerifier";
39+
import {EnvironmentDependenciesVerifier} from "./language/EnvironmentDependenciesVerifier";
4040
import {MsPythonExtensionWrapper} from "./language/MsPythonExtensionWrapper";
4141
import {DatabricksEnvFileManager} from "./file-managers/DatabricksEnvFileManager";
4242
import {getContextMetadata, Telemetry, toUserMetadata} from "./telemetry";
4343
import "./telemetry/commandExtensions";
4444
import {Events, Metadata} from "./telemetry/constants";
45-
import {DbConnectInstallPrompt} from "./language/DbConnectInstallPrompt";
45+
import {EnvironmentDependenciesInstaller} from "./language/EnvironmentDependenciesInstaller";
4646
import {setDbnbCellLimits} from "./language/notebooks/DatabricksNbCellLimits";
4747
import {DbConnectStatusBarButton} from "./language/DbConnectStatusBarButton";
48-
import {NotebookAccessVerifier} from "./language/notebooks/NotebookAccessVerifier";
4948
import {NotebookInitScriptManager} from "./language/notebooks/NotebookInitScriptManager";
5049
import {showRestartNotebookDialogue} from "./language/notebooks/restartNotebookDialogue";
5150
import {
@@ -72,6 +71,7 @@ import {BundleVariableModel} from "./bundle/models/BundleVariableModel";
7271
import {BundleVariableTreeDataProvider} from "./ui/bundle-variables/BundleVariableTreeDataProvider";
7372
import {ConfigurationTreeViewManager} from "./ui/configuration-view/ConfigurationTreeViewManager";
7473
import {getCLIDependenciesEnvVars} from "./utils/envVarGenerators";
74+
import {EnvironmentCommands} from "./language/EnvironmentCommands";
7575

7676
// eslint-disable-next-line @typescript-eslint/no-var-requires
7777
const packageJson = require("../package.json");
@@ -364,29 +364,46 @@ export async function activate(
364364

365365
const clusterModel = new ClusterModel(connectionManager);
366366

367-
const dbConnectInstallPrompt = new DbConnectInstallPrompt(
368-
stateStorage,
369-
pythonExtensionWrapper
370-
);
367+
const environmentDependenciesInstaller =
368+
new EnvironmentDependenciesInstaller(pythonExtensionWrapper);
371369
const featureManager = new FeatureManager<FeatureId>([]);
372370
featureManager.registerFeature(
373-
"debugging.dbconnect",
371+
"environment.dependencies",
374372
() =>
375-
new DbConnectAccessVerifier(
373+
new EnvironmentDependenciesVerifier(
376374
connectionManager,
377375
pythonExtensionWrapper,
378-
dbConnectInstallPrompt
376+
environmentDependenciesInstaller
379377
)
380378
);
381-
382-
featureManager.registerFeature(
383-
"notebooks.dbconnect",
384-
() =>
385-
new NotebookAccessVerifier(
386-
featureManager,
387-
pythonExtensionWrapper,
388-
stateStorage
389-
)
379+
const environmentCommands = new EnvironmentCommands(
380+
featureManager,
381+
pythonExtensionWrapper,
382+
environmentDependenciesInstaller
383+
);
384+
context.subscriptions.push(
385+
telemetry.registerCommand(
386+
"databricks.environment.setup",
387+
environmentCommands.setup,
388+
environmentCommands
389+
),
390+
telemetry.registerCommand(
391+
"databricks.environment.refresh",
392+
environmentCommands.refresh,
393+
environmentCommands
394+
),
395+
telemetry.registerCommand(
396+
"databricks.environment.selectPythonInterpreter",
397+
environmentCommands.selectPythonInterpreter,
398+
environmentCommands
399+
),
400+
telemetry.registerCommand(
401+
"databricks.environment.reinstallDBConnect",
402+
async () =>
403+
environmentCommands.reinstallDBConnect(
404+
connectionManager.cluster
405+
)
406+
)
390407
);
391408

392409
const dbConnectStatusBarButton = new DbConnectStatusBarButton(
@@ -406,8 +423,7 @@ export async function activate(
406423
connectionManager,
407424
featureManager,
408425
pythonExtensionWrapper,
409-
databricksEnvFileManager,
410-
configModel
426+
databricksEnvFileManager
411427
);
412428

413429
context.subscriptions.push(
@@ -417,30 +433,6 @@ export async function activate(
417433
"databricks.notebookInitScript.verify",
418434
notebookInitScriptManager.verifyInitScriptCommand,
419435
notebookInitScriptManager
420-
),
421-
workspace.onDidOpenNotebookDocument(() =>
422-
featureManager.isEnabled("notebooks.dbconnect")
423-
),
424-
featureManager.onDidChangeState(
425-
"notebooks.dbconnect",
426-
async (featureState) => {
427-
const dbconnectState = await featureManager.isEnabled(
428-
"debugging.dbconnect"
429-
);
430-
if (!dbconnectState.avaliable) {
431-
return; // Only take action of notebook errors, when dbconnect is avaliable
432-
}
433-
if (featureState.action) {
434-
featureState.action();
435-
} else if (
436-
!featureState.isDisabledByFf &&
437-
featureState.reason
438-
) {
439-
window.showErrorMessage(
440-
`Error while trying to initialize Databricks Notebooks. Some features may not work. Reason: ${featureState.reason}`
441-
);
442-
}
443-
}
444436
)
445437
);
446438

@@ -462,14 +454,14 @@ export async function activate(
462454
databricksEnvFileManager,
463455
showRestartNotebookDialogue(databricksEnvFileManager)
464456
);
465-
featureManager.isEnabled("debugging.dbconnect");
457+
featureManager.isEnabled("environment.dependencies");
466458

467459
const configureAutocomplete = new ConfigureAutocomplete(
468460
context,
469461
stateStorage,
470462
workspaceUri.fsPath,
471463
pythonExtensionWrapper,
472-
dbConnectInstallPrompt
464+
environmentDependenciesInstaller
473465
);
474466
context.subscriptions.push(
475467
configureAutocomplete,
@@ -484,7 +476,8 @@ export async function activate(
484476
connectionManager,
485477
bundleProjectManager,
486478
configModel,
487-
cli
479+
cli,
480+
featureManager
488481
);
489482
const configurationView = window.createTreeView("configurationView", {
490483
treeDataProvider: configurationDataProvider,

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

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,12 @@ export class DisabledFeature extends MultiStepAccessVerifier {
66
}
77

88
async check() {
9-
this.rejectStep("disabled", "feature is disabled", undefined, true);
9+
this.rejectStep(
10+
"disabled",
11+
"feature is disabled",
12+
undefined,
13+
undefined,
14+
true
15+
);
1016
}
1117
}

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

Lines changed: 20 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import {spy, verify} from "ts-mockito";
2-
import {FeatureManager, FeatureState} from "./FeatureManager";
2+
import {FeatureManager, FeatureStepState} from "./FeatureManager";
33
import {MultiStepAccessVerifier} from "./MultiStepAccessVerfier";
44
import * as assert from "assert";
55
class TestAccessVerifier extends MultiStepAccessVerifier {
@@ -27,11 +27,8 @@ class TestAccessVerifier extends MultiStepAccessVerifier {
2727
}
2828
}
2929

30-
function isAvailable(state: boolean | FeatureState) {
31-
if (typeof state === "boolean") {
32-
return state;
33-
}
34-
return state.avaliable;
30+
function isAvailable(state: FeatureStepState) {
31+
return state.available;
3532
}
3633

3734
describe(__filename, async () => {
@@ -41,8 +38,8 @@ describe(__filename, async () => {
4138
const spyTestVerifier = spy(testVerifier);
4239
fm.registerFeature("test", () => testVerifier);
4340

44-
assert.ok(!(await fm.isEnabled("test")).avaliable);
45-
assert.ok(!(await fm.isEnabled("test")).avaliable);
41+
assert.ok(!(await fm.isEnabled("test")).available);
42+
assert.ok(!(await fm.isEnabled("test")).available);
4643
verify(spyTestVerifier.check()).once();
4744
verify(spyTestVerifier.check1()).once();
4845
verify(spyTestVerifier.check2()).once();
@@ -55,26 +52,29 @@ describe(__filename, async () => {
5552
fm.registerFeature("test", () => testVerifier);
5653

5754
//check value is picked from cache
58-
assert.ok(!(await fm.isEnabled("test")).avaliable);
59-
assert.ok(!(await fm.isEnabled("test")).avaliable);
55+
assert.ok(!(await fm.isEnabled("test")).available);
56+
assert.ok(!(await fm.isEnabled("test")).available);
6057
verify(spyTestVerifier.check()).once();
6158
verify(spyTestVerifier.check1()).once();
6259
verify(spyTestVerifier.check2()).once();
6360

6461
//cache should be true only when both values are true
6562
assert.ok(isAvailable(await testVerifier.check1(true)));
66-
assert.ok(!(await fm.isEnabled("test")).avaliable);
63+
assert.ok(!(await fm.isEnabled("test")).available);
6764
assert.ok(isAvailable(await testVerifier.check2(true)));
68-
assert.ok((await fm.isEnabled("test")).avaliable);
65+
assert.ok((await fm.isEnabled("test")).available);
6966

7067
//cache should be false if even 1 value is false
7168
assert.ok(!isAvailable(await testVerifier.check2(false)));
72-
assert.ok(!(await fm.isEnabled("test")).avaliable);
73-
assert.ok((await fm.isEnabled("test")).reason === "reason2");
69+
assert.ok(!(await fm.isEnabled("test")).available);
70+
assert.strictEqual(
71+
(await fm.isEnabled("test")).steps.get("check2")?.title,
72+
"reason2"
73+
);
7474

7575
//cache should be reset to true if both values are true
7676
assert.ok(isAvailable(await testVerifier.check2(true)));
77-
assert.ok((await fm.isEnabled("test")).avaliable);
77+
assert.ok((await fm.isEnabled("test")).available);
7878
});
7979

8080
it("disabled features should always return false", async () => {
@@ -83,14 +83,15 @@ describe(__filename, async () => {
8383
const spyTestVerifier = spy(testVerifier);
8484
fm.registerFeature("test", () => testVerifier);
8585

86-
assert.ok(!(await fm.isEnabled("test")).avaliable);
87-
assert.ok(!(await fm.isEnabled("test")).avaliable);
86+
assert.ok(!(await fm.isEnabled("test")).available);
87+
assert.ok(!(await fm.isEnabled("test")).available);
8888
verify(spyTestVerifier.check()).never();
8989
verify(spyTestVerifier.check1()).never();
9090
verify(spyTestVerifier.check2()).never();
9191

92-
assert.ok(
93-
(await fm.isEnabled("test")).reason === "feature is disabled"
92+
assert.strictEqual(
93+
(await fm.isEnabled("test")).message,
94+
"Feature is disabled"
9495
);
9596
});
9697
});

0 commit comments

Comments
 (0)