Skip to content

Commit 93a7c36

Browse files
kartikgupta-dbkartikgupta-db
authored
Allow cancelling login, fix spaces in path for dbconnect and minor UI fixes (#1315)
## Changes * Allow cancelling login * Show confirmation prompt for dabs deploy in prod. * Fix spaces in path name for dbconnect paths. * Rename "workspace" in UI elements such as button "Migrate current workspace to Databricks Project" to "folder" to make it easier to understand. ## Tests <!-- How is this tested? -->
1 parent ded5382 commit 93a7c36

File tree

14 files changed

+303
-111
lines changed

14 files changed

+303
-111
lines changed

‎packages/databricks-vscode/package.json‎

Lines changed: 3 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -387,12 +387,12 @@
387387
"viewsWelcome": [
388388
{
389389
"view": "configurationView",
390-
"contents": "There are multiple Databricks projects in the workspace:\n[Open existing Databricks Project](command:databricks.bundle.openSubProject)",
390+
"contents": "There are multiple Databricks projects in the folder:\n[Open existing Databricks Project](command:databricks.bundle.openSubProject)",
391391
"when": "workspaceFolderCount > 0 && databricks.context.initialized && databricks.context.subProjectsAvailable"
392392
},
393393
{
394394
"view": "configurationView",
395-
"contents": "Migrate current workspace to a Databricks Project: \n[Migrate to a Databricks Project](command:databricks.bundle.startManualMigration)",
395+
"contents": "Migrate current folder to a Databricks Project: \n[Migrate current folder to a Databricks Project](command:databricks.bundle.startManualMigration)",
396396
"when": "workspaceFolderCount > 0 && databricks.context.initialized && databricks.context.pendingManualMigration"
397397
},
398398
{
@@ -412,7 +412,7 @@
412412
},
413413
{
414414
"view": "configurationView",
415-
"contents": "The workspace is empty.\n[Create a new Databricks Project](command:databricks.bundle.initNewProject)",
415+
"contents": "This folder is empty.\n[Create a new Databricks Project](command:databricks.bundle.initNewProject)",
416416
"when": "workspaceFolderCount == 0"
417417
}
418418
],
@@ -843,19 +843,6 @@
843843
"default": false,
844844
"description": "Enable/disable filtering for only accessible clusters (clusters on which the current user can run code)"
845845
},
846-
"databricks.sync.destinationType": {
847-
"type": "string",
848-
"default": "workspace",
849-
"enum": [
850-
"workspace",
851-
"repo"
852-
],
853-
"description": "Use a folder in Workspace or a Databricks Repo as sync destination (Reload for changes to take effect).",
854-
"enumDescriptions": [
855-
"Use a folder in Workspace as sync destination",
856-
"Use a Repo as sync destination"
857-
]
858-
},
859846
"databricks.overrideDatabricksConfigFile": {
860847
"type": "string",
861848
"default": "",
@@ -879,11 +866,6 @@
879866
"uniqueItems": true,
880867
"description": "Opt into experimental features."
881868
},
882-
"databricks.python.envFile": {
883-
"type": "string",
884-
"default": "",
885-
"description": "Similar to python.envFile. Absolute path to a file containing environment variable definitions."
886-
},
887869
"databricks.wsfs.rearrangeCells": {
888870
"type": "boolean",
889871
"default": true,

‎packages/databricks-vscode/src/cli/CliWrapper.ts‎

Lines changed: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,13 @@ import {
44
execFile as execFileCb,
55
spawn,
66
} from "child_process";
7-
import {ExtensionContext, window, Uri, commands} from "vscode";
7+
import {
8+
ExtensionContext,
9+
window,
10+
Uri,
11+
commands,
12+
CancellationToken,
13+
} from "vscode";
814
import {SyncDestinationMapper} from "../sync/SyncDestination";
915
import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
1016
import {promisify} from "node:util";
@@ -39,10 +45,31 @@ function getEscapedCommandAndAgrs(
3945
return {cmd, args, options};
4046
}
4147

48+
export async function cancellableExecFile(
49+
file: string,
50+
args: string[],
51+
options: Omit<SpawnOptionsWithoutStdio, "signal"> = {},
52+
cancellationToken?: CancellationToken
53+
): Promise<{
54+
stdout: string;
55+
stderr: string;
56+
}> {
57+
const abortController = new AbortController();
58+
cancellationToken?.onCancellationRequested(() => abortController.abort());
59+
const signal = abortController.signal;
60+
61+
const res = await promisify(execFileCb)(file, args, {
62+
...options,
63+
signal,
64+
});
65+
return {stdout: res.stdout.toString(), stderr: res.stderr.toString()};
66+
}
67+
4268
export const execFile = async (
4369
file: string,
4470
args: string[],
45-
options: any = {}
71+
options: Omit<SpawnOptionsWithoutStdio, "signal"> = {},
72+
cancellationToken?: CancellationToken
4673
): Promise<{
4774
stdout: string;
4875
stderr: string;
@@ -52,8 +79,13 @@ export const execFile = async (
5279
args: escapedArgs,
5380
options: escapedOptions,
5481
} = getEscapedCommandAndAgrs(file, args, options);
55-
const res = await promisify(execFileCb)(cmd, escapedArgs, escapedOptions);
56-
return {stdout: res.stdout.toString(), stderr: res.stderr.toString()};
82+
83+
return await cancellableExecFile(
84+
cmd,
85+
escapedArgs,
86+
escapedOptions,
87+
cancellationToken
88+
);
5789
};
5890

5991
export interface Command {

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

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ export class ConnectionManager implements Disposable {
4141
private _state: ConnectionState = "DISCONNECTED";
4242
private loginLogoutMutex: Mutex = new Mutex();
4343
private savedAuthMutex: Mutex = new Mutex();
44+
private configureLoginMutex: Mutex = new Mutex();
4445

4546
private _workspaceClient?: WorkspaceClient;
4647
private _syncDestinationMapper?: SyncDestinationMapper;
@@ -373,21 +374,29 @@ export class ConnectionManager implements Disposable {
373374
popup: {prefix: "Can't configure workspace. "},
374375
})
375376
async configureLogin(source: ManualLoginSource) {
376-
const recordEvent = this.telemetry.start(Events.MANUAL_LOGIN);
377-
try {
378-
const authProvider = await LoginWizard.run(
379-
this.cli,
380-
this.configModel.target,
381-
await this.configModel.get("host")
377+
if (this.configureLoginMutex.locked) {
378+
window.showErrorMessage(
379+
"Databricks: Already configuring workspace"
382380
);
383-
if (authProvider) {
384-
await this.connect(authProvider);
385-
}
386-
recordEvent({success: this.state === "CONNECTED", source});
387-
} catch (e) {
388-
recordEvent({success: false, source});
389-
throw e;
381+
return;
390382
}
383+
await this.configureLoginMutex.synchronise(async () => {
384+
const recordEvent = this.telemetry.start(Events.MANUAL_LOGIN);
385+
try {
386+
const authProvider = await LoginWizard.run(
387+
this.cli,
388+
this.configModel.target,
389+
await this.configModel.get("host")
390+
);
391+
if (authProvider) {
392+
await this.connect(authProvider);
393+
}
394+
recordEvent({success: this.state === "CONNECTED", source});
395+
} catch (e) {
396+
recordEvent({success: false, source});
397+
throw e;
398+
}
399+
});
391400
}
392401

393402
@onError({

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -122,10 +122,10 @@ export class LoginWizard {
122122
return;
123123
}
124124

125-
const choice = await window.showErrorMessage(
125+
const choice = await window.showInformationMessage(
126126
`Authentication using ${authProvider.describe()} failed.`,
127-
"Select a different auth method",
128-
"Cancel"
127+
{modal: true},
128+
"Select a different auth method"
129129
);
130130
if (choice === "Select a different auth method") {
131131
//Show input again with the select auth step.

‎packages/databricks-vscode/src/configuration/auth/AuthProvider.ts‎

Lines changed: 58 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ import {
55
WorkspaceClient,
66
logging,
77
} from "@databricks/databricks-sdk";
8-
import {ProgressLocation, window} from "vscode";
8+
import {CancellationToken, ProgressLocation, window} from "vscode";
99
import {normalizeHost} from "../../utils/urlUtils";
1010
import {workspaceConfigs} from "../../vscode-objs/WorkspaceConfigs";
1111

@@ -64,32 +64,62 @@ export abstract class AuthProvider {
6464
* This function should not throw an error and each implementing class must
6565
* handle it's own error messages and retry loops.
6666
*/
67-
protected abstract _check(): Promise<boolean>;
67+
protected abstract _check(token?: CancellationToken): Promise<boolean>;
6868

69-
public async check(force = false, showProgress = true): Promise<boolean> {
69+
public async check(
70+
force = false,
71+
showProgress = true,
72+
cancellationToken?: CancellationToken
73+
): Promise<boolean> {
7074
if (force) {
7175
this.checked = false;
7276
}
7377
if (this.checked) {
7478
return true;
7579
}
7680

77-
const checkFn = async () => {
78-
this.checked = await this._check();
81+
const checkFn = async (token?: CancellationToken) => {
82+
this.checked = await this._check(token);
7983
};
8084

8185
if (!showProgress) {
82-
await checkFn();
86+
await checkFn(cancellationToken);
8387
return this.checked;
8488
}
8589

90+
let cancellationRequested = false;
91+
let task: Promise<void> = Promise.resolve();
8692
await window.withProgress(
8793
{
8894
location: ProgressLocation.Notification,
8995
title: `Databricks: Logging in using ${this.describe()}`,
96+
cancellable: true,
9097
},
91-
async () => await checkFn()
98+
async (progress, token) => {
99+
task = checkFn(token);
100+
await Promise.race([
101+
task,
102+
new Promise((resolve) =>
103+
token.onCancellationRequested(resolve)
104+
),
105+
]);
106+
cancellationRequested = token.isCancellationRequested;
107+
}
92108
);
109+
if (cancellationRequested) {
110+
await window.withProgress(
111+
{
112+
location: ProgressLocation.Notification,
113+
title: `Databricks: Cancelling login using ${this.describe()}`,
114+
},
115+
async () => {
116+
await task;
117+
}
118+
);
119+
window.showErrorMessage("Databricks: Login cancelled");
120+
this.checked = false;
121+
}
122+
93123
if (this.checked) {
94124
window.showInformationMessage(
95125
`Databricks: Successfully logged in using ${this.describe()}`
@@ -230,8 +260,8 @@ export class ProfileAuthProvider extends AuthProvider {
230260
return ProfileAuthProvider.getSdkConfig(this.profile);
231261
}
232262

233-
protected async _check() {
234-
while (true) {
263+
protected async _check(cancellationToken?: CancellationToken) {
264+
while (cancellationToken?.isCancellationRequested !== true) {
235265
try {
236266
const sdkConfig = await this.getSdkConfig();
237267
const authProvider = AuthProvider.fromSdkConfig(
@@ -245,7 +275,11 @@ export class ProfileAuthProvider extends AuthProvider {
245275
return true;
246276
}
247277

248-
return await authProvider.check(false, false);
278+
return await authProvider.check(
279+
false,
280+
false,
281+
cancellationToken
282+
);
249283
} catch (e) {
250284
let message: string = `Can't login with config profile ${this.profile}`;
251285
if (e instanceof Error) {
@@ -266,6 +300,7 @@ export class ProfileAuthProvider extends AuthProvider {
266300
return false;
267301
}
268302
}
303+
return false;
269304
}
270305
}
271306

@@ -312,9 +347,11 @@ export class DatabricksCliAuthProvider extends AuthProvider {
312347
};
313348
}
314349

315-
protected async _check(): Promise<boolean> {
350+
protected async _check(
351+
cancellationToken?: CancellationToken
352+
): Promise<boolean> {
316353
const databricksCliCheck = new DatabricksCliCheck(this);
317-
return databricksCliCheck.check();
354+
return databricksCliCheck.check(cancellationToken);
318355
}
319356
}
320357

@@ -378,9 +415,11 @@ export class AzureCliAuthProvider extends AuthProvider {
378415
return envVars;
379416
}
380417

381-
protected async _check(): Promise<boolean> {
418+
protected async _check(
419+
cancellationToken?: CancellationToken
420+
): Promise<boolean> {
382421
const cliCheck = new AzureCliCheck(this);
383-
const result = await cliCheck.check();
422+
const result = await cliCheck.check(cancellationToken);
384423
this._tenantId = cliCheck.tenantId;
385424
this._appId = cliCheck.azureLoginAppId;
386425
return result;
@@ -419,8 +458,10 @@ export class PersonalAccessTokenAuthProvider extends AuthProvider {
419458
token: this.token,
420459
};
421460
}
422-
protected async _check(): Promise<boolean> {
423-
while (true) {
461+
protected async _check(
462+
cancellationToken?: CancellationToken
463+
): Promise<boolean> {
464+
while (cancellationToken?.isCancellationRequested !== true) {
424465
try {
425466
const workspaceClient = await this.getWorkspaceClient();
426467
await workspaceClient.currentUser.me();
@@ -445,6 +486,7 @@ export class PersonalAccessTokenAuthProvider extends AuthProvider {
445486
return false;
446487
}
447488
}
489+
return false;
448490
}
449491
protected _getSdkConfig(): Config {
450492
return new Config({

0 commit comments

Comments
 (0)