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

extensions/mssql/config.json

@@ -1,6 +1,6 @@
 {
 	"downloadUrl": "https://github.com/Microsoft/sqltoolsservice/releases/download/v{#version#}/microsoft.sqltools.servicelayer-{#fileName#}",
-	"version": "2.0.0-release.51",
+	"version": "2.0.0-release.53",
 	"downloadFileNames": {
 		"Windows_86": "win-x86-netcoreapp2.2.zip",
 		"Windows_64": "win-x64-netcoreapp2.2.zip",

src/sql/platform/accounts/common/accountPickerViewModel.ts

@@ -20,7 +20,7 @@ export class AccountPickerViewModel {
 	public selectedAccount: azdata.Account | undefined;
 
 	constructor(
-		private _providerId: string,
+		_providerId: string,
 		@IAccountManagementService private _accountManagementService: IAccountManagementService
 	) {
 		// Create our event emitters
@@ -37,7 +37,7 @@ export class AccountPickerViewModel {
 	 */
 	public initialize(): Thenable<azdata.Account[]> {
 		// Load a baseline of the accounts for the provider
-		return this._accountManagementService.getAccountsForProvider(this._providerId)
+		return this._accountManagementService.getAccounts()
 			.then(undefined, () => {
 				// In the event we failed to lookup accounts for the provider, just send
 				// back an empty collection

src/sql/platform/accounts/common/accountStore.ts

@@ -47,6 +47,7 @@ export default class AccountStore implements IAccountStore {
 		});
 	}
 
+
 	public remove(key: azdata.AccountKey): Thenable<boolean> {
 		return this.doOperation(() => {
 			return this.readFromMemento()

src/sql/platform/accounts/common/interfaces.ts

@@ -20,8 +20,10 @@ export interface IAccountManagementService {
 	addAccount(providerId: string): Thenable<void>;
 	getAccountProviderMetadata(): Thenable<azdata.AccountProviderMetadata[]>;
 	getAccountsForProvider(providerId: string): Thenable<azdata.Account[]>;
+	getAccounts(): Thenable<azdata.Account[]>;
 	getSecurityToken(account: azdata.Account, resource: azdata.AzureResource): Thenable<{ [key: string]: { token: string } }>;
 	removeAccount(accountKey: azdata.AccountKey): Thenable<boolean>;
+	removeAccounts(): Thenable<boolean>;
 	refreshAccount(account: azdata.Account): Thenable<azdata.Account>;
 
 	// UI METHODS //////////////////////////////////////////////////////////

src/sql/platform/accounts/test/common/accountPickerViewModel.test.ts

@@ -85,7 +85,7 @@ suite('Account picker view model tests', () => {
 				evUpdateAccounts.assertNotFired();
 
 				// ... The account management service should have been called
-				mockAccountManagementService.verify(x => x.getAccountsForProvider(TypeMoq.It.isAny()), TypeMoq.Times.once());
+				mockAccountManagementService.verify(x => x.getAccounts(), TypeMoq.Times.once());
 
 				// ... The results that were returned should be an array of account
 				assert.ok(Array.isArray(results));
@@ -108,7 +108,7 @@ suite('Account picker view model tests', () => {
 				evUpdateAccounts.assertNotFired();
 
 				// ... The account management service should have been called
-				mockAccountManagementService.verify(x => x.getAccountsForProvider(TypeMoq.It.isAny()), TypeMoq.Times.once());
+				mockAccountManagementService.verify(x => x.getAccounts(), TypeMoq.Times.once());
 
 				// ... The results should be an empty array
 				assert.ok(Array.isArray(result));
@@ -124,6 +124,8 @@ function getMockAccountManagementService(resolveProviders: boolean, resolveAccou
 		.returns(() => resolveProviders ? Promise.resolve(providers) : Promise.reject(null).then());
 	mockAccountManagementService.setup(x => x.getAccountsForProvider(TypeMoq.It.isAny()))
 		.returns(() => resolveAccounts ? Promise.resolve(accounts) : Promise.reject(null).then());
+	mockAccountManagementService.setup(x => x.getAccounts())
+		.returns(() => resolveAccounts ? Promise.resolve(accounts) : Promise.reject(null).then());
 
 	mockAccountManagementService.setup(x => x.updateAccountListEvent)
 		.returns(() => mockUpdateAccountEmitter.event);

src/sql/platform/accounts/test/common/testAccountManagementService.ts

@@ -43,6 +43,10 @@ export class TestAccountManagementService implements IAccountManagementService {
 		return Promise.resolve([]);
 	}
 
+	getAccounts(): Thenable<azdata.Account[]> {
+		return Promise.resolve([]);
+	}
+
 	getAccountsForProvider(providerId: string): Thenable<azdata.Account[]> {
 		return Promise.resolve([]);
 	}
@@ -55,6 +59,10 @@ export class TestAccountManagementService implements IAccountManagementService {
 		throw new Error('Method not implemented');
 	}
 
+	removeAccounts(): Thenable<boolean> {
+		throw new Error('Method not implemented');
+	}
+
 	refreshAccount(account: azdata.Account): Thenable<azdata.Account> {
 		throw new Error('Method not implemented');
 	}

src/sql/workbench/browser/actions.contribution.ts

@@ -11,6 +11,7 @@ import { Registry } from 'vs/platform/registry/common/platform';
 import { IConfigurationRegistry, Extensions as ConfigExtensions } from 'vs/platform/configuration/common/configurationRegistry';
 
 new Actions.ConfigureDashboardAction().registerTask();
+new Actions.ClearSavedAccountsAction().registerTask();
 
 Registry.as<IConfigurationRegistry>(ConfigExtensions.Configuration).registerConfiguration({
 	'id': 'previewFeatures',

src/sql/workbench/browser/actions.ts

@@ -17,6 +17,8 @@ import { ServicesAccessor } from 'vs/platform/instantiation/common/instantiation
 import { IInsightsConfig } from 'sql/platform/dashboard/browser/insightRegistry';
 import { IOpenerService } from 'vs/platform/opener/common/opener';
 import { URI } from 'vs/base/common/uri';
+import { IAccountManagementService } from 'sql/platform/accounts/common/interfaces';
+import { ILogService } from 'vs/platform/log/common/log';
 
 export interface BaseActionContext {
 	object?: ObjectMetadata;
@@ -88,3 +90,25 @@ export class ConfigureDashboardAction extends Task {
 		accessor.get(IOpenerService).open(URI.parse(ConfigureDashboardAction.configHelpUri));
 	}
 }
+
+export class ClearSavedAccountsAction extends Task {
+	public static readonly ID = 'clearSavedAccounts';
+	public static readonly LABEL = nls.localize('clearSavedAccounts', "Clear all saved accounts");
+
+	constructor() {
+		super({
+			id: ClearSavedAccountsAction.ID,
+			title: ClearSavedAccountsAction.LABEL,
+			iconPath: undefined
+		});
+	}
+
+	async runTask(accessor: ServicesAccessor): Promise<void> {
+		const logService = accessor.get(ILogService);
+		try {
+			await accessor.get(IAccountManagementService).removeAccounts();
+		} catch (ex) {
+			logService.error(ex);
+		}
+	}
+}

src/sql/workbench/contrib/accounts/test/browser/accountDialogController.test.ts

@@ -86,7 +86,7 @@ function createInstantiationService(addAccountFailureEmitter?: Emitter<string>):
 		.returns(() => undefined);
 
 	// Create a mock account dialog
-	let accountDialog = new AccountDialog(undefined!, undefined!, instantiationService.object, undefined!, undefined!, undefined!, undefined!, new MockContextKeyService(), undefined!, undefined!, undefined!, undefined!, undefined!, undefined!);
+	let accountDialog = new AccountDialog(undefined!, undefined!, instantiationService.object, undefined!, undefined!, undefined!, undefined!, new MockContextKeyService(), undefined!, undefined!, undefined!, undefined!, undefined!, undefined!, undefined!, undefined!);
 	let mockAccountDialog = TypeMoq.Mock.ofInstance(accountDialog);
 	mockAccountDialog.setup(x => x.onAddAccountErrorEvent)
 		.returns(() => { return addAccountFailureEmitter ? addAccountFailureEmitter.event : mockEvent.event; });

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(() => {
+		this._register(this._addAccountButton.onDidClick(async () => {
-			(<IProviderViewUiComponent>values(this._providerViewsMap)[0]).addAccountAction.run();
+			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.doWithProvider(accountKey.providerId, async provider => {
-			return this._accountStore.remove(accountKey)
+			const result = await this._accountStore.remove(accountKey);
-				.then(result => {
+			await provider.provider.clear(accountKey);
-					provider.provider.clear(accountKey);
+			if (!result) {
-					return result;
+				return result;
-				})
+			}
-				.then(result => {
-					if (!result) {
-						return result;
-					}
 
-					let indexToRemove: number = firstIndex(provider.accounts, account => {
+			let indexToRemove: number = firstIndex(provider.accounts, account => {
-						return account.key.accountId === accountKey.accountId;
+				return account.key.accountId === accountKey.accountId;
-					});
+			});
 
-					if (indexToRemove >= 0) {
+			if (indexToRemove >= 0) {
-						provider.accounts.splice(indexToRemove, 1);
+				provider.accounts.splice(indexToRemove, 1);
-						self.fireAccountListUpdate(provider, false);
+				this.fireAccountListUpdate(provider, false);
-					}
+			}
-					return result;
+			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);
 	});