Skip to content

Commit 8471909

Browse files
kartikgupta-dbkartikgupta-db
authored
Listen to changes to "host" in databricks.yml to retrigger login (#1057)
## Changes Auth resolution does not happen again when pre validate configs change (host, mode etc). We listen to specific changes to host to retrigger login. Also, check if the current host is not the same as the one in the saved profile (in case someone changes host in the bundle.yml). ## Tests <!-- How is this tested? -->
1 parent ab4e603 commit 8471909

File tree

6 files changed

+56
-53
lines changed

6 files changed

+56
-53
lines changed

‎packages/databricks-vscode/src/bundle/models/BundleRemoteStateModel.ts‎

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {logging} from "@databricks/databricks-sdk";
1111
import {Loggers} from "../../logger";
1212
import {Context, context} from "@databricks/databricks-sdk";
1313
import {BundleValidateModel} from "./BundleValidateModel";
14+
import {withOnErrorHandler} from "../../utils/onErrorDecorator";
1415

1516
/* eslint-disable @typescript-eslint/naming-convention */
1617
export type BundleResourceModifiedStatus = "created" | "deleted" | "updated";
@@ -43,9 +44,14 @@ export class BundleRemoteStateModel extends BaseModelWithStateCache<BundleRemote
4344
private readonly bundleValidateModel: BundleValidateModel
4445
) {
4546
super();
46-
this.bundleValidateModel.onDidChange(async () => {
47-
this.refresh();
48-
});
47+
this.bundleValidateModel.onDidChange(
48+
withOnErrorHandler(
49+
async () => {
50+
this.refresh();
51+
},
52+
{log: true, throw: false}
53+
)
54+
);
4955
}
5056

