Skip to content

Commit c6e9b0e

Browse files
kartikgupta-dbkartikgupta-db
authored
Pass CLI path to all auth providers and login only once during bundle init (#1114)
## Changes * Store overrides in a per target file instead of state store * Have a root overrides file for cases when we are not sure of the target. * Whenever a target is selected, and it does not already have an override file, it will MOVE the root override file to initialise it's own state json file. ## Tests <!-- How is this tested? -->
1 parent 12f306a commit c6e9b0e

File tree

13 files changed

+178
-65
lines changed

13 files changed

+178
-65
lines changed

‎packages/databricks-vscode/src/bundle/BundleInitWizard.ts‎

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,9 +14,13 @@ import {CliWrapper} from "../cli/CliWrapper";
1414
import {getSubProjects} from "./BundleFileSet";
1515
import {tmpdir} from "os";
1616
import {ShellUtils} from "../utils";
17+
import {OverrideableConfigModel} from "../configuration/models/OverrideableConfigModel";
18+
import {writeFile, mkdir} from "fs/promises";
19+
import path from "path";
1720

1821
export async function promptToOpenSubProjects(
19-
projects: {absolute: Uri; relative: Uri}[]
22+
projects: {absolute: Uri; relative: Uri}[],
23+
authProvider?: AuthProvider
2024
) {
2125
type OpenProjectItem = QuickPickItem & {uri?: Uri};
2226
const items: OpenProjectItem[] = projects.map((project) => {
@@ -34,9 +38,21 @@ export async function promptToOpenSubProjects(
3438
title: "Select the project you want to open",
3539
};
3640
const item = await window.showQuickPick<OpenProjectItem>(items, options);
37-
if (!item) {
41+
if (!item?.uri) {
3842
return;
3943
}
44+
45+
if (authProvider?.authType === "profile") {
46+
const rootOverrideFilePath =
47+
OverrideableConfigModel.getRootOverrideFile(item.uri);
48+
await mkdir(path.dirname(rootOverrideFilePath.fsPath), {
49+
recursive: true,
50+
});
51+
await writeFile(
52+
rootOverrideFilePath.fsPath,
53+
JSON.stringify({authProfile: authProvider.toJSON().profile})
54+
);
55+
}
4056
await commands.executeCommand("vscode.openFolder", item.uri);
4157
}
4258

@@ -71,7 +87,7 @@ export class BundleInitWizard {
7187
this.logger.debug(
7288
`Detected ${projects.length} sub projects after the init wizard, prompting to open one`
7389
);
74-
await promptToOpenSubProjects(projects);
90+
await promptToOpenSubProjects(projects, authProvider);
7591
} else {
7692
this.logger.debug(
7793
`No projects detected after the init wizard, showing notification to open a folder manually`

‎packages/databricks-vscode/src/bundle/BundleProjectManager.ts‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -186,7 +186,7 @@ export class BundleProjectManager {
186186
try {
187187
return await ProjectConfigFile.load(
188188
this.workspaceUri.fsPath,
189-
this.cli.cliPath
189+
this.cli
190190
);
191191
} catch (error) {
192192
this.logger.error("Failed to load legacy project config:", error);
@@ -214,7 +214,11 @@ export class BundleProjectManager {
214214
"Creating new profile before bundle migration",
215215
profileName
216216
);
217-
authProvider = await saveNewProfile(profileName, authProvider);
217+
authProvider = await saveNewProfile(
218+
profileName,
219+
authProvider,
220+
this.cli
221+
);
218222
}
219223
await this.migrateProjectJsonToBundle(
220224
authProvider as ProfileAuthProvider,

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ nothing = true
217217
const authProvider = new ProfileAuthProvider(
218218
new URL("https://test.com"),
219219
"PROFILE",
220+
cli,
220221
true
221222
);
222223
const workspaceFolder = Uri.file("/test/123");

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import {ConfigModel} from "./models/ConfigModel";
1515
import {saveNewProfile} from "./LoginWizard";
1616
import {PersonalAccessTokenAuthProvider} from "./auth/AuthProvider";
1717
import {normalizeHost} from "../utils/urlUtils";
18+
import {CliWrapper} from "../cli/CliWrapper";
1819
import {AUTH_TYPE_SWITCH_ID, AUTH_TYPE_LOGIN_ID} from "./ui/AuthTypeComponent";
1920
import {ManualLoginSource} from "../telemetry/constants";
2021

@@ -56,7 +57,8 @@ export class ConnectionCommands implements Disposable {
5657
private wsfsCommands: WorkspaceFsCommands,
5758
private connectionManager: ConnectionManager,
5859
private readonly clusterModel: ClusterModel,
59-
private readonly configModel: ConfigModel
60+
private readonly configModel: ConfigModel,
61+
private readonly cli: CliWrapper
6062
) {}
6163

6264
/**
@@ -94,7 +96,7 @@ export class ConnectionCommands implements Disposable {
9496
}
9597
const hostUrl = normalizeHost(host);
9698
const provider = new PersonalAccessTokenAuthProvider(hostUrl, token);
97-
await saveNewProfile(name, provider);
99+
await saveNewProfile(name, provider, this.cli);
98100
}
99101

100102
/**

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,10 @@ export class ConnectionManager implements Disposable {
226226
const savedProfile = (await this.configModel.get("overrides"))
227227
?.authProfile;
228228
if (savedProfile !== undefined) {
229-
const authProvider = await ProfileAuthProvider.from(savedProfile);
229+
const authProvider = await ProfileAuthProvider.from(
230+
savedProfile,
231+
this.cli
232+
);
230233
if (
231234
authProvider.host.toString() === host.toString() &&
232235
(await authProvider.check())
@@ -249,7 +252,10 @@ export class ConnectionManager implements Disposable {
249252
if (profiles.length !== 1) {
250253
return;
251254
}
252-
const authProvider = await ProfileAuthProvider.from(profiles[0].name);
255+
const authProvider = await ProfileAuthProvider.from(
256+
profiles[0].name,
257+
this.cli
258+
);
253259
if (await authProvider.check()) {
254260
return authProvider;
255261
}

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@ import {Uri} from "vscode";
44
import {ProfileAuthProvider} from "./auth/AuthProvider";
55
import {DatabricksWorkspace} from "./DatabricksWorkspace";
66
import {iam} from "@databricks/databricks-sdk";
7+
import {instance, mock} from "ts-mockito";
8+
import {CliWrapper} from "../cli/CliWrapper";
79

810
describe(__filename, () => {
911
it("create an instance", () => {
@@ -17,7 +19,8 @@ describe(__filename, () => {
1719
} as const;
1820
const authProvider = new ProfileAuthProvider(
1921
new URL("https://fabian.databricks.com"),
20-
"DEFAULT"
22+
"DEFAULT",
23+
instance(mock(CliWrapper))
2124
);
2225
const dbWorkspace: DatabricksWorkspace = new DatabricksWorkspace(
2326
authProvider,

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

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -243,7 +243,10 @@ export class LoginWizard {
243243
}
244244

245245
if (pick.profile !== undefined) {
246-
const authProvider = await ProfileAuthProvider.from(pick.profile);
246+
const authProvider = await ProfileAuthProvider.from(
247+
pick.profile,
248+
this.cliWrapper
249+
);
247250
const nextStep = await this.checkAuthProvider(authProvider, input);
248251
if (nextStep) {
249252
return nextStep;
@@ -307,7 +310,7 @@ export class LoginWizard {
307310
case "databricks-cli":
308311
authProvider = new DatabricksCliAuthProvider(
309312
this.state.host!,
310-
this.cliWrapper.cliPath
313+
this.cliWrapper
311314
);
312315
break;
313316

@@ -343,7 +346,8 @@ export class LoginWizard {
343346

344347
this.state.authProvider = await saveNewProfile(
345348
profileName,
346-
authProvider
349+
authProvider,
350+
this.cliWrapper
347351
);
348352
}
349353

@@ -367,7 +371,8 @@ export class LoginWizard {
367371

368372
export async function saveNewProfile(
369373
profileName: string,
370-
authProvider: AuthProvider
374+
authProvider: AuthProvider,
375+
cli: CliWrapper
371376
) {
372377
const iniData = authProvider.toIni();
373378
if (!iniData) {
@@ -397,7 +402,7 @@ export async function saveNewProfile(
397402
// Write the new profile to .databrickscfg
398403
await appendFile(configFilePath, finalStr);
399404

400-
return await ProfileAuthProvider.from(profileName, true);
405+
return await ProfileAuthProvider.from(profileName, cli, true);
401406
}
402407

403408
function humaniseSdkAuthType(sdkAuthType: string) {

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

Lines changed: 17 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ const extensionVersion = require("../../../package.json")
1616
import {AzureCliCheck} from "./AzureCliCheck";
1717
import {DatabricksCliCheck} from "./DatabricksCliCheck";
1818
import {Loggers} from "../../logger";
19+
import {CliWrapper} from "../../cli/CliWrapper";
1920

2021
// TODO: Resolve this with SDK's AuthType.
2122
export type AuthType =
@@ -98,10 +99,7 @@ export abstract class AuthProvider {
9899
}
99100
protected abstract getSdkConfig(): Config;
100101

101-
static fromJSON(
102-
json: Record<string, any>,
103-
databricksPath: string
104-
): AuthProvider {
102+
static fromJSON(json: Record<string, any>, cli: CliWrapper): AuthProvider {
105103
const host =
106104
json.host instanceof URL
107105
? json.host
@@ -123,20 +121,20 @@ export abstract class AuthProvider {
123121
);
124122

125123
case "databricks-cli":
126-
return new DatabricksCliAuthProvider(host, databricksPath);
124+
return new DatabricksCliAuthProvider(host, cli);
127125

128126
case "profile":
129127
if (!json.profile) {
130128
throw new Error("Missing profile");
131129
}
132-
return new ProfileAuthProvider(host, json.profile);
130+
return new ProfileAuthProvider(host, json.profile, cli);
133131

134132
default:
135133
throw new Error(`Unknown auth type: ${json.authType}`);
136134
}
137135
}
138136

139-
static fromSdkConfig(config: Config): AuthProvider {
137+
static fromSdkConfig(config: Config, cli: CliWrapper): AuthProvider {
140138
if (!config.host) {
141139
throw new Error("Missing host");
142140
}
@@ -151,33 +149,27 @@ export abstract class AuthProvider {
151149
);
152150

153151
case "databricks-cli":
154-
if (!config.databricksCliPath) {
155-
throw new Error("Missing path for databricks-cli");
156-
}
157-
158-
return new DatabricksCliAuthProvider(
159-
host,
160-
config.databricksCliPath
161-
);
152+
return new DatabricksCliAuthProvider(host, cli);
162153

163154
default:
164155
if (config.profile) {
165-
return new ProfileAuthProvider(host, config.profile);
156+
return new ProfileAuthProvider(host, config.profile, cli);
166157
}
167158
throw new Error(`Unknown auth type: ${config.authType}`);
168159
}
169160
}
170161
}
171162

172163
export class ProfileAuthProvider extends AuthProvider {
173-
static async from(profile: string, checked = false) {
164+
static async from(profile: string, cli: CliWrapper, checked = false) {
174165
const host = await ProfileAuthProvider.getSdkConfig(profile).getHost();
175-
return new ProfileAuthProvider(host, profile, checked);
166+
return new ProfileAuthProvider(host, profile, cli, checked);
176167
}
177168

178169
constructor(
179170
host: URL,
180171
private readonly profile: string,
172+
private readonly cli: CliWrapper,
181173
checked = false
182174
) {
183175
super(host, "profile", checked);
@@ -223,7 +215,10 @@ export class ProfileAuthProvider extends AuthProvider {
223215
try {
224216
const sdkConfig = this.getSdkConfig();
225217
await sdkConfig.ensureResolved();
226-
const authProvider = AuthProvider.fromSdkConfig(sdkConfig);
218+
const authProvider = AuthProvider.fromSdkConfig(
219+
sdkConfig,
220+
this.cli
221+
);
227222

228223
if (authProvider instanceof ProfileAuthProvider) {
229224
const workspaceClient = this.getWorkspaceClient();
@@ -257,7 +252,7 @@ export class ProfileAuthProvider extends AuthProvider {
257252
export class DatabricksCliAuthProvider extends AuthProvider {
258253
constructor(
259254
host: URL,
260-
readonly databricksPath: string
255+
readonly cli: CliWrapper
261256
) {
262257
super(host, "databricks-cli");
263258
}
@@ -270,15 +265,15 @@ export class DatabricksCliAuthProvider extends AuthProvider {
270265
return {
271266
host: this.host.toString(),
272267
authType: this.authType,
273-
databricksPath: this.databricksPath,
268+
databricksPath: this.cli.cliPath,
274269
};
275270
}
276271

277272
getSdkConfig(): Config {
278273
return new Config({
279274
host: this.host.toString(),
280275
authType: "databricks-cli",
281-
databricksCliPath: this.databricksPath,
276+
databricksCliPath: this.cli.cliPath,
282277
});
283278
}
284279

@@ -293,7 +288,6 @@ export class DatabricksCliAuthProvider extends AuthProvider {
293288
return {
294289
host: this.host.toString(),
295290
auth_type: "databricks-cli",
296-
databricks_cli_path: this.databricksPath,
297291
};
298292
}
299293

‎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.databricksPath,
76+
databricksCliPath: this.authProvider.cli.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.databricksPath, [
95+
await ExecUtils.execFile(this.authProvider.cli.cliPath, [
9696
"auth",
9797
"login",
9898
"--host",

0 commit comments

Comments
 (0)