Connection management service updates to support multiple providers (#9698)

* Connection management service work

* Fix tests

* Change how accounts are deleted

* Be consistent with names

* feedback

* Fix based on feedback

* Change sqltoolsservice version

src/sql/workbench/services/accountManagement/browser/accountDialog.ts

@@ -37,6 +37,8 @@ import { IAdsTelemetryService } from 'sql/platform/telemetry/common/telemetry';
 import { IViewPaneOptions, ViewPane } from 'vs/workbench/browser/parts/views/viewPaneContainer';
 import { attachModalDialogStyler, attachPanelStyler } from 'sql/workbench/common/styler';
 import { IViewDescriptorService } from 'vs/workbench/common/views';
+import { IQuickInputService, IQuickPickItem } from 'vs/platform/quickinput/common/quickInput';
+import { INotificationService } from 'vs/platform/notification/common/notification';
 import { IOpenerService } from 'vs/platform/opener/common/opener';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { ILayoutService } from 'vs/platform/layout/browser/layoutService';
@@ -135,8 +137,10 @@ export class AccountDialog extends Modal {
 		@ILogService logService: ILogService,
 		@IViewDescriptorService private viewDescriptorService: IViewDescriptorService,
 		@ITextResourcePropertiesService textResourcePropertiesService: ITextResourcePropertiesService,
+		@IQuickInputService private _quickInputService: IQuickInputService,
+		@INotificationService private _notificationService: INotificationService,
 		@IOpenerService protected readonly openerService: IOpenerService,
-		@ITelemetryService private readonly vstelemetryService: ITelemetryService,
+		@ITelemetryService private readonly vstelemetryService: ITelemetryService
 	) {
 		super(
 			localize('linkedAccounts', "Linked accounts"),
@@ -199,8 +203,34 @@ export class AccountDialog extends Modal {
 		this._addAccountButton = new Button(buttonSection);
 		this._addAccountButton.label = localize('accountDialog.addConnection', "Add an account");
 
-		this._register(this._addAccountButton.onDidClick(() => {
-			(<IProviderViewUiComponent>values(this._providerViewsMap)[0]).addAccountAction.run();
+		this._register(this._addAccountButton.onDidClick(async () => {
+			const vals = values(this._providerViewsMap);
+
+			let pickedValue: string;
+			if (vals.length === 0) {
+				this._notificationService.error(localize('accountDialog.noCloudsRegistered', "You have no clouds enabled. Go to Settings -> Search Azure Account Configuration -> Enable at least one cloud"));
+				return;
+			}
+			if (vals.length > 1) {
+				const buttons: IQuickPickItem[] = vals.map(v => {
+					return { label: v.view.title } as IQuickPickItem;
+				});
+
+				const picked = await this._quickInputService.pick(buttons, { canPickMany: false });
+
+				pickedValue = picked?.label;
+			} else {
+				pickedValue = vals[0].view.title;
+			}
+
+			const v = vals.filter(v => v.view.title === pickedValue)?.[0];
+
+			if (!v) {
+				this._notificationService.error(localize('accountDialog.didNotPickAuthProvider', "You didn't select any authentication provider. Please try again."));
+				return;
+			}
+
+			v.addAccountAction.run();
 		}));
 
 		DOM.append(container, this._noaccountViewContainer);

src/sql/workbench/services/accountManagement/browser/accountManagementService.ts

@@ -206,6 +206,13 @@ export class AccountManagementService implements IAccountManagementService {
 		});
 	}
 
+	/**
+	 * Retrieves all the accounts registered with ADS.
+	 */
+	public getAccounts(): Thenable<azdata.Account[]> {
+		return this._accountStore.getAllAccounts();
+	}
+
 	/**
 	 * Generates a security token by asking the account's provider
 	 * @param account Account to generate security token for
@@ -225,35 +232,49 @@ export class AccountManagementService implements IAccountManagementService {
 	 *                           removed, false otherwise.
 	 */
 	public removeAccount(accountKey: azdata.AccountKey): Thenable<boolean> {
-		let self = this;
 
 		// Step 1) Remove the account
 		// Step 2) Clear the sensitive data from the provider (regardless of whether the account was removed)
 		// Step 3) Update the account cache and fire an event
-		return this.doWithProvider(accountKey.providerId, provider => {
-			return this._accountStore.remove(accountKey)
-				.then(result => {
-					provider.provider.clear(accountKey);
-					return result;
-				})
-				.then(result => {
-					if (!result) {
-						return result;
-					}
+		return this.doWithProvider(accountKey.providerId, async provider => {
+			const result = await this._accountStore.remove(accountKey);
+			await provider.provider.clear(accountKey);
+			if (!result) {
+				return result;
+			}
 
-					let indexToRemove: number = firstIndex(provider.accounts, account => {
-						return account.key.accountId === accountKey.accountId;
-					});
+			let indexToRemove: number = firstIndex(provider.accounts, account => {
+				return account.key.accountId === accountKey.accountId;
+			});
 
-					if (indexToRemove >= 0) {
-						provider.accounts.splice(indexToRemove, 1);
-						self.fireAccountListUpdate(provider, false);
-					}
-					return result;
-				});
+			if (indexToRemove >= 0) {
+				provider.accounts.splice(indexToRemove, 1);
+				this.fireAccountListUpdate(provider, false);
+			}
+			return result;
 		});
 	}
 
+	/**
+	 * Removes all registered accounts
+	 */
+	public async removeAccounts(): Promise<boolean> {
+		const accounts = await this.getAccounts();
+		if (accounts.length === 0) {
+			return false;
+		}
+
+		let finalResult = true;
+		for (const account of accounts) {
+			const removeResult = await this.removeAccount(account.key);
+			if (removeResult === false) {
+				this.logService.info('Error when removing %s.', account.key);
+				finalResult = false;
+			}
+		}
+		return finalResult;
+	}
+
 	// UI METHODS //////////////////////////////////////////////////////////
 	/**
 	 * Opens the account list dialog

src/sql/workbench/services/accountManagement/browser/accountPickerService.ts

@@ -51,7 +51,7 @@ export class AccountPickerService implements IAccountPickerService {
 	public renderAccountPicker(container: HTMLElement): void {
 		if (!this._accountPicker) {
 			// TODO: expand support to multiple providers
-			const providerId: string = 'azurePublicCloud';
+			const providerId: string = 'azure_publicCloud';
 			this._accountPicker = this._instantiationService.createInstance(AccountPicker, providerId);
 			this._accountPicker.createAccountPickerComponent();
 		}

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

@@ -780,9 +780,9 @@ export class ConnectionManagementService extends Disposable implements IConnecti
 			return true;
 		}
 		let azureResource = this.getAzureResourceForConnection(connection);
-		let accounts = await this._accountManagementService.getAccountsForProvider('azurePublicCloud');
+		let accounts = (await this._accountManagementService.getAccounts()).filter(a => a.key.providerId.startsWith('azure'));
 		if (accounts && accounts.length > 0) {
-			let accountName = (connection.authenticationType !== Constants.azureMFA) ? connection.azureAccount : 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) {

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

@@ -56,7 +56,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 	private _azureTenantDropdown: SelectBox;
 	private _refreshCredentialsLink: HTMLLinkElement;
 	private _addAzureAccountMessage: string = localize('connectionWidget.AddAzureAccount', "Add an account...");
-	private readonly _azureProviderId = 'azurePublicCloud';
+	private readonly _azureProviderId = 'azure_publicCloud';
 	private _azureTenantId: string;
 	private _azureAccountList: azdata.Account[];
 	private _callbacks: IConnectionComponentCallbacks;
@@ -518,7 +518,8 @@ export class ConnectionWidget extends lifecycle.Disposable {
 
 	private async fillInAzureAccountOptions(): Promise<void> {
 		let oldSelection = this._azureAccountDropdown.value;
-		this._azureAccountList = await this._accountManagementService.getAccountsForProvider(this._azureProviderId);
+		const accounts = await this._accountManagementService.getAccounts();
+		this._azureAccountList = accounts.filter(a => a.key.providerId.startsWith('azure'));
 		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
@@ -821,7 +822,7 @@ export class ConnectionWidget extends lifecycle.Disposable {
 	}
 
 	public get azureAccount(): string {
-		return this.authenticationType === AuthenticationType.AzureMFAAndUser ? this._azureAccountDropdown.value : undefined;
+		return this.authenticationType === AuthenticationType.AzureMFAAndUser || this.authenticationType === AuthenticationType.AzureMFA ? this._azureAccountDropdown.value : undefined;
 	}
 
 	private validateAzureAccountSelection(showMessage: boolean = true): boolean {

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

@@ -796,9 +796,10 @@ suite('SQL ConnectionManagementService tests', () => {
 		let azureConnectionProfile = ConnectionProfile.fromIConnectionProfile(capabilitiesService, connectionProfile);
 		azureConnectionProfile.authenticationType = 'AzureMFA';
 		let username = 'testuser@microsoft.com';
-		azureConnectionProfile.userName = username;
+		azureConnectionProfile.azureAccount = username;
 		let servername = 'test-database.database.windows.net';
 		azureConnectionProfile.serverName = servername;
+		let providerId = 'azure_PublicCloud';
 
 		// Set up the account management service to return a token for the given user
 		accountManagementService.setup(x => x.getAccountsForProvider(TypeMoq.It.isAny())).returns(providerId => Promise.resolve<azdata.Account[]>([
@@ -812,9 +813,23 @@ suite('SQL ConnectionManagementService tests', () => {
 				properties: undefined
 			}
 		]));
+
+		accountManagementService.setup(x => x.getAccounts()).returns(() => {
+			return Promise.resolve<azdata.Account[]>([
+				{
+					key: {
+						accountId: username,
+						providerId: providerId
+					},
+					displayInfo: undefined,
+					isStale: false,
+					properties: undefined
+				}
+			]);
+		});
 		let testToken = 'testToken';
 		accountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve({
-			azurePublicCloud: {
+			azure_publicCloud: {
 				token: testToken
 			}
 		}));
@@ -827,7 +842,7 @@ suite('SQL ConnectionManagementService tests', () => {
 		let profileWithCredentials = await connectionManagementService.addSavedPassword(azureConnectionProfile);
 
 		// Then the returned profile has the account token set
-		assert.equal(profileWithCredentials.userName, username);
+		assert.equal(profileWithCredentials.userName, azureConnectionProfile.userName);
 		assert.equal(profileWithCredentials.options['azureAccountToken'], testToken);
 	});
 
@@ -836,11 +851,12 @@ suite('SQL ConnectionManagementService tests', () => {
 		let azureConnectionProfile = ConnectionProfile.fromIConnectionProfile(capabilitiesService, connectionProfile);
 		azureConnectionProfile.authenticationType = 'AzureMFA';
 		let username = 'testuser@microsoft.com';
-		azureConnectionProfile.userName = username;
+		azureConnectionProfile.azureAccount = username;
 		let servername = 'test-database.database.windows.net';
 		azureConnectionProfile.serverName = servername;
 		let azureTenantId = 'testTenant';
 		azureConnectionProfile.azureTenantId = azureTenantId;
+		let providerId = 'azure_PublicCloud';
 
 		// Set up the account management service to return a token for the given user
 		accountManagementService.setup(x => x.getAccountsForProvider(TypeMoq.It.isAny())).returns(providerId => Promise.resolve<azdata.Account[]>([
@@ -854,9 +870,24 @@ suite('SQL ConnectionManagementService tests', () => {
 				properties: undefined
 			}
 		]));
+
+		accountManagementService.setup(x => x.getAccounts()).returns(() => {
+			return Promise.resolve<azdata.Account[]>([
+				{
+					key: {
+						accountId: username,
+						providerId,
+					},
+					displayInfo: undefined,
+					isStale: false,
+					properties: undefined
+				}
+			]);
+		});
+
 		let testToken = 'testToken';
 		let returnedTokens = {};
-		returnedTokens['azurePublicCloud'] = { token: 'badToken' };
+		returnedTokens['azure_publicCloud'] = { token: 'badToken' };
 		returnedTokens[azureTenantId] = { token: testToken };
 		accountManagementService.setup(x => x.getSecurityToken(TypeMoq.It.isAny(), TypeMoq.It.isAny())).returns(() => Promise.resolve(returnedTokens));
 		connectionStore.setup(x => x.addSavedPassword(TypeMoq.It.is(profile => profile.authenticationType === 'AzureMFA'))).returns(profile => Promise.resolve({
@@ -868,7 +899,7 @@ suite('SQL ConnectionManagementService tests', () => {
 		let profileWithCredentials = await connectionManagementService.addSavedPassword(azureConnectionProfile);
 
 		// Then the returned profile has the account token set corresponding to the requested tenant
-		assert.equal(profileWithCredentials.userName, username);
+		assert.equal(profileWithCredentials.userName, azureConnectionProfile.userName);
 		assert.equal(profileWithCredentials.options['azureAccountToken'], testToken);
 	});