Skip to content

Commit 2cff68e

Browse files
kartikgupta-dbkartikgupta-db
authored
Use packaged CLI if a cli path is not specified in .databrickscfg profile (#1150)
For profiles which do not have `databricks_cli_path` specified, we should use the packaged CLI. ## Changes * Authproviders always try to resolve the sdk config. If no CLI is found after resolution (after loaders have loaded the profile data from .databrickscfg), we set the path to builtin databricks cli in the config. * Also, we now respect the `databricks_cli_path` config from profile. Earlier, we were always defaulting to packaged cli for `DatabricksCliAuthProvider`. https://github.com/databricks/databricks-vscode/assets/88345179/79213f70-e85a-42a2-98c9-531e546c0ab6 ## Tests <!-- How is this tested? -->
1 parent 450b002 commit 2cff68e

File tree

7 files changed

+65
-34
lines changed

7 files changed

+65
-34
lines changed

‎packages/databricks-vscode/src/bundle/run/JobRunStatus.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ export class JobRunStatus extends BundleRunStatus {
4242
throw new Error("No run id found");
4343
}
4444

45-
const client = this.authProvider.getWorkspaceClient();
45+
const client = await this.authProvider.getWorkspaceClient();
4646
try {
4747
this.runState = "running";
4848
await (
@@ -70,7 +70,7 @@ export class JobRunStatus extends BundleRunStatus {
7070
return;
7171
}
7272

73-
const client = this.authProvider.getWorkspaceClient();
73+
const client = await this.authProvider.getWorkspaceClient();
7474
await (
7575
await client.jobs.cancelRun({run_id: parseInt(this.runId)})
7676
).wait();

‎packages/databricks-vscode/src/bundle/run/PipelineRunStatus.ts‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ export class PipelineRunStatus extends BundleRunStatus {
4949
throw new Error("No update id");
5050
}
5151

52-
const client = this.authProvider.getWorkspaceClient();
52+
const client = await this.authProvider.getWorkspaceClient();
5353
this.runState = "running";
5454

5555
this.interval = setInterval(async () => {
@@ -99,7 +99,7 @@ export class PipelineRunStatus extends BundleRunStatus {
9999
return;
100100
}
101101

102-
const client = this.authProvider.getWorkspaceClient();
102+
const client = await this.authProvider.getWorkspaceClient();
103103
const update = await client.pipelines.getUpdate({
104104
pipeline_id: this.pipelineId,
105105
update_id: this.runId,

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -95,7 +95,11 @@ export class ConnectionCommands implements Disposable {
9595
throw new Error("Must login first");
9696
}
9797
const hostUrl = normalizeHost(host);
98-
const provider = new PersonalAccessTokenAuthProvider(hostUrl, token);
98+
const provider = new PersonalAccessTokenAuthProvider(
99+
hostUrl,
100+
token,
101+
this.cli
102+
);
99103
await saveNewProfile(name, provider, this.cli);
100104
}
101105

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -300,11 +300,11 @@ export class ConnectionManager implements Disposable {
300300
@Mutex.synchronise("loginLogoutMutex")
301301
private async _connect(authProvider: AuthProvider) {
302302
this.updateState("CONNECTING");
303+
this._workspaceClient = await authProvider.getWorkspaceClient();
303304
this._databricksWorkspace = await DatabricksWorkspace.load(
304-
authProvider.getWorkspaceClient(),
305+
this._workspaceClient,
305306
authProvider
306307
);
307-
this._workspaceClient = authProvider.getWorkspaceClient();
308308
await this.configModel.set(
309309
"authProfile",
310310
authProvider.toJSON().profile as string | undefined

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

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -304,12 +304,16 @@ export class LoginWizard {
304304
| PersonalAccessTokenAuthProvider;
305305
switch (pick.authType) {
306306
case "azure-cli":
307-
authProvider = new AzureCliAuthProvider(this.state.host!);
307+
authProvider = new AzureCliAuthProvider(
308+
this.state.host!,
309+
this.cliWrapper
310+
);
308311
break;
309312

310313
case "databricks-cli":
311314
authProvider = new DatabricksCliAuthProvider(
312315
this.state.host!,
316+
this.cliWrapper.cliPath,
313317
this.cliWrapper
314318
);
315319
break;
@@ -329,7 +333,8 @@ export class LoginWizard {
329333
}
330334
authProvider = new PersonalAccessTokenAuthProvider(
331335
this.state.host!,
332-
token
336+
token,
337+
this.cliWrapper
333338
);
334339
}
335340
break;

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

Lines changed: 45 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ export abstract class AuthProvider {
3030
constructor(
3131
private readonly _host: URL,
3232
private readonly _authType: AuthType,
33+
private readonly _cli: CliWrapper,
3334
private checked: boolean = false
3435
) {}
3536

@@ -49,8 +50,8 @@ export abstract class AuthProvider {
4950
abstract toEnv(): Record<string, string>;
5051
abstract toIni(): Record<string, string | undefined> | undefined;
5152

52-
getWorkspaceClient(): WorkspaceClient {
53-
const config = this.getSdkConfig();
53+
async getWorkspaceClient(): Promise<WorkspaceClient> {
54+
const config = await this.getSdkConfig();
5455

5556
return new WorkspaceClient(config, {
5657
product: "databricks-vscode",
@@ -97,7 +98,17 @@ export abstract class AuthProvider {
9798

9899
return this.checked;
99100
}
100-
protected abstract getSdkConfig(): Config;
101+
async getSdkConfig(): Promise<Config> {
102+
const config = this._getSdkConfig();
103+
await config.ensureResolved();
104+
if (config.databricksCliPath === undefined) {
105+
config.databricksCliPath = this._cli.cliPath;
106+
}
107+
108+
return config;
109+
}
110+
111+
protected abstract _getSdkConfig(): Config;
101112

102113
static fromJSON(json: Record<string, any>, cli: CliWrapper): AuthProvider {
103114
const host =
@@ -121,7 +132,11 @@ export abstract class AuthProvider {
121132
);
122133

123134
case "databricks-cli":
124-
return new DatabricksCliAuthProvider(host, cli);
135+
return new DatabricksCliAuthProvider(
136+
host,
137+
json.databricksPath ?? cli.cliPath,
138+
cli
139+
);
125140

126141
case "profile":
127142
if (!json.profile) {
@@ -144,12 +159,17 @@ export abstract class AuthProvider {
144159
case "azure-cli":
145160
return new AzureCliAuthProvider(
146161
host,
162+
cli,
147163
config.azureTenantId,
148164
config.azureLoginAppId
149165
);
150166

151167
case "databricks-cli":
152-
return new DatabricksCliAuthProvider(host, cli);
168+
return new DatabricksCliAuthProvider(
169+
host,
170+
config.databricksCliPath ?? cli.cliPath,
171+
cli
172+
);
153173

154174
default:
155175
if (config.profile) {
@@ -172,7 +192,7 @@ export class ProfileAuthProvider extends AuthProvider {
172192
private readonly cli: CliWrapper,
173193
checked = false
174194
) {
175-
super(host, "profile", checked);
195+
super(host, "profile", cli, checked);
176196
}
177197

178198
describe(): string {
@@ -198,33 +218,33 @@ export class ProfileAuthProvider extends AuthProvider {
198218
return undefined;
199219
}
200220

201-
public static getSdkConfig(profile: string): Config {
221+
private static getSdkConfig(profile: string): Config {
202222
return new Config({
203223
profile: profile,
204224
configFile: workspaceConfigs.databrickscfgLocation,
205225
env: {},
206226
});
207227
}
208228

209-
getSdkConfig(): Config {
229+
protected _getSdkConfig(): Config {
210230
return ProfileAuthProvider.getSdkConfig(this.profile);
211231
}
212232

213233
protected async _check() {
214234
while (true) {
215235
try {
216-
const sdkConfig = this.getSdkConfig();
217-
await sdkConfig.ensureResolved();
236+
const sdkConfig = await this.getSdkConfig();
218237
const authProvider = AuthProvider.fromSdkConfig(
219238
sdkConfig,
220239
this.cli
221240
);
222241

223242
if (authProvider instanceof ProfileAuthProvider) {
224-
const workspaceClient = this.getWorkspaceClient();
243+
const workspaceClient = await this.getWorkspaceClient();
225244
await workspaceClient.currentUser.me();
226245
return true;
227246
}
247+
228248
return await authProvider.check(false, false);
229249
} catch (e) {
230250
let message: string = `Can't login with config profile ${this.profile}`;
@@ -252,9 +272,10 @@ export class ProfileAuthProvider extends AuthProvider {
252272
export class DatabricksCliAuthProvider extends AuthProvider {
253273
constructor(
254274
host: URL,
255-
readonly cli: CliWrapper
275+
readonly cliPath: string,
276+
cli: CliWrapper
256277
) {
257-
super(host, "databricks-cli");
278+
super(host, "databricks-cli", cli);
258279
}
259280

260281
describe(): string {
@@ -265,15 +286,15 @@ export class DatabricksCliAuthProvider extends AuthProvider {
265286
return {
266287
host: this.host.toString(),
267288
authType: this.authType,
268-
databricksPath: this.cli.cliPath,
289+
databricksPath: this.cliPath,
269290
};
270291
}
271292

272-
getSdkConfig(): Config {
293+
_getSdkConfig(): Config {
273294
return new Config({
274295
host: this.host.toString(),
275296
authType: "databricks-cli",
276-
databricksCliPath: this.cli.cliPath,
297+
databricksCliPath: this.cliPath,
277298
});
278299
}
279300

@@ -301,8 +322,8 @@ export class AzureCliAuthProvider extends AuthProvider {
301322
private _tenantId?: string;
302323
private _appId?: string;
303324

304-
constructor(host: URL, tenantId?: string, appId?: string) {
305-
super(host, "azure-cli");
325+
constructor(host: URL, cli: CliWrapper, tenantId?: string, appId?: string) {
326+
super(host, "azure-cli", cli);
306327

307328
this._tenantId = tenantId;
308329
this._appId = appId;
@@ -329,7 +350,7 @@ export class AzureCliAuthProvider extends AuthProvider {
329350
};
330351
}
331352

332-
getSdkConfig(): Config {
353+
_getSdkConfig(): Config {
333354
return new Config({
334355
host: this.host.toString(),
335356
authType: "azure-cli",
@@ -369,9 +390,10 @@ export class AzureCliAuthProvider extends AuthProvider {
369390
export class PersonalAccessTokenAuthProvider extends AuthProvider {
370391
constructor(
371392
host: URL,
372-
private readonly token: string
393+
private readonly token: string,
394+
cli: CliWrapper
373395
) {
374-
super(host, "pat");
396+
super(host, "pat", cli);
375397
}
376398

377399
describe(): string {
@@ -400,7 +422,7 @@ export class PersonalAccessTokenAuthProvider extends AuthProvider {
400422
protected async _check(): Promise<boolean> {
401423
while (true) {
402424
try {
403-
const workspaceClient = this.getWorkspaceClient();
425+
const workspaceClient = await this.getWorkspaceClient();
404426
await workspaceClient.currentUser.me();
405427
return true;
406428
} catch (e) {
@@ -424,7 +446,7 @@ export class PersonalAccessTokenAuthProvider extends AuthProvider {
424446
}
425447
}
426448
}
427-
protected getSdkConfig(): Config {
449+
protected _getSdkConfig(): Config {
428450
return new Config({
429451
host: this.host.toString(),
430452
authType: "pat",

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ export class DatabricksCliCheck implements Disposable {
7373
{
7474
host: this.authProvider.host.toString(),
7575
authType: "databricks-cli",
76-
databricksCliPath: this.authProvider.cli.cliPath,
76+
databricksCliPath: this.authProvider.cliPath,
7777
},
7878
{
7979
product: "databricks-vscode",
@@ -92,7 +92,7 @@ export class DatabricksCliCheck implements Disposable {
9292

9393
private async login(): Promise<void> {
9494
try {
95-
await ExecUtils.execFile(this.authProvider.cli.cliPath, [
95+
await ExecUtils.execFile(this.authProvider.cliPath, [
9696
"auth",
9797
"login",
9898
"--host",

0 commit comments

Comments
 (0)