Interactive AAD setup (#308)

Logging in with `az login`. Including detecting the wrong tenant ID.


https://user-images.githubusercontent.com/40952/207889437-84b329c1-14d5-45b7-a9f8-1c96884ab79f.mov

Instructions for installing az CLI:


https://user-images.githubusercontent.com/40952/207889556-36ffd664-cd78-420d-aa37-e5fbac060e92.mov

Co-authored-by: Pieter Noordhuis <[email protected]>

Author: fjakobs

packages/databricks-sdk-js/src/api-client.test.ts

@@ -12,7 +12,7 @@ describe(__filename, () => {
     });
 
     it("should create proper user agent", () => {
-        const ua = new ApiClient("unit", "3.4.5").userAgent();
+        const ua = new ApiClient({extraUserAgent: {unit: "3.4.5"}}).userAgent();
         assert.equal(
             ua,
             `unit/3.4.5 databricks-sdk-js/${sdkVersion} nodejs/${process.version.slice(

packages/databricks-sdk-js/src/api-client.ts

@@ -9,6 +9,7 @@ import {context} from "./context";
 import {Context} from "./context";
 import retry, {RetriableError} from "./retries/retries";
 import Time, {TimeUnits} from "./retries/Time";
+import {CredentialProvider} from "./auth/types";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const sdkVersion = require("../package.json").version;
@@ -46,6 +47,9 @@ export class ApiClient {
     private agent: https.Agent;
     private _host?: URL;
 
+    private credentialProvider: CredentialProvider;
+    private readonly extraUserAgent: Record<string, string>;
+
     get host(): Promise<URL> {
         return (async () => {
             if (!this._host) {
@@ -56,24 +60,34 @@ export class ApiClient {
         })();
     }
 
-    constructor(
-        private readonly product: string,
-        private readonly productVersion: string,
-        private credentialProvider = fromDefaultChain
-    ) {
+    constructor({
+        credentialProvider = fromDefaultChain,
+        extraUserAgent = {},
+    }: {
+        credentialProvider?: CredentialProvider;
+        extraUserAgent?: Record<string, string>;
+    }) {
+        this.credentialProvider = credentialProvider;
+        this.extraUserAgent = extraUserAgent;
+
         this.agent = new https.Agent({
             keepAlive: true,
             keepAliveMsecs: 15_000,
         });
     }
 
     userAgent(): string {
-        const pairs = [
-            `${this.product}/${this.productVersion}`,
+        const pairs: Array<string> = [];
+        for (const [key, value] of Object.entries(this.extraUserAgent)) {
+            pairs.push(`${key}/${value}`);
+        }
+
+        pairs.push(
             `databricks-sdk-js/${sdkVersion}`,
             `nodejs/${process.version.slice(1)}`,
-            `os/${process.platform}`,
-        ];
+            `os/${process.platform}`
+        );
+
         // TODO: add ability of per-request extra-information,
         // so that we can track sub-functionality, like in Terraform
         return pairs.join(" ");
@@ -165,7 +179,7 @@ export class ApiClient {
 
         // throw error if the URL is incorrect and we get back an HTML page
         if (response.headers.get("content-type")?.match("text/html")) {
-            // When the AAD tenent is not configured correctly, the response is a HTML page with a title like this:
+            // When the AAD tenant is not configured correctly, the response is a HTML page with a title like this:
             // "Error 400 io.jsonwebtoken.IncorrectClaimException: Expected iss claim to be: https://sts.windows.net/aaaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaa/, but was: https://sts.windows.net/bbbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbb/."
             const m = responseText.match(/<title>(Error \d+.*?)<\/title>/);
             let error: HttpError;

packages/databricks-sdk-js/src/auth/fromAzureCli.integ.ts

@@ -8,7 +8,9 @@ describe.skip(__filename, function () {
     this.timeout(15_000);
 
     it("should login with Azure CLI", async () => {
-        const client = new ApiClient("test", "0.1", fromAzureCli());
+        const client = new ApiClient({
+            credentialProvider: fromAzureCli(),
+        });
 
         const scimApi = new CurrentUserService(client);
         await scimApi.me();

packages/databricks-sdk-js/src/test/IntegrationTestSetup.ts

@@ -14,7 +14,9 @@ export class IntegrationTestSetup {
     private static _instance: IntegrationTestSetup;
     static async getInstance(): Promise<IntegrationTestSetup> {
         if (!this._instance) {
-            const client = new ApiClient("integration-tests", "0.0.1");
+            const client = new ApiClient({
+                extraUserAgent: {"integration-tests": "0.0.1"},
+            });
 
             if (!process.env["TEST_DEFAULT_CLUSTER_ID"]) {
                 throw new Error(

packages/databricks-vscode/package.json

@@ -37,27 +37,7 @@
         "url": "https://github.com/databricks/databricks-vscode.git"
     },
     "activationEvents": [
-        "onCommand:databricks.connection.logout",
-        "onCommand:databricks.connection.configureWorkspace",
-        "onCommand:databricks.connection.openDatabricksConfigFile",
-        "onCommand:databricks.connection.attachCluster",
-        "onCommand:databricks.connection.attachClusterQuickPick",
-        "onCommand:databricks.connection.detachCluster",
-        "onCommand:databricks.connection.attachSyncDestination",
-        "onCommand:databricks.connection.detachSyncDestination",
-        "onCommand:databricks.sync.start",
-        "onCommand:databricks.sync.startFull",
-        "onCommand:databricks.sync.stop",
-        "onCommand:databricks.cluster.refresh",
-        "onCommand:databricks.cluster.filterByAll",
-        "onCommand:databricks.cluster.filterByMe",
-        "onCommand:databricks.cluster.filterByRunning",
-        "onCommand:databricks.run.getProgramName",
-        "onCommand:databricks.run.runEditorContentsAsWorkflow",
-        "onCommand:databricks.run.runEditorContents",
-        "onCommand:databricks.quickstart.open",
-        "onCommand:databricks.logs.openFolder",
-        "onCommand:databricks.autocomplete.configure",
+        "onCommand:databricks.*",
         "onView:configurationView",
         "onView:clusterView",
         "onTaskType:databricks",

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

@@ -171,6 +171,6 @@ export class AzureCliAuthProvider extends AuthProvider {
     }
 
     async check(silent: boolean): Promise<boolean> {
-        return await new AzureCliCheck().check(silent);
+        return await new AzureCliCheck(this).check(silent);
     }
 }

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

@@ -1,8 +1,9 @@
 import * as child_process from "node:child_process";
 import {promisify} from "node:util";
-import {CurrentUserService, fromAzureCli} from "@databricks/databricks-sdk";
+import {CurrentUserService} from "@databricks/databricks-sdk";
 import {commands, Disposable, Uri, window} from "vscode";
 import {ConnectionManager} from "./ConnectionManager";
+import {AuthProvider} from "./AuthProvider";
 
 const execFile = promisify(child_process.execFile);
 
@@ -59,24 +60,22 @@ type AzureStepName =
 export class AzureCliCheck implements Disposable {
     private disposables: Disposable[] = [];
 
-    constructor(private azBinPath: string = "az") {}
+    constructor(
+        private authProvider: AuthProvider,
+        private azBinPath: string = "az"
+    ) {}
 
     dispose() {
         this.disposables.forEach((i) => i.dispose());
         this.disposables = [];
     }
 
-    public async check(
-        silent = false,
-        host: URL = new URL(
-            "https://adb-309687753508875.15.azuredatabricks.net"
-        )
-    ): Promise<boolean> {
+    public async check(silent = false): Promise<boolean> {
         let tenant: string;
 
         const steps: Record<AzureStepName, Step<boolean, AzureStepName>> = {
             tryLogin: async () => {
-                const result = await this.tryLogin(host);
+                const result = await this.tryLogin();
                 if (typeof result === "string") {
                     tenant = result;
                     return {
@@ -162,8 +161,8 @@ export class AzureCliCheck implements Disposable {
         return result;
     }
 
-    private async tryLogin(host: URL): Promise<boolean | string> {
-        const apiClient = ConnectionManager.apiClientFrom(fromAzureCli(host));
+    private async tryLogin(): Promise<boolean | string> {
+        const apiClient = ConnectionManager.apiClientFrom(this.authProvider);
         try {
             await new CurrentUserService(apiClient).me();
         } catch (e: any) {

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

@@ -1,7 +1,6 @@
 import {
     ApiClient,
     Cluster,
-    CredentialProvider,
     WorkspaceService,
     HttpError,
 } from "@databricks/databricks-sdk";
@@ -15,12 +14,13 @@ import {
 import {CliWrapper} from "../cli/CliWrapper";
 import {SyncDestination} from "./SyncDestination";
 import {ProjectConfig, ProjectConfigFile} from "./ProjectConfigFile";
-import {configureWorkspaceWizard} from "./selectProfileWizard";
+import {configureWorkspaceWizard} from "./configureWorkspaceWizard";
 import {ClusterManager} from "../cluster/ClusterManager";
 import {workspace} from "@databricks/databricks-sdk";
 import {DatabricksWorkspace} from "./DatabricksWorkspace";
 import {NamedLogger} from "@databricks/databricks-sdk/dist/logging";
 import {Loggers} from "../logger";
+import {AuthProvider} from "./AuthProvider";
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const extensionVersion = require("../../package.json").version;
@@ -90,8 +90,16 @@ export class ConnectionManager {
         return this._apiClient;
     }
 
-    public static apiClientFrom(creds: CredentialProvider): ApiClient {
-        return new ApiClient("vscode-extension", extensionVersion, creds);
+    public static apiClientFrom(authProvider: AuthProvider): ApiClient {
+        return new ApiClient({
+            extraUserAgent: {
+                // eslint-disable-next-line @typescript-eslint/naming-convention
+                "vscode-extension": extensionVersion,
+                // eslint-disable-next-line @typescript-eslint/naming-convention
+                "auth-type": authProvider.authType,
+            },
+            credentialProvider: authProvider.getCredentialProvider(),
+        });
     }
 
     async login(interactive = false): Promise<void> {
@@ -118,18 +126,17 @@ export class ConnectionManager {
                 vscodeWorkspace.rootPath
             );
 
-            const credentialProvider =
-                projectConfigFile.authProvider.getCredentialProvider();
-
             if (!(await projectConfigFile.authProvider.check(true))) {
                 throw new Error(
                     `Can't login with ${projectConfigFile.authProvider.describe()}.`
                 );
             }
 
-            await credentialProvider();
+            await projectConfigFile.authProvider.getCredentialProvider();
 
-            apiClient = ConnectionManager.apiClientFrom(credentialProvider);
+            apiClient = ConnectionManager.apiClientFrom(
+                projectConfigFile.authProvider
+            );
             this._databricksWorkspace = await DatabricksWorkspace.load(
                 apiClient,
                 projectConfigFile.authProvider
@@ -238,12 +245,13 @@ export class ConnectionManager {
                 return;
             }
 
-            const credentialProvider =
-                config.authProvider.getCredentialProvider();
+            if (!(await config.authProvider.check(false))) {
+                return;
+            }
 
             try {
                 await DatabricksWorkspace.load(
-                    ConnectionManager.apiClientFrom(credentialProvider),
+                    ConnectionManager.apiClientFrom(config.authProvider),
                     config.authProvider
                 );
             } catch (e: any) {