AAD: Use device code flow on CodeSpaces (#392)

also store tenantID so users don't have to login twice

Author: fjakobs

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

@@ -1,6 +1,6 @@
 import {ExecUtils, WorkspaceClient} from "@databricks/databricks-sdk";
 import {commands, Disposable, Uri, window} from "vscode";
-import {AuthProvider} from "./auth/AuthProvider";
+import {AzureCliAuthProvider} from "./auth/AuthProvider";
 
 export type Step<S, N> = () => Promise<
     SuccessResult<S> | NextResult<N> | ErrorResult
@@ -54,25 +54,30 @@ type AzureStepName =
 
 export class AzureCliCheck implements Disposable {
     private disposables: Disposable[] = [];
+    private isCodeSpaces: boolean;
+
+    tenantId: string | undefined;
 
     constructor(
-        private authProvider: AuthProvider,
+        private authProvider: AzureCliAuthProvider,
         private azBinPath: string = "az"
-    ) {}
+    ) {
+        this.isCodeSpaces = process.env.CODESPACES === "true";
+    }
 
     dispose() {
         this.disposables.forEach((i) => i.dispose());
         this.disposables = [];
     }
 
     public async check(silent = false): Promise<boolean> {
-        let tenant: string;
+        this.tenantId = this.authProvider.tenantId;
 
         const steps: Record<AzureStepName, Step<boolean, AzureStepName>> = {
             tryLogin: async () => {
                 const result = await this.tryLogin(this.authProvider.host);
                 if (typeof result === "string") {
-                    tenant = result;
+                    this.tenantId = result;
                     return {
                         type: "next",
                         next: "loginAzureCli",
@@ -123,7 +128,7 @@ export class AzureCliCheck implements Disposable {
                 };
             },
             loginAzureCli: async () => {
-                if (await this.loginAzureCli(tenant)) {
+                if (await this.loginAzureCli(this.tenantId)) {
                     return {
                         type: "next",
                         next: "tryLogin",
@@ -233,7 +238,7 @@ export class AzureCliCheck implements Disposable {
     public async loginAzureCli(tenant = ""): Promise<boolean> {
         let message = 'You need to run "az login" to login with Azure.';
         if (tenant) {
-            message = `You are logged in with the wrong tenant ID. Run "az login -t ${tenant}" to login with Azure with the correct tenant.`;
+            message = `You need to tun "az login -t ${tenant}" to login with Azure.`;
         }
         const choice = await window.showInformationMessage(
             message,
@@ -245,8 +250,11 @@ export class AzureCliCheck implements Disposable {
             const terminal = window.createTerminal("az login");
             this.disposables.push(terminal);
             terminal.show();
+
+            const useDeviceCode = this.isCodeSpaces ? "--use-device-code" : "";
+
             terminal.sendText(
-                `${this.azBinPath} login ${
+                `${this.azBinPath} login ${useDeviceCode} ${
                     tenant ? "-t " + tenant : ""
                 }; echo "Press any key to close the terminal and continue ..."; read; exit`
             );

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

@@ -66,7 +66,7 @@ export abstract class AuthProvider {
 
         switch (json.authType as AuthType) {
             case "azure-cli":
-                return new AzureCliAuthProvider(host);
+                return new AzureCliAuthProvider(host, json.tenantId);
 
             case "profile":
                 if (!json.profile) {
@@ -114,8 +114,16 @@ export class ProfileAuthProvider extends AuthProvider {
 }
 
 export class AzureCliAuthProvider extends AuthProvider {
-    constructor(host: URL) {
+    private _tenantId?: string;
+
+    constructor(host: URL, tenantId?: string) {
         super(host, "azure-cli");
+
+        this._tenantId = tenantId;
+    }
+
+    get tenantId(): string | undefined {
+        return this._tenantId;
     }
 
     describe(): string {
@@ -126,6 +134,7 @@ export class AzureCliAuthProvider extends AuthProvider {
         return {
             host: this.host.toString(),
             authType: this.authType,
+            tenantId: this.tenantId,
         };
     }
 
@@ -145,6 +154,9 @@ export class AzureCliAuthProvider extends AuthProvider {
     }
 
     async check(silent: boolean): Promise<boolean> {
-        return await new AzureCliCheck(this).check(silent);
+        const cliCheck = new AzureCliCheck(this);
+        const result = await cliCheck.check(silent);
+        this._tenantId = cliCheck.tenantId;
+        return result;
     }
 }