Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 19 additions & 3 deletions packages/databricks-vscode/src/bundle/BundleInitWizard.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,13 @@ import {CliWrapper} from "../cli/CliWrapper";
import {getSubProjects} from "./BundleFileSet";
import {tmpdir} from "os";
import {ShellUtils} from "../utils";
import {OverrideableConfigModel} from "../configuration/models/OverrideableConfigModel";
import {writeFile, mkdir} from "fs/promises";
import path from "path";

export async function promptToOpenSubProjects(
projects: {absolute: Uri; relative: Uri}[]
projects: {absolute: Uri; relative: Uri}[],
authProvider?: AuthProvider
) {
type OpenProjectItem = QuickPickItem & {uri?: Uri};
const items: OpenProjectItem[] = projects.map((project) => {
Expand All @@ -34,9 +38,21 @@ export async function promptToOpenSubProjects(
title: "Select the project you want to open",
};
const item = await window.showQuickPick<OpenProjectItem>(items, options);
if (!item) {
if (!item?.uri) {
return;
}

if (authProvider?.authType === "profile") {
const rootOverrideFilePath =
OverrideableConfigModel.getRootOverrideFile(item.uri);
await mkdir(path.dirname(rootOverrideFilePath.fsPath), {
recursive: true,
});
await writeFile(
rootOverrideFilePath.fsPath,
JSON.stringify({authProfile: authProvider.toJSON().profile})
);
}
Comment on lines +45 to +55
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe move this logic to the OverrideableModel?

Although I'm not even sure we need a concept of the "root" override file (at least not yet). In the "normal" scenario with only one profile per host we won't really need this override, right? After opening that folder we will log-in automatically. Having all this logic around root override seems like a lot of effort for not much benefit, unless we already plan to use it for a bunch of other things.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't have a way to enforce 1 profile per host (and should not because we can have multiple auth methods). A lot of bricksfooders ended up creating multiple profiles for the same host when initializing multiple projects. The extension would then be in a login state until they manually reselected their already created profiles.

This allows us to establish some communication method between the pre bundle init extension instance and post bundle init extension instance. Right now it is just this, but in future with env management, I can imagine having to communicate that information, or cluster information as well.

await commands.executeCommand("vscode.openFolder", item.uri);
}

Expand Down Expand Up @@ -71,7 +87,7 @@ export class BundleInitWizard {
this.logger.debug(
`Detected ${projects.length} sub projects after the init wizard, prompting to open one`
);
await promptToOpenSubProjects(projects);
await promptToOpenSubProjects(projects, authProvider);
} else {
this.logger.debug(
`No projects detected after the init wizard, showing notification to open a folder manually`
Expand Down
8 changes: 6 additions & 2 deletions packages/databricks-vscode/src/bundle/BundleProjectManager.ts
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ export class BundleProjectManager {
try {
return await ProjectConfigFile.load(
this.workspaceUri.fsPath,
this.cli.cliPath
this.cli
);
} catch (error) {
this.logger.error("Failed to load legacy project config:", error);
Expand Down Expand Up @@ -214,7 +214,11 @@ export class BundleProjectManager {
"Creating new profile before bundle migration",
profileName
);
authProvider = await saveNewProfile(profileName, authProvider);
authProvider = await saveNewProfile(
profileName,
authProvider,
this.cli
);
}
await this.migrateProjectJsonToBundle(
authProvider as ProfileAuthProvider,
Expand Down
1 change: 1 addition & 0 deletions packages/databricks-vscode/src/cli/CliWrapper.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -217,6 +217,7 @@ nothing = true
const authProvider = new ProfileAuthProvider(
new URL("https://test.com"),
"PROFILE",
cli,
true
);
const workspaceFolder = Uri.file("/test/123");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import {ConfigModel} from "./models/ConfigModel";
import {saveNewProfile} from "./LoginWizard";
import {PersonalAccessTokenAuthProvider} from "./auth/AuthProvider";
import {normalizeHost} from "../utils/urlUtils";
import {CliWrapper} from "../cli/CliWrapper";
import {AUTH_TYPE_SWITCH_ID, AUTH_TYPE_LOGIN_ID} from "./ui/AuthTypeComponent";
import {ManualLoginSource} from "../telemetry/constants";

Expand Down Expand Up @@ -56,7 +57,8 @@ export class ConnectionCommands implements Disposable {
private wsfsCommands: WorkspaceFsCommands,
private connectionManager: ConnectionManager,
private readonly clusterModel: ClusterModel,
private readonly configModel: ConfigModel
private readonly configModel: ConfigModel,
private readonly cli: CliWrapper
) {}

/**
Expand Down Expand Up @@ -94,7 +96,7 @@ export class ConnectionCommands implements Disposable {
}
const hostUrl = normalizeHost(host);
const provider = new PersonalAccessTokenAuthProvider(hostUrl, token);
await saveNewProfile(name, provider);
await saveNewProfile(name, provider, this.cli);
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -226,7 +226,10 @@ export class ConnectionManager implements Disposable {
const savedProfile = (await this.configModel.get("overrides"))
?.authProfile;
if (savedProfile !== undefined) {
const authProvider = await ProfileAuthProvider.from(savedProfile);
const authProvider = await ProfileAuthProvider.from(
savedProfile,
this.cli
);
if (
authProvider.host.toString() === host.toString() &&
(await authProvider.check())
Expand All @@ -249,7 +252,10 @@ export class ConnectionManager implements Disposable {
if (profiles.length !== 1) {
return;
}
const authProvider = await ProfileAuthProvider.from(profiles[0].name);
const authProvider = await ProfileAuthProvider.from(
profiles[0].name,
this.cli
);
if (await authProvider.check()) {
return authProvider;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ import {Uri} from "vscode";
import {ProfileAuthProvider} from "./auth/AuthProvider";
import {DatabricksWorkspace} from "./DatabricksWorkspace";
import {iam} from "@databricks/databricks-sdk";
import {instance, mock} from "ts-mockito";
import {CliWrapper} from "../cli/CliWrapper";

describe(__filename, () => {
it("create an instance", () => {
Expand All @@ -17,7 +19,8 @@ describe(__filename, () => {
} as const;
const authProvider = new ProfileAuthProvider(
new URL("https://fabian.databricks.com"),
"DEFAULT"
"DEFAULT",
instance(mock(CliWrapper))
);
const dbWorkspace: DatabricksWorkspace = new DatabricksWorkspace(
authProvider,
Expand Down
15 changes: 10 additions & 5 deletions packages/databricks-vscode/src/configuration/LoginWizard.ts
Original file line number Diff line number Diff line change
Expand Up @@ -243,7 +243,10 @@ export class LoginWizard {
}

if (pick.profile !== undefined) {
const authProvider = await ProfileAuthProvider.from(pick.profile);
const authProvider = await ProfileAuthProvider.from(
pick.profile,
this.cliWrapper
);
const nextStep = await this.checkAuthProvider(authProvider, input);
if (nextStep) {
return nextStep;
Expand Down Expand Up @@ -307,7 +310,7 @@ export class LoginWizard {
case "databricks-cli":
authProvider = new DatabricksCliAuthProvider(
this.state.host!,
this.cliWrapper.cliPath
this.cliWrapper
);
break;

Expand Down Expand Up @@ -343,7 +346,8 @@ export class LoginWizard {

this.state.authProvider = await saveNewProfile(
profileName,
authProvider
authProvider,
this.cliWrapper
);
}

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

export async function saveNewProfile(
profileName: string,
authProvider: AuthProvider
authProvider: AuthProvider,
cli: CliWrapper
) {
const iniData = authProvider.toIni();
if (!iniData) {
Expand Down Expand Up @@ -397,7 +402,7 @@ export async function saveNewProfile(
// Write the new profile to .databrickscfg
await appendFile(configFilePath, finalStr);

return await ProfileAuthProvider.from(profileName, true);
return await ProfileAuthProvider.from(profileName, cli, true);
}

function humaniseSdkAuthType(sdkAuthType: string) {
Expand Down
40 changes: 17 additions & 23 deletions packages/databricks-vscode/src/configuration/auth/AuthProvider.ts
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ const extensionVersion = require("../../../package.json")
import {AzureCliCheck} from "./AzureCliCheck";
import {DatabricksCliCheck} from "./DatabricksCliCheck";
import {Loggers} from "../../logger";
import {CliWrapper} from "../../cli/CliWrapper";

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

static fromJSON(
json: Record<string, any>,
databricksPath: string
): AuthProvider {
static fromJSON(json: Record<string, any>, cli: CliWrapper): AuthProvider {
const host =
json.host instanceof URL
? json.host
Expand All @@ -123,20 +121,20 @@ export abstract class AuthProvider {
);

case "databricks-cli":
return new DatabricksCliAuthProvider(host, databricksPath);
return new DatabricksCliAuthProvider(host, cli);

case "profile":
if (!json.profile) {
throw new Error("Missing profile");
}
return new ProfileAuthProvider(host, json.profile);
return new ProfileAuthProvider(host, json.profile, cli);

default:
throw new Error(`Unknown auth type: ${json.authType}`);
}
}

static fromSdkConfig(config: Config): AuthProvider {
static fromSdkConfig(config: Config, cli: CliWrapper): AuthProvider {
if (!config.host) {
throw new Error("Missing host");
}
Expand All @@ -151,33 +149,27 @@ export abstract class AuthProvider {
);

case "databricks-cli":
if (!config.databricksCliPath) {
throw new Error("Missing path for databricks-cli");
}

return new DatabricksCliAuthProvider(
host,
config.databricksCliPath
);
return new DatabricksCliAuthProvider(host, cli);
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

With so many changes around passing cli to auth providers I'm not sure I can actually understand what exactly we are improving, can you explain it a bit?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If the profile does not have a databricksclipath variable, we would error out, https://github.com/databricks/databricks-vscode/pull/1114/files#diff-3e8d1317ab3ec8502ceebeae3f641a4e8f7388e442dc8804cdfdb5ea2db283b8L154-L156.

But, we can be using the internal databricks cli instead of relying on the profile to have the path. It is quiet a common profile pattern to not hard code the cli path. In fact, we should not be writing the cli path either because the profile will fail when we update the extension (hence breaking the path).

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was breaking oauth for users with existing profiles.


default:
if (config.profile) {
return new ProfileAuthProvider(host, config.profile);
return new ProfileAuthProvider(host, config.profile, cli);
}
throw new Error(`Unknown auth type: ${config.authType}`);
}
}
}

export class ProfileAuthProvider extends AuthProvider {
static async from(profile: string, checked = false) {
static async from(profile: string, cli: CliWrapper, checked = false) {
const host = await ProfileAuthProvider.getSdkConfig(profile).getHost();
return new ProfileAuthProvider(host, profile, checked);
return new ProfileAuthProvider(host, profile, cli, checked);
}

constructor(
host: URL,
private readonly profile: string,
private readonly cli: CliWrapper,
checked = false
) {
super(host, "profile", checked);
Expand Down Expand Up @@ -223,7 +215,10 @@ export class ProfileAuthProvider extends AuthProvider {
try {
const sdkConfig = this.getSdkConfig();
await sdkConfig.ensureResolved();
const authProvider = AuthProvider.fromSdkConfig(sdkConfig);
const authProvider = AuthProvider.fromSdkConfig(
sdkConfig,
this.cli
);

if (authProvider instanceof ProfileAuthProvider) {
const workspaceClient = this.getWorkspaceClient();
Expand Down Expand Up @@ -257,7 +252,7 @@ export class ProfileAuthProvider extends AuthProvider {
export class DatabricksCliAuthProvider extends AuthProvider {
constructor(
host: URL,
readonly databricksPath: string
readonly cli: CliWrapper
) {
super(host, "databricks-cli");
}
Expand All @@ -270,15 +265,15 @@ export class DatabricksCliAuthProvider extends AuthProvider {
return {
host: this.host.toString(),
authType: this.authType,
databricksPath: this.databricksPath,
databricksPath: this.cli.cliPath,
};
}

getSdkConfig(): Config {
return new Config({
host: this.host.toString(),
authType: "databricks-cli",
databricksCliPath: this.databricksPath,
databricksCliPath: this.cli.cliPath,
});
}

Expand All @@ -293,7 +288,6 @@ export class DatabricksCliAuthProvider extends AuthProvider {
return {
host: this.host.toString(),
auth_type: "databricks-cli",
databricks_cli_path: this.databricksPath,
};
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ export class DatabricksCliCheck implements Disposable {
{
host: this.authProvider.host.toString(),
authType: "databricks-cli",
databricksCliPath: this.authProvider.databricksPath,
databricksCliPath: this.authProvider.cli.cliPath,
},
{
product: "databricks-vscode",
Expand All @@ -92,7 +92,7 @@ export class DatabricksCliCheck implements Disposable {

private async login(): Promise<void> {
try {
await ExecUtils.execFile(this.authProvider.databricksPath, [
await ExecUtils.execFile(this.authProvider.cli.cliPath, [
"auth",
"login",
"--host",
Expand Down
Loading