Add UI to configure PAT auth (#1033)

## Changes
<!-- Summary of your changes that are easy to understand -->

## Tests
<!-- How is this tested? -->

Author: kartikgupta-db

packages/databricks-vscode/package.json

@@ -387,11 +387,6 @@
                 }
             ],
             "view/item/context": [
-                {
-                    "command": "databricks.utils.openExternal",
-                    "when": "viewItem =~ /^databricks.*\\.(has-url).*$/",
-                    "group": "inline@1"
-                },
                 {
                     "command": "databricks.utils.openExternal",
                     "when": "viewItem =~ /^databricks.*\\.(has-url).*$/",

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

@@ -11,13 +11,14 @@ import {
     MultiStepInput,
     ValidationMessageType,
 } from "../ui/MultiStepInputWizard";
-import {CliWrapper} from "../cli/CliWrapper";
+import {CliWrapper, ConfigEntry} from "../cli/CliWrapper";
 import {workspaceConfigs} from "../vscode-objs/WorkspaceConfigs";
 import {
     AuthProvider,
     AuthType,
     AzureCliAuthProvider,
     DatabricksCliAuthProvider,
+    PersonalAccessTokenAuthProvider,
     ProfileAuthProvider,
 } from "./auth/AuthProvider";
 import {UrlUtils} from "../utils";
@@ -45,6 +46,13 @@ interface State {
 export class LoginWizard {
     private state = {} as Partial<State>;
     private readonly title = "Configure Databricks Workspace";
+    private _profiles: Array<ConfigEntry> = [];
+    async getProfiles() {
+        if (this._profiles.length === 0) {
+            this._profiles = await listProfiles(this.cliWrapper);
+        }
+        return this._profiles;
+    }
     constructor(
         private readonly cliWrapper: CliWrapper,
         private readonly configModel: ConfigModel
@@ -64,9 +72,8 @@ export class LoginWizard {
             });
         }
 
-        const profiles = await listProfiles(this.cliWrapper);
         items.push(
-            ...profiles.map((profile) => {
+            ...(await this.getProfiles()).map((profile) => {
                 return {
                     label: profile.host!.toString(),
                     detail: `Profile: ${profile.name}`,
@@ -113,7 +120,10 @@ export class LoginWizard {
         input: MultiStepInput
     ): Promise<InputStep | void> {
         const items: Array<AuthTypeQuickPickItem> = [];
-
+        items.push({
+            label: "Create new Databricks CLI profile",
+            kind: QuickPickItemKind.Separator,
+        });
         for (const authMethod of authMethodsForHostname(this.state.host!)) {
             switch (authMethod) {
                 case "azure-cli":
@@ -131,10 +141,9 @@ export class LoginWizard {
                         authType: "databricks-cli",
                     });
                     break;
-
                 case "profile":
                     {
-                        const profiles = (await listProfiles(this.cliWrapper))
+                        const profiles = (await this.getProfiles())
                             .filter((profile) => {
                                 return (
                                     profile.host?.hostname ===
@@ -159,6 +168,11 @@ export class LoginWizard {
                                 };
                             });
 
+                        items.push({
+                            label: "Personal Access Token",
+                            detail: "Create a new profile and authenticate using the 'databricks' command line tool",
+                            authType: "pat",
+                        });
                         if (profiles.length !== 0) {
                             items.push(
                                 {
@@ -235,7 +249,7 @@ export class LoginWizard {
         let initialValue = this.configModel.target ?? "";
 
         // If the initialValue profile already exists, then create a unique name.
-        const profiles = await listProfiles(this.cliWrapper);
+        const profiles = await this.getProfiles();
         if (profiles.find((profile) => profile.name === initialValue)) {
             const suffix = randomUUID().slice(0, 8);
             initialValue = `${this.configModel.target}-${suffix}`;
@@ -254,7 +268,6 @@ export class LoginWizard {
                         type: "error",
                     };
                 }
-                const profiles = await listProfiles(this.cliWrapper);
                 if (profiles.find((profile) => profile.name === value)) {
                     return {
                         message: `Profile ${value} already exists`,
@@ -269,7 +282,10 @@ export class LoginWizard {
             return;
         }
 
-        let authProvider: AzureCliAuthProvider | DatabricksCliAuthProvider;
+        let authProvider:
+            | AzureCliAuthProvider
+            | DatabricksCliAuthProvider
+            | PersonalAccessTokenAuthProvider;
         switch (pick.authType) {
             case "azure-cli":
                 authProvider = new AzureCliAuthProvider(this.state.host!);
@@ -282,6 +298,20 @@ export class LoginWizard {
                 );
                 break;
 
+            case "pat":
+                {
+                    const token = await collectTokenForPatAuth(input, 4, 4);
+                    if (token === undefined) {
+                        // Token can never be undefined unless the users cancels the whole process.
+                        // Therefore, we can safely return here.
+                        return;
+                    }
+                    authProvider = new PersonalAccessTokenAuthProvider(
+                        this.state.host!,
+                        token
+                    );
+                }
+                break;
             default:
                 throw new Error(
                     `Unknown auth type: ${pick.authType} for profile creation`
@@ -430,3 +460,31 @@ function authMethodsForHostname(host: URL): Array<AuthType> {
 
     return ["profile"];
 }
+
+async function collectTokenForPatAuth(
+    input: MultiStepInput,
+    step: number,
+    totalSteps: number
+) {
+    const token = await input.showInputBox({
+        title: "Enter Personal Access Token",
+        step,
+        totalSteps,
+        validate: async (value) => {
+            if (value.length === 0) {
+                return {
+                    message: "Token cannot be empty",
+                    type: "error",
+                };
+            }
+        },
+        placeholder: "Enter Personal Access Token",
+        ignoreFocusOut: true,
+    });
+
+    if (token === undefined) {
+        return;
+    }
+
+    return token;
+}

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

@@ -312,3 +312,70 @@ export class AzureCliAuthProvider extends AuthProvider {
         return result;
     }
 }
+
+export class PersonalAccessTokenAuthProvider extends AuthProvider {
+    constructor(
+        host: URL,
+        private readonly token: string
+    ) {
+        super(host, "pat");
+    }
+
+    describe(): string {
+        return "Personal Access Token";
+    }
+    toJSON(): Record<string, string | undefined> {
+        return {
+            host: this.host.toString(),
+            authType: this.authType,
+            token: this.token,
+        };
+    }
+    toEnv(): Record<string, string> {
+        return {
+            DATABRICKS_HOST: this.host.toString(),
+            DATABRICKS_AUTH_TYPE: this.authType,
+            DATABRICKS_TOKEN: this.token,
+        };
+    }
+    toIni(): Record<string, string | undefined> | undefined {
+        return {
+            host: this.host.toString(),
+            token: this.token,
+        };
+    }
+    protected async _check(): Promise<boolean> {
+        while (true) {
+            try {
+                const workspaceClient = this.getWorkspaceClient();
+                await workspaceClient.currentUser.me();
+                return true;
+            } catch (e) {
+                let message: string = `Can't login with the provided Personal Access Token`;
+                if (e instanceof Error) {
+                    message = `Can't login with the provided Personal Access Token: ${e.message}`;
+                }
+                logging.NamedLogger.getOrCreate(Loggers.Extension).error(
+                    message,
+                    e
+                );
+                const choice = await window.showErrorMessage(
+                    message,
+                    "Retry",
+                    "Cancel"
+                );
+                if (choice === "Retry") {
+                    continue;
+                }
+                return false;
+            }
+        }
+    }
+    protected getSdkConfig(): Config {
+        return new Config({
+            host: this.host.toString(),
+            authType: "pat",
+            token: this.token,
+        });
+    }
+}

packages/databricks-vscode/src/extension.ts

@@ -539,7 +539,8 @@ export async function activate(
 
     const bundleCommands = new BundleCommands(
         bundleRemoteStateModel,
-        bundleRunStatusManager
+        bundleRunStatusManager,
+        bundleValidateModel
     );
     context.subscriptions.push(
         bundleResourceExplorerTreeDataProvider,

packages/databricks-vscode/src/ui/MultiStepInputWizard.ts

@@ -71,7 +71,7 @@ interface InputBoxParameters {
     totalSteps: number;
     placeholder: string;
     initialValue?: string;
-    validate: (
+    validate?: (
         value: string
     ) => Promise<string | undefined | ValidationMessageType>;
     buttons?: QuickInputButton[];

packages/databricks-vscode/src/ui/bundle-resource-explorer/BundleCommands.ts

@@ -4,6 +4,7 @@ import {onError} from "../../utils/onErrorDecorator";
 import {BundleResourceExplorerTreeNode} from "./types";
 import {BundleRunStatusManager} from "../../bundle/run/BundleRunStatusManager";
 import {Mutex} from "../../locking";
+import {BundleValidateModel} from "../../bundle/models/BundleValidateModel";
 import {PipelineTreeNode} from "./PipelineTreeNode";
 import {JobTreeNode} from "./JobTreeNode";
 
@@ -28,9 +29,15 @@ export class BundleCommands implements Disposable {
 
     constructor(
         private readonly bundleRemoteStateModel: BundleRemoteStateModel,
-        private readonly bundleRunStatusManager: BundleRunStatusManager
+        private readonly bundleRunStatusManager: BundleRunStatusManager,
+        private readonly bundleValidateModel: BundleValidateModel
     ) {
-        this.disposables.push(this.outputChannel);
+        this.disposables.push(
+            this.outputChannel,
+            this.bundleValidateModel.onDidChange(async () => {
+                await this.refreshRemoteState();
+            })
+        );
     }
 
     private refreshStateMutex = new Mutex();