Improve login wizard (#1045)

Ensure we can let users select host manually when they want to setup new auth for bundle init
databricks · Feb 7, 2024 · 25f964f · 25f964f
1 parent 0f85446
commit 25f964f
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 24 deletions.
diff --git a/packages/databricks-vscode/src/bundle/BundleInitWizard.ts b/packages/databricks-vscode/src/bundle/BundleInitWizard.ts
@@ -11,7 +11,6 @@ import {Loggers} from "../logger";
 import {AuthProvider} from "../configuration/auth/AuthProvider";
 import {LoginWizard} from "../configuration/LoginWizard";
 import {CliWrapper} from "../cli/CliWrapper";
-import {ConfigModel} from "../configuration/models/ConfigModel";
 import {getSubProjects} from "./BundleFileSet";
 import {tmpdir} from "os";
 
@@ -47,13 +46,10 @@ export class BundleInitWizard {
 
     public async initNewProject(
         workspaceUri?: Uri,
-        existingAuthProvider?: AuthProvider,
-        configModel?: ConfigModel
+        existingAuthProvider?: AuthProvider
     ) {
-        const authProvider = await this.configureAuthForBundleInit(
-            existingAuthProvider,
-            configModel
-        );
+        const authProvider =
+            await this.configureAuthForBundleInit(existingAuthProvider);
         if (!authProvider) {
             this.logger.debug(
                 "No valid auth providers, can't proceed with bundle init wizard"
@@ -91,8 +87,7 @@ export class BundleInitWizard {
     }
 
     private async configureAuthForBundleInit(
-        authProvider?: AuthProvider,
-        configModel?: ConfigModel
+        authProvider?: AuthProvider
     ): Promise<AuthProvider | undefined> {
         if (authProvider) {
             const response = await this.promptToUseExistingAuth(authProvider);
@@ -103,7 +98,7 @@ export class BundleInitWizard {
             }
         }
         if (!authProvider) {
-            authProvider = await LoginWizard.run(this.cli, configModel);
+            authProvider = await LoginWizard.run(this.cli);
         }
         if (authProvider && (await authProvider.check())) {
             return authProvider;
@@ -147,6 +142,9 @@ export class BundleInitWizard {
             isTransient: true,
             location: TerminalLocation.Editor,
             env: this.cli.getBundleInitEnvVars(authProvider),
+            // Without strict env we will inherit our environmentVariableCollection
+            // which will override auth env vars we provide in this call.
+            strictEnv: true,
             // Setting CWD avoids a possibility of the CLI picking up unrelated bundle configuration
             // in the current workspace root or while traversing up the folder structure.
             cwd: tmpdir(),

diff --git a/packages/databricks-vscode/src/bundle/BundleProjectManager.ts b/packages/databricks-vscode/src/bundle/BundleProjectManager.ts
@@ -217,7 +217,7 @@ export class BundleProjectManager {
         },
     })
     public async startManualMigration() {
-        const authProvider = await LoginWizard.run(this.cli, this.configModel);
+        const authProvider = await LoginWizard.run(this.cli);
         if (
             authProvider instanceof ProfileAuthProvider &&
             (await authProvider.check())
@@ -275,8 +275,7 @@ export class BundleProjectManager {
             this.connectionManager.databricksWorkspace?.authProvider;
         const parentFolder = await bundleInitWizard.initNewProject(
             this.workspaceUri,
-            authProvider,
-            this.configModel
+            authProvider
         );
         if (parentFolder) {
             await this.isBundleProjectCache.refresh();

diff --git a/packages/databricks-vscode/src/cli/CliWrapper.ts b/packages/databricks-vscode/src/cli/CliWrapper.ts
@@ -303,7 +303,7 @@ export class CliWrapper {
     }
 
     getBundleInitEnvVars(authProvider: AuthProvider) {
-        return {
+        return removeUndefinedKeys({
             ...EnvVarGenerators.getEnvVarsForCli(
                 workspaceConfigs.databrickscfgLocation
             ),
@@ -312,7 +312,7 @@ export class CliWrapper {
             ...authProvider.toEnv(),
             // eslint-disable-next-line @typescript-eslint/naming-convention
             DATABRICKS_OUTPUT_FORMAT: "text",
-        };
+        });
     }
 
     async bundleInit(

diff --git a/packages/databricks-vscode/src/configuration/ConnectionManager.ts b/packages/databricks-vscode/src/configuration/ConnectionManager.ts
@@ -286,7 +286,11 @@ export class ConnectionManager implements Disposable {
         popup: {prefix: "Can't configure workspace. "},
     })
     async configureLogin() {
-        const authProvider = await LoginWizard.run(this.cli, this.configModel);
+        const authProvider = await LoginWizard.run(
+            this.cli,
+            this.configModel.target,
+            await this.configModel.get("host")
+        );
         if (!authProvider) {
             return;
         }

diff --git a/packages/databricks-vscode/src/configuration/LoginWizard.ts b/packages/databricks-vscode/src/configuration/LoginWizard.ts
@@ -26,7 +26,6 @@ import {
     loadConfigFile,
     AuthType as SdkAuthType,
 } from "@databricks/databricks-sdk";
-import {ConfigModel} from "./models/ConfigModel";
 import {randomUUID} from "crypto";
 import ini from "ini";
 import {copyFile, writeFile} from "fs/promises";
@@ -55,7 +54,7 @@ export class LoginWizard {
     }
     constructor(
         private readonly cliWrapper: CliWrapper,
-        private readonly configModel?: ConfigModel
+        private readonly target?: string
     ) {}
 
     private async inputHost(input: MultiStepInput) {
@@ -253,13 +252,13 @@ export class LoginWizard {
         input: MultiStepInput,
         pick: AuthTypeQuickPickItem
     ) {
-        let initialValue = this.configModel?.target ?? "";
+        let initialValue = this.target ?? "";
 
         // If the initialValue profile already exists, then create a unique name.
         const profiles = await this.getProfiles();
         if (profiles.find((profile) => profile.name === initialValue)) {
             const suffix = randomUUID().slice(0, 8);
-            initialValue = `${this.configModel?.target ?? "dev"}-${suffix}`;
+            initialValue = `${this.target ?? "dev"}-${suffix}`;
         }
 
         const profileName = await input.showInputBox({
@@ -342,11 +341,12 @@ export class LoginWizard {
 
     static async run(
         cliWrapper: CliWrapper,
-        configModel?: ConfigModel
+        target?: string,
+        host?: URL
     ): Promise<AuthProvider | undefined> {
-        const wizard = new LoginWizard(cliWrapper, configModel);
-        if (configModel) {
-            wizard.state.host = await configModel.get("host");
+        const wizard = new LoginWizard(cliWrapper, target);
+        if (host) {
+            wizard.state.host = host;
         }
         await MultiStepInput.run(wizard.inputHost.bind(wizard));
         if (!wizard.state.host || !wizard.state.authProvider) {