Enable Azure Active Directory MFA authentication (#3125)

src/sql/base/browser/ui/modal/dialogHelper.ts

@@ -12,9 +12,10 @@ import * as types from 'vs/base/common/types';
 
 import * as sqlops from 'sqlops';
 
-export function appendRow(container: Builder, label: string, labelClass: string, cellContainerClass: string): Builder {
+export function appendRow(container: Builder, label: string, labelClass: string, cellContainerClass: string, rowContainerClass?: string): Builder {
 	let cellContainer: Builder;
-	container.element('tr', {}, (rowContainer) => {
+	let rowAttributes = rowContainerClass ? { class: rowContainerClass } : {};
+	container.element('tr', rowAttributes, (rowContainer) => {
 		rowContainer.element('td', { class: labelClass }, (labelCellContainer) => {
 			labelCellContainer.div({}, (labelContainer) => {
 				labelContainer.text(label);

src/sql/base/browser/ui/selectBox/selectBox.ts

@@ -75,7 +75,12 @@ export class SelectBox extends vsSelectBox {
 
 		// explicitly set the accessible role so that the screen readers can read the control type properly
 		this.selectElement.setAttribute('role', 'combobox');
+
 		this._selectBoxOptions = selectBoxOptions;
+		var focusTracker = dom.trackFocus(this.selectElement);
+		this._register(focusTracker);
+		this._register(focusTracker.onDidBlur(() => this._hideMessage()));
+		this._register(focusTracker.onDidFocus(() => this._showMessage()));
 	}
 
 	public style(styles: ISelectBoxStyles): void {
@@ -142,6 +147,10 @@ export class SelectBox extends vsSelectBox {
 		this.applyStyles();
 	}
 
+	public hasFocus(): boolean {
+		return document.activeElement === this.selectElement;
+	}
+
 	public showMessage(message: IMessage): void {
 		this.message = message;
 
@@ -163,7 +172,9 @@ export class SelectBox extends vsSelectBox {
 
 		aria.alert(alertText);
 
-		this._showMessage();
+		if (this.hasFocus()) {
+			this._showMessage();
+		}
 	}
 
 	public _showMessage(): void {

src/sql/parts/accountManagement/firewallRuleDialog/firewallRuleDialogController.ts

@@ -12,7 +12,7 @@ import * as sqlops from 'sqlops';
 import { IConnectionProfile } from 'sql/parts/connection/common/interfaces';
 import { IErrorMessageService } from 'sql/parts/connection/common/connectionManagement';
 import { FirewallRuleDialog } from 'sql/parts/accountManagement/firewallRuleDialog/firewallRuleDialog';
-import { IAccountManagementService } from 'sql/services/accountManagement/interfaces';
+import { IAccountManagementService, AzureResource } from 'sql/services/accountManagement/interfaces';
 import { IResourceProviderService } from 'sql/parts/accountManagement/common/interfaces';
 import { Deferred } from 'sql/base/common/promise';
 
@@ -61,7 +61,7 @@ export class FirewallRuleDialogController {
 	private handleOnCreateFirewallRule(): void {
 		let resourceProviderId = this._resourceProviderId;
 
-		this._accountManagementService.getSecurityToken(this._firewallRuleDialog.viewModel.selectedAccount).then(tokenMappings => {
+		this._accountManagementService.getSecurityToken(this._firewallRuleDialog.viewModel.selectedAccount, AzureResource.ResourceManagement).then(tokenMappings => {
 			let firewallRuleInfo: sqlops.FirewallRuleInfo = {
 				startIpAddress: this._firewallRuleDialog.viewModel.isIPAddressSelected ? this._firewallRuleDialog.viewModel.defaultIPAddress : this._firewallRuleDialog.viewModel.fromSubnetIPRange,
 				endIpAddress: this._firewallRuleDialog.viewModel.isIPAddressSelected ? this._firewallRuleDialog.viewModel.defaultIPAddress : this._firewallRuleDialog.viewModel.toSubnetIPRange,

src/sql/parts/connection/common/connectionManagementService.ts

@@ -34,12 +34,13 @@ import { Deferred } from 'sql/base/common/promise';
 import { ConnectionOptionSpecialType } from 'sql/workbench/api/common/sqlExtHostTypes';
 import { values } from 'sql/base/common/objects';
 import { ConnectionProviderProperties, IConnectionProviderRegistry, Extensions as ConnectionProviderExtensions } from 'sql/workbench/parts/connection/common/connectionProviderExtension';
+import { IAccountManagementService, AzureResource } from 'sql/services/accountManagement/interfaces';
 
 import * as sqlops from 'sqlops';
 
 import * as nls from 'vs/nls';
 import * as errors from 'vs/base/common/errors';
-import { IDisposable, dispose, Disposable } from 'vs/base/common/lifecycle';
+import { Disposable } from 'vs/base/common/lifecycle';
 import { IInstantiationService } from 'vs/platform/instantiation/common/instantiation';
 import { IEditorService, ACTIVE_GROUP } from 'vs/workbench/services/editor/common/editorService';
 import * as platform from 'vs/platform/registry/common/platform';
@@ -58,7 +59,6 @@ import * as statusbar from 'vs/workbench/browser/parts/statusbar/statusbar';
 import { IViewletService } from 'vs/workbench/services/viewlet/browser/viewlet';
 import { IStatusbarService } from 'vs/platform/statusbar/common/statusbar';
 import { ICommandService } from 'vs/platform/commands/common/commands';
-import { EditorGroup } from 'vs/workbench/common/editor/editorGroup';
 
 export class ConnectionManagementService extends Disposable implements IConnectionManagementService {
 
@@ -100,7 +100,8 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		@IStatusbarService private _statusBarService: IStatusbarService,
 		@IResourceProviderService private _resourceProviderService: IResourceProviderService,
 		@IViewletService private _viewletService: IViewletService,
-		@IAngularEventingService private _angularEventing: IAngularEventingService
+		@IAngularEventingService private _angularEventing: IAngularEventingService,
+		@IAccountManagementService private _accountManagementService: IAccountManagementService
 	) {
 		super();
 		if (this._instantiationService) {
@@ -248,7 +249,8 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	 * Load the password for the profile
 	 * @param connectionProfile Connection Profile
 	 */
-	public addSavedPassword(connectionProfile: IConnectionProfile): Promise<IConnectionProfile> {
+	public async addSavedPassword(connectionProfile: IConnectionProfile): Promise<IConnectionProfile> {
+		await this.fillInAzureTokenIfNeeded(connectionProfile);
 		return this._connectionStore.addSavedPassword(connectionProfile).then(result => result.profile);
 	}
 
@@ -274,7 +276,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		let self = this;
 		return new Promise<IConnectionResult>((resolve, reject) => {
 			// Load the password if it's not already loaded
-			self._connectionStore.addSavedPassword(connection).then(result => {
+			self._connectionStore.addSavedPassword(connection).then(async result => {
 				let newConnection = result.profile;
 				let foundPassword = result.savedCred;
 
@@ -286,8 +288,12 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 						foundPassword = true;
 					}
 				}
+
+				// Fill in the Azure account token if needed and open the connection dialog if it fails
+				let tokenFillSuccess = await self.fillInAzureTokenIfNeeded(newConnection);
+
 				// If the password is required and still not loaded show the dialog
-				if (!foundPassword && self._connectionStore.isPasswordRequired(newConnection) && !newConnection.password) {
+				if ((!foundPassword && self._connectionStore.isPasswordRequired(newConnection) && !newConnection.password) || !tokenFillSuccess) {
 					resolve(self.showConnectionDialogOnError(connection, owner, { connected: false, errorMessage: undefined, callStack: undefined, errorCode: undefined }, options));
 				} else {
 					// Try to connect
@@ -449,10 +455,14 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 				showFirewallRuleOnError: true
 			};
 		}
-		return new Promise<IConnectionResult>((resolve, reject) => {
+		return new Promise<IConnectionResult>(async (resolve, reject) => {
 			if (callbacks.onConnectStart) {
 				callbacks.onConnectStart();
 			}
+			let tokenFillSuccess = await this.fillInAzureTokenIfNeeded(connection);
+			if (!tokenFillSuccess) {
+				throw new Error(nls.localize('connection.noAzureAccount', 'Failed to get Azure account token for connection'));
+			}
 			this.createNewConnection(uri, connection).then(connectionResult => {
 				if (connectionResult && connectionResult.connected) {
 					if (callbacks.onConnectSuccess) {
@@ -743,8 +753,33 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		}
 	}
 
+	private async fillInAzureTokenIfNeeded(connection: IConnectionProfile): Promise<boolean> {
+		if (connection.authenticationType !== Constants.azureMFA || connection.options['azureAccountToken']) {
+			return true;
+		}
+		let accounts = await this._accountManagementService.getAccountsForProvider('azurePublicCloud');
+		if (accounts && accounts.length > 0) {
+			let account = accounts.find(account => account.key.accountId === connection.userName);
+			if (account) {
+				if (account.isStale) {
+					try {
+						account = await this._accountManagementService.refreshAccount(account);
+					} catch {
+						// refreshAccount throws an error if the user cancels the dialog
+						return false;
+					}
+				}
+				let tokens = await this._accountManagementService.getSecurityToken(account, AzureResource.Sql);
+				connection.options['azureAccountToken'] = Object.values(tokens)[0].token;
+				connection.options['password'] = '';
+				return true;
+			}
+		}
+		return false;
+	}
+
 	// Request Senders
-	private sendConnectRequest(connection: IConnectionProfile, uri: string): Thenable<boolean> {
+	private async sendConnectRequest(connection: IConnectionProfile, uri: string): Promise<boolean> {
 		let connectionInfo = Object.assign({}, {
 			options: connection.options
 		});

src/sql/parts/connection/common/constants.ts

@@ -33,4 +33,5 @@ export const passwordChars = '***************';
 /* authentication types */
 export const sqlLogin = 'SqlLogin';
 export const integrated = 'Integrated';
+export const azureMFA = 'AzureMFA';
 

src/sql/parts/connection/connectionDialog/connectionWidget.ts

@@ -22,6 +22,8 @@ import { Dropdown } from 'sql/base/browser/ui/editableDropdown/dropdown';
 import { IConnectionManagementService } from 'sql/parts/connection/common/connectionManagement';
 import { ICapabilitiesService } from 'sql/services/capabilities/capabilitiesService';
 import { ConnectionProfile } from '../common/connectionProfile';
+import * as styler from 'sql/common/theme/styler';
+import { IAccountManagementService } from 'sql/services/accountManagement/interfaces';
 
 import * as sqlops from 'sqlops';
 
@@ -30,7 +32,6 @@ import { IContextViewService } from 'vs/platform/contextview/browser/contextView
 import { localize } from 'vs/nls';
 import * as DOM from 'vs/base/browser/dom';
 import { IThemeService } from 'vs/platform/theme/common/themeService';
-import * as styler from 'vs/platform/theme/common/styler';
 import { OS, OperatingSystem } from 'vs/base/common/platform';
 import { Builder, $ } from 'vs/base/browser/builder';
 import { MessageType } from 'vs/base/browser/ui/inputbox/inputBox';
@@ -50,6 +51,11 @@ export class ConnectionWidget {
 	private _passwordInputBox: InputBox;
 	private _password: string;
 	private _rememberPasswordCheckBox: Checkbox;
+	private _azureAccountDropdown: SelectBox;
+	private _refreshCredentialsLinkBuilder: Builder;
+	private _addAzureAccountMessage: string = localize('connectionWidget.AddAzureAccount', 'Add an account...');
+	private readonly _azureProviderId = 'azurePublicCloud';
+	private _azureAccountList: sqlops.Account[];
 	private _advancedButton: Button;
 	private _callbacks: IConnectionComponentCallbacks;
 	private _authTypeSelectBox: SelectBox;
@@ -59,7 +65,7 @@ export class ConnectionWidget {
 	private _focusedBeforeHandleOnConnection: HTMLElement;
 	private _providerName: string;
 	private _authTypeMap: { [providerName: string]: AuthenticationType[] } = {
-		[Constants.mssqlProviderName]: [new AuthenticationType(Constants.integrated, false), new AuthenticationType(Constants.sqlLogin, true)]
+		[Constants.mssqlProviderName]: [AuthenticationType.SqlLogin, AuthenticationType.Integrated, AuthenticationType.AzureMFA]
 	};
 	private _saveProfile: boolean;
 	private _databaseDropdownExpanded: boolean = false;
@@ -96,7 +102,8 @@ export class ConnectionWidget {
 		@IConnectionManagementService private _connectionManagementService: IConnectionManagementService,
 		@ICapabilitiesService private _capabilitiesService: ICapabilitiesService,
 		@IClipboardService private _clipboardService: IClipboardService,
-		@IConfigurationService private _configurationService: IConfigurationService
+		@IConfigurationService private _configurationService: IConfigurationService,
+		@IAccountManagementService private _accountManagementService: IAccountManagementService
 	) {
 		this._callbacks = callbacks;
 		this._toDispose = [];
@@ -109,9 +116,9 @@ export class ConnectionWidget {
 		var authTypeOption = this._optionsMaps[ConnectionOptionSpecialType.authType];
 		if (authTypeOption) {
 			if (OS === OperatingSystem.Windows) {
-				authTypeOption.defaultValue = this.getAuthTypeDisplayName(Constants.integrated);
+				authTypeOption.defaultValue = this.getAuthTypeDisplayName(AuthenticationType.Integrated);
 			} else {
-				authTypeOption.defaultValue = this.getAuthTypeDisplayName(Constants.sqlLogin);
+				authTypeOption.defaultValue = this.getAuthTypeDisplayName(AuthenticationType.SqlLogin);
 			}
 			this._authTypeSelectBox = new SelectBox(authTypeOption.categoryValues.map(c => c.displayName), authTypeOption.defaultValue, this._contextViewService, undefined, { ariaLabel: authTypeOption.displayName });
 		}
@@ -182,7 +189,7 @@ export class ConnectionWidget {
 		// Username
 		let self = this;
 		let userNameOption = this._optionsMaps[ConnectionOptionSpecialType.userName];
-		let userNameBuilder = DialogHelper.appendRow(this._tableContainer, userNameOption.displayName, 'connection-label', 'connection-input');
+		let userNameBuilder = DialogHelper.appendRow(this._tableContainer, userNameOption.displayName, 'connection-label', 'connection-input', 'username-password-row');
 		this._userNameInputBox = new InputBox(userNameBuilder.getHTMLElement(), this._contextViewService, {
 			validationOptions: {
 				validation: (value: string) => self.validateUsername(value, userNameOption.isRequired) ? ({ type: MessageType.ERROR, content: localize('connectionWidget.missingRequireField', '{0} is required.', userNameOption.displayName) }) : null
@@ -191,14 +198,22 @@ export class ConnectionWidget {
 		});
 		// Password
 		let passwordOption = this._optionsMaps[ConnectionOptionSpecialType.password];
-		let passwordBuilder = DialogHelper.appendRow(this._tableContainer, passwordOption.displayName, 'connection-label', 'connection-input');
+		let passwordBuilder = DialogHelper.appendRow(this._tableContainer, passwordOption.displayName, 'connection-label', 'connection-input', 'username-password-row');
 		this._passwordInputBox = new InputBox(passwordBuilder.getHTMLElement(), this._contextViewService, { ariaLabel: passwordOption.displayName });
 		this._passwordInputBox.inputElement.type = 'password';
 		this._password = '';
 
 		// Remember password
 		let rememberPasswordLabel = localize('rememberPassword', 'Remember password');
-		this._rememberPasswordCheckBox = this.appendCheckbox(this._tableContainer, rememberPasswordLabel, 'connection-checkbox', 'connection-input', false);
+		this._rememberPasswordCheckBox = this.appendCheckbox(this._tableContainer, rememberPasswordLabel, 'connection-checkbox', 'connection-input', 'username-password-row', false);
+
+		// Azure account picker
+		let accountLabel = localize('connection.azureAccountDropdownLabel', 'Account');
+		let accountDropdownBuilder = DialogHelper.appendRow(this._tableContainer, accountLabel, 'connection-label', 'connection-input', 'azure-account-row');
+		this._azureAccountDropdown = new SelectBox([], undefined, this._contextViewService, accountDropdownBuilder.getContainer(), { ariaLabel: accountLabel });
+		DialogHelper.appendInputSelectBox(accountDropdownBuilder, this._azureAccountDropdown);
+		let refreshCredentialsBuilder = DialogHelper.appendRow(this._tableContainer, '', 'connection-label', 'connection-input', 'azure-account-row refresh-credentials-link');
+		this._refreshCredentialsLinkBuilder = refreshCredentialsBuilder.a({ href: '#' }).text(localize('connectionWidget.refreshAzureCredentials', 'Refresh account credentials'));
 
 		// Database
 		let databaseOption = this._optionsMaps[ConnectionOptionSpecialType.databaseName];
@@ -228,7 +243,7 @@ export class ConnectionWidget {
 
 	private validateUsername(value: string, isOptionRequired: boolean): boolean {
 		let currentAuthType = this._authTypeSelectBox ? this.getMatchingAuthType(this._authTypeSelectBox.value) : undefined;
-		if (!currentAuthType || currentAuthType.showUsernameAndPassword) {
+		if (!currentAuthType || currentAuthType === AuthenticationType.SqlLogin) {
 			if (!value && isOptionRequired) {
 				return true;
 			}
@@ -254,9 +269,9 @@ export class ConnectionWidget {
 		return button;
 	}
 
-	private appendCheckbox(container: Builder, label: string, checkboxClass: string, cellContainerClass: string, isChecked: boolean): Checkbox {
+	private appendCheckbox(container: Builder, label: string, checkboxClass: string, cellContainerClass: string, rowContainerClass: string, isChecked: boolean): Checkbox {
 		let checkbox: Checkbox;
-		container.element('tr', {}, (rowContainer) => {
+		container.element('tr', { class: rowContainerClass }, (rowContainer) => {
 			rowContainer.element('td');
 			rowContainer.element('td', { class: cellContainerClass }, (inputCellContainer) => {
 				checkbox = new Checkbox(inputCellContainer.getHTMLElement(), { label, checked: isChecked, ariaLabel: label });
@@ -275,6 +290,7 @@ export class ConnectionWidget {
 		this._toDispose.push(styler.attachSelectBoxStyler(this._serverGroupSelectBox, this._themeService));
 		this._toDispose.push(attachButtonStyler(this._advancedButton, this._themeService));
 		this._toDispose.push(attachCheckboxStyler(this._rememberPasswordCheckBox, this._themeService));
+		this._toDispose.push(styler.attachSelectBoxStyler(this._azureAccountDropdown, this._themeService));
 
 		if (this._authTypeSelectBox) {
 			// Theme styler
@@ -285,6 +301,23 @@ export class ConnectionWidget {
 			}));
 		}
 
+		if (this._azureAccountDropdown) {
+			this._toDispose.push(styler.attachSelectBoxStyler(this._azureAccountDropdown, this._themeService));
+			this._toDispose.push(this._azureAccountDropdown.onDidSelect(() => {
+				this.onAzureAccountSelected();
+			}));
+		}
+
+		if (this._refreshCredentialsLinkBuilder) {
+			this._toDispose.push(this._refreshCredentialsLinkBuilder.on(DOM.EventType.CLICK, async () => {
+				let account = this._azureAccountList.find(account => account.key.accountId === this._azureAccountDropdown.value);
+				if (account) {
+					await this._accountManagementService.refreshAccount(account);
+					this.fillInAzureAccountOptions();
+				}
+			}));
+		}
+
 		this._toDispose.push(this._serverGroupSelectBox.onDidSelect(selectedGroup => {
 			this.onGroupSelected(selectedGroup.selected);
 		}));
@@ -342,7 +375,7 @@ export class ConnectionWidget {
 	private setConnectButton(): void {
 		let showUsernameAndPassword: boolean = true;
 		if (this.authType) {
-			showUsernameAndPassword = this.authType.showUsernameAndPassword;
+			showUsernameAndPassword = this.authType === AuthenticationType.SqlLogin;
 		}
 		showUsernameAndPassword ? this._callbacks.onSetConnectButton(!!this.serverName && !!this.userName) :
 			this._callbacks.onSetConnectButton(!!this.serverName);
@@ -350,7 +383,7 @@ export class ConnectionWidget {
 
 	private onAuthTypeSelected(selectedAuthType: string) {
 		let currentAuthType = this.getMatchingAuthType(selectedAuthType);
-		if (!currentAuthType.showUsernameAndPassword) {
+		if (currentAuthType !== AuthenticationType.SqlLogin) {
 			this._userNameInputBox.disable();
 			this._passwordInputBox.disable();
 			this._userNameInputBox.hideMessage();
@@ -366,6 +399,68 @@ export class ConnectionWidget {
 			this._passwordInputBox.enable();
 			this._rememberPasswordCheckBox.enabled = true;
 		}
+
+		if (currentAuthType === AuthenticationType.AzureMFA) {
+			this.fillInAzureAccountOptions();
+			this._azureAccountDropdown.enable();
+			let tableContainer = this._tableContainer.getContainer();
+			tableContainer.classList.add('hide-username-password');
+			tableContainer.classList.remove('hide-azure-accounts');
+		} else {
+			this._azureAccountDropdown.disable();
+			let tableContainer = this._tableContainer.getContainer();
+			tableContainer.classList.remove('hide-username-password');
+			tableContainer.classList.add('hide-azure-accounts');
+			this._azureAccountDropdown.hideMessage();
+		}
+	}
+
+	private async fillInAzureAccountOptions(): Promise<void> {
+		let oldSelection = this._azureAccountDropdown.value;
+		this._azureAccountList = await this._accountManagementService.getAccountsForProvider(this._azureProviderId);
+		let accountDropdownOptions = this._azureAccountList.map(account => account.key.accountId);
+		if (accountDropdownOptions.length === 0) {
+			// If there are no accounts add a blank option so that add account isn't automatically selected
+			accountDropdownOptions.unshift('');
+		}
+		accountDropdownOptions.push(this._addAzureAccountMessage);
+		this._azureAccountDropdown.setOptions(accountDropdownOptions);
+		this._azureAccountDropdown.selectWithOptionName(oldSelection);
+		this.updateRefreshCredentialsLink();
+	}
+
+	private async updateRefreshCredentialsLink(): Promise<void> {
+		let chosenAccount = this._azureAccountList.find(account => account.key.accountId === this._azureAccountDropdown.value);
+		if (chosenAccount && chosenAccount.isStale) {
+			this._tableContainer.getContainer().classList.remove('hide-refresh-link');
+		} else {
+			this._tableContainer.getContainer().classList.add('hide-refresh-link');
+		}
+	}
+
+	private async onAzureAccountSelected(): Promise<void> {
+		// Reset the dropdown's validation message if the old selection was not valid but the new one is
+		this.validateAzureAccountSelection(false);
+		this._refreshCredentialsLinkBuilder.display('none');
+
+		// Open the add account dialog if needed, then select the added account
+		if (this._azureAccountDropdown.value === this._addAzureAccountMessage) {
+			let oldAccountIds = this._azureAccountList.map(account => account.key.accountId);
+			await this._accountManagementService.addAccount(this._azureProviderId);
+
+			// Refresh the dropdown's list to include the added account
+			await this.fillInAzureAccountOptions();
+
+			// If a new account was added find it and select it, otherwise select the first account
+			let newAccount = this._azureAccountList.find(option => !oldAccountIds.some(oldId => oldId === option.key.accountId));
+			if (newAccount) {
+				this._azureAccountDropdown.selectWithOptionName(newAccount.key.accountId);
+			} else {
+				this._azureAccountDropdown.select(0);
+			}
+		}
+
+		this.updateRefreshCredentialsLink();
 	}
 
 	private serverNameChanged(serverName: string) {
@@ -407,6 +502,7 @@ export class ConnectionWidget {
 	private clearValidationMessages(): void {
 		this._serverNameInputBox.hideMessage();
 		this._userNameInputBox.hideMessage();
+		this._azureAccountDropdown.hideMessage();
 	}
 
 	private getModelValue(value: string): string {
@@ -449,8 +545,8 @@ export class ConnectionWidget {
 
 			if (this._authTypeSelectBox) {
 				this.onAuthTypeSelected(this._authTypeSelectBox.value);
-
 			}
+
 			// Disable connect button if -
 			// 1. Authentication type is SQL Login and no username is provided
 			// 2. No server name is provided
@@ -513,7 +609,7 @@ export class ConnectionWidget {
 			currentAuthType = this.getMatchingAuthType(this._authTypeSelectBox.value);
 		}
 
-		if (!currentAuthType || currentAuthType.showUsernameAndPassword) {
+		if (!currentAuthType || currentAuthType === AuthenticationType.SqlLogin) {
 			this._userNameInputBox.enable();
 			this._passwordInputBox.enable();
 			this._rememberPasswordCheckBox.enabled = true;
@@ -537,7 +633,7 @@ export class ConnectionWidget {
 	}
 
 	public get userName(): string {
-		return this._userNameInputBox.value;
+		return this.authenticationType === AuthenticationType.AzureMFA ? this._azureAccountDropdown.value : this._userNameInputBox.value;
 	}
 
 	public get password(): string {
@@ -548,6 +644,27 @@ export class ConnectionWidget {
 		return this._authTypeSelectBox ? this.getAuthTypeName(this._authTypeSelectBox.value) : undefined;
 	}
 
+	private validateAzureAccountSelection(showMessage: boolean = true): boolean {
+		if (this.authType !== AuthenticationType.AzureMFA) {
+			return true;
+		}
+
+		let selected = this._azureAccountDropdown.value;
+		if (selected === '' || selected === this._addAzureAccountMessage) {
+			if (showMessage) {
+				this._azureAccountDropdown.showMessage({
+					content: localize('connectionWidget.invalidAzureAccount', 'You must select an account'),
+					type: MessageType.ERROR
+				});
+			}
+			return false;
+		} else {
+			this._azureAccountDropdown.hideMessage();
+		}
+
+		return true;
+	}
+
 	private validateInputs(): boolean {
 		let isFocused = false;
 		let validateServerName = this._serverNameInputBox.validate();
@@ -565,7 +682,12 @@ export class ConnectionWidget {
 			this._passwordInputBox.focus();
 			isFocused = true;
 		}
-		return validateServerName && validateUserName && validatePassword;
+		let validateAzureAccount = this.validateAzureAccountSelection();
+		if (!validateAzureAccount && !isFocused) {
+			this._azureAccountDropdown.focus();
+			isFocused = true;
+		}
+		return validateServerName && validateUserName && validatePassword && validateAzureAccount;
 	}
 
 	public connect(model: IConnectionProfile): boolean {
@@ -613,7 +735,7 @@ export class ConnectionWidget {
 
 	private getMatchingAuthType(displayName: string): AuthenticationType {
 		const authType = this._authTypeMap[this._providerName];
-		return authType ? authType.find(authType => this.getAuthTypeDisplayName(authType.name) === displayName) : undefined;
+		return authType ? authType.find(authType => this.getAuthTypeDisplayName(authType) === displayName) : undefined;
 	}
 
 	public closeDatabaseDropdown(): void {
@@ -634,18 +756,14 @@ export class ConnectionWidget {
 	}
 
 	private focusPasswordIfNeeded(): void {
-		if (this.authType && this.authType.showUsernameAndPassword && this.userName && !this.password) {
+		if (this.authType && this.authType === AuthenticationType.SqlLogin && this.userName && !this.password) {
 			this._passwordInputBox.focus();
 		}
 	}
 }
 
-class AuthenticationType {
-	public name: string;
-	public showUsernameAndPassword: boolean;
-
-	constructor(name: string, showUsernameAndPassword: boolean) {
-		this.name = name;
-		this.showUsernameAndPassword = showUsernameAndPassword;
-	}
+enum AuthenticationType {
+	SqlLogin = 'SqlLogin',
+	Integrated = 'Integrated',
+	AzureMFA = 'AzureMFA'
 }

src/sql/parts/connection/connectionDialog/media/connectionDialog.css

@@ -28,11 +28,12 @@
 	overflow: hidden;
 	margin: 0px 11px;
 }
+
 .connection-dialog .tabBody {
 	overflow: hidden;
-    flex: 1 1;
-    display: flex;
-    flex-direction: column;
+	flex: 1 1;
+	display: flex;
+	flex-direction: column;
 }
 
 .connection-recent, .connection-saved {
@@ -114,4 +115,16 @@
 	margin: 5px 0px;
 	padding: 5px 15px;
 	font-weight: 600;
-}
+}
+
+.hide-azure-accounts .azure-account-row {
+	display: none;
+}
+
+.hide-username-password .username-password-row {
+	display: none;
+}
+
+.hide-refresh-link .azure-account-row.refresh-credentials-link {
+	display: none;
+}

src/sql/parts/objectExplorer/common/objectExplorerService.ts

@@ -261,7 +261,7 @@ export class ObjectExplorerService implements IObjectExplorerService {
 		return this._activeObjectExplorerNodes[connection.id];
 	}
 
-	public createNewSession(providerId: string, connection: ConnectionProfile): Thenable<sqlops.ObjectExplorerSessionResponse> {
+	public async createNewSession(providerId: string, connection: ConnectionProfile): Promise<sqlops.ObjectExplorerSessionResponse> {
 		let self = this;
 		return new Promise<sqlops.ObjectExplorerSessionResponse>((resolve, reject) => {
 			let provider = this._providers[providerId];

src/sql/services/accountManagement/accountManagementService.ts

@@ -20,7 +20,7 @@ import { AccountDialogController } from 'sql/parts/accountManagement/accountDial
 import { AutoOAuthDialogController } from 'sql/parts/accountManagement/autoOAuthDialog/autoOAuthDialogController';
 import { AccountListStatusbarItem } from 'sql/parts/accountManagement/accountListStatusbar/accountListStatusbarItem';
 import { AccountProviderAddedEventParams, UpdateAccountListEventParams } from 'sql/services/accountManagement/eventTypes';
-import { IAccountManagementService } from 'sql/services/accountManagement/interfaces';
+import { IAccountManagementService, AzureResource } from 'sql/services/accountManagement/interfaces';
 import { IClipboardService } from 'vs/platform/clipboard/common/clipboardService';
 
 export class AccountManagementService implements IAccountManagementService {
@@ -217,11 +217,12 @@ export class AccountManagementService implements IAccountManagementService {
 	/**
 	 * Generates a security token by asking the account's provider
 	 * @param {Account} account Account to generate security token for
+	 * @param {AzureResource} resource The resource to get the security token for
 	 * @return {Thenable<{}>} Promise to return the security token
 	 */
-	public getSecurityToken(account: sqlops.Account): Thenable<{}> {
+	public getSecurityToken(account: sqlops.Account, resource: sqlops.AzureResource): Thenable<{}> {
 		return this.doWithProvider(account.key.providerId, provider => {
-			return provider.provider.getSecurityToken(account);
+			return provider.provider.getSecurityToken(account, resource);
 		});
 	}
 

src/sql/services/accountManagement/interfaces.ts

@@ -22,7 +22,7 @@ export interface IAccountManagementService {
 	addAccount(providerId: string): Thenable<void>;
 	getAccountProviderMetadata(): Thenable<sqlops.AccountProviderMetadata[]>;
 	getAccountsForProvider(providerId: string): Thenable<sqlops.Account[]>;
-	getSecurityToken(account: sqlops.Account): Thenable<{}>;
+	getSecurityToken(account: sqlops.Account, resource: sqlops.AzureResource): Thenable<{}>;
 	removeAccount(accountKey: sqlops.AccountKey): Thenable<boolean>;
 	refreshAccount(account: sqlops.Account): Thenable<sqlops.Account>;
 
@@ -44,6 +44,12 @@ export interface IAccountManagementService {
 	readonly updateAccountListEvent: Event<UpdateAccountListEventParams>;
 }
 
+// Enum matching the AzureResource enum from sqlops.d.ts
+export enum AzureResource {
+	ResourceManagement = 0,
+	Sql = 1
+}
+
 export interface IAccountStore {
 	/**
 	 * Adds the provided account if the account doesn't exist. Updates the account if it already exists

src/sql/sqlops.d.ts

@@ -1915,7 +1915,7 @@ declare module 'sqlops' {
 		 * @param {Account} account Account to generate security token for
 		 * @return {Thenable<{}>} Promise to return the security token
 		 */
-		export function getSecurityToken(account: Account): Thenable<{}>;
+		export function getSecurityToken(account: Account, resource: AzureResource): Thenable<{}>;
 
 		/**
 		 * An [event](#Event) which fires when the accounts have changed.
@@ -1988,6 +1988,11 @@ declare module 'sqlops' {
 		isStale: boolean;
 	}
 
+	export enum AzureResource {
+		ResourceManagement = 0,
+		Sql = 1
+	}
+
 	export interface DidChangeAccountsParams {
 		// Updated accounts
 		accounts: Account[];
@@ -2045,9 +2050,10 @@ declare module 'sqlops' {
 		/**
 		 * Generates a security token for the provided account
 		 * @param {Account} account The account to generate a security token for
+		 * @param {AzureResource} resource The resource to get the token for
 		 * @return {Thenable<{}>} Promise to return a security token object
 		 */
-		getSecurityToken(account: Account): Thenable<{}>;
+		getSecurityToken(account: Account, resource: AzureResource): Thenable<{}>;
 
 		/**
 		 * Prompts the user to enter account information.

src/sql/workbench/api/common/sqlExtHostTypes.ts

@@ -313,6 +313,11 @@ export class TreeComponentItem extends TreeItem {
 	checked?: boolean;
 }
 
+export enum AzureResource {
+	ResourceManagement = 0,
+	Sql = 1
+}
+
 export class SqlThemeIcon {
 
 	static readonly Folder = new SqlThemeIcon('Folder');

src/sql/workbench/api/node/extHostAccountManagement.ts

@@ -89,12 +89,12 @@ export class ExtHostAccountManagement extends ExtHostAccountManagementShape {
 		return Promise.all(promises).then(() => resultAccounts);
 	}
 
-	public $getSecurityToken(account: sqlops.Account): Thenable<{}> {
+	public $getSecurityToken(account: sqlops.Account, resource: sqlops.AzureResource): Thenable<{}> {
 		return this.$getAllAccounts().then(() => {
 			for (const handle in this._accounts) {
 				const providerHandle = parseInt(handle);
 				if (this._accounts[handle].findIndex((acct) => acct.key.accountId === account.key.accountId) !== -1) {
-					return this._withProvider(providerHandle, (provider: sqlops.AccountProvider) => provider.getSecurityToken(account));
+					return this._withProvider(providerHandle, (provider: sqlops.AccountProvider) => provider.getSecurityToken(account, resource));
 				}
 			}
 

src/sql/workbench/api/node/mainThreadAccountManagement.ts

@@ -76,8 +76,8 @@ export class MainThreadAccountManagement implements MainThreadAccountManagementS
 			clear(accountKey: sqlops.AccountKey): Thenable<void> {
 				return self._proxy.$clear(handle, accountKey);
 			},
-			getSecurityToken(account: sqlops.Account): Thenable<{}> {
-				return self._proxy.$getSecurityToken(account);
+			getSecurityToken(account: sqlops.Account, resource: sqlops.AzureResource): Thenable<{}> {
+				return self._proxy.$getSecurityToken(account, resource);
 			},
 			initialize(restoredAccounts: sqlops.Account[]): Thenable<sqlops.Account[]> {
 				return self._proxy.$initialize(handle, restoredAccounts);

src/sql/workbench/api/node/sqlExtHost.api.impl.ts

@@ -97,8 +97,8 @@ export function createApiFactory(
 				getAllAccounts(): Thenable<sqlops.Account[]> {
 					return extHostAccountManagement.$getAllAccounts();
 				},
-				getSecurityToken(account: sqlops.Account): Thenable<{}> {
-					return extHostAccountManagement.$getSecurityToken(account);
+				getSecurityToken(account: sqlops.Account, resource: sqlops.AzureResource): Thenable<{}> {
+					return extHostAccountManagement.$getSecurityToken(account, resource);
 				},
 				onDidChangeAccounts(listener: (e: sqlops.DidChangeAccountsParams) => void, thisArgs?: any, disposables?: extHostTypes.Disposable[]) {
 					return extHostAccountManagement.onDidChangeAccounts(listener, thisArgs, disposables);
@@ -452,7 +452,8 @@ export function createApiFactory(
 				Orientation: sqlExtHostTypes.Orientation,
 				SqlThemeIcon: sqlExtHostTypes.SqlThemeIcon,
 				TreeComponentItem: sqlExtHostTypes.TreeComponentItem,
-				nb: nb
+				nb: nb,
+				AzureResource: sqlExtHostTypes.AzureResource
 			};
 		}
 	};

src/sql/workbench/api/node/sqlExtHost.protocol.ts

@@ -27,7 +27,7 @@ import {
 export abstract class ExtHostAccountManagementShape {
 	$autoOAuthCancelled(handle: number): Thenable<void> { throw ni(); }
 	$clear(handle: number, accountKey: sqlops.AccountKey): Thenable<void> { throw ni(); }
-	$getSecurityToken(account: sqlops.Account): Thenable<{}> { throw ni(); }
+	$getSecurityToken(account: sqlops.Account, resource?: sqlops.AzureResource): Thenable<{}> { throw ni(); }
 	$initialize(handle: number, restoredAccounts: sqlops.Account[]): Thenable<sqlops.Account[]> { throw ni(); }
 	$prompt(handle: number): Thenable<sqlops.Account> { throw ni(); }
 	$refresh(handle: number, account: sqlops.Account): Thenable<sqlops.Account> { throw ni(); }