Modify connection management to support Active Directory authType for non-SQL DBs (#8434)

* First attempt to add Azure MFA login for PostgreSQL

* Finish merge with master

* Fix auth type default selection

* Add AzureMFAAndUser auth type for Orcas

* Fix formatting

* Update change log

* Incorporate some review comments

* Missed an occurrence of AzureResource

* Try to move all changes out of azdata.d.ts and sqlops.d.ts

* Concrete implementation of ConnectionProfile in azdata no longer has azureAccount

* Use enum names instead of numbers in config files

src/sql/workbench/api/browser/mainThreadConnectionManagement.ts

@@ -60,6 +60,7 @@ export class MainThreadConnectionManagement extends Disposable implements MainTh
 				saveProfile: inputProfile.saveProfile,
 				id: inputProfile.id,
 				azureTenantId: inputProfile.azureTenantId,
+				azureAccount: inputProfile.azureAccount,
 				options: inputProfile.options
 			};
 			return outputProfile;

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

@@ -390,7 +390,8 @@ export class TreeComponentItem extends vsExtTypes.TreeItem {
 
 export enum AzureResource {
 	ResourceManagement = 0,
-	Sql = 1
+	Sql = 1,
+	OssRdbms = 2
 }
 
 export class TreeItem extends vsExtTypes.TreeItem {
@@ -697,6 +698,14 @@ export class ConnectionProfile {
 		this.options['azureTenantId'] = value;
 	}
 
+	get azureAccount(): string {
+		return this.options['azureAccount'];
+	}
+
+	set azureAccount(value: string) {
+		this.options['azureAccount'] = value;
+	}
+
 	options: { [key: string]: any } = {};
 
 	static createFrom(options: { [key: string]: any }): ConnectionProfile {

src/sql/workbench/contrib/connection/common/connectionProviderExtension.ts

@@ -132,6 +132,21 @@ const ConnectionProviderContrib: IJSONSchema = {
 					defaultValue: {
 						type: 'any'
 					},
+					defaultValueOsOverrides: {
+						type: 'array',
+						items: {
+							type: 'object',
+							properties: {
+								os: {
+									type: 'string',
+									enum: ['Windows', 'Macintosh', 'Linux']
+								},
+								defaultValueOverride: {
+									type: 'any'
+								}
+							}
+						}
+					},
 					objectType: {
 						type: 'any'
 					},

src/sql/workbench/contrib/notebook/test/electron-browser/notebookUtils.test.ts

@@ -26,7 +26,8 @@ suite('notebookUtils', function (): void {
 		saveProfile: true,
 		id: '',
 		options: {},
-		azureTenantId: undefined
+		azureTenantId: undefined,
+		azureAccount: undefined
 	};
 
 	test('Should format server and database name correctly for attach to', async function (): Promise<void> {

src/sql/workbench/services/connection/browser/cmsConnectionWidget.ts

@@ -56,12 +56,9 @@ export class CmsConnectionWidget extends ConnectionWidget {
 			_clipboardService, _configurationService, _accountManagementService, _logService);
 		let authTypeOption = this._optionsMaps[ConnectionOptionSpecialType.authType];
 		if (authTypeOption) {
-			if (OS === OperatingSystem.Windows) {
-				authTypeOption.defaultValue = this.getAuthTypeDisplayName(AuthenticationType.Integrated);
-			} else {
-				authTypeOption.defaultValue = this.getAuthTypeDisplayName(AuthenticationType.SqlLogin);
-			}
-			this._authTypeSelectBox = new SelectBox(authTypeOption.categoryValues.map(c => c.displayName), authTypeOption.defaultValue, this._contextViewService, undefined, { ariaLabel: authTypeOption.displayName });
+			let authTypeDefault = this.getAuthTypeDefault(authTypeOption, OS);
+			let authTypeDefaultDisplay = this.getAuthTypeDisplayName(authTypeDefault);
+			this._authTypeSelectBox = new SelectBox(authTypeOption.categoryValues.map(c => c.displayName), authTypeDefaultDisplay, this._contextViewService, undefined, { ariaLabel: authTypeOption.displayName });
 		}
 	}
 

src/sql/workbench/services/connection/browser/connectionManagementService.ts

@@ -72,6 +72,8 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 	private _mementoContext: Memento;
 	private _mementoObj: any;
 	private static readonly CONNECTION_MEMENTO = 'ConnectionManagement';
+	private static readonly _azureResources: AzureResource[] =
+		[AzureResource.ResourceManagement, AzureResource.Sql, AzureResource.OssRdbms];
 
 	constructor(
 		private _connectionStore: ConnectionStore,
@@ -700,22 +702,40 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 		return defaultProvider && this._providers.has(defaultProvider) ? defaultProvider : undefined;
 	}
 
+	/**
+	 * Previously, the only resource available for AAD access tokens was for Azure SQL / SQL Server.
+	 * Use that as a default if the provider extension does not configure a different one. If one is
+	 * configured, then use it.
+	 * @param connection The connection to fill in or update
+	 */
+	private getAzureResourceForConnection(connection: interfaces.IConnectionProfile): azdata.AzureResource {
+		let provider = this._providers.get(connection.providerName);
+		if (!provider || !provider.properties || !provider.properties.azureResource) {
+			return AzureResource.Sql;
+		}
+
+		let result = find(ConnectionManagementService._azureResources, r => AzureResource[r] === provider.properties.azureResource);
+		return result ? result : AzureResource.Sql;
+	}
+
 	/**
 	 * Fills in the Azure account token if it's needed for this connection and doesn't already have one
 	 * and clears it if it isn't.
 	 * @param connection The connection to fill in or update
 	 */
 	private async fillInOrClearAzureToken(connection: interfaces.IConnectionProfile): Promise<boolean> {
-		if (connection.authenticationType !== Constants.azureMFA) {
+		if (connection.authenticationType !== Constants.azureMFA && connection.authenticationType !== Constants.azureMFAAndUser) {
 			connection.options['azureAccountToken'] = undefined;
 			return true;
 		}
 		if (connection.options['azureAccountToken']) {
 			return true;
 		}
+		let azureResource = this.getAzureResourceForConnection(connection);
 		let accounts = await this._accountManagementService.getAccountsForProvider('azurePublicCloud');
 		if (accounts && accounts.length > 0) {
-			let account = find(accounts, account => account.key.accountId === connection.userName);
+			let accountName = (connection.authenticationType !== Constants.azureMFA) ? connection.azureAccount : connection.userName;
+			let account = find(accounts, account => account.key.accountId === accountName);
 			if (account) {
 				if (account.isStale) {
 					try {
@@ -725,7 +745,7 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 						return false;
 					}
 				}
-				let tokensByTenant = await this._accountManagementService.getSecurityToken(account, AzureResource.Sql);
+				let tokensByTenant = await this._accountManagementService.getSecurityToken(account, azureResource);
 				let token: string;
 				let tenantId = connection.azureTenantId;
 				if (tenantId && tokensByTenant[tenantId]) {

src/sql/workbench/services/connection/browser/connectionWidget.ts

@@ -37,6 +37,13 @@ import { IConfigurationService } from 'vs/platform/configuration/common/configur
 import { ILogService } from 'vs/platform/log/common/log';
 import { find } from 'vs/base/common/arrays';
 
+export enum AuthenticationType {
+	SqlLogin = 'SqlLogin',
+	Integrated = 'Integrated',
+	AzureMFA = 'AzureMFA',
+	AzureMFAAndUser = 'AzureMFAAndUser'
+}
+
 export class ConnectionWidget extends lifecycle.Disposable {
 	private _previousGroupOption: string;
 	private _serverGroupOptions: IConnectionProfileGroup[];
@@ -65,12 +72,16 @@ export class ConnectionWidget extends lifecycle.Disposable {
 	protected _optionsMaps: { [optionType: number]: azdata.ConnectionOption };
 	protected _tableContainer: HTMLElement;
 	protected _providerName: string;
-	protected _authTypeMap: { [providerName: string]: AuthenticationType[] } = {
-		[Constants.mssqlProviderName]: [AuthenticationType.SqlLogin, AuthenticationType.Integrated, AuthenticationType.AzureMFA]
-	};
 	protected _connectionNameInputBox: InputBox;
 	protected _databaseNameInputBox: Dropdown;
 	protected _advancedButton: Button;
+	private static readonly _authTypes: AuthenticationType[] =
+		[AuthenticationType.AzureMFA, AuthenticationType.AzureMFAAndUser, AuthenticationType.Integrated, AuthenticationType.SqlLogin];
+	private static readonly _osByName = {
+		Windows: OperatingSystem.Windows,
+		Macintosh: OperatingSystem.Macintosh,
+		Linux: OperatingSystem.Linux
+	};
 	public DefaultServerGroup: IConnectionProfileGroup = {
 		id: '',
 		name: localize('defaultServerGroup', "<Default>"),
@@ -115,16 +126,36 @@ export class ConnectionWidget extends lifecycle.Disposable {
 
 		let authTypeOption = this._optionsMaps[ConnectionOptionSpecialType.authType];
 		if (authTypeOption) {
-			if (OS === OperatingSystem.Windows) {
-				authTypeOption.defaultValue = this.getAuthTypeDisplayName(AuthenticationType.Integrated);
-			} else {
-				authTypeOption.defaultValue = this.getAuthTypeDisplayName(AuthenticationType.SqlLogin);
-			}
-			this._authTypeSelectBox = new SelectBox(authTypeOption.categoryValues.map(c => c.displayName), authTypeOption.defaultValue, this._contextViewService, undefined, { ariaLabel: authTypeOption.displayName });
+			let authTypeDefault = this.getAuthTypeDefault(authTypeOption, OS);
+			let authTypeDefaultDisplay = this.getAuthTypeDisplayName(authTypeDefault);
+			this._authTypeSelectBox = new SelectBox(authTypeOption.categoryValues.map(c => c.displayName), authTypeDefaultDisplay, this._contextViewService, undefined, { ariaLabel: authTypeOption.displayName });
 		}
 		this._providerName = providerName;
 	}
 
+	protected getAuthTypeDefault(option: azdata.ConnectionOption, os: OperatingSystem): string {
+		// Check for OS-specific default value
+		if (option.defaultValueOsOverrides) {
+			let result = find(option.defaultValueOsOverrides, d => ConnectionWidget._osByName[d.os] === os);
+			if (result) {
+				return result.defaultValueOverride;
+			}
+		}
+
+		// No OS-specific default, and so return global default, if any.
+		if (option.defaultValue) {
+			return option.defaultValue;
+		}
+
+		// No default value specified at all. Return first category value.
+		if (option.categoryValues.length > 0) {
+			return option.categoryValues[0].name;
+		}
+
+		// Give up.
+		return undefined;
+	}
+
 	public createConnectionWidget(container: HTMLElement, authTypeChanged: boolean = false): void {
 		this._serverGroupOptions = [this.DefaultServerGroup];
 		this._serverGroupSelectBox = new SelectBox(this._serverGroupOptions.map(g => g.name), this.DefaultServerGroup.name, this._contextViewService, undefined, { ariaLabel: this._serverGroupDisplayString });
@@ -210,7 +241,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 		// Username
 		let self = this;
 		let userNameOption = this._optionsMaps[ConnectionOptionSpecialType.userName];
-		let userName = DialogHelper.appendRow(this._tableContainer, userNameOption.displayName, 'connection-label', 'connection-input', 'username-password-row');
+		let userName = DialogHelper.appendRow(this._tableContainer, userNameOption.displayName, 'connection-label', 'connection-input', 'username-row');
 		this._userNameInputBox = new InputBox(userName, this._contextViewService, {
 			validationOptions: {
 				validation: (value: string) => self.validateUsername(value, userNameOption.isRequired) ? ({ type: MessageType.ERROR, content: localize('connectionWidget.missingRequireField', "{0} is required.", userNameOption.displayName) }) : null
@@ -219,14 +250,14 @@ export class ConnectionWidget extends lifecycle.Disposable {
 		});
 		// Password
 		let passwordOption = this._optionsMaps[ConnectionOptionSpecialType.password];
-		let password = DialogHelper.appendRow(this._tableContainer, passwordOption.displayName, 'connection-label', 'connection-input', 'username-password-row');
+		let password = DialogHelper.appendRow(this._tableContainer, passwordOption.displayName, 'connection-label', 'connection-input', 'password-row');
 		this._passwordInputBox = new InputBox(password, 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-input', 'username-password-row', false);
+		this._rememberPasswordCheckBox = this.appendCheckbox(this._tableContainer, rememberPasswordLabel, 'connection-input', 'password-row', false);
 
 		// Azure account picker
 		let accountLabel = localize('connection.azureAccountDropdownLabel', "Account");
@@ -414,17 +445,26 @@ export class ConnectionWidget extends lifecycle.Disposable {
 	}
 
 	private setConnectButton(): void {
-		let showUsernameAndPassword: boolean;
+		let showUsername: boolean;
 		if (this.authType) {
-			showUsernameAndPassword = this.authType === AuthenticationType.SqlLogin;
+			showUsername = this.authType === AuthenticationType.SqlLogin || this.authType === AuthenticationType.AzureMFAAndUser;
 		}
-		showUsernameAndPassword ? this._callbacks.onSetConnectButton(!!this.serverName && !!this.userName) :
+		showUsername ? this._callbacks.onSetConnectButton(!!this.serverName && !!this.userName) :
 			this._callbacks.onSetConnectButton(!!this.serverName);
 	}
 
 	protected onAuthTypeSelected(selectedAuthType: string) {
 		let currentAuthType = this.getMatchingAuthType(selectedAuthType);
-		if (currentAuthType !== AuthenticationType.SqlLogin) {
+		if (currentAuthType === AuthenticationType.AzureMFAAndUser) {
+			this._userNameInputBox.enable();
+			this._passwordInputBox.disable();
+			this._passwordInputBox.hideMessage();
+			this._passwordInputBox.value = '';
+			this._password = '';
+
+			this._rememberPasswordCheckBox.checked = false;
+			this._rememberPasswordCheckBox.enabled = false;
+		} else if (currentAuthType !== AuthenticationType.SqlLogin) {
 			this._userNameInputBox.disable();
 			this._passwordInputBox.disable();
 			this._userNameInputBox.hideMessage();
@@ -442,18 +482,37 @@ export class ConnectionWidget extends lifecycle.Disposable {
 		}
 
 		if (currentAuthType === AuthenticationType.AzureMFA) {
-			this.fillInAzureAccountOptions().then(() => {
+			this.fillInAzureAccountOptions().then(async () => {
 				// Don't enable the control until we've populated it
 				this._azureAccountDropdown.enable();
+				// Populate tenants
+				await this.onAzureAccountSelected();
+				this._azureTenantDropdown.enable();
 			}).catch(err => this._logService.error(`Unexpected error populating Azure Account dropdown : ${err}`));
 			// Immediately show/hide appropriate elements though so user gets immediate feedback while we load accounts
-			DOM.addClass(this._tableContainer, 'hide-username-password');
+			DOM.addClass(this._tableContainer, 'hide-username');
+			DOM.addClass(this._tableContainer, 'hide-password');
+			DOM.removeClass(this._tableContainer, 'hide-azure-accounts');
+		} else if (currentAuthType === AuthenticationType.AzureMFAAndUser) {
+			this.fillInAzureAccountOptions().then(async () => {
+				// Don't enable the control until we've populated it
+				this._azureAccountDropdown.enable();
+				// Populate tenants
+				await this.onAzureAccountSelected();
+				this._azureTenantDropdown.enable();
+			}).catch(err => this._logService.error(`Unexpected error populating Azure Account dropdown : ${err}`));
+			// Immediately show/hide appropriate elements though so user gets immediate feedback while we load accounts
+			DOM.removeClass(this._tableContainer, 'hide-username');
+			DOM.addClass(this._tableContainer, 'hide-password');
 			DOM.removeClass(this._tableContainer, 'hide-azure-accounts');
 		} else {
 			this._azureAccountDropdown.disable();
-			DOM.removeClass(this._tableContainer, 'hide-username-password');
-			DOM.addClass(this._tableContainer, 'hide-azure-accounts');
 			this._azureAccountDropdown.hideMessage();
+			this._azureTenantDropdown.disable();
+			this._azureTenantDropdown.hideMessage();
+			DOM.removeClass(this._tableContainer, 'hide-username');
+			DOM.removeClass(this._tableContainer, 'hide-password');
+			DOM.addClass(this._tableContainer, 'hide-azure-accounts');
 		}
 	}
 
@@ -629,13 +688,16 @@ export class ConnectionWidget extends lifecycle.Disposable {
 			if (this._authTypeSelectBox) {
 				this.onAuthTypeSelected(this._authTypeSelectBox.value);
 			} else {
-				DOM.removeClass(this._tableContainer, 'hide-username-password');
+				DOM.removeClass(this._tableContainer, 'hide-username');
+				DOM.removeClass(this._tableContainer, 'hide-password');
 				DOM.addClass(this._tableContainer, 'hide-azure-accounts');
 			}
 
-			if (this.authType === AuthenticationType.AzureMFA) {
+			if (this.authType === AuthenticationType.AzureMFA || this.authType === AuthenticationType.AzureMFAAndUser) {
 				this.fillInAzureAccountOptions().then(async () => {
-					this._azureAccountDropdown.selectWithOptionName(this.getModelValue(connectionInfo.userName));
+					let accountName = (this.authType === AuthenticationType.AzureMFA)
+						? connectionInfo.userName : connectionInfo.azureAccount;
+					this._azureAccountDropdown.selectWithOptionName(this.getModelValue(accountName));
 					await this.onAzureAccountSelected();
 					let tenantId = connectionInfo.azureTenantId;
 					let account = find(this._azureAccountList, account => account.key.accountId === this._azureAccountDropdown.value);
@@ -659,7 +721,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 
 	protected getAuthTypeDisplayName(authTypeName: string) {
 		let displayName: string;
-		let authTypeOption = this._optionsMaps[ConnectionOptionSpecialType.authType];
+		let authTypeOption: azdata.ConnectionOption = this._optionsMaps[ConnectionOptionSpecialType.authType];
 
 		if (authTypeOption) {
 			authTypeOption.categoryValues.forEach(c => {
@@ -673,7 +735,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 
 	private getAuthTypeName(authTypeDisplayName: string) {
 		let authTypeName: string;
-		let authTypeOption = this._optionsMaps[ConnectionOptionSpecialType.authType];
+		let authTypeOption: azdata.ConnectionOption = this._optionsMaps[ConnectionOptionSpecialType.authType];
 		authTypeOption.categoryValues.forEach(c => {
 			if (c.displayName === authTypeDisplayName) {
 				authTypeName = c.name;
@@ -716,6 +778,8 @@ export class ConnectionWidget extends lifecycle.Disposable {
 			this._userNameInputBox.enable();
 			this._passwordInputBox.enable();
 			this._rememberPasswordCheckBox.enabled = true;
+		} else if (currentAuthType === AuthenticationType.AzureMFAAndUser) {
+			this._userNameInputBox.enable();
 		}
 
 		if (this._focusedBeforeHandleOnConnection) {
@@ -755,8 +819,12 @@ export class ConnectionWidget extends lifecycle.Disposable {
 		return this._authTypeSelectBox ? this.getAuthTypeName(this._authTypeSelectBox.value) : undefined;
 	}
 
+	public get azureAccount(): string {
+		return this.authenticationType === AuthenticationType.AzureMFAAndUser ? this._azureAccountDropdown.value : undefined;
+	}
+
 	private validateAzureAccountSelection(showMessage: boolean = true): boolean {
-		if (this.authType !== AuthenticationType.AzureMFA) {
+		if (this.authType !== AuthenticationType.AzureMFA && this.authType !== AuthenticationType.AzureMFAAndUser) {
 			return true;
 		}
 
@@ -808,6 +876,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 			model.userName = this.userName;
 			model.password = this.password;
 			model.authenticationType = this.authenticationType;
+			model.azureAccount = this.azureAccount;
 			model.savePassword = this._rememberPasswordCheckBox.checked;
 			model.connectionName = this.connectionName;
 			model.databaseName = this.databaseName;
@@ -825,7 +894,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 					model.groupId = this.findGroupId(model.groupFullName);
 				}
 			}
-			if (this.authType === AuthenticationType.AzureMFA) {
+			if (this.authType === AuthenticationType.AzureMFA || this.authType === AuthenticationType.AzureMFAAndUser) {
 				model.azureTenantId = this._azureTenantId;
 			}
 		}
@@ -846,8 +915,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 	}
 
 	private getMatchingAuthType(displayName: string): AuthenticationType {
-		const authType = this._authTypeMap[this._providerName];
-		return authType ? find(authType, authType => this.getAuthTypeDisplayName(authType) === displayName) : undefined;
+		return find(ConnectionWidget._authTypes, authType => this.getAuthTypeDisplayName(authType) === displayName);
 	}
 
 	public closeDatabaseDropdown(): void {
@@ -873,9 +941,3 @@ export class ConnectionWidget extends lifecycle.Disposable {
 		}
 	}
 }
-
-export enum AuthenticationType {
-	SqlLogin = 'SqlLogin',
-	Integrated = 'Integrated',
-	AzureMFA = 'AzureMFA'
-}

src/sql/workbench/services/connection/browser/media/connectionDialog.css

@@ -138,7 +138,11 @@
 	display: none;
 }
 
-.hide-username-password .username-password-row {
+.hide-username .username-row {
+	display: none;
+}
+
+.hide-password .password-row {
 	display: none;
 }