From c18a54bc1d13cd3494e1c5595b048f79015e36cc Mon Sep 17 00:00:00 2001
From: Justin M <63619224+JustinMDotNet@users.noreply.github.com>
Date: Mon, 16 Nov 2020 11:24:22 -0800
Subject: [PATCH] 12666 Passed azureAccount into onFetchDatabases and set on
 tempProfile. (#13239)

---
 src/sql/platform/connection/common/connectionStore.ts      | 2 ++
 .../services/connection/browser/connectionController.ts    | 7 ++++---
 .../services/connection/browser/connectionDialogService.ts | 2 +-
 .../services/connection/browser/connectionWidget.ts        | 2 +-
 4 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/src/sql/platform/connection/common/connectionStore.ts b/src/sql/platform/connection/common/connectionStore.ts
index 1a632eac62..944dfab6cd 100644
--- a/src/sql/platform/connection/common/connectionStore.ts
+++ b/src/sql/platform/connection/common/connectionStore.ts
@@ -88,6 +88,8 @@ export class ConnectionStore {
 					}
 					return { profile: credentialsItem, savedCred: !!savedCred };
 				});
+		} else if (credentialsItem.authenticationType === 'AzureMFA' && credentialsItem.azureAccount) {
+			return Promise.resolve({ profile: credentialsItem, savedCred: true });
 		} else {
 			// No need to look up the password
 			return Promise.resolve({ profile: credentialsItem, savedCred: credentialsItem.savePassword });
diff --git a/src/sql/workbench/services/connection/browser/connectionController.ts b/src/sql/workbench/services/connection/browser/connectionController.ts
index bd92aeca4f..659f671050 100644
--- a/src/sql/workbench/services/connection/browser/connectionController.ts
+++ b/src/sql/workbench/services/connection/browser/connectionController.ts
@@ -47,8 +47,8 @@ export class ConnectionController implements IConnectionComponentController {
 			onCreateNewServerGroup: () => this.onCreateNewServerGroup(),
 			onAdvancedProperties: () => this.handleOnAdvancedProperties(),
 			onSetAzureTimeOut: () => this.handleonSetAzureTimeOut(),
-			onFetchDatabases: (serverName: string, authenticationType: string, userName?: string, password?: string) => this.onFetchDatabases(
-				serverName, authenticationType, userName, password).then(result => {
+			onFetchDatabases: (serverName: string, authenticationType: string, userName?: string, password?: string, azureAccount?: string) => this.onFetchDatabases(
+				serverName, authenticationType, userName, password, azureAccount).then(result => {
 					return result;
 				}),
 			onAzureTenantSelection: (azureTenantId?: string) => this.onAzureTenantSelection(azureTenantId),
@@ -56,7 +56,7 @@ export class ConnectionController implements IConnectionComponentController {
 		this._providerName = providerName;
 	}
 
-	protected async onFetchDatabases(serverName: string, authenticationType: string, userName?: string, password?: string): Promise<string[]> {
+	protected async onFetchDatabases(serverName: string, authenticationType: string, userName?: string, password?: string, azureAccount?: string): Promise<string[]> {
 		let tempProfile = this._model;
 		tempProfile.serverName = serverName;
 		tempProfile.authenticationType = authenticationType;
@@ -64,6 +64,7 @@ export class ConnectionController implements IConnectionComponentController {
 		tempProfile.password = password;
 		tempProfile.groupFullName = '';
 		tempProfile.saveProfile = false;
+		tempProfile.azureAccount = azureAccount;
 		let uri = this._connectionManagementService.getConnectionUri(tempProfile);
 		if (this._databaseCache.has(uri)) {
 			let cachedDatabases: string[] = this._databaseCache.get(uri);
diff --git a/src/sql/workbench/services/connection/browser/connectionDialogService.ts b/src/sql/workbench/services/connection/browser/connectionDialogService.ts
index 3aa51c390f..b256edff42 100644
--- a/src/sql/workbench/services/connection/browser/connectionDialogService.ts
+++ b/src/sql/workbench/services/connection/browser/connectionDialogService.ts
@@ -42,7 +42,7 @@ export interface IConnectionComponentCallbacks {
 	onCreateNewServerGroup?: () => void;
 	onAdvancedProperties?: () => void;
 	onSetAzureTimeOut?: () => void;
-	onFetchDatabases?: (serverName: string, authenticationType: string, userName?: string, password?: string) => Promise<string[]>;
+	onFetchDatabases?: (serverName: string, authenticationType: string, userName?: string, password?: string, azureAccount?: string) => Promise<string[]>;
 	onAzureTenantSelection?: (azureTenantId?: string) => void;
 }
 
diff --git a/src/sql/workbench/services/connection/browser/connectionWidget.ts b/src/sql/workbench/services/connection/browser/connectionWidget.ts
index fad6bad70a..d28ae54097 100644
--- a/src/sql/workbench/services/connection/browser/connectionWidget.ts
+++ b/src/sql/workbench/services/connection/browser/connectionWidget.ts
@@ -364,7 +364,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 				this._databaseDropdownExpanded = true;
 				if (this.serverName) {
 					this._databaseNameInputBox.values = [this._loadingDatabaseName];
-					this._callbacks.onFetchDatabases(this.serverName, this.authenticationType, this.userName, this._password).then(databases => {
+					this._callbacks.onFetchDatabases(this.serverName, this.authenticationType, this.userName, this._password, this.azureAccount).then(databases => {
 						if (databases) {
 							this._databaseNameInputBox.values = databases.sort((a, b) => a.localeCompare(b));
 						} else {