5157
public async refresh() {

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

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -194,10 +194,7 @@ export class ConfigureWorkspaceWizard {
194194
if (pick.profile !== undefined) {
195195
// TODO: If the profile has an auth type with deeper support in the Extension, (azure-cli, databricks-cli),
196196
// then run the necessary checks to validate the profile and help user fix any errors.
197-
const authProvider = new ProfileAuthProvider(
198-
this.state.host!,
199-
pick.profile
200-
);
197+
const authProvider = await ProfileAuthProvider.from(pick.profile);
201198
const checkResult = await this.checkAuthProvider(
202199
authProvider,
203200
`profile '${pick.profile}'`
@@ -303,8 +300,7 @@ export class ConfigureWorkspaceWizard {
303300
// Write the new profile to .databrickscfg
304301
await writeFile(configFilePath, ini.stringify(iniFile));
305302

306-
this.state.authProvider = new ProfileAuthProvider(
307-
this.state.host!,
303+
this.state.authProvider = await ProfileAuthProvider.from(
308304
profileName,
309305
true
310306
);

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

Lines changed: 31 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -123,22 +123,31 @@ export class ConnectionManager implements Disposable {
123123
return this._metadataService.url;
124124
}
125125

126+
private _host: URL | undefined;
127+
126128
public async init() {
127-
await this.loginWithSavedAuth();
128-
129-
this.disposables.push(
130-
this.configModel.onDidChangeKey("remoteRootPath")(
131-
this.updateSyncDestinationMapper,
132-
this
133-
),
134-
this.configModel.onDidChangeKey("clusterId")(
135-
this.updateClusterManager,
136-
this
137-
),
138-
this.configModel.onDidChangeTarget(this.loginWithSavedAuth, this)
139-
// TODO: We don't react to changes in authProfile from the config model. We instead listen to changes
140-
// in individual models to react to (Overrides and BundlePreValidate)
141-
);
129+
try {
130+
await this.loginWithSavedAuth();
131+
} finally {
132+
this.disposables.push(
133+
this.configModel.onDidChangeKey("remoteRootPath")(
134+
this.updateSyncDestinationMapper,
135+
this
136+
),
137+
this.configModel.onDidChangeKey("clusterId")(
138+
this.updateClusterManager,
139+
this
140+
),
141+
// Don't listen to target change for logging in. Explictly listen for changes in the keys we care about.
142+
// We don't have to listen to changes in authProfile as it's set by the login method and we don't respect other
143+
// user changes.
144+
// TODO: start listening to changes in authParams
145+
this.configModel.onDidChangeKey("host")(
146+
this.loginWithSavedAuth,
147+
this
148+
)
149+
);
150+
}
142151
}
143152

144153
get state(): ConnectionState {
@@ -198,15 +207,16 @@ export class ConnectionManager implements Disposable {
198207
const savedProfile = (await this.configModel.get("overrides"))
199208
?.authProfile;
200209
if (savedProfile !== undefined) {
201-
const authProvider = new ProfileAuthProvider(host, savedProfile);
202-
if (await authProvider.check()) {
210+
const authProvider = await ProfileAuthProvider.from(savedProfile);
211+
if (authProvider.host === host && (await authProvider.check())) {
203212
return authProvider;
204213
}
205214
}
206215

207216
// Try to load any parameters that are hard coded in the bundle
208-
const bundleAuthParams =
209-
await this.configModel.get("preValidateConfig");
217+
const bundleAuthParams = await this.configModel.get(
218+
"preValidateConfig"
219+
);
210220
if (bundleAuthParams?.authParams !== undefined) {
211221
throw new Error("Bundle auth params not implemented");
212222
}
@@ -218,7 +228,7 @@ export class ConnectionManager implements Disposable {
218228
if (profiles.length !== 1) {
219229
return;
220230
}
221-
const authProvider = new ProfileAuthProvider(host, profiles[0].name);
231+
const authProvider = await ProfileAuthProvider.from(profiles[0].name);
222232
if (await authProvider.check()) {
223233
return authProvider;
224234
}
@@ -239,7 +249,7 @@ export class ConnectionManager implements Disposable {
239249
e
240250
);
241251
if (e instanceof Error) {
242-
await window.showWarningMessage(
252+
window.showWarningMessage(
243253
`Can't connect to the workspace: "${e.message}"."`
244254
);
245255
}

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

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -230,10 +230,7 @@ export class LoginWizard {
230230
if (pick.profile !== undefined) {
231231
// We assume that the profile is setup correctly (even for the auth types that have a deeper integration with vscode such as azure-cli).
232232
// To fix errors, users can create a new profile.
233-
const authProvider = new ProfileAuthProvider(
234-
this.state.host!,
235-
pick.profile
236-
);
233+
const authProvider = await ProfileAuthProvider.from(pick.profile);
237234
const checkResult = await this.checkAuthProvider(
238235
authProvider,
239236
`profile '${pick.profile}'`,
@@ -395,7 +392,7 @@ export async function saveNewProfile(
395392
// Write the new profile to .databrickscfg
396393
await appendFile(configFilePath, finalStr);
397394

398-
return new ProfileAuthProvider(authProvider.host, profileName, true);
395+
return await ProfileAuthProvider.from(profileName, true);
399396
}
400397

401398
function humaniseSdkAuthType(sdkAuthType: string) {

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

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,11 @@ export abstract class AuthProvider {
125125
}
126126

127127
export class ProfileAuthProvider extends AuthProvider {
128+
static async from(profile: string, checked = false) {
129+
const host = await ProfileAuthProvider._getSdkConfig(profile).getHost();
130+
return new ProfileAuthProvider(host, profile, checked);
131+
}
132+
128133
constructor(
129134
host: URL,
130135
private readonly profile: string,
@@ -156,16 +161,20 @@ export class ProfileAuthProvider extends AuthProvider {
156161
return undefined;
157162
}
158163

159-
getSdkConfig(): Config {
164+
private static _getSdkConfig(profile: string): Config {
160165
return new Config({
161-
profile: this.profile,
166+
profile: profile,
162167
configFile:
163168
workspaceConfigs.databrickscfgLocation ??
164169
process.env.DATABRICKS_CONFIG_FILE,
165170
env: {},
166171
});
167172
}
168173

174+
getSdkConfig(): Config {
175+
return ProfileAuthProvider._getSdkConfig(this.profile);
176+
}
177+
169178
protected async _check() {
170179
while (true) {
171180
try {

‎packages/databricks-vscode/src/file-managers/ProjectConfigFile.ts‎

Lines changed: 1 addition & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,6 @@ import {
55
ProfileAuthProvider,
66
} from "../configuration/auth/AuthProvider";
77
import {Uri} from "vscode";
8-
import {Config} from "@databricks/databricks-sdk";
9-
import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
108

119
export interface ProjectConfig {
1210
authProvider: AuthProvider;
@@ -48,20 +46,7 @@ export class ProjectConfigFile {
4846
}
4947

5048
static async importOldConfig(config: any): Promise<ProfileAuthProvider> {
51-
const sdkConfig = new Config({
52-
profile: config.profile,
53-
configFile:
54-
workspaceConfigs.databrickscfgLocation ??
55-
process.env.DATABRICKS_CONFIG_FILE,
56-
env: {},
57-
});
58-
59-
await sdkConfig.ensureResolved();
60-
61-
return new ProfileAuthProvider(
62-
new URL(sdkConfig.host!),
63-
sdkConfig.profile!
64-
);
49+
return await ProfileAuthProvider.from(config.profile);
6550
}
6651

6752
static async load(

0 commit comments

Comments
 (0)